diff --git a/docs/dqx/docs/guide/data_profiling.mdx b/docs/dqx/docs/guide/data_profiling.mdx index ad5b00878..a410ecbcc 100644 --- a/docs/dqx/docs/guide/data_profiling.mdx +++ b/docs/dqx/docs/guide/data_profiling.mdx @@ -262,7 +262,7 @@ When running the profiler workflow using Databricks API or UI, you have the same - If the `checks_location` in the run config points to a table, the checks will be saved to that table. If the `checks_location` in the run config points to a file, file name is replaced with "<input_table>.yml". In addition, if the location is specified as a relative path, it is prefixed with the workspace installation folder. For example: - - If "checks_location=catalog.schema.table", the location will be resolved to "catalog.schema.table". + - If "checks_location=catalog.schema.table", the location will be resolved to "catalog.schema.table" or "database.schema.table" in case of using Lakebase to store checks. - If "checks_location=folder/checks.yml", the location will be resolved to "install_folder/folder/<input_table>.yml". - If "checks_location=/App/checks.yml", the location will be resolved to "/App/<input_table>.yml". - If "checks_location=/Volume/catalog/schema/folder/checks.yml", the location will be resolved to "/Volume/catalog/schema/folder/<input_table>.yml". diff --git a/docs/dqx/docs/guide/quality_checks_apply.mdx b/docs/dqx/docs/guide/quality_checks_apply.mdx index b98ec403c..d66d9fb69 100644 --- a/docs/dqx/docs/guide/quality_checks_apply.mdx +++ b/docs/dqx/docs/guide/quality_checks_apply.mdx @@ -575,7 +575,7 @@ When running the quality checker workflow using Databricks API or UI, you have t - If the `checks_location` in the run config points to a table, the checks will be directly loaded from that table. If the `checks_location` in the run config points to a file, file name is replaced with "<input_table>.yml". In addition, if the location is specified as a relative path, it is prefixed with the workspace installation folder. For example: - - If "checks_location=catalog.schema.table", the location will be resolved to "catalog.schema.table". + - If "checks_location=catalog.schema.table", the location will be resolved to "catalog.schema.table" or "database.schema.table" in case of using Lakebase to store checks. - If "checks_location=folder/checks.yml", the location will be resolved to "install_folder/folder/<input_table>.yml". - If "checks_location=/App/checks.yml", the location will be resolved to "/App/<input_table>.yml". - If "checks_location=/Volume/catalog/schema/folder/checks.yml", the location will be resolved to "/Volume/catalog/schema/folder/<input_table>.yml". @@ -690,7 +690,7 @@ When running the e2e workflow using Databricks API or UI, you have the same exec - If the `checks_location` in the run config points to a table, the checks will be directly loaded from that table. If the `checks_location` in the run config points to a file, file name is replaced with <input_table>.yml. In addition, if the location is specified as a relative path, it is prefixed with the workspace installation folder. For example: - - If "checks_location=catalog.schema.table", the location will be resolved to "catalog.schema.table". + - If "checks_location=catalog.schema.table", the location will be resolved to "catalog.schema.table" or "database.schema.table" in case of using Lakebase to store checks. - If "checks_location=folder/checks.yml", the location will be resolved to "install_folder/folder/<input_table>.yml". - If "checks_location=/App/checks.yml", the location will be resolved to "/App/<input_table>.yml". - If "checks_location=/Volume/catalog/schema/folder/checks.yml", the location will be resolved to "/Volume/catalog/schema/folder/<input_table>.yml". diff --git a/docs/dqx/docs/guide/quality_checks_storage.mdx b/docs/dqx/docs/guide/quality_checks_storage.mdx index 44018aff5..ab32aaebe 100644 --- a/docs/dqx/docs/guide/quality_checks_storage.mdx +++ b/docs/dqx/docs/guide/quality_checks_storage.mdx @@ -22,12 +22,20 @@ Saving and loading methods accept a storage backend configuration as input. The * `mode`: (optional) write mode for saving checks (`overwrite` or `append`, default is `overwrite`). The `overwrite` mode will only replace checks for the specific run config and not all checks in the table. * `VolumeFileChecksStorageConfig`: Unity Catalog Volume (JSON/YAML file). Containing fields: * `location`: Unity Catalog Volume file path (JSON or YAML). +* `LakebaseChecksStorageConfig`: Lakebase table. Containing fields: + * `instance_name`: name of the Lakebase instance, e.g., "my-instance". + * `user`: user to connect to the Lakebase instance, e.g., "user@domain.com" or Databricks service principal client ID. + * `location`: fully-qualified table name in the format "database.schema.table". + * `port`: (optional) port on which to connect to the Lakebase instance (use 5432 if not provided). + * `run_config_name`: (optional) run configuration name to load (use "default" if not provided). + * `mode`: (optional) write mode for saving checks (`overwrite` or `append`, default is `overwrite`). The `overwrite` mode will only replace checks for the specific run config and not all checks in the table. * `InstallationChecksStorageConfig`: installation-managed location from the run config, ignores location and infers it from `checks_location` in the run config. Containing fields: * `location` (optional): automatically set based on the `checks_location` field from the run configuration. * `install_folder`: (optional) installation folder where DQX is installed, only required when custom installation folder is used. * `run_config_name` (optional) - run configuration name to load (it can be any string), e.g. input table or job name (use "default" if not provided). * `product_name`: (optional) name of the product (use "dqx" if not provided). * `assume_user`: (optional) if True, assume user installation, otherwise global installation (skipped if `install_folder` is provided). + * the config inherits from the specific configs such as `WorkspaceFileChecksStorageConfig`, `TableChecksStorageConfig`, `VolumeFileChecksStorageConfig`, and `LakebaseChecksStorageConfig` so relevant fields from these specific configs can be provided (e.g. instance_name and user for lakebase). You can find details on how to define checks [here](/docs/guide/quality_checks_definition). @@ -49,7 +57,8 @@ If you create checks as a list of `DQRule` objects, you can convert them to meta WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, - VolumeFileChecksStorageConfig + VolumeFileChecksStorageConfig, + LakebaseChecksStorageConfig, ) from databricks.sdk import WorkspaceClient @@ -81,6 +90,9 @@ If you create checks as a list of `DQRule` objects, you can convert them to meta # save checks as a YAML in a Unity Catalog Volume location (overwrite the file) dq_engine.save_checks(checks, config=VolumeFileChecksStorageConfig(location="/Volumes/dq/config/checks_volume/App1/checks.yml")) + # save checks as a Lakebase table using a Databricks service principal + dq_engine.save_checks(checks, config=LakebaseChecksStorageConfig(instance_name="my-instance", user="00000000-0000-0000-0000-000000000000", location="dqx.config.checks")) + # save checks as a YAML file or table defined in 'checks_location' of the run config # only works if DQX is installed in the workspace # the run config name can be any string, e.g. input table or job name @@ -195,7 +207,8 @@ If you create checks as a list of DQRule objects, you can convert them using the WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, - VolumeFileChecksStorageConfig + VolumeFileChecksStorageConfig, + LakebaseChecksStorageConfig, ) from databricks.sdk import WorkspaceClient @@ -217,6 +230,9 @@ If you create checks as a list of DQRule objects, you can convert them using the # load checks from a Unity Catalog Volume checks: list[dict] = dq_engine.load_checks(config=VolumeFileChecksStorageConfig(location="/Volumes/dq/config/checks_volume/App1/checks.yml")) + # load checks from a Lakebase table using a Databricks service principal + checks: list[dict] = dq_engine.load_checks(config=LakebaseChecksStorageConfig(instance_name="my-instance", user="00000000-0000-0000-0000-000000000000", location="dqx.config.checks")) + # load checks from a file or table defined in the run config ('checks_location' field) # only works if DQX is installed in the workspace # the run config name is a string (e.g. input table or job name) diff --git a/docs/dqx/docs/installation.mdx b/docs/dqx/docs/installation.mdx index 8894f868a..31dbf7c05 100644 --- a/docs/dqx/docs/installation.mdx +++ b/docs/dqx/docs/installation.mdx @@ -212,6 +212,12 @@ run_configs: # <- list of run configurations, each run co checks_location: iot_checks.yml # <- Quality rules (checks) can be stored in a table or defined in JSON or YAML files, located at absolute or relative path within the installation folder or volume file path. + # if wanting to store checks in lakebase table + # checks_location: dqx.config.checks # <- fully qualified Lakebase table for storing quality rules (checks) + # lakebase_instance_name: my-lakebase-instance # <- the name of the lakebase instance to use for storing checks + # lakebase_user: 00000000-0000-0000-0000-000000000000 # <- the user to connect to the lakebase, e.g., user@domain.com or a Databricks service principal client ID + # lakebase_port: 5432 # <- optional port to connect to Lakebase, default is 5432 + custom_check_functions: # <- optional mapping of custom check function name to Python file (module) containing check function definition my_func: custom_checks/my_funcs.py # relative workspace path (installation folder prefix applied) my_other: /Workspace/Shared/MyApp/my_funcs.py # or absolute workspace path @@ -236,7 +242,7 @@ run_configs: # <- list of run configurations, each run co warehouse_id: your-warehouse-id # <- warehouse id for refreshing dashboard -- name: another_run_config # <- unique name of the run config +- name: another_run_config # <- unique name of the run config ... ``` diff --git a/pyproject.toml b/pyproject.toml index 95649f2bb..1a8143dd5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,8 +29,9 @@ classifiers = [ "Topic :: Utilities", ] dependencies = ["databricks-labs-blueprint>=0.9.1,<0.10", - "databricks-sdk~=0.57", + "databricks-sdk~=0.67", "databricks-labs-lsql>=0.5,<=0.16", + "sqlalchemy>=1.4,<3.0", ] [project.optional-dependencies] @@ -101,6 +102,7 @@ dependencies = [ "dbldatagen~=0.4.0", "pyparsing~=3.2.3", "jmespath~=1.0.1", + "psycopg2-binary~=2.9.10", ] python="3.12" # must match the version required by databricks-connect and python version on the test clusters @@ -111,7 +113,7 @@ path = ".venv" [tool.hatch.envs.default.scripts] test = "pytest tests/unit/ -n 10 --cov --cov-report=xml:coverage-unit.xml --timeout 30 --durations 20" coverage = "pytest tests/ -n 10 --cov --cov-report=html --timeout 600 --durations 20" -integration = "pytest tests/integration/ -n 10 --cov --cov-report=xml --timeout 1200 --durations 20" +integration = "pytest tests/integration/ -n 5 --cov --cov-report=xml --timeout 1200 --durations 20" e2e = "pytest tests/e2e/ -n 10 --cov --cov-report=xml --timeout 600 --durations 20" perf = "pytest tests/perf/ -n 10 --cov --cov-report=xml --timeout 600 --durations 20" fmt = ["black .", diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 7541d97ab..5fa940d95 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -1,10 +1,27 @@ import json import logging +import uuid import os from io import StringIO, BytesIO from pathlib import Path from abc import ABC, abstractmethod -from typing import Generic, TypeVar +from typing import Generic, TypeVar, NoReturn +from sqlalchemy import ( + Engine, + create_engine, + MetaData, + Table, + Column, + String, + Text, + insert, + select, + delete, + null, +) +from sqlalchemy.schema import CreateSchema +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.exc import DatabaseError, ProgrammingError, OperationalError, IntegrityError import yaml @@ -14,11 +31,13 @@ from databricks.labs.dqx.config import ( TableChecksStorageConfig, + LakebaseChecksStorageConfig, FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, BaseChecksStorageConfig, VolumeFileChecksStorageConfig, + RunConfig, ) from databricks.labs.dqx.errors import InvalidCheckError, InvalidConfigError, CheckDownloadError from databricks.sdk import WorkspaceClient @@ -110,6 +129,294 @@ def save(self, checks: list[dict], config: TableChecksStorageConfig) -> None: ) +class LakebaseChecksStorageHandler(ChecksStorageHandler[LakebaseChecksStorageConfig]): + """ + Handler for storing dq rules (checks) in a Lakebase table. + """ + + def __init__(self, ws: WorkspaceClient, spark: SparkSession, engine: Engine | None = None): + self.ws = ws + self.spark = spark + self.engine = engine + + def _get_connection_url(self, config: LakebaseChecksStorageConfig) -> str: + """ + Generate a Lakebase connection URL. + + Args: + config: Configuration for saving and loading checks to Lakebase. + + Returns: + Lakebase connection URL. + """ + if not config.instance_name: + raise InvalidConfigError("instance_name must be provided for Lakebase storage") + if not config.user: + raise InvalidConfigError("user must be provided for Lakebase storage") + if not config.location: + raise InvalidConfigError("location must be provided for Lakebase storage") + + instance = self.ws.database.get_database_instance(config.instance_name) + cred = self.ws.database.generate_database_credential( + request_id=str(uuid.uuid4()), instance_names=[config.instance_name] + ) + host = instance.read_write_dns + password = cred.token + + return f"postgresql://{config.user}:{password}@{host}:{config.port}/{config.database_name}?sslmode=require" + + def _get_engine(self, config: LakebaseChecksStorageConfig) -> Engine: + """ + Create a SQLAlchemy engine for the Lakebase instance. + + Args: + config: Configuration for saving and loading checks to Lakebase. + + Returns: + SQLAlchemy engine for the Lakebase instance. + """ + connection_url = self._get_connection_url(config) + return create_engine(connection_url) + + @staticmethod + def get_table_definition(schema_name: str, table_name: str) -> Table: + """ + Create a SQLAlchemy table definition for storing DQ rules (checks) in Lakebase. + + Args: + schema_name: The schema where the checks table is located. + table_name: The table where the checks are stored. + + Returns: + SQLAlchemy table definition for the Lakebase instance. + """ + return Table( + table_name, + MetaData(schema=schema_name), + Column("name", String(255)), + Column("criticality", String(50), server_default="error"), + Column("check", JSONB), + Column("filter", Text), + Column("run_config_name", String(255), server_default="default"), + Column("user_metadata", JSONB), + ) + + @staticmethod + def _normalize_checks(checks: list[dict], config: LakebaseChecksStorageConfig) -> list[dict]: + """ + Normalize the checks to be compatible with the Lakebase table. + + Args: + checks: List of dq rules (checks) to normalize. + config: Configuration for saving and loading checks to Lakebase. + + Returns: + List of normalized dq rules (checks). + """ + normalized_checks = [] + for check in checks: + user_metadata = check.get("user_metadata") + normalized_check = { + "name": check.get("name"), + "criticality": check.get("criticality", "error"), + "check": check.get("check"), + "filter": check.get("filter"), + "run_config_name": check.get("run_config_name", config.run_config_name), + "user_metadata": null() if user_metadata is None else user_metadata, + } + normalized_checks.append(normalized_check) + return normalized_checks + + def _save_checks_to_lakebase(self, checks: list[dict], config: LakebaseChecksStorageConfig, engine: Engine) -> None: + """ + Save dq rules (checks) to a Lakebase table. + + Args: + checks: List of dq rules (checks) to save. + config: Configuration for saving and loading checks to Lakebase. + engine: SQLAlchemy engine for the Lakebase instance. + + Returns: + None + + Raises: + OperationalError: If connecting to the database fails. + """ + try: + with engine.connect() as conn: + pass + except OperationalError as e: + logger.error( + f"Failed to connect to database '{config.database_name}'. " + f"Please verify the database exists and the user has access: {e}" + ) + raise + + with engine.begin() as conn: + if not conn.dialect.has_schema(conn, config.schema_name): + conn.execute(CreateSchema(config.schema_name)) + logger.info(f"Successfully created schema '{config.schema_name}'.") + + with engine.begin() as conn: + table = self.get_table_definition(config.schema_name, config.table_name) + table.metadata.create_all(engine, checkfirst=True) + logger.info( + f"Successfully created or verified table '{config.database_name}.{config.schema_name}.{config.table_name}'." + ) + + if config.mode == "overwrite": + delete_stmt = delete(table).where(table.c.run_config_name == config.run_config_name) + result = conn.execute(delete_stmt) + logger.info(f"Deleted {result.rowcount} existing checks for run_config_name '{config.run_config_name}'") + + normalized_checks = self._normalize_checks(checks, config) + insert_stmt = insert(table) + conn.execute(insert_stmt, normalized_checks) + logger.info( + f"Inserted {len(normalized_checks)} checks to {config.database_name}.{config.schema_name}.{config.table_name} " + f"with run_config_name='{config.run_config_name}'" + ) + + def _load_checks_from_lakebase(self, config: LakebaseChecksStorageConfig, engine: Engine) -> list[dict]: + """ + Load dq rules (checks) from a Lakebase table. + + Args: + config: Configuration for saving and loading checks to Lakebase. + engine: SQLAlchemy engine for the Lakebase instance. + + Returns: + List of dq rules. + """ + table = self.get_table_definition(config.schema_name, config.table_name) + + stmt = select(table) + if config.run_config_name: + logger.info(f"Filtering checks by run_config_name='{config.run_config_name}'") + stmt = stmt.where(table.c.run_config_name == config.run_config_name) + else: + logger.info("Loading all checks (no run_config_name filter)") + + with engine.connect() as conn: + result = conn.execute(stmt) + checks = result.mappings().all() + logger.info( + f"Successfully loaded {len(checks)} checks from {config.database_name}.{config.schema_name}.{config.table_name} " + f"for run_config_name='{config.run_config_name}'" + ) + if len(checks) == 0: + logger.warning( + f"No checks found in {config.database_name}.{config.schema_name}.{config.table_name} " + f"for run_config_name='{config.run_config_name}'. " + f"Make sure the profiler has run successfully and saved checks to this location." + ) + return [dict(check) for check in checks] + + def _check_for_undefined_table_error(self, e: ProgrammingError, config: LakebaseChecksStorageConfig) -> NoReturn: + """ + Check if the error is an undefined table error and raise an appropriate exception. + + This method always raises an exception and never returns normally. + + Args: + e: Programming error to check. + config: Configuration for saving and loading checks to Lakebase. + + Raises: + NotFound: If the table does not exist in the Lakebase instance (pgcode 42P01). + ProgrammingError: Re-raises the original error if it's not an undefined table error. + """ + pgcode = getattr(getattr(e, 'orig', None), 'pgcode', None) + postgres_undefined_table_error = '42P01' + if pgcode == postgres_undefined_table_error: + raise NotFound(f"Table '{config.location}' does not exist in the Lakebase instance") from e + raise e + + @telemetry_logger("load_checks", "lakebase") + def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: + """ + Load dq rules (checks) from a Lakebase table. + + Args: + config: Configuration for saving and loading checks to Lakebase. + + Returns: + List of dq rules or error if loading checks fails. + + Raises: + NotFound: If the table does not exist in the Lakebase instance. + ProgrammingError: If SQL syntax errors or missing objects (converted to NotFound for missing tables). + DatabaseError: If other database operations fail (includes OperationalError, IntegrityError, etc.). + """ + engine = self.engine + engine_created_internally = False + if not engine: + engine = self._get_engine(config) + engine_created_internally = True + + try: + return self._load_checks_from_lakebase(config, engine) + + except ProgrammingError as e: + logger.error(f"Programming error while loading checks from Lakebase: {e}") + self._check_for_undefined_table_error(e, config) + + except DatabaseError as e: + logger.error(f"Database error while loading checks from Lakebase: {e}") + raise + + finally: + if engine_created_internally: + engine.dispose() + + @telemetry_logger("save_checks", "lakebase") + def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: + """ + Save dq rules (checks) to a Lakebase table. + + Args: + checks: List of dq rules (checks) to save. + config: Configuration for saving and loading checks to Lakebase. + + Returns: + None + + Raises: + InvalidCheckError: If any check is invalid or unsupported. + IntegrityError: If constraint violations occur (e.g., duplicate keys). + ProgrammingError: If SQL syntax errors or missing objects. + DatabaseError: If other database operations fail (includes OperationalError, DataError, etc.). + """ + if not checks: + raise InvalidCheckError("Checks cannot be empty or None.") + + engine = self.engine + engine_created_internally = False + if not engine: + engine = self._get_engine(config) + engine_created_internally = True + + try: + self._save_checks_to_lakebase(checks, config, engine) + logger.info(f"Successfully saved {len(checks)} checks to Lakebase.") + + except IntegrityError as e: + logger.error(f"Integrity constraint violation while saving checks to Lakebase: {e}") + raise + + except ProgrammingError as e: + logger.error(f"Programming error while saving checks to Lakebase: {e}") + raise + + except DatabaseError as e: + logger.error(f"Database error while saving checks to Lakebase: {e}") + raise + + finally: + if engine_created_internally: + engine.dispose() + + class WorkspaceFileChecksStorageHandler(ChecksStorageHandler[WorkspaceFileChecksStorageConfig]): """ Handler for storing quality rules (checks) in a file (json or yaml) in the workspace. @@ -234,6 +541,7 @@ def __init__(self, ws: WorkspaceClient, spark: SparkSession, run_config_loader: self.workspace_file_handler = WorkspaceFileChecksStorageHandler(ws) self.table_handler = TableChecksStorageHandler(ws, spark) self.volume_handler = VolumeFileChecksStorageHandler(ws) + self.lakebase_handler = LakebaseChecksStorageHandler(ws, spark, None) @telemetry_logger("load_checks", "installation") def load(self, config: InstallationChecksStorageConfig) -> list[dict]: @@ -269,6 +577,7 @@ def save(self, checks: list[dict], config: InstallationChecksStorageConfig) -> N def _get_storage_handler_and_config( self, config: InstallationChecksStorageConfig ) -> tuple[ChecksStorageHandler, InstallationChecksStorageConfig]: + # Overwrite location if overwrite_location is set if config.overwrite_location: checks_location = config.location else: @@ -278,24 +587,44 @@ def _get_storage_handler_and_config( product_name=config.product_name, install_folder=config.install_folder, ) + checks_location = run_config.checks_location + # transfer lakebase fields from run config to storage config if not already set + if run_config.lakebase_instance_name and not config.instance_name: + config.instance_name = run_config.lakebase_instance_name + if run_config.lakebase_user and not config.user: + config.user = run_config.lakebase_user + # replace port if non-default is specified in the run config + if run_config.lakebase_port and config.port != run_config.lakebase_port: + config.port = run_config.lakebase_port + installation = self._run_config_loader.get_installation( config.assume_user, config.product_name, config.install_folder ) config.location = checks_location - if is_table_location(config.location): + matches_table_pattern = is_table_location(config.location) + is_lakebase_storage = config.instance_name is not None + + if matches_table_pattern and is_lakebase_storage: + logger.debug(f"Using LakebaseChecksStorageHandler for location '{config.location}'") + return self.lakebase_handler, config + + if matches_table_pattern: + logger.debug(f"Using TableChecksStorageHandler for location '{config.location}'") return self.table_handler, config if config.location.startswith("/Volumes/"): + logger.debug(f"Using VolumeChecksStorageHandler for location '{config.location}'") return self.volume_handler, config if not config.location.startswith("/"): # if absolute path is not provided, the location should be set relative to the installation folder config.location = f"{installation.install_folder()}/{config.location}" + logger.debug(f"Using WorkspaceFileChecksStorageHandler for location '{config.location}'") return self.workspace_file_handler, config @@ -395,6 +724,23 @@ def create_for_location( An instance of the corresponding BaseChecksStorageHandler. """ + @abstractmethod + def create_for_run_config(self, run_config: RunConfig) -> tuple[ChecksStorageHandler, BaseChecksStorageConfig]: + """ + Abstract method to create a handler and config based on a RunConfig. + + This method inspects the RunConfig to determine the appropriate storage handler. + If Lakebase connection parameters are present (lakebase_instance_name), it creates + a LakebaseChecksStorageHandler. Otherwise, it delegates to create_for_location + to infer the handler from the checks location string. + + Args: + run_config: RunConfig containing checks location and optional Lakebase parameters. + + Returns: + A tuple of (ChecksStorageHandler, BaseChecksStorageConfig). + """ + class ChecksStorageHandlerFactory(BaseChecksStorageHandlerFactory): def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession): @@ -422,6 +768,8 @@ def create(self, config: BaseChecksStorageConfig) -> ChecksStorageHandler: return WorkspaceFileChecksStorageHandler(self.workspace_client) if isinstance(config, TableChecksStorageConfig): return TableChecksStorageHandler(self.workspace_client, self.spark) + if isinstance(config, LakebaseChecksStorageConfig): + return LakebaseChecksStorageHandler(self.workspace_client, self.spark, None) if isinstance(config, VolumeFileChecksStorageConfig): return VolumeFileChecksStorageHandler(self.workspace_client) @@ -448,6 +796,57 @@ def create_for_location( return FileChecksStorageHandler(), FileChecksStorageConfig(location=location) + def create_for_run_config(self, run_config: RunConfig) -> tuple[ChecksStorageHandler, BaseChecksStorageConfig]: + """ + Factory method to create a handler and config based on a RunConfig. + + This method inspects the RunConfig to determine the appropriate storage handler. + If Lakebase connection parameters are present (lakebase_instance_name), it creates + a LakebaseChecksStorageHandler. Otherwise, it delegates to create_for_location + to infer the handler from the checks location string. + + Args: + run_config: RunConfig containing checks location and optional Lakebase parameters. + + Returns: + A tuple of (ChecksStorageHandler, BaseChecksStorageConfig). + + Raises: + InvalidConfigError: If the configuration is invalid or unsupported. + """ + if run_config.lakebase_instance_name: + if not run_config.lakebase_user: + raise InvalidConfigError( + f"Lakebase user must be specified in run config '{run_config.name}' when " + f"lakebase_instance_name is set. Please add 'lakebase_user' to your run configuration." + ) + + if not run_config.checks_location: + raise InvalidConfigError( + f"checks_location must be specified in run config '{run_config.name}' when using Lakebase. " + f"Expected format: 'database.schema.table'." + ) + + if len(run_config.checks_location.split(".")) != 3: + raise InvalidConfigError( + f"Invalid Lakebase table name '{run_config.checks_location}' in run config '{run_config.name}'. " + f"Must be in the format 'database.schema.table'. " + f"Example: 'my_database.my_schema.my_table'." + ) + + return ( + LakebaseChecksStorageHandler(self.workspace_client, self.spark, None), + LakebaseChecksStorageConfig( + location=run_config.checks_location, + instance_name=run_config.lakebase_instance_name, + user=run_config.lakebase_user, + port=run_config.lakebase_port or "5432", + run_config_name=run_config.name, + ), + ) + + return self.create_for_location(run_config.checks_location, run_config.name) + def is_table_location(location: str) -> bool: """ diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index fd62cafba..d1df0b93f 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -1,7 +1,8 @@ import abc +from functools import cached_property from datetime import datetime, timezone from dataclasses import dataclass, field -from databricks.labs.dqx.errors import InvalidConfigError +from databricks.labs.dqx.errors import InvalidConfigError, InvalidParameterError __all__ = [ "WorkspaceConfig", @@ -14,6 +15,7 @@ "FileChecksStorageConfig", "WorkspaceFileChecksStorageConfig", "TableChecksStorageConfig", + "LakebaseChecksStorageConfig", "InstallationChecksStorageConfig", "VolumeFileChecksStorageConfig", ] @@ -60,16 +62,24 @@ class RunConfig: input_config: InputConfig | None = None output_config: OutputConfig | None = None quarantine_config: OutputConfig | None = None # quarantined data table + profiler_config: ProfilerConfig = field(default_factory=ProfilerConfig) + checks_location: str = ( "checks.yml" # absolute or relative workspace file path or table containing quality rules / checks ) + warehouse_id: str | None = None # warehouse id to use in the dashboard - profiler_config: ProfilerConfig = field(default_factory=ProfilerConfig) + reference_tables: dict[str, InputConfig] = field(default_factory=dict) # reference tables to use in the checks # mapping of fully qualified custom check function (e.g. my_func) to the module location in the workspace # (e.g. {"my_func": "/Workspace/my_repo/my_module.py"}) custom_check_functions: dict[str, str] = field(default_factory=dict) + # Lakebase connection parameters, if wanting to store checks in lakebase database + lakebase_instance_name: str | None = None + lakebase_user: str | None = None + lakebase_port: str | None = None + def _default_run_time() -> str: return datetime.now(timezone.utc).isoformat() @@ -195,6 +205,65 @@ def __post_init__(self): raise InvalidConfigError("The table name ('location' field) must not be empty or None.") +@dataclass +class LakebaseChecksStorageConfig(BaseChecksStorageConfig): + """ + Configuration class for storing checks in a Lakebase table. + + Args: + instance_name: Name of the Lakebase instance. + user: Name of the user for the Lakebase connection. + location: Fully qualified name of the Lakebase table to store checks in the format 'database.schema.table'. + port: The Lakebase port (default is '5432'). + run_config_name: Name of the run configuration to use for checks (default is 'default'). + mode: The mode for writing checks to a table (e.g., 'append' or 'overwrite'). The *overwrite* mode + only replaces checks for the specific run config and not all checks in the table (default is 'overwrite'). + """ + + instance_name: str | None = None + user: str | None = None + location: str | None = None + port: str = "5432" + run_config_name: str = "default" + mode: str = "overwrite" + + def __post_init__(self): + if not self.location or self.location == "": + raise InvalidParameterError("Location must not be empty or None.") + + if not self.instance_name or self.instance_name == "": + raise InvalidParameterError("Instance name must not be empty or None.") + + if not self.user or self.user == "": + raise InvalidParameterError("User must not be empty or None.") + + if len(self.location.split(".")) != 3: + raise InvalidConfigError( + f"Invalid Lakebase table name '{self.location}'. Must be in the format 'database.schema.table'." + ) + + if self.mode not in ("append", "overwrite"): + raise InvalidConfigError(f"Invalid mode '{self.mode}'. Must be 'append' or 'overwrite'.") + + def _split_location(self) -> tuple[str, ...]: + """Splits 'database.schema.table' into components.""" + if not self.location: + raise InvalidConfigError("location must be set before accessing database components.") + return tuple(self.location.split(".")) + + @cached_property + def database_name(self) -> str: + return self._split_location()[0] + + @cached_property + def schema_name(self) -> str: + return self._split_location()[1] + + @cached_property + def table_name(self) -> str: + return self._split_location()[2] + + @dataclass class VolumeFileChecksStorageConfig(BaseChecksStorageConfig): """ @@ -213,7 +282,10 @@ def __post_init__(self): @dataclass class InstallationChecksStorageConfig( - WorkspaceFileChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig + WorkspaceFileChecksStorageConfig, + TableChecksStorageConfig, + VolumeFileChecksStorageConfig, + LakebaseChecksStorageConfig, ): """ Configuration class for storing checks in an installation. diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py index 0aa7cc119..1daff6fc2 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -802,6 +802,7 @@ def load_checks(self, config: BaseChecksStorageConfig) -> list[dict]: - *FileChecksStorageConfig* (local file); - *WorkspaceFileChecksStorageConfig* (Databricks workspace file); - *TableChecksStorageConfig* (table-backed storage); + - *LakebaseChecksStorageConfig* (Lakebase table); - *InstallationChecksStorageConfig* (installation directory); - *VolumeFileChecksStorageConfig* (Unity Catalog volume file); @@ -829,6 +830,7 @@ def save_checks(self, checks: list[dict], config: BaseChecksStorageConfig) -> No - *FileChecksStorageConfig* (local file); - *WorkspaceFileChecksStorageConfig* (Databricks workspace file); - *TableChecksStorageConfig* (table-backed storage); + - *LakebaseChecksStorageConfig* (Lakebase table); - *InstallationChecksStorageConfig* (installation directory); - *VolumeFileChecksStorageConfig* (Unity Catalog volume file); @@ -852,6 +854,10 @@ def _apply_checks_for_run_config(self, run_config: RunConfig) -> None: This method loads checks from the specified location, reads input data using the input config, and writes results using the output and optionally quarantine configs. + The storage handler is determined by the factory based on the RunConfig. If Lakebase + connection parameters are present (lakebase_instance_name), checks will be loaded from + a Lakebase table. Otherwise, the checks location will be inferred from the checks_location string. + Args: run_config (RunConfig): Specifies the inputs, outputs, and checks file. """ @@ -863,9 +869,7 @@ def _apply_checks_for_run_config(self, run_config: RunConfig) -> None: logger.info(f"Applying checks from: {run_config.checks_location}") - storage_handler, storage_config = self._checks_handler_factory.create_for_location( - run_config.checks_location, run_config.name - ) + storage_handler, storage_config = self._checks_handler_factory.create_for_run_config(run_config) # if checks are not found, return empty list # raise an error if checks location not found checks = storage_handler.load(storage_config) diff --git a/tests/conftest.py b/tests/conftest.py index 3dade4e14..f3ed47c16 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,10 +1,13 @@ import os +from typing import Any import re +import logging from collections.abc import Callable, Generator -from dataclasses import replace +from dataclasses import replace, dataclass from functools import cached_property import pytest +from pyspark.sql.types import StructType, StructField, StringType, IntegerType from databricks.labs.blueprint.installation import Installation, MockInstallation from databricks.labs.blueprint.tui import MockPrompts from databricks.labs.blueprint.wheels import ProductInfo, WheelsV2 @@ -19,9 +22,12 @@ from databricks.labs.pytester.fixtures.baseline import factory from databricks.sdk import WorkspaceClient from databricks.sdk.service.workspace import ImportFormat +from databricks.sdk.service.database import DatabaseInstance, DatabaseCatalog +logger = logging.getLogger(__name__) -@pytest.fixture + +@pytest.fixture(scope="session") def debug_env_name(): return "ws" # Specify the name of the debug environment from ~/.databricks/debug-env.json @@ -406,6 +412,35 @@ def expected_checks(): ] +@pytest.fixture +def df(spark): + return spark.createDataFrame( + data=[ + (1, None, "2006-04-09", "ymason@example.net", "APAC", "France", "High"), + (2, "Mark Brooks", "1992-07-27", "johnthomas@example.net", "LATAM", "Trinidad and Tobago", None), + (3, "Lori Gardner", "2001-01-22", "heather68@example.com", None, None, None), + (4, None, "1968-10-24", "hollandfrank@example.com", None, "Palau", "Low"), + (5, "Laura Mitchell DDS", "1968-01-08", "paynebrett@example.org", "NA", "Qatar", None), + (6, None, "1951-09-11", "williamstracy@example.org", "EU", "Benin", "High"), + (7, "Benjamin Parrish", "1971-08-17", "hmcpherson@example.net", "EU", "Kazakhstan", "Medium"), + (8, "April Hamilton", "1989-09-04", "adamrichards@example.net", "EU", "Saint Lucia", None), + (9, "Stephanie Price", "1975-03-01", "ktrujillo@example.com", "NA", "Togo", "High"), + (10, "Jonathan Sherman", "1976-04-13", "charles93@example.org", "NA", "Japan", "Low"), + ], + schema=StructType( + [ + StructField("id", IntegerType(), True), + StructField("name", StringType(), True), + StructField("birthdate", StringType(), True), + StructField("email", StringType(), True), + StructField("region", StringType(), True), + StructField("state", StringType(), True), + StructField("tier", StringType(), True), + ] + ), + ) + + @pytest.fixture def make_local_check_file_as_yaml(checks_yaml_content, make_random): file_path = f"checks_{make_random(10).lower()}.yml" @@ -640,3 +675,110 @@ def delete(volume_file_path: str) -> None: ws.files.delete(volume_file_path) yield from factory("file", create, delete) + + +@pytest.fixture(scope="session") +def lakebase_user(): + """ + Get the Lakebase user (service principal client ID) from environment variable. + + This fixture reads ARM_CLIENT_ID which is set in the CI/CD pipeline. + For local development, you can set this environment variable to your service principal's client ID. + + Returns: + The client ID to use as the Lakebase user. + + Raises: + pytest.skip: If ARM_CLIENT_ID is not set in the environment. + """ + client_id = os.environ.get("ARM_CLIENT_ID") + if not client_id: + pytest.skip("ARM_CLIENT_ID environment variable not set. Required for Lakebase tests.") + return client_id + + +@dataclass +class LakebaseInstance: + name: str + catalog: str + location: str + + +@pytest.fixture +def make_lakebase_instance(ws, make_random): + + def create() -> LakebaseInstance: + instance_name = f"dqxtest-{make_random(10).lower()}" + database_name = "dqx" # does not need to be random + catalog_name = f"dqxtest-{make_random(10).lower()}" + capacity = "CU_2" + schema_name = "config" # auto-created when saving checks, can be static since each test create new db + table_name = "checks" # auto-created when saving checks, can be static since each test create new db + + _ = ws.database.create_database_instance_and_wait( + database_instance=DatabaseInstance(name=instance_name, capacity=capacity) + ) + logger.info(f"Successfully created database instance: {instance_name}") + + ws.database.create_database_catalog( + DatabaseCatalog( + name=catalog_name, + database_name=database_name, + database_instance_name=instance_name, + create_database_if_not_exists=True, + ) + ) + logger.info(f"Successfully created database catalog: {catalog_name}") + + return LakebaseInstance( + name=instance_name, catalog=catalog_name, location=f"{database_name}.{schema_name}.{table_name}" + ) + + def delete(instance: LakebaseInstance) -> None: + try: + ws.database.delete_database_catalog(name=instance.catalog) + logger.info(f"Successfully deleted database catalog: {instance.catalog}") + except Exception as e: + logger.warning(f"Failed to delete database catalog {instance.catalog}: {e}") + + try: + ws.database.delete_database_instance(name=instance.name) + logger.info(f"Successfully deleted database instance: {instance.name}") + except Exception as e: + logger.error(f"Failed to delete database instance {instance.name}: {e}") + raise + + yield from factory("lakebase", create, delete) + + +def compare_checks(result: list[dict[str, Any]], expected: list[dict[str, Any]]) -> None: + """ + Compares two lists of checks dictionaries for equality, ensuring + they contain the same checks with identical fields. + + Args: + result: The result checks list. + expected: The expected checks list. + + Returns: + None + """ + assert len(result) == len(expected), f"Expected {len(expected)} checks, got {len(result)}" + sorted_result = sorted(result, key=_sort_key) + sorted_expected = sorted(expected, key=_sort_key) + for res, exp in zip(sorted_result, sorted_expected): + for key in exp: + assert res.get(key) == exp[key], f"Mismatch for key '{key}': {res.get(key)} != {exp[key]}" + + +def _sort_key(check: dict[str, Any]) -> str: + """ + Sorts a checks dictionary by the 'name' field. + + Args: + check: The check dictionary. + + Returns: + The name of the check as a string, or an empty string if not found. + """ + return str(check.get("name", "")) diff --git a/tests/integration/test_save_and_load_checks_from_lakebase_table.py b/tests/integration/test_save_and_load_checks_from_lakebase_table.py new file mode 100644 index 000000000..4ce99d908 --- /dev/null +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -0,0 +1,149 @@ +import pytest + +from databricks.labs.dqx.config import InstallationChecksStorageConfig, LakebaseChecksStorageConfig +from databricks.labs.dqx.engine import DQEngine +from databricks.sdk.errors import NotFound + +from tests.conftest import compare_checks +from tests.integration.test_save_and_load_checks_from_table import EXPECTED_CHECKS as TEST_CHECKS + + +def test_load_checks_when_lakebase_table_does_not_exist(ws, spark, make_lakebase_instance, lakebase_user): + instance = make_lakebase_instance() + dq_engine = DQEngine(ws, spark) + config = LakebaseChecksStorageConfig(location=instance.location, user=lakebase_user, instance_name=instance.name) + + with pytest.raises(NotFound, match=f"Table '{config.location}' does not exist in the Lakebase instance"): + dq_engine.load_checks(config=config) + + +def test_save_and_load_checks_from_lakebase_table(ws, spark, make_lakebase_instance, lakebase_user): + instance = make_lakebase_instance() + dq_engine = DQEngine(ws, spark) + config = LakebaseChecksStorageConfig(location=instance.location, user=lakebase_user, instance_name=instance.name) + + dq_engine.save_checks(checks=TEST_CHECKS, config=config) + checks = dq_engine.load_checks(config=config) + + compare_checks(checks, TEST_CHECKS) + + +def test_save_and_load_checks_from_lakebase_table_with_run_config(ws, spark, make_lakebase_instance, lakebase_user): + instance = make_lakebase_instance() + dq_engine = DQEngine(ws, spark) + + # test first run config + config = LakebaseChecksStorageConfig( + location=instance.location, user=lakebase_user, instance_name=instance.name, run_config_name="workflow_001" + ) + dq_engine.save_checks(TEST_CHECKS[:1], config=config) + checks = dq_engine.load_checks(config=config) + + compare_checks(checks, TEST_CHECKS[:1]) + + # test second run config + second_config = LakebaseChecksStorageConfig( + location=instance.location, user=lakebase_user, instance_name=instance.name, run_config_name="workflow_002" + ) + dq_engine.save_checks(TEST_CHECKS[1:], config=second_config) + checks = dq_engine.load_checks(config=second_config) + + compare_checks(checks, TEST_CHECKS[1:]) + + +def test_save_and_load_checks_from_lakebase_table_with_output_modes( + ws, + spark, + make_lakebase_instance, + lakebase_user, +): + instance = make_lakebase_instance() + dq_engine = DQEngine(ws, spark) + + run_config_name = "workflow_003" + dq_engine.save_checks( + TEST_CHECKS[:1], + config=LakebaseChecksStorageConfig( + location=instance.location, + instance_name=instance.name, + user=lakebase_user, + run_config_name=run_config_name, + mode="append", + ), + ) + checks = dq_engine.load_checks( + config=LakebaseChecksStorageConfig( + location=instance.location, instance_name=instance.name, user=lakebase_user, run_config_name=run_config_name + ) + ) + compare_checks(checks, TEST_CHECKS[:1]) + + run_config_name = "workflow_004" + dq_engine.save_checks( + TEST_CHECKS[1:], + config=LakebaseChecksStorageConfig( + location=instance.location, + instance_name=instance.name, + user=lakebase_user, + run_config_name=run_config_name, + mode="overwrite", + ), + ) + checks = dq_engine.load_checks( + config=LakebaseChecksStorageConfig( + location=instance.location, instance_name=instance.name, user=lakebase_user, run_config_name=run_config_name + ) + ) + compare_checks(checks, TEST_CHECKS[1:]) + + +def test_save_and_load_checks_from_lakebase_table_with_user_installation( + ws, spark, installation_ctx, make_lakebase_instance, lakebase_user +): + instance = make_lakebase_instance() + + config = installation_ctx.config + run_config = config.get_run_config() + run_config.checks_location = instance.location + installation_ctx.installation.save(installation_ctx.config) + product_name = installation_ctx.product_info.product_name() + + config = InstallationChecksStorageConfig( + instance_name=instance.name, + user=lakebase_user, + run_config_name=run_config.name, + product_name=product_name, + assume_user=True, + ) + + dq_engine = DQEngine(ws, spark) + + dq_engine.save_checks(TEST_CHECKS, config=config) + checks = dq_engine.load_checks(config=config) + + compare_checks(checks, TEST_CHECKS) + + +def test_profiler_workflow_save_to_lakebase(ws, spark, setup_workflows, make_lakebase_instance, lakebase_user): + installation_ctx, run_config = setup_workflows() + + instance = make_lakebase_instance() + + config = installation_ctx.config + run_config = config.get_run_config() + run_config.checks_location = instance.location + run_config.lakebase_user = lakebase_user + run_config.lakebase_instance_name = instance.name + run_config.lakebase_port = "5432" + + installation_ctx.installation.save(installation_ctx.config) + + installation_ctx.deployed_workflows.run_workflow("profiler", run_config.name) + + dq_engine = DQEngine(ws, spark) + checks = dq_engine.load_checks( + config=LakebaseChecksStorageConfig( + location=instance.location, instance_name=instance.name, user=lakebase_user, run_config_name=run_config.name + ) + ) + assert checks, "Checks are missing" diff --git a/tests/unit/test_lakebase_config.py b/tests/unit/test_lakebase_config.py new file mode 100644 index 000000000..7787db2c0 --- /dev/null +++ b/tests/unit/test_lakebase_config.py @@ -0,0 +1,199 @@ +"""Unit tests for LakebaseConnectionConfig and LakebaseChecksStorageConfig.""" + +import pytest +from databricks.labs.dqx.config import LakebaseChecksStorageConfig +from databricks.labs.dqx.errors import InvalidParameterError, InvalidConfigError + + +def test_create_valid_lakebase_check_storage_config(): + """Test creating LakebaseConnectionConfig with valid parameters.""" + config = LakebaseChecksStorageConfig( + instance_name="testhost", location="testdb.testschema.testtable", user="testuser", port="5432" + ) + + assert config.instance_name == "testhost" + assert config.database_name == "testdb" + assert config.user == "testuser" + assert config.port == "5432" + + +def test_create_lakebase_check_storage_config_default_port(): + """Test creating LakebaseConnectionConfig with default port.""" + config = LakebaseChecksStorageConfig( + instance_name="testhost", location="testdb.testschema.testtable", user="testuser" + ) + + assert config.instance_name == "testhost" + assert config.database_name == "testdb" + assert config.user == "testuser" + assert config.port == "5432" # Default port + + +@pytest.mark.parametrize( + "instance_name,location,user,port,expected_error", + [ + # Empty/None instance_name + ("", "testdb.testschema.testtable", "user", "5432", "Instance name must not be empty or None."), + (None, "testdb.testschema.testtable", "user", "5432", "Instance name must not be empty or None."), + # Empty/None database + ("instance", "", "user", "5432", "Location must not be empty or None."), + ("instance", None, "user", "5432", "Location must not be empty or None."), + # Empty/None user + ("instance", "testdb.testschema.testtable", "", "5432", "User must not be empty or None."), + ("instance", "testdb.testschema.testtable", None, "5432", "User must not be empty or None."), + ], +) +def test_create_invalid_lakebase_check_storage_config(instance_name, location, user, port, expected_error): + """Test LakebaseChecksStorageConfig validation with invalid parameters.""" + with pytest.raises(InvalidParameterError, match=expected_error): + LakebaseChecksStorageConfig(instance_name=instance_name, location=location, user=user, port=port) + + +def test_valid_lakebase_checks_storage_config(): + """Test creating LakebaseChecksStorageConfig with valid parameters.""" + config = LakebaseChecksStorageConfig( + instance_name="testinstance", + location="testdb.testschema.testtable", + user="testuser", + port="5432", + run_config_name="test_run", + mode="append", + ) + + assert config.instance_name == "testinstance" + assert config.location == "testdb.testschema.testtable" + assert config.user == "testuser" + assert config.port == "5432" + assert config.run_config_name == "test_run" + assert config.mode == "append" + assert config.database_name == "testdb" + assert config.schema_name == "testschema" + assert config.table_name == "testtable" + + +def test_valid_lakebase_checks_storage_config_defaults(): + """Test creating LakebaseChecksStorageConfig with default values.""" + config = LakebaseChecksStorageConfig( + instance_name="testinstance", location="testdb.testschema.testtable", user="testuser" + ) + + assert config.run_config_name == "default" + assert config.mode == "overwrite" + assert config.port == "5432" # Default port + + +@pytest.mark.parametrize( + "location,expected_error", + [ + # Empty/None location + ("", "Location must not be empty or None."), + (None, "Location must not be empty or None."), + ], +) +def test_invalid_location_empty(location, expected_error): + """Test LakebaseChecksStorageConfig with empty/None location.""" + with pytest.raises(InvalidParameterError, match=expected_error): + LakebaseChecksStorageConfig(instance_name="testinstance", location=location, user="testuser") + + +@pytest.mark.parametrize( + "location,expected_error", + [ + # Invalid location format (not 3 parts) + ("db", "Invalid Lakebase table name 'db'. Must be in the format 'database.schema.table'."), + ("db.schema", "Invalid Lakebase table name 'db.schema'. Must be in the format 'database.schema.table'."), + ( + "db.schema.table.extra", + "Invalid Lakebase table name 'db.schema.table.extra'. Must be in the format 'database.schema.table'.", + ), + ], +) +def test_invalid_location_format(location, expected_error): + """Test LakebaseChecksStorageConfig with invalid location format.""" + with pytest.raises(InvalidConfigError, match=expected_error): + LakebaseChecksStorageConfig(instance_name="testinstance", location=location, user="testuser") + + +@pytest.mark.parametrize( + "instance_name,expected_error", + [ + ("", "Instance name must not be empty or None."), + (None, "Instance name must not be empty or None."), + ], +) +def test_invalid_instance_name(instance_name, expected_error): + """Test LakebaseChecksStorageConfig with invalid instance name.""" + with pytest.raises(InvalidParameterError, match=expected_error): + LakebaseChecksStorageConfig( + instance_name=instance_name, location="testdb.testschema.testtable", user="testuser" + ) + + +@pytest.mark.parametrize( + "mode,expected_error", + [ + # Invalid modes + ("invalid", "Invalid mode 'invalid'. Must be 'append' or 'overwrite'."), + ("INSERT", "Invalid mode 'INSERT'. Must be 'append' or 'overwrite'."), + ("update", "Invalid mode 'update'. Must be 'append' or 'overwrite'."), + ("delete", "Invalid mode 'delete'. Must be 'append' or 'overwrite'."), + ("", "Invalid mode ''. Must be 'append' or 'overwrite'."), + ], +) +def test_invalid_mode(mode, expected_error): + """Test LakebaseChecksStorageConfig with invalid mode.""" + with pytest.raises(InvalidConfigError, match=expected_error): + LakebaseChecksStorageConfig( + instance_name="testinstance", location="testdb.testschema.testtable", user="testuser", mode=mode + ) + + +def test_valid_modes(): + """Test LakebaseChecksStorageConfig with valid modes.""" + # Test 'append' mode + config_append = LakebaseChecksStorageConfig( + instance_name="testinstance", location="testdb.testschema.testtable", user="testuser", mode="append" + ) + assert config_append.mode == "append" + + # Test 'overwrite' mode + config_overwrite = LakebaseChecksStorageConfig( + instance_name="testinstance", location="testdb.testschema.testtable", user="testuser", mode="overwrite" + ) + assert config_overwrite.mode == "overwrite" + + +def test_location_properties(): + """Test database_name, schema_name, and table_name properties.""" + config = LakebaseChecksStorageConfig( + instance_name="testinstance", location="testdb.testschema.testtable", user="testuser" + ) + + assert config.database_name == "testdb" + assert config.schema_name == "testschema" + assert config.table_name == "testtable" + + +def test_complex_location_with_special_characters(): + """Test location parsing with special characters in names.""" + config = LakebaseChecksStorageConfig( + instance_name="testinstance", location="test_db.test_schema.test_table", user="testuser" + ) + + assert config.database_name == "test_db" + assert config.schema_name == "test_schema" + assert config.table_name == "test_table" + + +def test_eager_validation_on_init(): + """Test that validation happens eagerly during __post_init__.""" + # This should fail immediately during construction, not later + with pytest.raises(InvalidParameterError, match="Instance name must not be empty or None"): + LakebaseChecksStorageConfig(instance_name="", location="testdb.testschema.testtable", user="testuser") + + +def test_multiple_validation_errors_precedence(): + """Test that the first validation error is raised when multiple issues exist.""" + # Empty location should be caught first, before instance name validation + with pytest.raises(InvalidParameterError, match="Location must not be empty or None"): + LakebaseChecksStorageConfig(instance_name="", location="", user="testuser") diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index 136c131aa..777528294 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -1,8 +1,9 @@ from unittest.mock import create_autospec + import pytest -from databricks.sdk.service.files import DownloadResponse from databricks.sdk import WorkspaceClient from databricks.sdk.errors import NotFound +from databricks.sdk.service.files import DownloadResponse from databricks.labs.dqx.checks_storage import VolumeFileChecksStorageHandler from databricks.labs.dqx.config import VolumeFileChecksStorageConfig diff --git a/tests/unit/test_save_checks.py b/tests/unit/test_save_checks.py index 376eae87b..4b0fd0593 100644 --- a/tests/unit/test_save_checks.py +++ b/tests/unit/test_save_checks.py @@ -1,18 +1,21 @@ import os import json +import yaml import pytest -import yaml from databricks.labs.dqx.engine import DQEngineCore from databricks.labs.dqx.errors import InvalidConfigError + TEST_CHECKS = [ { + "name": "id_is_null", "criticality": "error", "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, }, { + "name": "name_is_null", "criticality": "error", "check": { "function": "is_not_null",