conftest.py

import logging
import os
import threading
from datetime import datetime, timezone
from io import BytesIO
from typing import Any
from unittest.mock import patch
from pyspark.sql import DataFrame, functions as F
from pyspark.sql.types import ArrayType, StructType
from pyspark.testing.utils import assertDataFrameEqual

import pytest
from databricks.sdk.service.workspace import ImportFormat
from databricks.labs.blueprint.installation import Installation
from databricks.labs.pytester.fixtures.baseline import factory
from databricks.labs.dqx.checks_serializer import serialize_checks
from databricks.labs.dqx.rule_fingerprint import compute_rule_fingerprint, compute_rule_set_fingerprint_by_metadata
from databricks.labs.dqx.checks_storage import InstallationChecksStorageHandler
from databricks.labs.dqx.config import InputConfig, OutputConfig, InstallationChecksStorageConfig, ExtraParams
from databricks.labs.dqx.engine import DQEngine
from databricks.labs.dqx.installer.mixins import InstallationMixin
from databricks.labs.dqx.rule import DQRule
from databricks.labs.dqx.schema import dq_result_schema
from databricks.sdk.service.compute import DataSecurityMode, Kind

from tests.constants import TEST_CATALOG


logging.getLogger("tests").setLevel("DEBUG")
logging.getLogger("databricks.labs.dqx").setLevel("DEBUG")

logger = logging.getLogger(__name__)


REPORTING_COLUMNS = f", _errors: {dq_result_schema.simpleString()}, _warnings: {dq_result_schema.simpleString()}"
RUN_TIME = datetime(2025, 1, 1, 0, 0, 0, 0, tzinfo=timezone.utc)
RUN_ID = "2f9120cf-e9f2-446a-8278-12d508b00639"
EXTRA_PARAMS = ExtraParams(run_time_overwrite=RUN_TIME.isoformat(), run_id_overwrite=RUN_ID)


FINGERPRINT_FIELDS = ["rule_fingerprint", "rule_set_fingerprint"]


def _strip_fingerprints_from_result_column(df: DataFrame, col_name: str) -> DataFrame:
    """Remove fingerprint fields from a result array column (_errors or _warnings)."""

    schema = df.schema
    if col_name not in df.columns:
        return df

    field = schema[col_name]
    if not isinstance(field.dataType, ArrayType) or not isinstance(field.dataType.elementType, StructType):
        return df

    struct_fields = field.dataType.elementType.fieldNames()
    keep_fields = [f for f in struct_fields if f not in FINGERPRINT_FIELDS]

    return df.withColumn(
        col_name,
        F.when(
            F.col(col_name).isNotNull(),
            F.transform(
                F.col(col_name),
                lambda x: F.struct(*[x[f].alias(f) for f in keep_fields]),
            ),
        ),
    )


_FINGERPRINT_COLUMNS = ("_errors", "dq_errors", "_warnings", "dq_warnings")


def _strip_all_fingerprints(df: DataFrame) -> DataFrame:
    """Remove fingerprint fields from all known result columns."""
    for col in _FINGERPRINT_COLUMNS:
        df = _strip_fingerprints_from_result_column(df, col)
    return df


def assert_df_equality_ignore_fingerprints(
    df1: DataFrame,
    df2: DataFrame,
    **kwargs,
):
    """Assert DataFrame equality after stripping fingerprint fields from result columns."""
    df1_clean = _strip_all_fingerprints(df1)
    df2_clean = _strip_all_fingerprints(df2)

    check_row_order = not kwargs.pop("ignore_row_order", False)
    ignore_column_order = kwargs.pop("ignore_column_order", False)
    kwargs.pop("ignore_nullable", None)
    assertDataFrameEqual(
        df1_clean, df2_clean, checkRowOrder=check_row_order, ignoreColumnOrder=ignore_column_order, **kwargs
    )


def build_quality_violation(
    name: str,
    message: str,
    columns: list[str] | None,
    *,
    function: str = "is_not_null_and_not_empty",
    filter_expr: str | None = None,
    user_metadata: dict | None = None,
    rule_fingerprint: str | None = None,
    rule_set_fingerprint: str | None = None,
) -> dict[str, Any]:
    """Helper for constructing expected violation entries with shared metadata."""

    return {
        "name": name,
        "message": message,
        "columns": columns,
        "filter": filter_expr,
        "function": function,
        "run_time": RUN_TIME,
        "run_id": RUN_ID,
        "user_metadata": user_metadata or {},
        "rule_fingerprint": rule_fingerprint,
        "rule_set_fingerprint": rule_set_fingerprint,
    }


def build_skipped_violation(
    name: str,
    message: str,
    columns: list[str] | None,
    *,
    function: str = "is_not_null",
    filter_expr: str | None = None,
    user_metadata: dict | None = None,
) -> dict[str, Any]:
    """Helper for constructing expected entries for checks that were skipped during evaluation."""

    return {
        "name": name,
        "message": message,
        "columns": columns,
        "filter": filter_expr,
        "function": function,
        "run_time": RUN_TIME,
        "run_id": RUN_ID,
        "user_metadata": user_metadata or {},
        "skipped": True,
    }


def assert_check_and_split_results(
    checked: DataFrame,
    good_df: DataFrame,
    bad_df: DataFrame,
    expected: DataFrame,
    columns: list[str],
    *,
    ignore_column_order: bool = False,
) -> None:
    """Assert equality of checked, bad, and good DataFrames against expected results."""
    assert_df_equality_ignore_fingerprints(
        checked, expected, ignore_nullable=True, ignore_column_order=ignore_column_order
    )
    assert_df_equality_ignore_fingerprints(
        bad_df, expected.where(F.col("_errors").isNotNull() | F.col("_warnings").isNotNull()), ignore_nullable=True
    )
    assert_df_equality_ignore_fingerprints(
        good_df, expected.where(F.col("_errors").isNull()).select(*columns), ignore_nullable=True
    )


class SparkKeepAlive:
    """
    Utility to keep Spark Connect session alive during long-running operations.

    Runs lightweight queries periodically to prevent INACTIVITY_TIMEOUT.
    """

    def __init__(self, spark, interval_seconds=60, join_timeout=5):
        """
        Args:
            spark: SparkSession to keep alive
            interval_seconds: How often to run keep-alive query (default 60s)
            join_timeout: Seconds to wait for the background thread to stop (default 5s)
        """
        self.spark = spark
        self.interval = interval_seconds
        self._join_timeout = join_timeout
        self._stop_flag = threading.Event()
        self._thread = None

    def _keep_alive_loop(self):
        """Background thread that runs periodic queries."""
        logger.debug(f"SparkKeepAlive started (interval={self.interval}s)")

        while not self._stop_flag.is_set():
            try:
                # Lightweight query to keep session active
                self.spark.sql("SELECT 1").collect()
                logger.debug("SparkKeepAlive: sent keep-alive query")
            except Exception as e:
                logger.warning(f"SparkKeepAlive: query failed: {e}")

            # Wait for interval or stop signal
            self._stop_flag.wait(self.interval)

        logger.debug("SparkKeepAlive stopped")

    def start(self):
        """Start the keep-alive background thread."""
        if self._thread and self._thread.is_alive():
            logger.warning("SparkKeepAlive already running")
            return

        self._stop_flag.clear()
        self._thread = threading.Thread(target=self._keep_alive_loop, daemon=True)
        self._thread.start()

    def stop(self):
        """Stop the keep-alive background thread."""
        if self._thread and self._thread.is_alive():
            self._stop_flag.set()
            self._thread.join(timeout=self._join_timeout)
            if self._thread.is_alive():
                logger.warning(f"SparkKeepAlive thread did not stop within {self._join_timeout}s")


@pytest.fixture
def spark_keep_alive(spark):
    """
    Fixture that provides a SparkKeepAlive utility to prevent session timeouts during long-running tests (e.g. workflows).
    """
    keep_alive = SparkKeepAlive(spark, interval_seconds=60)
    keep_alive.start()
    yield keep_alive
    keep_alive.stop()


@pytest.fixture
def webbrowser_open():
    with patch("webbrowser.open") as mock_open:
        yield mock_open


@pytest.fixture
def setup_workflows(ws, spark, installation_ctx, make_schema, make_table, make_random):
    """
    Set up the workflows with serverless cluster for the tests in the workspace.
    """

    if os.getenv("DATABRICKS_SERVERLESS_COMPUTE_ID"):
        pytest.skip()

    def create(_spark, **kwargs):
        installation_ctx.installation_service.run()

        quarantine = kwargs.get("quarantine", False)
        quarantine_only = kwargs.get("quarantine_only", False)

        checks_location = None
        if "checks" in kwargs and kwargs["checks"]:
            checks_location = _setup_quality_checks(installation_ctx, _spark, ws)

        run_config = _setup_workflows_deps(
            installation_ctx,
            make_schema,
            make_table,
            make_random,
            checks_location,
            quarantine,
            quarantine_only=quarantine_only,
        )
        return installation_ctx, run_config

    def delete(resource) -> None:
        ctx, run_config = resource
        checks_location = f"{ctx.installation.install_folder()}/{run_config.checks_location}"
        ws.workspace.delete(checks_location)

    yield from factory("workflows", lambda **kw: create(spark, **kw), delete)


@pytest.fixture
def setup_serverless_workflows(ws, spark, serverless_installation_ctx, make_schema, make_table, make_random):
    """
    Set up the workflows with serverless cluster for the tests in the workspace.
    """

    if not os.getenv("DATABRICKS_SERVERLESS_COMPUTE_ID"):
        pytest.skip()

    def create(_spark, **kwargs):
        serverless_installation_ctx.installation_service.run()

        quarantine = kwargs.get("quarantine", False)
        quarantine_only = kwargs.get("quarantine_only", False)

        checks_location = None
        if "checks" in kwargs and kwargs["checks"]:
            checks_location = _setup_quality_checks(serverless_installation_ctx, _spark, ws)

        run_config = _setup_workflows_deps(
            serverless_installation_ctx,
            make_schema,
            make_table,
            make_random,
            checks_location,
            quarantine,
            is_streaming=kwargs.get("is_streaming", False),
            quarantine_only=quarantine_only,
        )
        return serverless_installation_ctx, run_config

    def delete(resource) -> None:
        ctx, run_config = resource
        checks_location = f"{ctx.installation.install_folder()}/{run_config.checks_location}"
        ws.workspace.delete(checks_location)

    yield from factory("workflows", lambda **kw: create(spark, **kw), delete)


@pytest.fixture
def setup_workflows_with_metrics(ws, spark, installation_ctx, make_schema, make_table, make_cluster, make_random):
    """Set up workflows with metrics configuration for testing."""

    if os.getenv("DATABRICKS_SERVERLESS_COMPUTE_ID"):
        pytest.skip()

    def create(_spark, **kwargs):
        cluster = make_cluster(
            single_node=True,
            kind=Kind.CLASSIC_PREVIEW,
            data_security_mode=DataSecurityMode.DATA_SECURITY_MODE_DEDICATED,
        )
        cluster_id = cluster.cluster_id
        installation_ctx.config.serverless_clusters = False
        installation_ctx.config.profiler_override_clusters["default"] = cluster_id
        installation_ctx.config.quality_checker_override_clusters["default"] = cluster_id
        installation_ctx.config.e2e_override_clusters["default"] = cluster_id
        installation_ctx.installation_service.run()

        quarantine = False
        if "quarantine" in kwargs and kwargs["quarantine"]:
            quarantine = True

        checks_location = _setup_quality_checks(installation_ctx, _spark, ws)

        run_config = _setup_workflows_deps(
            installation_ctx,
            make_schema,
            make_table,
            make_random,
            checks_location,
            quarantine,
            is_streaming=kwargs.get("is_streaming", False),
            is_continuous_streaming=kwargs.get("is_continuous_streaming", False),
        )

        config = installation_ctx.config
        run_config = config.get_run_config()

        catalog_name = TEST_CATALOG
        schema_name = run_config.output_config.location.split(".")[1]
        metrics_table_name = f"{catalog_name}.{schema_name}.metrics_{make_random(6).lower()}"
        run_config.metrics_config = OutputConfig(location=metrics_table_name)

        custom_metrics = kwargs.get("custom_metrics")
        if custom_metrics:
            config.custom_metrics = custom_metrics

        installation_ctx.installation.save(config)

        return installation_ctx, run_config

    def delete(resource):
        ctx, run_config = resource
        checks_location = f"{ctx.installation.install_folder()}/{run_config.checks_location}"
        try:
            ws.workspace.delete(checks_location)
        except Exception:
            pass

    yield from factory("workflows_with_metrics", lambda **kw: create(spark, **kw), delete)


@pytest.fixture
def setup_workflows_with_custom_folder(
    ws, spark, installation_ctx_custom_install_folder, make_schema, make_table, make_random
):
    """
    Set up the workflows with installation in the custom install folder.
    """

    if os.getenv("DATABRICKS_SERVERLESS_COMPUTE_ID"):
        pytest.skip()

    def create(_spark, **kwargs):
        installation_ctx_custom_install_folder.installation_service.run()

        quarantine = kwargs.get("quarantine", False)
        quarantine_only = kwargs.get("quarantine_only", False)

        checks_location = None
        if "checks" in kwargs and kwargs["checks"]:
            checks_location = _setup_quality_checks(installation_ctx_custom_install_folder, _spark, ws)

        run_config = _setup_workflows_deps(
            installation_ctx_custom_install_folder,
            make_schema,
            make_table,
            make_random,
            checks_location,
            quarantine,
            quarantine_only=quarantine_only,
        )
        return installation_ctx_custom_install_folder, run_config

    def delete(resource) -> None:
        ctx, run_config = resource
        checks_location = f"{ctx.installation.install_folder()}/{run_config.checks_location}"
        ws.workspace.delete(checks_location)

    yield from factory("workflows", lambda **kw: create(spark, **kw), delete)


class TestInstallationMixin(InstallationMixin):
    def get_my_username(self):
        return self._my_username

    def get_me(self):
        return self._me

    def get_installation(
        self, product_name: str, assume_user: bool = True, install_folder: str | None = None
    ) -> Installation:
        return self._get_installation(product_name, assume_user, install_folder)


def _setup_workflows_deps(
    ctx,
    make_schema,
    make_table,
    make_random,
    checks_location: str | None = None,
    quarantine: bool = False,
    is_streaming: bool = False,
    is_continuous_streaming: bool = False,
    quarantine_only: bool = False,
):
    # prepare test data
    catalog_name = TEST_CATALOG
    schema = make_schema(catalog_name=catalog_name)

    input_table = make_table(
        catalog_name=catalog_name,
        schema_name=schema.name,
        # sample data
        ctas="SELECT * FROM VALUES "
        "(1, 'a'), (2, 'b'), (3, NULL), (NULL, 'c'), (3, NULL), (1, 'a'), (6, 'a'), (2, 'c'), (4, 'a'), (5, 'd') "
        "AS data(id, name)",
    )

    # update input and output locations
    config = ctx.config
    config.extra_params = EXTRA_PARAMS

    run_config = config.get_run_config()
    run_config.input_config = InputConfig(
        location=input_table.full_name,
        options={"versionAsOf": "0"} if not is_streaming else {},
        is_streaming=is_streaming,
    )

    trigger: dict[str, Any] = {}
    if is_streaming:
        if is_continuous_streaming:
            trigger = {"processingTime": "60 seconds"}
        else:
            trigger = {"availableNow": True}

    # quarantine_only writes only invalid records (no output table); quarantine writes both.
    if quarantine_only:
        run_config.output_config = None
    else:
        output_table = f"{catalog_name}.{schema.name}.{make_random(10).lower()}"
        run_config.output_config = OutputConfig(
            location=output_table,
            trigger=trigger,
            options=({"checkpointLocation": f"/tmp/dqx_tests/{make_random(10)}_out_ckpt"} if is_streaming else {}),
        )

    if quarantine or quarantine_only:
        quarantine_table = f"{catalog_name}.{schema.name}.{make_random(10).lower()}_quarantine"
        run_config.quarantine_config = OutputConfig(
            location=quarantine_table,
            trigger=trigger,
            options=({"checkpointLocation": f"/tmp/dqx_tests/{make_random(10)}_qr_ckpt"} if is_streaming else {}),
        )

    if checks_location:
        run_config.checks_location = checks_location

    # ensure tests are deterministic
    run_config.profiler_config.sample_fraction = 1.0
    run_config.profiler_config.sample_seed = 100
    # relax null/empty thresholds so the profiler generates is_not_null / is_not_null_or_empty
    run_config.profiler_config.max_null_ratio = 0.25
    run_config.profiler_config.max_empty_ratio = 0.25

    ctx.installation.save(ctx.config)

    return run_config


@pytest.fixture
def expected_quality_checking_output(spark) -> DataFrame:
    return spark.createDataFrame(
        [
            [1, "a", None, None],
            [2, "b", None, None],
            [
                3,
                None,
                [
                    build_quality_violation(
                        "name_is_not_null_and_not_empty", "Column 'name' value is null or empty", ["name"]
                    )
                ],
                None,
            ],
            [
                None,
                "c",
                [
                    build_quality_violation(
                        "id_is_not_null", "Column 'id' value is null", ["id"], function="is_not_null"
                    )
                ],
                None,
            ],
            [
                3,
                None,
                [
                    build_quality_violation(
                        "name_is_not_null_and_not_empty", "Column 'name' value is null or empty", ["name"]
                    )
                ],
                None,
            ],
            [1, "a", None, None],
            [6, "a", None, None],
            [2, "c", None, None],
            [4, "a", None, None],
            [5, "d", None, None],
        ],
        f"id int, name string {REPORTING_COLUMNS}",
    )


WORKFLOW_CHECKS = [
    {
        "name": "id_is_not_null",
        "criticality": "error",
        "check": {"function": "is_not_null", "arguments": {"column": "id"}},
    },
    {
        "name": "name_is_not_null_and_not_empty",
        "criticality": "error",
        "check": {"function": "is_not_null_and_not_empty", "arguments": {"column": "name"}},
    },
]


def _setup_quality_checks(ctx, spark, ws):
    config = ctx.config
    checks_location = config.get_run_config().checks_location
    checks = WORKFLOW_CHECKS
    config = InstallationChecksStorageConfig(
        location=checks_location,
        product_name=ctx.installation.product(),
        install_folder=ctx.installation.install_folder(),
    )

    InstallationChecksStorageHandler(ws, spark).save(checks=checks, config=config)
    return checks_location


def setup_custom_check_func(ws, installation_ctx, custom_checks_funcs_location):
    content = '''from databricks.labs.dqx.check_funcs import make_condition, register_rule
from pyspark.sql import functions as F

@register_rule("row")
def not_ends_with_suffix(column: str, suffix: str):
    """
    Example of custom python row-level check function.
    """
    return make_condition(
        F.col(column).endswith(suffix), f"Column '{column}' ends with '{suffix}'", f"{column}_ends_with_{suffix}"
    )
'''
    if custom_checks_funcs_location.startswith("/Workspace/"):
        ws.workspace.upload(
            path=custom_checks_funcs_location, format=ImportFormat.AUTO, content=content.encode(), overwrite=True
        )
    elif custom_checks_funcs_location.startswith("/Volumes/"):
        binary_data = BytesIO(content.encode("utf-8"))
        ws.files.upload(custom_checks_funcs_location, binary_data, overwrite=True)
    else:  # relative workspace path
        installation_dir = installation_ctx.installation.install_folder()
        ws.workspace.upload(
            path=f"{installation_dir}/{custom_checks_funcs_location}",
            format=ImportFormat.AUTO,
            content=content.encode(),
            overwrite=True,
        )

    config = installation_ctx.config
    run_config = config.get_run_config()
    run_config.custom_check_functions = {"not_ends_with_suffix": custom_checks_funcs_location}
    installation_ctx.installation.save(config)


def contains_expected_workflows(workflows, state):
    for workflow in workflows:
        if all(item in workflow.items() for item in state.items()):
            return True
    return False


def assert_quarantine_and_output_dfs(ws, spark, expected_output, output_config, quarantine_config):
    dq_engine = DQEngine(ws, spark)
    expected_output_df = dq_engine.get_valid(expected_output)
    expected_quarantine_df = dq_engine.get_invalid(expected_output)

    output_df = spark.table(output_config.location)
    assert_df_equality_ignore_fingerprints(output_df, expected_output_df, ignore_nullable=True)

    quarantine_df = spark.table(quarantine_config.location)
    assert_df_equality_ignore_fingerprints(quarantine_df, expected_quarantine_df, ignore_nullable=True)


def assert_quarantine_df_only(ws, spark, expected_output, quarantine_config):
    dq_engine = DQEngine(ws, spark)
    expected_quarantine_df = dq_engine.get_invalid(expected_output)

    quarantine_df = spark.table(quarantine_config.location)
    assert_df_equality_ignore_fingerprints(quarantine_df, expected_quarantine_df, ignore_nullable=True)


def assert_output_df(spark, expected_output, output_config):
    checked_df = spark.table(output_config.location)
    assert_df_equality_ignore_fingerprints(checked_df, expected_output, ignore_nullable=True)


def generate_checks_with_rule_and_set_fingerprint_from_rules(rules: list[DQRule]) -> list[dict]:
    """Generate check dicts with rule_fingerprint and rule_set_fingerprint from DQRule instances."""
    checks_dict = serialize_checks(rules)
    rule_set_fingerprint = compute_rule_set_fingerprint_by_metadata(checks_dict)
    return [
        {**check, "rule_fingerprint": compute_rule_fingerprint(check), "rule_set_fingerprint": rule_set_fingerprint}
        for check in checks_dict
    ]


def generate_checks_with_rule_and_set_fingerprint_from_dicts(checks: list[dict]) -> list[dict]:
    """Generate check dicts with rule_fingerprint and rule_set_fingerprint from check dicts."""
    checks_dict = [dict(check) for check in checks]
    rule_set_fingerprint = compute_rule_set_fingerprint_by_metadata(checks_dict)
    return [
        {**check, "rule_fingerprint": compute_rule_fingerprint(check), "rule_set_fingerprint": rule_set_fingerprint}
        for check in checks_dict
    ]


def get_rule_fingerprint_from_checks(
    versioning_rules_checks: list[dict] | None, check_name: str, criticality: str
) -> str | None:
    """Extract the rule_fingerprint for the first check matching name and criticality."""
    if not versioning_rules_checks:
        return None
    match = next(
        (c for c in versioning_rules_checks if c.get("name") == check_name and c.get("criticality") == criticality),
        None,
    )
    return match.get("rule_fingerprint") if match else None


def get_rule_set_fingerprint_from_checks(versioning_rules_checks: list[dict] | None) -> str | None:
    """
    Helper function to extract the rule_set_fingerprint from the versioning rules checks
    based on the check name, function, criticality and column (if applicable).
    versioning_rules_checks: list of versioning rules checks
    check_name: name of the check
    """

    if not versioning_rules_checks:
        return None
    return versioning_rules_checks[0].get("rule_set_fingerprint", None)