From b3334136270fb1e32138a264ea8f7a9ee3607e2a Mon Sep 17 00:00:00 2001 From: Manuel Tilgner <58051108+tlgnr@users.noreply.github.com> Date: Wed, 10 Sep 2025 12:09:03 +0200 Subject: [PATCH 001/125] refactor: delete only records for specific run_config_name Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/databricks/labs/dqx/checks_storage.py | 209 ++++++++++++++++++++++ src/databricks/labs/dqx/config.py | 22 +++ src/databricks/labs/dqx/engine.py | 1 + 3 files changed, 232 insertions(+) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 65b65fc0b..78aee0d4c 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -1,10 +1,12 @@ 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 psycopg2.pool import ThreadedConnectionPool import yaml from pyspark.sql import SparkSession @@ -13,6 +15,7 @@ from databricks.labs.dqx.config import ( TableChecksStorageConfig, + LakebaseChecksStorageConfig, FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, @@ -34,6 +37,7 @@ logger = logging.getLogger(__name__) T = TypeVar("T", bound=BaseChecksStorageConfig) +connection_pool: ThreadedConnectionPool | None = None class ChecksStorageHandler(ABC, Generic[T]): @@ -103,6 +107,209 @@ def save(self, checks: list[dict], config: TableChecksStorageConfig) -> None: ) +class LakebaseChecksStorageHandler(ChecksStorageHandler[LakebaseChecksStorageConfig]): + """ + Handler for storing quality rules (checks) in a Lakebase table in the workspace. + """ + + def __init__(self, ws: WorkspaceClient, spark: SparkSession): + self.ws = ws + self.spark = spark + + def get_connection_pool(self, config: LakebaseChecksStorageConfig) -> ThreadedConnectionPool: + """ + Get a connection from a threaded connection pool for a Lakebase instance. + + Args: + config: configuration for loading checks, including the table location, instance name, and run configuration name. + + Returns: + A connection pool. + """ + global connection_pool + + if not connection_pool: + 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] + ) + + user = self.ws.current_user.me().user_name + password = cred.token + + connection_pool = ThreadedConnectionPool( + minconn=1, + maxconn=10, + user=user, + password=password, + host=instance.read_write_dns, + port="5432", + database=config.location, + ) + + logger.info(f"Successfully created connection pool for instance {config.instance_name}") + + return connection_pool + + def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: + """ + Load checks from a Lakebase table. + + Args: + config: configuration for loading checks, including the table location and run configuration name. + + Returns: + list of dq rules or raise an error if checks table is missing or is invalid. + """ + logger.info(f"Loading quality rules (checks) from Lakebase table '{config.location}'") + connection_pool = self.get_connection_pool(config) + connection = None + cursor = None + + try: + connection = connection_pool.getconn() + if not connection: + raise RuntimeError(f"Failed to get connection from pool for instance {config.instance_name}") + + cursor = connection.cursor() + + check_table_query = f""" + SELECT EXISTS ( + SELECT FROM information_schema.tables + WHERE table_schema = '{config.location}' + AND table_name = 'checks' + ); + """ + cursor.execute(check_table_query) + table_exists = cursor.fetchone()[0] + + if not table_exists: + logger.warning(f"Table {config.location}.checks does not exist") + return [] + + select_query = f""" + SELECT name, criticality, "check", filter, run_config_name, user_metadata + FROM {config.location}.checks + WHERE run_config_name = %s + """ + cursor.execute(select_query, (config.run_config_name,)) + rows = cursor.fetchall() + + checks = [] + for row in rows: + field_name, field_criticality, field_check, field_filter, field_run_config_name, field_user_metadata = ( + row + ) + + check_dict = { + "name": field_name, + "criticality": field_criticality, + "check": field_check, + "filter": field_filter, + "run_config_name": field_run_config_name, + "user_metadata": field_user_metadata, + } + + checks.append(check_dict) + + logger.info(f"Successfully loaded checks from {config.location}") + return checks + except Exception as e: + logger.error(f"Failed to load checks from {config.instance_name}: {e}") + raise + finally: + if cursor: + cursor.close() + if connection: + connection_pool.putconn(connection) + + def save(self, checks: list[dict], config: LakebaseChecksStorageConfig): + """ + Save checks to a Lakebase table. + + Args: + config: configuration for loading checks, including the table location, instance name, and run configuration name. + + Returns: + + """ + connection_pool = self.get_connection_pool(config) + connection = None + cursor = None + + try: + connection = connection_pool.getconn() + if not connection: + raise RuntimeError(f"Failed to get connection from pool for instance {config.instance_name}") + + cursor = connection.cursor() + + create_dataset_query = f"CREATE SCHEMA IF NOT EXISTS {config.location};" + cursor.execute(create_dataset_query) + + create_table_query = f""" + CREATE TABLE IF NOT EXISTS {config.location}.checks ( + name VARCHAR(255), + criticality VARCHAR(50), + "check" JSONB, + filter TEXT, + run_config_name VARCHAR(255), + user_metadata JSONB + ) + """ + cursor.execute(create_table_query) + + if config.mode == "overwrite": + delete_query = f"DELETE FROM {config.location}.checks WHERE run_config_name = %s" + cursor.execute(delete_query, (config.run_config_name,)) + + values = [] + placeholders = [] + + for check in checks: + field_name = check.get("name") + field_criticality = check.get("criticality", "error") + field_filter = check.get("filter") + field_run_config_name = check.get("run_config_name", config.run_config_name) + field_check = check.get("check") + field_user_metadata = check.get("user_metadata") + + field_check = json.dumps(field_check) if field_check is not None else None + field_user_metadata = json.dumps(field_user_metadata) if field_user_metadata is not None else None + + placeholders.append("(%s, %s, %s, %s, %s, %s)") + values.extend( + [ + field_name, + field_criticality, + field_check, + field_filter, + field_run_config_name, + field_user_metadata, + ] + ) + + if values: + insert_values_query = f""" + INSERT INTO {config.location}.checks (name, criticality, "check", filter, run_config_name, user_metadata) + VALUES {", ".join(placeholders)} + """ + cursor.execute(insert_values_query, values) + + connection.commit() + logger.info(f"Successfully saved checks to {config.location}") + except Exception as e: + logger.error(f"Failed to save checks to {config.instance_name}: {e}") + if connection: + connection.rollback() + raise + finally: + if cursor: + cursor.close() + if connection: + connection_pool.putconn(connection) + + class WorkspaceFileChecksStorageHandler(ChecksStorageHandler[WorkspaceFileChecksStorageConfig]): """ Handler for storing quality rules (checks) in a file (json or yaml) in the workspace. @@ -378,6 +585,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) if isinstance(config, VolumeFileChecksStorageConfig): return VolumeFileChecksStorageHandler(self.workspace_client) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 4eaf67a83..2f047a124 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -13,6 +13,7 @@ "FileChecksStorageConfig", "WorkspaceFileChecksStorageConfig", "TableChecksStorageConfig", + "LakebaseChecksStorageConfig", "InstallationChecksStorageConfig", "VolumeFileChecksStorageConfig", ] @@ -189,6 +190,27 @@ def __post_init__(self): if not self.location: raise ValueError("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: + location: The table name where the checks are stored. + run_config_name: The 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 will only replace checks for the specific run config and not all checks in the table. + """ + + location: str + instance_name: str + run_config_name: str = "default" # to filter checks by run config + mode: str = "overwrite" + + def __post_init__(self): + if not self.location: + raise ValueError("The table name ('location' field) must not be empty or None.") + @dataclass class VolumeFileChecksStorageConfig(BaseChecksStorageConfig): diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py index 5734f99ac..2d7c499ef 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -671,6 +671,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); From f61ef7e17465c0cd43a80c0753d3a083e04c0709 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 22 Sep 2025 15:04:02 +0200 Subject: [PATCH 002/125] refactor: use sqlalchemy for saving and loading checks to Lakebase --- pyproject.toml | 1 + src/databricks/labs/dqx/checks_storage.py | 281 ++++++++++------------ src/databricks/labs/dqx/config.py | 49 +++- 3 files changed, 173 insertions(+), 158 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 22651cb63..450035c06 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ classifiers = [ dependencies = ["databricks-labs-blueprint>=0.9.1,<0.10", "databricks-sdk~=0.57", "databricks-labs-lsql>=0.5,<=0.16", + "sqlalchemy>=1.4,<3.0", ] [project.optional-dependencies] diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 6e261726a..4c9c2cdf0 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -7,6 +7,22 @@ from abc import ABC, abstractmethod from typing import Generic, TypeVar from psycopg2.pool import ThreadedConnectionPool +from sqlalchemy import ( + Engine, + create_engine, + text, + MetaData, + Table, + Column, + String, + Text, + insert, + select, + delete, +) +from sqlalchemy.schema import CreateSchema +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.exc import OperationalError import yaml from pyspark.sql import SparkSession @@ -21,6 +37,7 @@ InstallationChecksStorageConfig, BaseChecksStorageConfig, VolumeFileChecksStorageConfig, + ConnectionInfo, ) from databricks.sdk import WorkspaceClient @@ -112,205 +129,167 @@ def save(self, checks: list[dict], config: TableChecksStorageConfig) -> None: class LakebaseChecksStorageHandler(ChecksStorageHandler[LakebaseChecksStorageConfig]): """ - Handler for storing quality rules (checks) in a Lakebase table in the workspace. + Handler for storing quality rules (checks) in a Lakebase table. """ def __init__(self, ws: WorkspaceClient, spark: SparkSession): self.ws = ws self.spark = spark - def get_connection_pool(self, config: LakebaseChecksStorageConfig) -> ThreadedConnectionPool: + def _get_connection_info(self, config: LakebaseChecksStorageConfig) -> ConnectionInfo: """ - Get a connection from a threaded connection pool for a Lakebase instance. + Get connection information for the Lakebase instance. Args: - config: configuration for loading checks, including the table location, instance name, and run configuration name. + config: configuration for saving and loading checks, including instance name, database, schema, table, and port. Returns: - A connection pool. + Connection information for the Lakebase instance. """ - global connection_pool + 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 + user = self.ws.current_user.me().user_name + password = cred.token + + return ConnectionInfo( + user=user, + password=password, + host=host, + port=config.port, + database=config.database, + ) - if not connection_pool: - 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] - ) + def _get_engine(self, connection_info: ConnectionInfo) -> Engine: + """ + Create a SQLAlchemy engine for the Lakebase instance. - user = self.ws.current_user.me().user_name - password = cred.token + Args: + connection_info: connection information for the Lakebase instance. - connection_pool = ThreadedConnectionPool( - minconn=1, - maxconn=10, - user=user, - password=password, - host=instance.read_write_dns, - port="5432", - database=config.location, - ) + Returns: + SQLAlchemy engine for the Lakebase instance. + """ + return create_engine(connection_info.to_url()) - logger.info(f"Successfully created connection pool for instance {config.instance_name}") + def _get_table_definition(self, config: LakebaseChecksStorageConfig) -> Table: + """ + Create table definition for consistency between load and save methods. - return connection_pool + Args: + config: configuration for saving and loading checks, including instance name, database, schema, table, and port. + Returns: + SQLAlchemy table definition for the Lakebase instance. + """ + metadata = MetaData(schema=config.schema) + return Table( + config.table, + metadata, + Column("name", String(255)), + Column("criticality", String(50), default="error"), + Column("check", JSONB), + Column("filter", Text), + Column("run_config_name", String(255), default=config.run_config_name), + Column("user_metadata", JSONB), + ) + + @telemetry_logger("load_checks", "lakebase") def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: """ - Load checks from a Lakebase table. + Load dq rules (checks) from a Lakebase table. Args: - config: configuration for loading checks, including the table location and run configuration name. + config: configuration for saving and loading checks, including instance name, database, schema, table, and port. Returns: - list of dq rules or raise an error if checks table is missing or is invalid. + list of dq rules or error if checks table is missing or is invalid. """ - logger.info(f"Loading quality rules (checks) from Lakebase table '{config.location}'") - connection_pool = self.get_connection_pool(config) - connection = None - cursor = None + connection_info = self._get_connection_info(config) + engine = self._get_engine(connection_info) try: - connection = connection_pool.getconn() - if not connection: - raise RuntimeError(f"Failed to get connection from pool for instance {config.instance_name}") - - cursor = connection.cursor() - - check_table_query = f""" - SELECT EXISTS ( - SELECT FROM information_schema.tables - WHERE table_schema = '{config.location}' - AND table_name = 'checks' - ); - """ - cursor.execute(check_table_query) - table_exists = cursor.fetchone()[0] - - if not table_exists: - logger.warning(f"Table {config.location}.checks does not exist") - return [] - - select_query = f""" - SELECT name, criticality, "check", filter, run_config_name, user_metadata - FROM {config.location}.checks - WHERE run_config_name = %s - """ - cursor.execute(select_query, (config.run_config_name,)) - rows = cursor.fetchall() - - checks = [] - for row in rows: - field_name, field_criticality, field_check, field_filter, field_run_config_name, field_user_metadata = ( - row - ) + logger.info(f"Loading checks from Lakebase instance {config.instance_name}") - check_dict = { - "name": field_name, - "criticality": field_criticality, - "check": field_check, - "filter": field_filter, - "run_config_name": field_run_config_name, - "user_metadata": field_user_metadata, - } + table = self._get_table_definition(config) + table.metadata.create_all(engine) - checks.append(check_dict) + stmt = select(table) + if config.run_config_name: + stmt = stmt.where(table.c.run_config_name == config.run_config_name) - logger.info(f"Successfully loaded checks from {config.location}") - return checks + with engine.connect() as conn: + result = conn.execute(stmt) + checks = result.mappings().all() + logger.info(f"Successfully loaded {len(checks)} checks from {config.schema}.{config.table}") + return [dict(check) for check in checks] except Exception as e: - logger.error(f"Failed to load checks from {config.instance_name}: {e}") + logger.error(f"Failed to load checks from Lakebase instance {config.instance_name}: {e}") raise finally: - if cursor: - cursor.close() - if connection: - connection_pool.putconn(connection) + engine.dispose() - def save(self, checks: list[dict], config: LakebaseChecksStorageConfig): + @telemetry_logger("save_checks", "lakebase") + def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: """ - Save checks to a Lakebase table. + Save dq rules (checks) to a Lakebase table. Args: - config: configuration for loading checks, including the table location, instance name, and run configuration name. + checks: list of dq rules (checks) to save. + config: configuration for saving and loading checks, including instance name, database, schema, table, and port. Returns: + None + Raises: + OperationalError: If connection to Lakebase instance fails. + ValueError: If invalid mode is specified. + Exception: If saving checks fails. """ - connection_pool = self.get_connection_pool(config) - connection = None - cursor = None + if not checks: + logger.warning("No checks provided to save") + return + + if config.mode not in ("append", "overwrite"): + raise ValueError(f"Invalid mode '{config.mode}'. Must be 'append' or 'overwrite'") + + connection_info = self._get_connection_info(config) + engine = self._get_engine(connection_info) try: - connection = connection_pool.getconn() - if not connection: - raise RuntimeError(f"Failed to get connection from pool for instance {config.instance_name}") - - cursor = connection.cursor() - - create_dataset_query = f"CREATE SCHEMA IF NOT EXISTS {config.location};" - cursor.execute(create_dataset_query) - - create_table_query = f""" - CREATE TABLE IF NOT EXISTS {config.location}.checks ( - name VARCHAR(255), - criticality VARCHAR(50), - "check" JSONB, - filter TEXT, - run_config_name VARCHAR(255), - user_metadata JSONB - ) - """ - cursor.execute(create_table_query) - - if config.mode == "overwrite": - delete_query = f"DELETE FROM {config.location}.checks WHERE run_config_name = %s" - cursor.execute(delete_query, (config.run_config_name,)) - - values = [] - placeholders = [] - - for check in checks: - field_name = check.get("name") - field_criticality = check.get("criticality", "error") - field_filter = check.get("filter") - field_run_config_name = check.get("run_config_name", config.run_config_name) - field_check = check.get("check") - field_user_metadata = check.get("user_metadata") - - field_check = json.dumps(field_check) if field_check is not None else None - field_user_metadata = json.dumps(field_user_metadata) if field_user_metadata is not None else None - - placeholders.append("(%s, %s, %s, %s, %s, %s)") - values.extend( - [ - field_name, - field_criticality, - field_check, - field_filter, - field_run_config_name, - field_user_metadata, - ] - ) + logger.info(f"Saving {len(checks)} checks to Lakebase instance {config.instance_name}") - if values: - insert_values_query = f""" - INSERT INTO {config.location}.checks (name, criticality, "check", filter, run_config_name, user_metadata) - VALUES {", ".join(placeholders)} - """ - cursor.execute(insert_values_query, values) + with engine.begin() as conn: + conn.execute(text("SELECT 1")) + logger.info(f"Connected to Lakebase instance {config.instance_name}") + + conn.execute(CreateSchema(config.schema, if_not_exists=True)) + logger.info( + f"Successfully verified or created schema '{config.schema}' in database '{config.database}'." + ) - connection.commit() - logger.info(f"Successfully saved checks to {config.location}") + table = self._get_table_definition(config) + table.metadata.create_all(engine) + + 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}'" + ) + + conn.execute(insert(table), checks) + logger.info(f"Successfully saved {len(checks)} checks to {config.schema}.{config.table}") + except OperationalError as e: + logger.error(f"Failed to connect to Lakebase instance {config.instance_name}: {e}") + raise OperationalError(f"Failed to connect to Lakebase instance {config.instance_name}: {e}") from e except Exception as e: - logger.error(f"Failed to save checks to {config.instance_name}: {e}") - if connection: - connection.rollback() + logger.error(f"Failed to save checks to Lakebase: {e}") raise finally: - if cursor: - cursor.close() - if connection: - connection_pool.putconn(connection) + engine.dispose() class WorkspaceFileChecksStorageHandler(ChecksStorageHandler[WorkspaceFileChecksStorageConfig]): diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 2f047a124..9494c36cf 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -190,26 +190,34 @@ def __post_init__(self): if not self.location: raise ValueError("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: - location: The table name where the checks are stored. - run_config_name: The 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 will only replace checks for the specific run config and not all checks in the table. + instance_name: Name of the Lakebase instance. + database: Name of the database to use for checks. + schema: Name of the schema to use for checks. + table: Name of the table to use for checks. + port: Port on which to connect to the Lakebase instance. + 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 + will only replace checks for the specific run config and not all checks in the table. Default is 'overwrite'. """ - location: str instance_name: str + database: str = "dqx" + schema: str = "config" + table: str = "checks" + port: str = "5432" run_config_name: str = "default" # to filter checks by run config mode: str = "overwrite" def __post_init__(self): - if not self.location: - raise ValueError("The table name ('location' field) must not be empty or None.") + if not self.instance_name: + raise ValueError("The table name ('instance_name' field) must not be empty or None.") @dataclass @@ -247,3 +255,30 @@ class InstallationChecksStorageConfig( run_config_name: str = "default" # to retrieve run config product_name: str = "dqx" assume_user: bool = True + +@dataclass +class ConnectionInfo: + """ + Configuration class for a PostgreSQL connection. + + Args: + user: user name for the connection. + password: password for the connection. + host: host for the connection. + port: port for the connection. + database: database for the connection. + """ + + user: str + password: str + host: str + port: str + database: str + + def to_url(self) -> str: + """Generate PostgreSQL connection URL. + + Returns: + PostgreSQL connection URL. + """ + return f"postgresql://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}?sslmode=require" From 269720b2cceaf179ef31bdd8513cf76987cc2c79 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 22 Sep 2025 15:31:10 +0200 Subject: [PATCH 003/125] refactor: make load checks exception more specific --- src/databricks/labs/dqx/checks_storage.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 14479b04c..86b2f7112 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -22,7 +22,7 @@ ) from sqlalchemy.schema import CreateSchema from sqlalchemy.dialects.postgresql import JSONB -from sqlalchemy.exc import OperationalError +from sqlalchemy.exc import OperationalError, DatabaseError, ProgrammingError import yaml from pyspark.sql import SparkSession @@ -225,7 +225,7 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: checks = result.mappings().all() logger.info(f"Successfully loaded {len(checks)} checks from {config.schema}.{config.table}") return [dict(check) for check in checks] - except Exception as e: + except (OperationalError, DatabaseError, ProgrammingError) as e: logger.error(f"Failed to load checks from Lakebase instance {config.instance_name}: {e}") raise finally: From 0683135590d641e070268a484c569882dd601f0f Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 22 Sep 2025 16:47:36 +0200 Subject: [PATCH 004/125] refactor: remove connection info class and provide external engine for unit testing --- src/databricks/labs/dqx/checks_storage.py | 45 ++++++++++------------- src/databricks/labs/dqx/config.py | 30 +-------------- 2 files changed, 21 insertions(+), 54 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 86b2f7112..ddeee9a97 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -37,7 +37,6 @@ InstallationChecksStorageConfig, BaseChecksStorageConfig, VolumeFileChecksStorageConfig, - ConnectionInfo, ) from databricks.sdk import WorkspaceClient @@ -132,19 +131,19 @@ class LakebaseChecksStorageHandler(ChecksStorageHandler[LakebaseChecksStorageCon Handler for storing quality rules (checks) in a Lakebase table. """ - def __init__(self, ws: WorkspaceClient, spark: SparkSession): + def __init__(self, ws: WorkspaceClient, spark: SparkSession, engine: Engine | None = None): self.ws = ws self.spark = spark + self.engine = engine - def _get_connection_info(self, config: LakebaseChecksStorageConfig) -> ConnectionInfo: - """ - Get connection information for the Lakebase instance. + def _get_connection_url(self, config: LakebaseChecksStorageConfig) -> str: + """Generate PostgreSQL connection URL. Args: config: configuration for saving and loading checks, including instance name, database, schema, table, and port. Returns: - Connection information for the Lakebase instance. + PostgreSQL connection URL. """ instance = self.ws.database.get_database_instance(config.instance_name) cred = self.ws.database.generate_database_credential( @@ -154,25 +153,20 @@ def _get_connection_info(self, config: LakebaseChecksStorageConfig) -> Connectio user = self.ws.current_user.me().user_name password = cred.token - return ConnectionInfo( - user=user, - password=password, - host=host, - port=config.port, - database=config.database, - ) + return f"postgresql://{user}:{password}@{host}:{config.port}/{config.database}?sslmode=require" - def _get_engine(self, connection_info: ConnectionInfo) -> Engine: + def _get_engine(self, config: LakebaseChecksStorageConfig) -> Engine: """ Create a SQLAlchemy engine for the Lakebase instance. Args: - connection_info: connection information for the Lakebase instance. + config: configuration for saving and loading checks, including instance name, database, schema, table, and port. Returns: SQLAlchemy engine for the Lakebase instance. """ - return create_engine(connection_info.to_url()) + connection_url = self._get_connection_url(config) + return create_engine(connection_url) def _get_table_definition(self, config: LakebaseChecksStorageConfig) -> Table: """ @@ -207,8 +201,10 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: Returns: list of dq rules or error if checks table is missing or is invalid. """ - connection_info = self._get_connection_info(config) - engine = self._get_engine(connection_info) + if not self.engine: + engine = self._get_engine(config) + else: + engine = self.engine try: logger.info(f"Loading checks from Lakebase instance {config.instance_name}") @@ -255,8 +251,10 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: if config.mode not in ("append", "overwrite"): raise ValueError(f"Invalid mode '{config.mode}'. Must be 'append' or 'overwrite'") - connection_info = self._get_connection_info(config) - engine = self._get_engine(connection_info) + if not self.engine: + engine = self._get_engine(config) + else: + engine = self.engine try: logger.info(f"Saving {len(checks)} checks to Lakebase instance {config.instance_name}") @@ -282,10 +280,7 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: conn.execute(insert(table), checks) logger.info(f"Successfully saved {len(checks)} checks to {config.schema}.{config.table}") - except OperationalError as e: - logger.error(f"Failed to connect to Lakebase instance {config.instance_name}: {e}") - raise OperationalError(f"Failed to connect to Lakebase instance {config.instance_name}: {e}") from e - except Exception as e: + except (OperationalError, DatabaseError, ProgrammingError) as e: logger.error(f"Failed to save checks to Lakebase: {e}") raise finally: @@ -580,7 +575,7 @@ def create(self, config: BaseChecksStorageConfig) -> ChecksStorageHandler: if isinstance(config, TableChecksStorageConfig): return TableChecksStorageHandler(self.workspace_client, self.spark) if isinstance(config, LakebaseChecksStorageConfig): - return LakebaseChecksStorageHandler(self.workspace_client, self.spark) + return LakebaseChecksStorageHandler(self.workspace_client, self.spark, None) if isinstance(config, VolumeFileChecksStorageConfig): return VolumeFileChecksStorageHandler(self.workspace_client) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index c3b970a43..c99d7daf1 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -217,7 +217,7 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig): def __post_init__(self): if not self.instance_name: - raise ValueError("The table name ('instance_name' field) must not be empty or None.") + raise ValueError("The instance name ('instance_name' field) must not be empty or None.") @dataclass @@ -260,31 +260,3 @@ class InstallationChecksStorageConfig( product_name: str = "dqx" assume_user: bool = True install_folder: str | None = None - - -@dataclass -class ConnectionInfo: - """ - Configuration class for a PostgreSQL connection. - - Args: - user: user name for the connection. - password: password for the connection. - host: host for the connection. - port: port for the connection. - database: database for the connection. - """ - - user: str - password: str - host: str - port: str - database: str - - def to_url(self) -> str: - """Generate PostgreSQL connection URL. - - Returns: - PostgreSQL connection URL. - """ - return f"postgresql://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}?sslmode=require" From 8445740c5c52ce5b0a5772bcfab51b8cd70f799a Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 22 Sep 2025 16:53:18 +0200 Subject: [PATCH 005/125] test: add test for lakebase load checks functionality --- tests/unit/test_load_checks.py | 82 +++++++++++++++++++++++++++++++++- 1 file changed, 80 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index d88256144..283c9c756 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -4,10 +4,28 @@ from databricks.sdk import WorkspaceClient from databricks.sdk.errors import NotFound -from databricks.labs.dqx.checks_storage import VolumeFileChecksStorageHandler -from databricks.labs.dqx.config import VolumeFileChecksStorageConfig +from databricks.labs.dqx.checks_storage import VolumeFileChecksStorageHandler, LakebaseChecksStorageHandler +from databricks.labs.dqx.config import LakebaseChecksStorageConfig, VolumeFileChecksStorageConfig from databricks.labs.dqx.engine import DQEngineCore +import testing.postgresql +from sqlalchemy import ( + Engine, + create_engine, + text, + MetaData, + Table, + Column, + String, + Text, + insert, + select, + delete, +) +from sqlalchemy.schema import CreateSchema +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.exc import OperationalError, DatabaseError, ProgrammingError + def test_load_checks_from_local_file_json(make_local_check_file_as_json, expected_checks): file = make_local_check_file_as_json @@ -82,3 +100,63 @@ def test_file_download_contents_read_none(): with pytest.raises(NotFound, match="No contents at Unity Catalog volume path"): handler.load(VolumeFileChecksStorageConfig(location="test_path")) + + +def test_lakebase_checks_storage_handler_load(): + # Simulate successful loading of checks from Lakebase instance + ws = create_autospec(WorkspaceClient) + spark = create_autospec("pyspark.sql.SparkSession") + + schema_name = "public" + table_name = "checks" + metadata = MetaData(schema=schema_name) + + table = Table( + table_name, + metadata, + Column("name", String(255)), + Column("criticality", String(50), default="error"), + Column("check", JSONB), + Column("filter", Text), + Column("run_config_name", String(255), default="default"), + Column("user_metadata", JSONB), + ) + + expected_checks = [ + { + 'name': 'id_is_null', + 'criticality': 'error', + 'check': {'function': 'is_not_null', 'arguments': {'column': 'id'}}, + 'filter': None, + 'run_config_name': 'default', + 'user_metadata': None, + }, + { + 'name': 'name_is_null', + 'criticality': 'error', + 'check': {'function': 'is_not_null', 'arguments': {'column': 'name'}}, + 'filter': None, + 'run_config_name': 'default', + 'user_metadata': None, + }, + ] + + with testing.postgresql.Postgresql() as postgresql: + engine = create_engine(postgresql.url()) + table.metadata.create_all(engine) + + with engine.begin() as conn: + conn.execute(insert(table), expected_checks) + + handler = LakebaseChecksStorageHandler(ws, spark, engine) + config = LakebaseChecksStorageConfig(instance_name="test", schema=schema_name) + + result = handler.load(config) + + assert result == expected_checks, "Loaded checks do not match expected checks" + assert len(result) == 2, "Expected exactly 2 checks" + assert all('name' in check for check in result), "All checks should have 'name' field" + assert all('criticality' in check for check in result), "All checks should have 'criticality' field" + assert all('check' in check for check in result), "All checks should have 'check' field" + assert result[0]['name'] == 'id_is_null', "First check name should be 'id_is_null'" + assert result[1]['name'] == 'name_is_null', "Second check name should be 'name_is_null'" From 8d97cd60d64f97aba465faf1aea0c5bf3f3cf8a8 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 24 Sep 2025 13:29:30 +0200 Subject: [PATCH 006/125] refactor: remove threaded connection pool --- src/databricks/labs/dqx/checks_storage.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index ddeee9a97..235067ae7 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -6,7 +6,6 @@ from pathlib import Path from abc import ABC, abstractmethod from typing import Generic, TypeVar -from psycopg2.pool import ThreadedConnectionPool from sqlalchemy import ( Engine, create_engine, @@ -54,7 +53,6 @@ logger = logging.getLogger(__name__) T = TypeVar("T", bound=BaseChecksStorageConfig) -connection_pool: ThreadedConnectionPool | None = None class ChecksStorageHandler(ABC, Generic[T]): From b7c0803f0e899a7cc882a4830769426e376225ee Mon Sep 17 00:00:00 2001 From: Manuel Tilgner <58051108+tlgnr@users.noreply.github.com> Date: Wed, 24 Sep 2025 13:30:34 +0200 Subject: [PATCH 007/125] style: remove blank line Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/unit/test_load_checks.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index 283c9c756..493d68641 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -150,7 +150,6 @@ def test_lakebase_checks_storage_handler_load(): handler = LakebaseChecksStorageHandler(ws, spark, engine) config = LakebaseChecksStorageConfig(instance_name="test", schema=schema_name) - result = handler.load(config) assert result == expected_checks, "Loaded checks do not match expected checks" From 18bd6bb51f0802e74563f0f21db48b67e1436713 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 24 Sep 2025 13:36:51 +0200 Subject: [PATCH 008/125] refactor: only dispose engines that were created internally --- src/databricks/labs/dqx/checks_storage.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 235067ae7..eb3fb7717 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -201,8 +201,10 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: """ if not self.engine: engine = self._get_engine(config) + engine_created_internally = True else: engine = self.engine + engine_created_internally = False try: logger.info(f"Loading checks from Lakebase instance {config.instance_name}") @@ -223,7 +225,8 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: logger.error(f"Failed to load checks from Lakebase instance {config.instance_name}: {e}") raise finally: - engine.dispose() + if engine_created_internally: + engine.dispose() @telemetry_logger("save_checks", "lakebase") def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: @@ -251,8 +254,10 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: if not self.engine: engine = self._get_engine(config) + engine_created_internally = True else: engine = self.engine + engine_created_internally = False try: logger.info(f"Saving {len(checks)} checks to Lakebase instance {config.instance_name}") @@ -282,7 +287,8 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: logger.error(f"Failed to save checks to Lakebase: {e}") raise finally: - engine.dispose() + if engine_created_internally: + engine.dispose() class WorkspaceFileChecksStorageHandler(ChecksStorageHandler[WorkspaceFileChecksStorageConfig]): From 10528f79da0559ff8d23b98f0d91d19a73dfbaa9 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 24 Sep 2025 13:57:12 +0200 Subject: [PATCH 009/125] refactor: enable providing a custom user for connecting to lakebase --- src/databricks/labs/dqx/checks_storage.py | 2 +- src/databricks/labs/dqx/config.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index eb3fb7717..b70ce23ea 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -148,7 +148,7 @@ def _get_connection_url(self, config: LakebaseChecksStorageConfig) -> str: request_id=str(uuid.uuid4()), instance_names=[config.instance_name] ) host = instance.read_write_dns - user = self.ws.current_user.me().user_name + user = config.user if config.user else self.ws.current_user.me().user_name password = cred.token return f"postgresql://{user}:{password}@{host}:{config.port}/{config.database}?sslmode=require" diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index c99d7daf1..690595f43 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -202,6 +202,7 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig): schema: Name of the schema to use for checks. table: Name of the table to use for checks. port: Port on which to connect to the Lakebase instance. + user: User to use for the connection to the Lakebase instance. 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 will only replace checks for the specific run config and not all checks in the table. Default is 'overwrite'. @@ -212,6 +213,7 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig): schema: str = "config" table: str = "checks" port: str = "5432" + user: str | None = None run_config_name: str = "default" # to filter checks by run config mode: str = "overwrite" From 4ba265f1e7e311db92dfb9a71502f858088dd09a Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 24 Sep 2025 15:04:55 +0200 Subject: [PATCH 010/125] refactor: supply lakebase user only from config --- src/databricks/labs/dqx/checks_storage.py | 4 ++-- src/databricks/labs/dqx/config.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index b70ce23ea..429baf5f1 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -148,10 +148,9 @@ def _get_connection_url(self, config: LakebaseChecksStorageConfig) -> str: request_id=str(uuid.uuid4()), instance_names=[config.instance_name] ) host = instance.read_write_dns - user = config.user if config.user else self.ws.current_user.me().user_name password = cred.token - return f"postgresql://{user}:{password}@{host}:{config.port}/{config.database}?sslmode=require" + return f"postgresql://{config.user}:{password}@{host}:{config.port}/{config.database}?sslmode=require" def _get_engine(self, config: LakebaseChecksStorageConfig) -> Engine: """ @@ -462,6 +461,7 @@ def _get_storage_handler_and_config( tuple(FILE_SERIALIZERS.keys()) ): return self.table_handler, config + if config.location.startswith("/Volumes/"): return self.volume_handler, config diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 690595f43..596629054 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -202,7 +202,7 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig): schema: Name of the schema to use for checks. table: Name of the table to use for checks. port: Port on which to connect to the Lakebase instance. - user: User to use for the connection to the Lakebase instance. + user: User for the connection to the Lakebase instance. 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 will only replace checks for the specific run config and not all checks in the table. Default is 'overwrite'. @@ -214,7 +214,7 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig): table: str = "checks" port: str = "5432" user: str | None = None - run_config_name: str = "default" # to filter checks by run config + run_config_name: str = "default" mode: str = "overwrite" def __post_init__(self): From c59aa7f278a57172ac824cc396753fc83c39360b Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 24 Sep 2025 16:45:51 +0200 Subject: [PATCH 011/125] feature: add lakebase implementation for installationchecksstoragehandler --- src/databricks/labs/dqx/checks_storage.py | 79 +++++++++++++++++++++++ src/databricks/labs/dqx/config.py | 12 +++- 2 files changed, 90 insertions(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 429baf5f1..6ce9a8d74 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -49,6 +49,7 @@ from databricks.labs.dqx.telemetry import telemetry_logger from databricks.labs.dqx.utils import TABLE_PATTERN from databricks.labs.dqx.checks_serializer import FILE_SERIALIZERS +from urllib.parse import urlparse, unquote logger = logging.getLogger(__name__) @@ -411,6 +412,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]: @@ -442,6 +444,74 @@ def save(self, checks: list[dict], config: InstallationChecksStorageConfig) -> N handler, config = self._get_storage_handler_and_config(config) return handler.save(checks, config) + def _parse_lakebase_config( + self, config: InstallationChecksStorageConfig + ) -> tuple[str | None, str | None, str | None, str | None]: + """ + Parse PostgreSQL connection string to extract connection parameters. + + Expected format: postgresql://user:password@instance_name:port/database?params + Examples: + User: postgresql://user@domain.com:${PGPASSWORD}@instance-1234.database.azuredatabricks.net:5432/databricks_postgres?sslmode=require + Service Principal: postgresql://1234567890:${PGPASSWORD}@instance-1234.database.azuredatabricks.net:5432/databricks_postgres?sslmode=require + + Args: + config: Installation checks storage configuration containing the location URL + + Returns: + Tuple of (user, instance_name, port, database) - any may be None if parsing fails + + Raises: + ValueError: If the URL format is invalid or required components are missing + """ + if not config.location: + raise ValueError("Location field is empty or None - cannot parse Lakebase configuration") + + try: + parsed = urlparse(config.location) + except Exception as e: + raise ValueError(f"Failed to parse URL '{config.location}': {e}") from e + + if parsed.scheme != "postgresql": + raise ValueError( + f"Invalid URL scheme '{parsed.scheme}'. Expected 'postgresql' for Lakebase connections. " + f"URL: {config.location}" + ) + + try: + user = unquote(parsed.username) if parsed.username else None + except Exception as e: + raise ValueError(f"Failed to decode username from URL: {e}") from e + + instance_name = parsed.hostname + if not instance_name: + raise ValueError(f"Missing hostname in URL: {config.location}") + + port = None + if parsed.port: + try: + port = str(parsed.port) + except (ValueError, TypeError) as e: + raise ValueError(f"Invalid port '{parsed.port}' in URL: {e}") from e + + database = None + if parsed.path: + try: + database = parsed.path.lstrip("/") + if not database: + raise ValueError("Database name is missing") + except Exception as e: + raise ValueError(f"Failed to extract database name from connection string '{parsed.path}': {e}") from e + + if not database: + raise ValueError(f"Missing required database name in connection string: {config.location}") + + logger.debug( + f"Parsed Lakebase config - User: {user}, Instance name: {instance_name}, Port: {port}, Database: {database}" + ) + + return user, instance_name, port, database + def _get_storage_handler_and_config( self, config: InstallationChecksStorageConfig ) -> tuple[ChecksStorageHandler, InstallationChecksStorageConfig]: @@ -457,6 +527,7 @@ def _get_storage_handler_and_config( config.location = run_config.checks_location + if TABLE_PATTERN.match(config.location) and not config.location.lower().endswith( tuple(FILE_SERIALIZERS.keys()) ): @@ -464,6 +535,14 @@ def _get_storage_handler_and_config( if config.location.startswith("/Volumes/"): return self.volume_handler, config + + if config.location.startswith("postgresql://"): + user, instance_name, port, database = self._parse_lakebase_config(config) + config.user = user + config.instance_name = instance_name + config.port = port + config.database = database + return self.lakebase_handler, config if not config.location.startswith("/"): # if absolute path is not provided, the location should be set relative to the installation folder diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 596629054..b6b6298a0 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -240,7 +240,10 @@ def __post_init__(self): @dataclass class InstallationChecksStorageConfig( - WorkspaceFileChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig + WorkspaceFileChecksStorageConfig, + TableChecksStorageConfig, + VolumeFileChecksStorageConfig, + LakebaseChecksStorageConfig, ): """ Configuration class for storing checks in an installation. @@ -262,3 +265,10 @@ class InstallationChecksStorageConfig( product_name: str = "dqx" assume_user: bool = True install_folder: str | None = None + instance_name: str | None = None + database: str = "dqx" + schema: str = "config" + table: str = "checks" + port: str = "5432" + user: str | None = None + mode: str = "overwrite" From c697d515dce199460ed01d6f5286304fa6d67907 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 24 Sep 2025 17:30:42 +0200 Subject: [PATCH 012/125] test: add and refine tests for saving and loading checks --- tests/unit/test_load_checks.py | 89 ++++++++++++++---------- tests/unit/test_save_checks.py | 121 +++++++++++++++++++++++++++++++++ 2 files changed, 175 insertions(+), 35 deletions(-) diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index 493d68641..07a2247a9 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -10,21 +10,15 @@ import testing.postgresql from sqlalchemy import ( - Engine, create_engine, - text, MetaData, Table, Column, String, Text, insert, - select, - delete, ) -from sqlalchemy.schema import CreateSchema from sqlalchemy.dialects.postgresql import JSONB -from sqlalchemy.exc import OperationalError, DatabaseError, ProgrammingError def test_load_checks_from_local_file_json(make_local_check_file_as_json, expected_checks): @@ -103,24 +97,11 @@ def test_file_download_contents_read_none(): def test_lakebase_checks_storage_handler_load(): - # Simulate successful loading of checks from Lakebase instance ws = create_autospec(WorkspaceClient) spark = create_autospec("pyspark.sql.SparkSession") - schema_name = "public" - table_name = "checks" - metadata = MetaData(schema=schema_name) - - table = Table( - table_name, - metadata, - Column("name", String(255)), - Column("criticality", String(50), default="error"), - Column("check", JSONB), - Column("filter", Text), - Column("run_config_name", String(255), default="default"), - Column("user_metadata", JSONB), - ) + schema_name = "public" + table_name = "checks" expected_checks = [ { @@ -133,29 +114,67 @@ def test_lakebase_checks_storage_handler_load(): }, { 'name': 'name_is_null', - 'criticality': 'error', + 'criticality': 'warning', 'check': {'function': 'is_not_null', 'arguments': {'column': 'name'}}, - 'filter': None, - 'run_config_name': 'default', - 'user_metadata': None, + 'filter': "col1 < 3", + 'run_config_name': 'default', + 'user_metadata': {'team': 'data-engineers'}, }, ] with testing.postgresql.Postgresql() as postgresql: engine = create_engine(postgresql.url()) - table.metadata.create_all(engine) + + metadata = MetaData(schema=schema_name) + + table = Table( + table_name, + metadata, + Column("name", String(255)), + Column("criticality", String(50), default="error"), + Column("check", JSONB), + Column("filter", Text), + Column("run_config_name", String(255), default="default"), + Column("user_metadata", JSONB), + ) + + metadata.create_all(engine) with engine.begin() as conn: - conn.execute(insert(table), expected_checks) + conn.execute(insert(table), expected_checks) handler = LakebaseChecksStorageHandler(ws, spark, engine) - config = LakebaseChecksStorageConfig(instance_name="test", schema=schema_name) + + config = LakebaseChecksStorageConfig( + instance_name="test", + schema=schema_name, + ) + result = handler.load(config) - assert result == expected_checks, "Loaded checks do not match expected checks" - assert len(result) == 2, "Expected exactly 2 checks" - assert all('name' in check for check in result), "All checks should have 'name' field" - assert all('criticality' in check for check in result), "All checks should have 'criticality' field" - assert all('check' in check for check in result), "All checks should have 'check' field" - assert result[0]['name'] == 'id_is_null', "First check name should be 'id_is_null'" - assert result[1]['name'] == 'name_is_null', "Second check name should be 'name_is_null'" + assert len(result) == 2, f"Expected 2 checks, got {len(result)}" + + for check in result: + assert 'name' in check, "Missing 'name' field" + assert 'criticality' in check, "Missing 'criticality' field" + assert 'check' in check, "Missing 'check' field" + assert 'run_config_name' in check, "Missing 'run_config_name' field" + assert check['criticality'] in ['error', 'warning', 'info'], f"Invalid criticality: {check['criticality']}" + + id_check = next((c for c in result if c['name'] == 'id_is_null'), None) + name_check = next((c for c in result if c['name'] == 'name_is_null'), None) + + assert id_check is not None, "Missing 'id_is_null' check" + assert name_check is not None, "Missing 'name_is_null' check" + + assert id_check['criticality'] == 'error' + assert id_check['check'] == {'function': 'is_not_null', 'arguments': {'column': 'id'}} + assert id_check['filter'] is None + assert id_check['run_config_name'] == 'default' + assert id_check['user_metadata'] is None + + assert name_check['criticality'] == 'warning' + assert name_check['check'] == {'function': 'is_not_null', 'arguments': {'column': 'name'}} + assert name_check['filter'] == "col1 < 3" + assert name_check['run_config_name'] == 'default' + assert name_check['user_metadata'] == {'team': 'data-engineers'} diff --git a/tests/unit/test_save_checks.py b/tests/unit/test_save_checks.py index 02e4a85de..c6dd8b9d8 100644 --- a/tests/unit/test_save_checks.py +++ b/tests/unit/test_save_checks.py @@ -5,6 +5,24 @@ import yaml from databricks.labs.dqx.engine import DQEngineCore +from unittest.mock import create_autospec +from databricks.sdk import WorkspaceClient +from databricks.labs.dqx.checks_storage import LakebaseChecksStorageHandler +from databricks.labs.dqx.config import LakebaseChecksStorageConfig + +import testing.postgresql +from sqlalchemy import ( + create_engine, + MetaData, + Table, + Column, + String, + Text, + select, +) +from sqlalchemy.dialects.postgresql import JSONB +from databricks.labs.dqx.checks_storage import InstallationChecksStorageHandler +from databricks.labs.dqx.config import InstallationChecksStorageConfig TEST_CHECKS = [ @@ -62,3 +80,106 @@ def _validate_file(file_path: str, file_format: str = "yaml") -> None: if file_format == "json": json.load(file) yaml.safe_load(file) + +def test_lakebase_checks_storage_handler_save(): + ws = create_autospec(WorkspaceClient) + spark = create_autospec("pyspark.sql.SparkSession") + + schema_name = "public" + table_name = "checks" + + expected_checks = [ + { + 'name': 'id_is_null', + 'criticality': 'error', + 'check': {'function': 'is_not_null', 'arguments': {'column': 'id'}}, + 'filter': None, + 'run_config_name': 'default', + 'user_metadata': None, + }, + { + 'name': 'name_is_null', + 'criticality': 'warning', + 'check': {'function': 'is_not_null', 'arguments': {'column': 'name'}}, + 'filter': "col1 < 3", + 'run_config_name': 'extended', + 'user_metadata': {'team': 'data-engineers'}, + }, + ] + + with testing.postgresql.Postgresql() as postgresql: + engine = create_engine(postgresql.url()) + + metadata = MetaData(schema=schema_name) + + table = Table( + table_name, + metadata, + Column("name", String(255)), + Column("criticality", String(50), default="error"), + Column("check", JSONB), + Column("filter", Text), + Column("run_config_name", String(255), default="default"), + Column("user_metadata", JSONB), + ) + + metadata.create_all(engine) + + handler = LakebaseChecksStorageHandler(ws, spark, engine) + + config = LakebaseChecksStorageConfig( + instance_name="test", + schema=schema_name, + ) + + handler.save(expected_checks, config) + + with engine.connect() as conn: + result = conn.execute(select(table)).mappings().all() + result = [dict(check) for check in result] + + assert len(result) == 2, f"Expected 2 checks, got {len(result)}" + + for check in result: + assert 'name' in check, "Missing 'name' field" + assert 'criticality' in check, "Missing 'criticality' field" + assert 'check' in check, "Missing 'check' field" + assert 'run_config_name' in check, "Missing 'run_config_name' field" + assert check['criticality'] in ['error', 'warning', 'info'], f"Invalid criticality: {check['criticality']}" + + id_check = next((c for c in result if c['name'] == 'id_is_null'), None) + name_check = next((c for c in result if c['name'] == 'name_is_null'), None) + + assert id_check is not None, "Missing 'id_is_null' check" + assert name_check is not None, "Missing 'name_is_null' check" + + assert id_check['criticality'] == 'error' + assert id_check['check'] == {'function': 'is_not_null', 'arguments': {'column': 'id'}} + assert id_check['filter'] is None + assert id_check['run_config_name'] == 'default' + assert id_check['user_metadata'] is None + + assert name_check['criticality'] == 'warning' + assert name_check['check'] == {'function': 'is_not_null', 'arguments': {'column': 'name'}} + assert name_check['filter'] == "col1 < 3" + assert name_check['run_config_name'] == 'extended' + assert name_check['user_metadata'] == {'team': 'data-engineers'} + +def test_installation_checks_storage_handler_postgresql_parsing(): + ws = create_autospec(WorkspaceClient) + spark = create_autospec("pyspark.sql.SparkSession") + + handler = InstallationChecksStorageHandler(ws, spark) + config = InstallationChecksStorageConfig() + config.location = "postgresql://user@databricks.com:password@instance-test.database.azuredatabricks.net:5432/dqx?sslmode=require" + + user, instance_name, port, database = handler._parse_lakebase_config(config) + config.user = user + config.instance_name = instance_name + config.port = port + config.database = database + + assert config.user == "user@databricks.com" + assert config.instance_name == "instance-test.database.azuredatabricks.net" + assert config.port == "5432" + assert config.database == "dqx" \ No newline at end of file From d51c88696bf291feef2877d6503ff1f4022a26a6 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner <58051108+tlgnr@users.noreply.github.com> Date: Thu, 25 Sep 2025 16:33:26 +0200 Subject: [PATCH 013/125] style: remove trailing whitespace Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/unit/test_save_checks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_save_checks.py b/tests/unit/test_save_checks.py index c6dd8b9d8..08da85f4b 100644 --- a/tests/unit/test_save_checks.py +++ b/tests/unit/test_save_checks.py @@ -86,7 +86,7 @@ def test_lakebase_checks_storage_handler_save(): spark = create_autospec("pyspark.sql.SparkSession") schema_name = "public" - table_name = "checks" + table_name = "checks" expected_checks = [ { From 8cf705a7eb41c1f1618d97b19ebaaa6ec1d0b0d1 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner <58051108+tlgnr@users.noreply.github.com> Date: Thu, 25 Sep 2025 16:33:44 +0200 Subject: [PATCH 014/125] style: remove trailing whitespace Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/unit/test_load_checks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index 07a2247a9..dba8ff750 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -101,7 +101,7 @@ def test_lakebase_checks_storage_handler_load(): spark = create_autospec("pyspark.sql.SparkSession") schema_name = "public" - table_name = "checks" + table_name = "checks" expected_checks = [ { From 457cb06a42685dfda674be509610708e3f3fe382 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 25 Sep 2025 16:38:38 +0200 Subject: [PATCH 015/125] refactor: remove redundant database query --- src/databricks/labs/dqx/checks_storage.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 6ce9a8d74..ddba09c3d 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -263,9 +263,6 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: logger.info(f"Saving {len(checks)} checks to Lakebase instance {config.instance_name}") with engine.begin() as conn: - conn.execute(text("SELECT 1")) - logger.info(f"Connected to Lakebase instance {config.instance_name}") - conn.execute(CreateSchema(config.schema, if_not_exists=True)) logger.info( f"Successfully verified or created schema '{config.schema}' in database '{config.database}'." From cb194dce4d404fc15b1d25dafa95f80ef9521454 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 25 Sep 2025 17:07:54 +0200 Subject: [PATCH 016/125] refactor: move parse_lakebase_config to lakebasechecksstorageconfig --- src/databricks/labs/dqx/checks_storage.py | 74 +---------------------- src/databricks/labs/dqx/config.py | 72 +++++++++++++++++++--- 2 files changed, 66 insertions(+), 80 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index ddba09c3d..0fe2af6bb 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -49,7 +49,6 @@ from databricks.labs.dqx.telemetry import telemetry_logger from databricks.labs.dqx.utils import TABLE_PATTERN from databricks.labs.dqx.checks_serializer import FILE_SERIALIZERS -from urllib.parse import urlparse, unquote logger = logging.getLogger(__name__) @@ -441,74 +440,6 @@ def save(self, checks: list[dict], config: InstallationChecksStorageConfig) -> N handler, config = self._get_storage_handler_and_config(config) return handler.save(checks, config) - def _parse_lakebase_config( - self, config: InstallationChecksStorageConfig - ) -> tuple[str | None, str | None, str | None, str | None]: - """ - Parse PostgreSQL connection string to extract connection parameters. - - Expected format: postgresql://user:password@instance_name:port/database?params - Examples: - User: postgresql://user@domain.com:${PGPASSWORD}@instance-1234.database.azuredatabricks.net:5432/databricks_postgres?sslmode=require - Service Principal: postgresql://1234567890:${PGPASSWORD}@instance-1234.database.azuredatabricks.net:5432/databricks_postgres?sslmode=require - - Args: - config: Installation checks storage configuration containing the location URL - - Returns: - Tuple of (user, instance_name, port, database) - any may be None if parsing fails - - Raises: - ValueError: If the URL format is invalid or required components are missing - """ - if not config.location: - raise ValueError("Location field is empty or None - cannot parse Lakebase configuration") - - try: - parsed = urlparse(config.location) - except Exception as e: - raise ValueError(f"Failed to parse URL '{config.location}': {e}") from e - - if parsed.scheme != "postgresql": - raise ValueError( - f"Invalid URL scheme '{parsed.scheme}'. Expected 'postgresql' for Lakebase connections. " - f"URL: {config.location}" - ) - - try: - user = unquote(parsed.username) if parsed.username else None - except Exception as e: - raise ValueError(f"Failed to decode username from URL: {e}") from e - - instance_name = parsed.hostname - if not instance_name: - raise ValueError(f"Missing hostname in URL: {config.location}") - - port = None - if parsed.port: - try: - port = str(parsed.port) - except (ValueError, TypeError) as e: - raise ValueError(f"Invalid port '{parsed.port}' in URL: {e}") from e - - database = None - if parsed.path: - try: - database = parsed.path.lstrip("/") - if not database: - raise ValueError("Database name is missing") - except Exception as e: - raise ValueError(f"Failed to extract database name from connection string '{parsed.path}': {e}") from e - - if not database: - raise ValueError(f"Missing required database name in connection string: {config.location}") - - logger.debug( - f"Parsed Lakebase config - User: {user}, Instance name: {instance_name}, Port: {port}, Database: {database}" - ) - - return user, instance_name, port, database - def _get_storage_handler_and_config( self, config: InstallationChecksStorageConfig ) -> tuple[ChecksStorageHandler, InstallationChecksStorageConfig]: @@ -524,7 +455,6 @@ def _get_storage_handler_and_config( config.location = run_config.checks_location - if TABLE_PATTERN.match(config.location) and not config.location.lower().endswith( tuple(FILE_SERIALIZERS.keys()) ): @@ -532,9 +462,9 @@ def _get_storage_handler_and_config( if config.location.startswith("/Volumes/"): return self.volume_handler, config - + if config.location.startswith("postgresql://"): - user, instance_name, port, database = self._parse_lakebase_config(config) + user, instance_name, port, database = config._parse_lakebase_config(config.location) config.user = user config.instance_name = instance_name config.port = port diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index b6b6298a0..cdb85fae6 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -1,6 +1,7 @@ import abc from datetime import datetime, timezone from dataclasses import dataclass, field +from urllib.parse import urlparse, unquote __all__ = [ "WorkspaceConfig", @@ -213,10 +214,72 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig): schema: str = "config" table: str = "checks" port: str = "5432" - user: str | None = None + user: str run_config_name: str = "default" mode: str = "overwrite" + def _parse_lakebase_config(self, location: str) -> tuple[str | None, str | None, str | None, str | None]: + """ + Parse PostgreSQL connection string to extract connection parameters. + + Expected format: postgresql://user:password@instance_name:port/database?params + Examples: + User: postgresql://user@domain.com:${PGPASSWORD}@instance-1234.database.azuredatabricks.net:5432/databricks_postgres?sslmode=require + Service Principal: postgresql://1234567890:${PGPASSWORD}@instance-1234.database.azuredatabricks.net:5432/databricks_postgres?sslmode=require + + Args: + config: Installation checks storage configuration containing the location URL + + Returns: + Tuple of (user, instance_name, port, database) - any may be None if parsing fails + + Raises: + ValueError: If the URL format is invalid or required components are missing + """ + if not location: + raise ValueError("Location field is empty or None - cannot parse Lakebase configuration") + + try: + parsed = urlparse(location) + except Exception as e: + raise ValueError(f"Failed to parse URL '{location}': {e}") from e + + if parsed.scheme != "postgresql": + raise ValueError( + f"Invalid URL scheme '{parsed.scheme}'. Expected 'postgresql' for Lakebase connections. " + f"URL: {location}" + ) + + try: + user = unquote(parsed.username) if parsed.username else None + except Exception as e: + raise ValueError(f"Failed to decode username from URL: {e}") from e + + instance_name = parsed.hostname + if not instance_name: + raise ValueError(f"Missing hostname in URL: {location}") + + port = None + if parsed.port: + try: + port = str(parsed.port) + except (ValueError, TypeError) as e: + raise ValueError(f"Invalid port '{parsed.port}' in URL: {e}") from e + + database = None + if parsed.path: + try: + database = parsed.path.lstrip("/") + if not database: + raise ValueError("Database name is missing") + except Exception as e: + raise ValueError(f"Failed to extract database name from connection string '{parsed.path}': {e}") from e + + if not database: + raise ValueError(f"Missing required database name in connection string: {location}") + + return user, instance_name, port, database + def __post_init__(self): if not self.instance_name: raise ValueError("The instance name ('instance_name' field) must not be empty or None.") @@ -265,10 +328,3 @@ class InstallationChecksStorageConfig( product_name: str = "dqx" assume_user: bool = True install_folder: str | None = None - instance_name: str | None = None - database: str = "dqx" - schema: str = "config" - table: str = "checks" - port: str = "5432" - user: str | None = None - mode: str = "overwrite" From 694f86a48452af455bb467bb49cc83bad0206712 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 26 Sep 2025 17:23:33 +0200 Subject: [PATCH 017/125] refactor: change structure of lakebasechecksstoragehandler --- src/databricks/labs/dqx/checks_storage.py | 152 ++++++++++++++-------- src/databricks/labs/dqx/config.py | 113 ++++++++++------ 2 files changed, 168 insertions(+), 97 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 0fe2af6bb..ab9c58101 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -5,11 +5,10 @@ from io import StringIO, BytesIO from pathlib import Path from abc import ABC, abstractmethod -from typing import Generic, TypeVar +from typing import Dict, Generic, List, TypeVar from sqlalchemy import ( Engine, create_engine, - text, MetaData, Table, Column, @@ -29,6 +28,7 @@ from databricks.sdk.service.workspace import ImportFormat from databricks.labs.dqx.config import ( + LakebaseConnectionConfig, TableChecksStorageConfig, LakebaseChecksStorageConfig, FileChecksStorageConfig, @@ -126,7 +126,7 @@ def save(self, checks: list[dict], config: TableChecksStorageConfig) -> None: class LakebaseChecksStorageHandler(ChecksStorageHandler[LakebaseChecksStorageConfig]): """ - Handler for storing quality rules (checks) in a Lakebase table. + Handler for storing dq rules (checks) in a Lakebase table. """ def __init__(self, ws: WorkspaceClient, spark: SparkSession, engine: Engine | None = None): @@ -134,141 +134,178 @@ def __init__(self, ws: WorkspaceClient, spark: SparkSession, engine: Engine | No self.spark = spark self.engine = engine - def _get_connection_url(self, config: LakebaseChecksStorageConfig) -> str: - """Generate PostgreSQL connection URL. + def _get_connection_url(self, connection_config: LakebaseConnectionConfig) -> str: + """ + Generate a Lakebase connection URL. Args: - config: configuration for saving and loading checks, including instance name, database, schema, table, and port. + config: Configuration for saving and loading checks to Lakebase. Returns: - PostgreSQL connection URL. + Lakebase connection URL. """ - instance = self.ws.database.get_database_instance(config.instance_name) + instance = self.ws.database.get_database_instance(connection_config.instance_name) cred = self.ws.database.generate_database_credential( - request_id=str(uuid.uuid4()), instance_names=[config.instance_name] + request_id=str(uuid.uuid4()), instance_names=[connection_config.instance_name] ) host = instance.read_write_dns password = cred.token - return f"postgresql://{config.user}:{password}@{host}:{config.port}/{config.database}?sslmode=require" + return f"postgresql://{connection_config.user}:{password}@{host}:{connection_config.port}/{connection_config.database}?sslmode=require" - def _get_engine(self, config: LakebaseChecksStorageConfig) -> Engine: + def _get_engine(self, connection_config: LakebaseConnectionConfig) -> Engine: """ Create a SQLAlchemy engine for the Lakebase instance. Args: - config: configuration for saving and loading checks, including instance name, database, schema, table, and port. + config: Configuration for saving and loading checks to Lakebase. Returns: SQLAlchemy engine for the Lakebase instance. """ - connection_url = self._get_connection_url(config) + connection_url = self._get_connection_url(connection_config) return create_engine(connection_url) - def _get_table_definition(self, config: LakebaseChecksStorageConfig) -> Table: + def _get_schema_and_table_name(self, config: LakebaseChecksStorageConfig) -> tuple[str, str]: + """ + Extract the schema and table name from the fully qualified table name. + + Args: + config: Configuration for saving and loading checks to Lakebase. + + Returns: + A tuple containing the schema and table name. + + Raises: + ValueError: If the fully qualified table name is not in the correct format. """ - Create table definition for consistency between load and save methods. + location_components = config.location.split(".") + + if len(location_components) != 3: + raise ValueError( + f"Invalid Lakebase table name '{config.location}'. Must be in the format 'database.schema.table'." + ) + + _, schema_name, table_name = location_components + return schema_name, table_name + + def _get_table_definition(self, schema_name: str, table_name: str) -> Table: + """ + Create a table definition for consistency between the load and save methods. Args: - config: configuration for saving and loading checks, including instance name, database, schema, table, and port. + schema: The schema name where the checks table is located. + table: The table name where the checks are stored. Returns: SQLAlchemy table definition for the Lakebase instance. """ - metadata = MetaData(schema=config.schema) return Table( - config.table, - metadata, + table_name, + MetaData(schema=schema_name), Column("name", String(255)), - Column("criticality", String(50), default="error"), + Column("criticality", String(50), server_default="error"), Column("check", JSONB), Column("filter", Text), - Column("run_config_name", String(255), default=config.run_config_name), + Column("run_config_name", String(255), server_default="default"), Column("user_metadata", JSONB), ) @telemetry_logger("load_checks", "lakebase") - def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: + def load(self, config: LakebaseChecksStorageConfig) -> List[Dict]: """ Load dq rules (checks) from a Lakebase table. Args: - config: configuration for saving and loading checks, including instance name, database, schema, table, and port. + config: Configuration for saving and loading checks to Lakebase. Returns: - list of dq rules or error if checks table is missing or is invalid. + List of dq rules or error if loading checks fails. + + Raises: + OperationalError: If connection to Lakebase instance fails. + DatabaseError: If loading checks fails. + ProgrammingError: If loading checks fails. """ + connection_config = LakebaseConnectionConfig._parse_connection_string(config.connection_string) + + engine_created_internally = False if not self.engine: - engine = self._get_engine(config) + engine = self._get_engine(connection_config) engine_created_internally = True else: engine = self.engine - engine_created_internally = False try: - logger.info(f"Loading checks from Lakebase instance {config.instance_name}") - - table = self._get_table_definition(config) - table.metadata.create_all(engine) + logger.info(f"Loading checks from Lakebase instance {connection_config.instance_name}") + schema_name, table_name = self._get_schema_and_table_name(config) + table = self._get_table_definition(schema_name, table_name) stmt = select(table) + if config.run_config_name: stmt = stmt.where(table.c.run_config_name == config.run_config_name) with engine.connect() as conn: result = conn.execute(stmt) checks = result.mappings().all() - logger.info(f"Successfully loaded {len(checks)} checks from {config.schema}.{config.table}") + logger.info(f"Successfully loaded {len(checks)} checks from {schema_name}.{table_name}") return [dict(check) for check in checks] except (OperationalError, DatabaseError, ProgrammingError) as e: - logger.error(f"Failed to load checks from Lakebase instance {config.instance_name}: {e}") + logger.error(f"Failed to load checks from Lakebase instance {connection_config.instance_name}: {e}") raise finally: if engine_created_internally: engine.dispose() @telemetry_logger("save_checks", "lakebase") - def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: + 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, including instance name, database, schema, table, and port. + checks: List of dq rules (checks) to save. + config: Configuration for saving and loading checks to Lakebase. Returns: None Raises: OperationalError: If connection to Lakebase instance fails. + DatabaseError: If saving checks fails. + ProgrammingError: If saving checks fails. ValueError: If invalid mode is specified. - Exception: If saving checks fails. """ if not checks: - logger.warning("No checks provided to save") + logger.warning("No checks provided to save.") return if config.mode not in ("append", "overwrite"): - raise ValueError(f"Invalid mode '{config.mode}'. Must be 'append' or 'overwrite'") + raise ValueError(f"Invalid mode '{config.mode}'. Must be 'append' or 'overwrite'.") + connection_config = LakebaseConnectionConfig._parse_connection_string(config.connection_string) + + engine_created_internally = False if not self.engine: - engine = self._get_engine(config) + engine = self._get_engine(connection_config) engine_created_internally = True else: engine = self.engine - engine_created_internally = False + + schema_name, table_name = self._get_schema_and_table_name(config) try: - logger.info(f"Saving {len(checks)} checks to Lakebase instance {config.instance_name}") + logger.info(f"Saving {len(checks)} checks to Lakebase instance {connection_config.instance_name}") with engine.begin() as conn: - conn.execute(CreateSchema(config.schema, if_not_exists=True)) - logger.info( - f"Successfully verified or created schema '{config.schema}' in database '{config.database}'." - ) + if not conn.dialect.has_schema(conn, schema_name): + conn.execute(CreateSchema(schema_name)) + logger.info( + f"Successfully created schema '{schema_name}' in database '{connection_config.database}'." + ) - table = self._get_table_definition(config) - table.metadata.create_all(engine) + table = self._get_table_definition(schema_name, table_name) + table.metadata.create_all(engine, checkfirst=True) if config.mode == "overwrite": delete_stmt = delete(table).where(table.c.run_config_name == config.run_config_name) @@ -277,8 +314,9 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: f"Deleted {result.rowcount} existing checks for run_config_name '{config.run_config_name}'" ) - conn.execute(insert(table), checks) - logger.info(f"Successfully saved {len(checks)} checks to {config.schema}.{config.table}") + insert_stmt = insert(table) + conn.execute(insert_stmt, checks) + logger.info(f"Inserted {len(checks)} new checks into {schema_name}.{table_name}") except (OperationalError, DatabaseError, ProgrammingError) as e: logger.error(f"Failed to save checks to Lakebase: {e}") raise @@ -460,17 +498,17 @@ def _get_storage_handler_and_config( ): return self.table_handler, config + if ( + TABLE_PATTERN.match(config.location) + and not config.location.lower().endswith(tuple(FILE_SERIALIZERS.keys())) + and config.connection_string + and config.connection_string.startswith("postgresql://") + ): + return self.lakebase_handler, config + if config.location.startswith("/Volumes/"): return self.volume_handler, config - if config.location.startswith("postgresql://"): - user, instance_name, port, database = config._parse_lakebase_config(config.location) - config.user = user - config.instance_name = instance_name - config.port = port - config.database = database - return self.lakebase_handler, config - if not config.location.startswith("/"): # if absolute path is not provided, the location should be set relative to the installation folder workspace_path = f"{installation.install_folder()}/{run_config.checks_location}" diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index cdb85fae6..1a1af6475 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -19,6 +19,8 @@ "VolumeFileChecksStorageConfig", ] +LAKEBASE_DEFAULT_PORT = 5432 + @dataclass class InputConfig: @@ -193,76 +195,70 @@ def __post_init__(self): @dataclass -class LakebaseChecksStorageConfig(BaseChecksStorageConfig): +class LakebaseConnectionConfig: """ - Configuration class for storing checks in a Lakebase table. + Configuration class for a Lakebase connection. + + Note: + The connection string format includes a placeholder for a password, + e.g., postgresql://user:password@host:port/database), but the password + is not stored or parsed by this class. Args: - instance_name: Name of the Lakebase instance. - database: Name of the database to use for checks. - schema: Name of the schema to use for checks. - table: Name of the table to use for checks. - port: Port on which to connect to the Lakebase instance. - user: User for the connection to the Lakebase instance. - 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 - will only replace checks for the specific run config and not all checks in the table. Default is 'overwrite'. + instance_name: The Lakebase instance name (hostname). + database: The Lakebase database name. + user: The Lakebase username. + port: The Lakebase port (default is '5432'). """ instance_name: str - database: str = "dqx" - schema: str = "config" - table: str = "checks" - port: str = "5432" + database: str user: str - run_config_name: str = "default" - mode: str = "overwrite" + port: int = LAKEBASE_DEFAULT_PORT - def _parse_lakebase_config(self, location: str) -> tuple[str | None, str | None, str | None, str | None]: + @staticmethod + def _parse_connection_string(connection_string: str) -> "LakebaseConnectionConfig": """ Parse PostgreSQL connection string to extract connection parameters. - Expected format: postgresql://user:password@instance_name:port/database?params - Examples: - User: postgresql://user@domain.com:${PGPASSWORD}@instance-1234.database.azuredatabricks.net:5432/databricks_postgres?sslmode=require - Service Principal: postgresql://1234567890:${PGPASSWORD}@instance-1234.database.azuredatabricks.net:5432/databricks_postgres?sslmode=require + Expected format: postgresql://user:password@instance_name:port/database?params. Args: - config: Installation checks storage configuration containing the location URL + connection_string: Lakebase (SQLAlchemy) connection string. Returns: - Tuple of (user, instance_name, port, database) - any may be None if parsing fails + Instance of LakebaseConnectionConfig with extracted parameters. Raises: - ValueError: If the URL format is invalid or required components are missing + ValueError: If the URL format is invalid or required components are missing. """ - if not location: - raise ValueError("Location field is empty or None - cannot parse Lakebase configuration") + if not connection_string: + raise ValueError("Connection string cannot be empty or None.") try: - parsed = urlparse(location) + parsed = urlparse(connection_string) except Exception as e: - raise ValueError(f"Failed to parse URL '{location}': {e}") from e + raise ValueError(f"Failed to parse URL '{connection_string}': {e}") from e if parsed.scheme != "postgresql": - raise ValueError( - f"Invalid URL scheme '{parsed.scheme}'. Expected 'postgresql' for Lakebase connections. " - f"URL: {location}" - ) + raise ValueError(f"Invalid URL scheme '{parsed.scheme}'. Expected 'postgresql' for Lakebase connections.") try: user = unquote(parsed.username) if parsed.username else None except Exception as e: raise ValueError(f"Failed to decode username from URL: {e}") from e + if not user: + raise ValueError(f"Missing username in URL: {connection_string}") + instance_name = parsed.hostname if not instance_name: - raise ValueError(f"Missing hostname in URL: {location}") + raise ValueError(f"Missing hostname in URL: {connection_string}") - port = None + port = LAKEBASE_DEFAULT_PORT if parsed.port: try: - port = str(parsed.port) + port = int(parsed.port) except (ValueError, TypeError) as e: raise ValueError(f"Invalid port '{parsed.port}' in URL: {e}") from e @@ -276,13 +272,50 @@ def _parse_lakebase_config(self, location: str) -> tuple[str | None, str | None, raise ValueError(f"Failed to extract database name from connection string '{parsed.path}': {e}") from e if not database: - raise ValueError(f"Missing required database name in connection string: {location}") + raise ValueError(f"Missing required database name in connection string: {connection_string}") + + return LakebaseConnectionConfig( + instance_name=instance_name, + database=database, + user=user, + port=port, + ) + + +@dataclass +class LakebaseChecksStorageConfig(BaseChecksStorageConfig): + """ + Configuration class for storing checks in a Lakebase table. + + Note: + The `schema` and `table` fields are not parsed from the connection string, but + must be provided in the `location` field in the format 'database.schema.table'. - return user, instance_name, port, database + Args: + location: Fully qualified Lakebase table name in the format 'database.schema.table'. + connection_string: (SQLAlchemy) Lakebase connection string, e.g., postgresql://user:password@host:port/database. + 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'). + """ + + location: str + connection_string: str + run_config_name: str = "default" + mode: str = "overwrite" def __post_init__(self): - if not self.instance_name: - raise ValueError("The instance name ('instance_name' field) must not be empty or None.") + if not self.location: + raise ValueError("The location ('location' field) must not be empty or None.") + + if not self.connection_string: + raise ValueError("The connection string ('connection_string' field) must not be empty or None.") + + if self.connection_string and self.connection_string.startswith("postgresql://"): + try: + LakebaseConnectionConfig._parse_connection_string(self.connection_string) + except Exception as e: + raise ValueError(f"Failed to parse connection string '{self.connection_string}': {e}") from e @dataclass From 6d02f17de358a9fa9259044e51e8b414b6c50f12 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 26 Sep 2025 17:27:43 +0200 Subject: [PATCH 018/125] refactor: make lakebase port a string --- src/databricks/labs/dqx/config.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 1a1af6475..bf3cdbca2 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -19,7 +19,7 @@ "VolumeFileChecksStorageConfig", ] -LAKEBASE_DEFAULT_PORT = 5432 +LAKEBASE_DEFAULT_PORT = "5432" @dataclass @@ -214,7 +214,7 @@ class LakebaseConnectionConfig: instance_name: str database: str user: str - port: int = LAKEBASE_DEFAULT_PORT + port: str = LAKEBASE_DEFAULT_PORT @staticmethod def _parse_connection_string(connection_string: str) -> "LakebaseConnectionConfig": @@ -258,7 +258,7 @@ def _parse_connection_string(connection_string: str) -> "LakebaseConnectionConfi port = LAKEBASE_DEFAULT_PORT if parsed.port: try: - port = int(parsed.port) + port = str(parsed.port) except (ValueError, TypeError) as e: raise ValueError(f"Invalid port '{parsed.port}' in URL: {e}") from e From 2d7c011ea4feb47071134941cebe0fc658b8c29a Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 26 Sep 2025 17:31:37 +0200 Subject: [PATCH 019/125] refactor: remove code duplication --- src/databricks/labs/dqx/checks_storage.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index ab9c58101..482ae9c04 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -493,17 +493,15 @@ def _get_storage_handler_and_config( config.location = run_config.checks_location - if TABLE_PATTERN.match(config.location) and not config.location.lower().endswith( + matches_table_pattern = TABLE_PATTERN.match(config.location) and not config.location.lower().endswith( tuple(FILE_SERIALIZERS.keys()) - ): + ) + matches_lakebase_pattern = config.connection_string and config.connection_string.startswith("postgresql://") + + if matches_table_pattern: return self.table_handler, config - if ( - TABLE_PATTERN.match(config.location) - and not config.location.lower().endswith(tuple(FILE_SERIALIZERS.keys())) - and config.connection_string - and config.connection_string.startswith("postgresql://") - ): + if matches_table_pattern and matches_lakebase_pattern: return self.lakebase_handler, config if config.location.startswith("/Volumes/"): From 7e1de5ae2e3ca9a40abd767a1746f12f9ec46528 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner <58051108+tlgnr@users.noreply.github.com> Date: Fri, 26 Sep 2025 17:41:34 +0200 Subject: [PATCH 020/125] refactor: change order of if statements Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/databricks/labs/dqx/checks_storage.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 482ae9c04..e81b2b9b7 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -498,12 +498,11 @@ def _get_storage_handler_and_config( ) matches_lakebase_pattern = config.connection_string and config.connection_string.startswith("postgresql://") - if matches_table_pattern: - return self.table_handler, config - if matches_table_pattern and matches_lakebase_pattern: return self.lakebase_handler, config + if matches_table_pattern: + return self.table_handler, config if config.location.startswith("/Volumes/"): return self.volume_handler, config From 3d0c803aea471524241e888f36c8ffb1f78206d4 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 26 Sep 2025 17:58:02 +0200 Subject: [PATCH 021/125] refactor: add default value in installationchecksstorageconfig for connection string field --- src/databricks/labs/dqx/config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index bf3cdbca2..357b7a86f 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -356,6 +356,7 @@ class InstallationChecksStorageConfig( * Global directory if `DQX_FORCE_INSTALL=global`: "/Applications/dqx" """ + connection_string: str | None = None # retrieved from the installation config location: str = "installation" # retrieved from the installation config run_config_name: str = "default" # to retrieve run config product_name: str = "dqx" From cc21466947b6d8dcdf613e9bfab489420b8da8b4 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 26 Sep 2025 18:33:31 +0200 Subject: [PATCH 022/125] refactor: streamline unit test for lakebasechecksstoragehandler --- tests/unit/test_save_checks.py | 99 ++++++++++++---------------------- 1 file changed, 33 insertions(+), 66 deletions(-) diff --git a/tests/unit/test_save_checks.py b/tests/unit/test_save_checks.py index 08da85f4b..195de6337 100644 --- a/tests/unit/test_save_checks.py +++ b/tests/unit/test_save_checks.py @@ -1,28 +1,19 @@ import os import json - -import pytest import yaml -from databricks.labs.dqx.engine import DQEngineCore +import pytest from unittest.mock import create_autospec +from sqlalchemy import create_engine, select + from databricks.sdk import WorkspaceClient + +from databricks.labs.dqx.engine import DQEngineCore from databricks.labs.dqx.checks_storage import LakebaseChecksStorageHandler from databricks.labs.dqx.config import LakebaseChecksStorageConfig +from databricks.labs.dqx.config import LakebaseConnectionConfig import testing.postgresql -from sqlalchemy import ( - create_engine, - MetaData, - Table, - Column, - String, - Text, - select, -) -from sqlalchemy.dialects.postgresql import JSONB -from databricks.labs.dqx.checks_storage import InstallationChecksStorageHandler -from databricks.labs.dqx.config import InstallationChecksStorageConfig TEST_CHECKS = [ @@ -81,12 +72,11 @@ def _validate_file(file_path: str, file_format: str = "yaml") -> None: json.load(file) yaml.safe_load(file) + def test_lakebase_checks_storage_handler_save(): ws = create_autospec(WorkspaceClient) spark = create_autospec("pyspark.sql.SparkSession") - - schema_name = "public" - table_name = "checks" + location = "test.public.checks" expected_checks = [ { @@ -99,39 +89,23 @@ def test_lakebase_checks_storage_handler_save(): }, { 'name': 'name_is_null', - 'criticality': 'warning', + 'criticality': 'warning', 'check': {'function': 'is_not_null', 'arguments': {'column': 'name'}}, - 'filter': "col1 < 3", - 'run_config_name': 'extended', - 'user_metadata': {'team': 'data-engineers'}, + 'filter': "col1 < 3", + 'run_config_name': 'extended', + 'user_metadata': {'team': 'data-engineers'}, }, ] with testing.postgresql.Postgresql() as postgresql: - engine = create_engine(postgresql.url()) - - metadata = MetaData(schema=schema_name) - - table = Table( - table_name, - metadata, - Column("name", String(255)), - Column("criticality", String(50), default="error"), - Column("check", JSONB), - Column("filter", Text), - Column("run_config_name", String(255), default="default"), - Column("user_metadata", JSONB), - ) - - metadata.create_all(engine) + connection_string = postgresql.url() + engine = create_engine(connection_string) handler = LakebaseChecksStorageHandler(ws, spark, engine) - - config = LakebaseChecksStorageConfig( - instance_name="test", - schema=schema_name, - ) - + config = LakebaseChecksStorageConfig(location, connection_string) + schema_name, table_name = handler._get_schema_and_table_name(config) + table = handler._get_table_definition(schema_name, table_name) + handler.save(expected_checks, config) with engine.connect() as conn: @@ -139,47 +113,40 @@ def test_lakebase_checks_storage_handler_save(): result = [dict(check) for check in result] assert len(result) == 2, f"Expected 2 checks, got {len(result)}" - + for check in result: assert 'name' in check, "Missing 'name' field" assert 'criticality' in check, "Missing 'criticality' field" assert 'check' in check, "Missing 'check' field" assert 'run_config_name' in check, "Missing 'run_config_name' field" assert check['criticality'] in ['error', 'warning', 'info'], f"Invalid criticality: {check['criticality']}" - + id_check = next((c for c in result if c['name'] == 'id_is_null'), None) name_check = next((c for c in result if c['name'] == 'name_is_null'), None) - + assert id_check is not None, "Missing 'id_is_null' check" assert name_check is not None, "Missing 'name_is_null' check" - + assert id_check['criticality'] == 'error' assert id_check['check'] == {'function': 'is_not_null', 'arguments': {'column': 'id'}} assert id_check['filter'] is None assert id_check['run_config_name'] == 'default' assert id_check['user_metadata'] is None - + assert name_check['criticality'] == 'warning' assert name_check['check'] == {'function': 'is_not_null', 'arguments': {'column': 'name'}} assert name_check['filter'] == "col1 < 3" assert name_check['run_config_name'] == 'extended' assert name_check['user_metadata'] == {'team': 'data-engineers'} + def test_installation_checks_storage_handler_postgresql_parsing(): - ws = create_autospec(WorkspaceClient) - spark = create_autospec("pyspark.sql.SparkSession") - - handler = InstallationChecksStorageHandler(ws, spark) - config = InstallationChecksStorageConfig() - config.location = "postgresql://user@databricks.com:password@instance-test.database.azuredatabricks.net:5432/dqx?sslmode=require" - - user, instance_name, port, database = handler._parse_lakebase_config(config) - config.user = user - config.instance_name = instance_name - config.port = port - config.database = database - - assert config.user == "user@databricks.com" - assert config.instance_name == "instance-test.database.azuredatabricks.net" - assert config.port == "5432" - assert config.database == "dqx" \ No newline at end of file + connection_string = ( + "postgresql://user@databricks.com:password@instance-test.database.azuredatabricks.net:5432/dqx?sslmode=require" + ) + connection_config = LakebaseConnectionConfig._parse_connection_string(connection_string) + + assert connection_config.user == "user@databricks.com" + assert connection_config.instance_name == "instance-test.database.azuredatabricks.net" + assert connection_config.port == "5432" + assert connection_config.database == "dqx" From 0b88f5a10fea2f3acadc8689dc0b588b800b6a56 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 26 Sep 2025 18:45:29 +0200 Subject: [PATCH 023/125] style: insert line break --- src/databricks/labs/dqx/checks_storage.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 482ae9c04..036ad0dbd 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -241,8 +241,8 @@ def load(self, config: LakebaseChecksStorageConfig) -> List[Dict]: schema_name, table_name = self._get_schema_and_table_name(config) table = self._get_table_definition(schema_name, table_name) + stmt = select(table) - if config.run_config_name: stmt = stmt.where(table.c.run_config_name == config.run_config_name) From 2af23a232ec24f11cc3bcd574c159f3c288efcf0 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 26 Sep 2025 18:45:56 +0200 Subject: [PATCH 024/125] refactor: revise load lakebasechecksstoragehandler test --- tests/unit/test_load_checks.py | 69 +++++++++++++++------------------- 1 file changed, 30 insertions(+), 39 deletions(-) diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index dba8ff750..75271aa1c 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -1,5 +1,6 @@ -from unittest.mock import create_autospec import pytest +from unittest.mock import create_autospec + from databricks.sdk.service.files import DownloadResponse from databricks.sdk import WorkspaceClient from databricks.sdk.errors import NotFound @@ -9,16 +10,9 @@ from databricks.labs.dqx.engine import DQEngineCore import testing.postgresql -from sqlalchemy import ( - create_engine, - MetaData, - Table, - Column, - String, - Text, - insert, -) -from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy import Column, MetaData, create_engine, insert +from sqlalchemy.schema import Table +from sqlalchemy.types import JSON as JSONB, String, Text def test_load_checks_from_local_file_json(make_local_check_file_as_json, expected_checks): @@ -99,9 +93,7 @@ def test_file_download_contents_read_none(): def test_lakebase_checks_storage_handler_load(): ws = create_autospec(WorkspaceClient) spark = create_autospec("pyspark.sql.SparkSession") - - schema_name = "public" - table_name = "checks" + location = "test.public.checks" expected_checks = [ { @@ -114,65 +106,64 @@ def test_lakebase_checks_storage_handler_load(): }, { 'name': 'name_is_null', - 'criticality': 'warning', + 'criticality': 'warning', 'check': {'function': 'is_not_null', 'arguments': {'column': 'name'}}, - 'filter': "col1 < 3", - 'run_config_name': 'default', - 'user_metadata': {'team': 'data-engineers'}, + 'filter': "col1 < 3", + 'run_config_name': 'default', + 'user_metadata': {'team': 'data-engineers'}, }, ] with testing.postgresql.Postgresql() as postgresql: - engine = create_engine(postgresql.url()) - - metadata = MetaData(schema=schema_name) + connection_string = postgresql.url() + engine = create_engine(connection_string) + handler = LakebaseChecksStorageHandler(ws, spark, engine) + config = LakebaseChecksStorageConfig(location, connection_string) + schema_name, table_name = handler._get_schema_and_table_name(config) + metadata = MetaData(schema=schema_name) table = Table( table_name, metadata, Column("name", String(255)), - Column("criticality", String(50), default="error"), + Column("criticality", String(50), server_default="error"), Column("check", JSONB), Column("filter", Text), - Column("run_config_name", String(255), default="default"), + Column("run_config_name", String(255), server_default="default"), Column("user_metadata", JSONB), ) - - metadata.create_all(engine) + metadata.create_all(engine, checkfirst=True) with engine.begin() as conn: - conn.execute(insert(table), expected_checks) + conn.execute(insert(table), expected_checks) - handler = LakebaseChecksStorageHandler(ws, spark, engine) - - config = LakebaseChecksStorageConfig( - instance_name="test", - schema=schema_name, - ) - result = handler.load(config) assert len(result) == 2, f"Expected 2 checks, got {len(result)}" - + for check in result: assert 'name' in check, "Missing 'name' field" assert 'criticality' in check, "Missing 'criticality' field" assert 'check' in check, "Missing 'check' field" assert 'run_config_name' in check, "Missing 'run_config_name' field" - assert check['criticality'] in ['error', 'warning', 'info'], f"Invalid criticality: {check['criticality']}" - + assert check['criticality'] in [ + 'error', + 'warning', + 'info', + ], f"Invalid criticality: {check['criticality']}" + id_check = next((c for c in result if c['name'] == 'id_is_null'), None) name_check = next((c for c in result if c['name'] == 'name_is_null'), None) - + assert id_check is not None, "Missing 'id_is_null' check" assert name_check is not None, "Missing 'name_is_null' check" - + assert id_check['criticality'] == 'error' assert id_check['check'] == {'function': 'is_not_null', 'arguments': {'column': 'id'}} assert id_check['filter'] is None assert id_check['run_config_name'] == 'default' assert id_check['user_metadata'] is None - + assert name_check['criticality'] == 'warning' assert name_check['check'] == {'function': 'is_not_null', 'arguments': {'column': 'name'}} assert name_check['filter'] == "col1 < 3" From 31036d9f794044ae42c925fe0c6fe361236eb0fa Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 29 Sep 2025 11:41:56 +0200 Subject: [PATCH 025/125] chore: add testing library to project.toml --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index ae839f960..17388f141 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,6 +97,8 @@ dependencies = [ "dbldatagen~=0.4.0", "pyparsing~=3.2.3", "jmespath~=1.0.1", + "testing.postgresql~=1.3.0", + "testing.common.database~=2.0.3", ] python="3.12" # must match the version required by databricks-connect and python version on the test clusters From c60637f619bda1a0d8d5766b86853bc911922847 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 29 Sep 2025 11:47:31 +0200 Subject: [PATCH 026/125] chore: add psycopg2-binary library to projec.toml --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 17388f141..e3601b3de 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -99,6 +99,7 @@ dependencies = [ "jmespath~=1.0.1", "testing.postgresql~=1.3.0", "testing.common.database~=2.0.3", + psycopg2-binary~=2.9.10, # required by testing.postgresql ] python="3.12" # must match the version required by databricks-connect and python version on the test clusters From 43c3cfb10d0e8caed57d0ceb43294a289d0a599f Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 29 Sep 2025 11:53:00 +0200 Subject: [PATCH 027/125] fix: allow connection_string to be optional in LakebaseChecksStorageConfig --- src/databricks/labs/dqx/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 357b7a86f..4cd835c15 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -300,7 +300,7 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig): """ location: str - connection_string: str + connection_string: str | None = None run_config_name: str = "default" mode: str = "overwrite" From b28443d5071fc0764712cd8162f6c1e6ae49a61e Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 29 Sep 2025 11:58:54 +0200 Subject: [PATCH 028/125] fix: add quotation marks to library name --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index e3601b3de..9af39a2c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -99,7 +99,7 @@ dependencies = [ "jmespath~=1.0.1", "testing.postgresql~=1.3.0", "testing.common.database~=2.0.3", - psycopg2-binary~=2.9.10, # required by testing.postgresql + "psycopg2-binary~=2.9.10", ] python="3.12" # must match the version required by databricks-connect and python version on the test clusters From 166f60f297e1f40de940f9f638ae24cc660e592b Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 29 Sep 2025 12:04:57 +0200 Subject: [PATCH 029/125] refactor: simplify unit test test_lakebase_checks_storage_handler_load --- tests/unit/test_load_checks.py | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index 75271aa1c..26c7dceff 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -121,18 +121,8 @@ def test_lakebase_checks_storage_handler_load(): config = LakebaseChecksStorageConfig(location, connection_string) schema_name, table_name = handler._get_schema_and_table_name(config) - metadata = MetaData(schema=schema_name) - table = Table( - table_name, - metadata, - 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), - ) - metadata.create_all(engine, checkfirst=True) + table = handler._get_table_definition(schema_name=schema_name, table_name=table_name) + table.metadata.create_all(engine, checkfirst=True) with engine.begin() as conn: conn.execute(insert(table), expected_checks) From d8108ee5629f6c80cb990198faa983ff140a59b3 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 29 Sep 2025 12:06:47 +0200 Subject: [PATCH 030/125] docs: remove errant paranthesis --- src/databricks/labs/dqx/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 4cd835c15..95aa99104 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -201,7 +201,7 @@ class LakebaseConnectionConfig: Note: The connection string format includes a placeholder for a password, - e.g., postgresql://user:password@host:port/database), but the password + e.g., postgresql://user:password@host:port/database, but the password is not stored or parsed by this class. Args: From edc0fa3fe5af785eee98d7fa4c09cbbf12e18f29 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 29 Sep 2025 13:03:14 +0200 Subject: [PATCH 031/125] refactor: allow connection_string in _parse_connection_string method to be None --- src/databricks/labs/dqx/config.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 95aa99104..48823715f 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -201,7 +201,7 @@ class LakebaseConnectionConfig: Note: The connection string format includes a placeholder for a password, - e.g., postgresql://user:password@host:port/database, but the password + e.g., postgresql://user:password@host:port/database, but the password is not stored or parsed by this class. Args: @@ -217,7 +217,7 @@ class LakebaseConnectionConfig: port: str = LAKEBASE_DEFAULT_PORT @staticmethod - def _parse_connection_string(connection_string: str) -> "LakebaseConnectionConfig": + def _parse_connection_string(connection_string: str | None) -> "LakebaseConnectionConfig": """ Parse PostgreSQL connection string to extract connection parameters. @@ -287,8 +287,8 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig): """ Configuration class for storing checks in a Lakebase table. - Note: - The `schema` and `table` fields are not parsed from the connection string, but + Note: + The `schema` and `table` fields are not parsed from the connection string, but must be provided in the `location` field in the format 'database.schema.table'. Args: From 49774427fe9c7749f8e42e10e4c20633d44c3b0c Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 29 Sep 2025 13:10:12 +0200 Subject: [PATCH 032/125] chore: exclude testing packet from type checking --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 9af39a2c5..504090b8d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -162,7 +162,7 @@ profile = "black" exclude = ['venv', '.venv', 'demos/*', 'tests/e2e/*'] [[tool.mypy.overrides]] -module = ["google", "google.*"] # Google libraries are not type annotated +module = ["google", "google.*", "testing.postgresql", "testing.common.database"] # Google and testing libraries are not type annotated ignore_missing_imports = true [tool.pytest.ini_options] From 010166d88c208221e0b72831da7754608098dbb9 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 29 Sep 2025 13:16:02 +0200 Subject: [PATCH 033/125] chore: exclude testing.* library from type checking --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 504090b8d..7b5b6da3a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -162,7 +162,7 @@ profile = "black" exclude = ['venv', '.venv', 'demos/*', 'tests/e2e/*'] [[tool.mypy.overrides]] -module = ["google", "google.*", "testing.postgresql", "testing.common.database"] # Google and testing libraries are not type annotated +module = ["google", "google.*", "testing.*"] # Google and testing libraries are not type annotated ignore_missing_imports = true [tool.pytest.ini_options] From 01ba12371c2efbb9704958191937f3377471793d Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 29 Sep 2025 14:39:06 +0200 Subject: [PATCH 034/125] refactor: simplified load and save Lakebase checks test cases --- src/databricks/labs/dqx/utils.py | 12 ++++ tests/unit/test_load_checks.py | 101 +++++++++++++------------------ tests/unit/test_save_checks.py | 92 +++++++++------------------- 3 files changed, 83 insertions(+), 122 deletions(-) diff --git a/src/databricks/labs/dqx/utils.py b/src/databricks/labs/dqx/utils.py index fd81e54c9..73d11946a 100644 --- a/src/databricks/labs/dqx/utils.py +++ b/src/databricks/labs/dqx/utils.py @@ -339,3 +339,15 @@ def safe_json_load(value: str): return json.loads(value) # load as json if possible except json.JSONDecodeError: return value + +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/unit/test_load_checks.py b/tests/unit/test_load_checks.py index 26c7dceff..9fce84116 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -1,18 +1,37 @@ -import pytest from unittest.mock import create_autospec -from databricks.sdk.service.files import DownloadResponse +import pytest +import testing.postgresql +from pyspark.sql import SparkSession +from sqlalchemy import create_engine, insert + 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, LakebaseChecksStorageHandler from databricks.labs.dqx.config import LakebaseChecksStorageConfig, VolumeFileChecksStorageConfig from databricks.labs.dqx.engine import DQEngineCore - -import testing.postgresql -from sqlalchemy import Column, MetaData, create_engine, insert -from sqlalchemy.schema import Table -from sqlalchemy.types import JSON as JSONB, String, Text +from databricks.labs.dqx.utils import sort_key + + +TEST_CHECKS = [ + { + "name": "id_is_null", + "criticality": "error", + "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, + }, + { + "name": "id_is_null", + "criticality": "error", + "check": { + "function": "is_not_null", + "for_each_column": ["col1", "col2"], + "arguments": {}, + "user_metadata": {"rule_type": "completeness"}, + }, + }, +] def test_load_checks_from_local_file_json(make_local_check_file_as_json, expected_checks): @@ -92,70 +111,32 @@ def test_file_download_contents_read_none(): def test_lakebase_checks_storage_handler_load(): ws = create_autospec(WorkspaceClient) - spark = create_autospec("pyspark.sql.SparkSession") + spark = create_autospec(SparkSession) location = "test.public.checks" - expected_checks = [ - { - 'name': 'id_is_null', - 'criticality': 'error', - 'check': {'function': 'is_not_null', 'arguments': {'column': 'id'}}, - 'filter': None, - 'run_config_name': 'default', - 'user_metadata': None, - }, - { - 'name': 'name_is_null', - 'criticality': 'warning', - 'check': {'function': 'is_not_null', 'arguments': {'column': 'name'}}, - 'filter': "col1 < 3", - 'run_config_name': 'default', - 'user_metadata': {'team': 'data-engineers'}, - }, - ] - with testing.postgresql.Postgresql() as postgresql: connection_string = postgresql.url() engine = create_engine(connection_string) handler = LakebaseChecksStorageHandler(ws, spark, engine) config = LakebaseChecksStorageConfig(location, connection_string) - schema_name, table_name = handler._get_schema_and_table_name(config) - table = handler._get_table_definition(schema_name=schema_name, table_name=table_name) + schema_name, table_name = handler.get_schema_and_table_name(config) + table = handler.get_table_definition(schema_name=schema_name, table_name=table_name) table.metadata.create_all(engine, checkfirst=True) with engine.begin() as conn: - conn.execute(insert(table), expected_checks) + conn.execute(insert(table), TEST_CHECKS) result = handler.load(config) - assert len(result) == 2, f"Expected 2 checks, got {len(result)}" - - for check in result: - assert 'name' in check, "Missing 'name' field" - assert 'criticality' in check, "Missing 'criticality' field" - assert 'check' in check, "Missing 'check' field" - assert 'run_config_name' in check, "Missing 'run_config_name' field" - assert check['criticality'] in [ - 'error', - 'warning', - 'info', - ], f"Invalid criticality: {check['criticality']}" - - id_check = next((c for c in result if c['name'] == 'id_is_null'), None) - name_check = next((c for c in result if c['name'] == 'name_is_null'), None) - - assert id_check is not None, "Missing 'id_is_null' check" - assert name_check is not None, "Missing 'name_is_null' check" - - assert id_check['criticality'] == 'error' - assert id_check['check'] == {'function': 'is_not_null', 'arguments': {'column': 'id'}} - assert id_check['filter'] is None - assert id_check['run_config_name'] == 'default' - assert id_check['user_metadata'] is None - - assert name_check['criticality'] == 'warning' - assert name_check['check'] == {'function': 'is_not_null', 'arguments': {'column': 'name'}} - assert name_check['filter'] == "col1 < 3" - assert name_check['run_config_name'] == 'default' - assert name_check['user_metadata'] == {'team': 'data-engineers'} + assert len(result) == len(TEST_CHECKS), f"Expected {len(TEST_CHECKS)} checks, got {len(result)}" + + sorted_result = sorted(result, key=sort_key) + sorted_expected = sorted(TEST_CHECKS, key=sort_key) + + for result, expected in zip(sorted_result, sorted_expected): + for key in expected: + print(result.get(key), expected[key]) + assert ( + result.get(key) == expected[key] + ), f"Mismatch for key '{key}': {result.get(key)} != {expected[key]}" diff --git a/tests/unit/test_save_checks.py b/tests/unit/test_save_checks.py index 195de6337..51ce75cf0 100644 --- a/tests/unit/test_save_checks.py +++ b/tests/unit/test_save_checks.py @@ -1,34 +1,37 @@ import os import json import yaml - -import pytest from unittest.mock import create_autospec -from sqlalchemy import create_engine, select +from sqlalchemy import create_engine, select +from pyspark.sql import SparkSession from databricks.sdk import WorkspaceClient +import pytest +import testing.postgresql from databricks.labs.dqx.engine import DQEngineCore from databricks.labs.dqx.checks_storage import LakebaseChecksStorageHandler from databricks.labs.dqx.config import LakebaseChecksStorageConfig from databricks.labs.dqx.config import LakebaseConnectionConfig - -import testing.postgresql +from databricks.labs.dqx.utils import sort_key TEST_CHECKS = [ { - "criticality": "error", - "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, + 'name': 'id_is_null', + 'criticality': 'error', + 'check': {'function': 'is_not_null', 'arguments': {'column': 'id'}}, + 'filter': None, + 'run_config_name': 'default', + 'user_metadata': None, }, { - "criticality": "error", - "check": { - "function": "is_not_null", - "for_each_column": ["col1", "col2"], - "arguments": {}, - "user_metadata": {"rule_type": "completeness"}, - }, + 'name': 'name_is_null', + 'criticality': 'warning', + 'check': {'function': 'is_not_null', 'arguments': {'column': 'name'}}, + 'filter': "col1 < 3", + 'run_config_name': 'default', + 'user_metadata': {'team': 'data-engineers'}, }, ] @@ -75,76 +78,41 @@ def _validate_file(file_path: str, file_format: str = "yaml") -> None: def test_lakebase_checks_storage_handler_save(): ws = create_autospec(WorkspaceClient) - spark = create_autospec("pyspark.sql.SparkSession") + spark = create_autospec(SparkSession) location = "test.public.checks" - expected_checks = [ - { - 'name': 'id_is_null', - 'criticality': 'error', - 'check': {'function': 'is_not_null', 'arguments': {'column': 'id'}}, - 'filter': None, - 'run_config_name': 'default', - 'user_metadata': None, - }, - { - 'name': 'name_is_null', - 'criticality': 'warning', - 'check': {'function': 'is_not_null', 'arguments': {'column': 'name'}}, - 'filter': "col1 < 3", - 'run_config_name': 'extended', - 'user_metadata': {'team': 'data-engineers'}, - }, - ] - with testing.postgresql.Postgresql() as postgresql: connection_string = postgresql.url() engine = create_engine(connection_string) handler = LakebaseChecksStorageHandler(ws, spark, engine) config = LakebaseChecksStorageConfig(location, connection_string) - schema_name, table_name = handler._get_schema_and_table_name(config) - table = handler._get_table_definition(schema_name, table_name) + schema_name, table_name = handler.get_schema_and_table_name(config) + table = handler.get_table_definition(schema_name, table_name) - handler.save(expected_checks, config) + handler.save(TEST_CHECKS, config) with engine.connect() as conn: result = conn.execute(select(table)).mappings().all() result = [dict(check) for check in result] - assert len(result) == 2, f"Expected 2 checks, got {len(result)}" - - for check in result: - assert 'name' in check, "Missing 'name' field" - assert 'criticality' in check, "Missing 'criticality' field" - assert 'check' in check, "Missing 'check' field" - assert 'run_config_name' in check, "Missing 'run_config_name' field" - assert check['criticality'] in ['error', 'warning', 'info'], f"Invalid criticality: {check['criticality']}" - - id_check = next((c for c in result if c['name'] == 'id_is_null'), None) - name_check = next((c for c in result if c['name'] == 'name_is_null'), None) - - assert id_check is not None, "Missing 'id_is_null' check" - assert name_check is not None, "Missing 'name_is_null' check" + assert len(result) == len(TEST_CHECKS), f"Expected {len(TEST_CHECKS)} checks, got {len(result)}" - assert id_check['criticality'] == 'error' - assert id_check['check'] == {'function': 'is_not_null', 'arguments': {'column': 'id'}} - assert id_check['filter'] is None - assert id_check['run_config_name'] == 'default' - assert id_check['user_metadata'] is None + sorted_result = sorted(result, key=sort_key) + sorted_expected = sorted(TEST_CHECKS, key=sort_key) - assert name_check['criticality'] == 'warning' - assert name_check['check'] == {'function': 'is_not_null', 'arguments': {'column': 'name'}} - assert name_check['filter'] == "col1 < 3" - assert name_check['run_config_name'] == 'extended' - assert name_check['user_metadata'] == {'team': 'data-engineers'} + for result, expected in zip(sorted_result, sorted_expected): + for key in expected: + assert ( + result.get(key) == expected[key] + ), f"Mismatch for key '{key}': {result.get(key)} != {expected[key]}" def test_installation_checks_storage_handler_postgresql_parsing(): connection_string = ( "postgresql://user@databricks.com:password@instance-test.database.azuredatabricks.net:5432/dqx?sslmode=require" ) - connection_config = LakebaseConnectionConfig._parse_connection_string(connection_string) + connection_config = LakebaseConnectionConfig.parse_connection_string(connection_string) assert connection_config.user == "user@databricks.com" assert connection_config.instance_name == "instance-test.database.azuredatabricks.net" From 08ba30265cac81a32c2ce3e0cd46511f92c7cfaf Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 29 Sep 2025 14:40:52 +0200 Subject: [PATCH 035/125] refactor: apply code review suggestions --- src/databricks/labs/dqx/checks_storage.py | 43 ++++++++++----------- src/databricks/labs/dqx/config.py | 46 +++-------------------- 2 files changed, 25 insertions(+), 64 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index a7341e678..89265c9b8 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -5,7 +5,7 @@ from io import StringIO, BytesIO from pathlib import Path from abc import ABC, abstractmethod -from typing import Dict, Generic, List, TypeVar +from typing import Generic, TypeVar from sqlalchemy import ( Engine, create_engine, @@ -20,7 +20,7 @@ ) from sqlalchemy.schema import CreateSchema from sqlalchemy.dialects.postgresql import JSONB -from sqlalchemy.exc import OperationalError, DatabaseError, ProgrammingError +from sqlalchemy.exc import DatabaseError import yaml from pyspark.sql import SparkSession @@ -139,7 +139,7 @@ def _get_connection_url(self, connection_config: LakebaseConnectionConfig) -> st Generate a Lakebase connection URL. Args: - config: Configuration for saving and loading checks to Lakebase. + connection_config: Configuration for a Lakebase connection. Returns: Lakebase connection URL. @@ -158,7 +158,7 @@ def _get_engine(self, connection_config: LakebaseConnectionConfig) -> Engine: Create a SQLAlchemy engine for the Lakebase instance. Args: - config: Configuration for saving and loading checks to Lakebase. + connection_config: Configuration for a Lakebase connection. Returns: SQLAlchemy engine for the Lakebase instance. @@ -166,7 +166,7 @@ def _get_engine(self, connection_config: LakebaseConnectionConfig) -> Engine: connection_url = self._get_connection_url(connection_config) return create_engine(connection_url) - def _get_schema_and_table_name(self, config: LakebaseChecksStorageConfig) -> tuple[str, str]: + def get_schema_and_table_name(self, config: LakebaseChecksStorageConfig) -> tuple[str, str]: """ Extract the schema and table name from the fully qualified table name. @@ -189,13 +189,13 @@ def _get_schema_and_table_name(self, config: LakebaseChecksStorageConfig) -> tup _, schema_name, table_name = location_components return schema_name, table_name - def _get_table_definition(self, schema_name: str, table_name: str) -> Table: + def get_table_definition(self, schema_name: str, table_name: str) -> Table: """ Create a table definition for consistency between the load and save methods. Args: - schema: The schema name where the checks table is located. - table: The table name where the checks are stored. + 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. @@ -212,7 +212,7 @@ def _get_table_definition(self, schema_name: str, table_name: str) -> Table: ) @telemetry_logger("load_checks", "lakebase") - def load(self, config: LakebaseChecksStorageConfig) -> List[Dict]: + def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: """ Load dq rules (checks) from a Lakebase table. @@ -223,11 +223,9 @@ def load(self, config: LakebaseChecksStorageConfig) -> List[Dict]: List of dq rules or error if loading checks fails. Raises: - OperationalError: If connection to Lakebase instance fails. DatabaseError: If loading checks fails. - ProgrammingError: If loading checks fails. """ - connection_config = LakebaseConnectionConfig._parse_connection_string(config.connection_string) + connection_config = LakebaseConnectionConfig.parse_connection_string(config.connection_string) engine_created_internally = False if not self.engine: @@ -239,9 +237,9 @@ def load(self, config: LakebaseChecksStorageConfig) -> List[Dict]: try: logger.info(f"Loading checks from Lakebase instance {connection_config.instance_name}") - schema_name, table_name = self._get_schema_and_table_name(config) - table = self._get_table_definition(schema_name, table_name) - + schema_name, table_name = self.get_schema_and_table_name(config) + table = self.get_table_definition(schema_name, table_name) + stmt = select(table) if config.run_config_name: stmt = stmt.where(table.c.run_config_name == config.run_config_name) @@ -251,7 +249,7 @@ def load(self, config: LakebaseChecksStorageConfig) -> List[Dict]: checks = result.mappings().all() logger.info(f"Successfully loaded {len(checks)} checks from {schema_name}.{table_name}") return [dict(check) for check in checks] - except (OperationalError, DatabaseError, ProgrammingError) as e: + except DatabaseError as e: logger.error(f"Failed to load checks from Lakebase instance {connection_config.instance_name}: {e}") raise finally: @@ -259,7 +257,7 @@ def load(self, config: LakebaseChecksStorageConfig) -> List[Dict]: engine.dispose() @telemetry_logger("save_checks", "lakebase") - def save(self, checks: List[Dict], config: LakebaseChecksStorageConfig) -> None: + def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: """ Save dq rules (checks) to a Lakebase table. @@ -271,10 +269,7 @@ def save(self, checks: List[Dict], config: LakebaseChecksStorageConfig) -> None: None Raises: - OperationalError: If connection to Lakebase instance fails. DatabaseError: If saving checks fails. - ProgrammingError: If saving checks fails. - ValueError: If invalid mode is specified. """ if not checks: logger.warning("No checks provided to save.") @@ -283,7 +278,7 @@ def save(self, checks: List[Dict], config: LakebaseChecksStorageConfig) -> None: if config.mode not in ("append", "overwrite"): raise ValueError(f"Invalid mode '{config.mode}'. Must be 'append' or 'overwrite'.") - connection_config = LakebaseConnectionConfig._parse_connection_string(config.connection_string) + connection_config = LakebaseConnectionConfig.parse_connection_string(config.connection_string) engine_created_internally = False if not self.engine: @@ -292,7 +287,7 @@ def save(self, checks: List[Dict], config: LakebaseChecksStorageConfig) -> None: else: engine = self.engine - schema_name, table_name = self._get_schema_and_table_name(config) + schema_name, table_name = self.get_schema_and_table_name(config) try: logger.info(f"Saving {len(checks)} checks to Lakebase instance {connection_config.instance_name}") @@ -304,7 +299,7 @@ def save(self, checks: List[Dict], config: LakebaseChecksStorageConfig) -> None: f"Successfully created schema '{schema_name}' in database '{connection_config.database}'." ) - table = self._get_table_definition(schema_name, table_name) + table = self.get_table_definition(schema_name, table_name) table.metadata.create_all(engine, checkfirst=True) if config.mode == "overwrite": @@ -317,7 +312,7 @@ def save(self, checks: List[Dict], config: LakebaseChecksStorageConfig) -> None: insert_stmt = insert(table) conn.execute(insert_stmt, checks) logger.info(f"Inserted {len(checks)} new checks into {schema_name}.{table_name}") - except (OperationalError, DatabaseError, ProgrammingError) as e: + except DatabaseError as e: logger.error(f"Failed to save checks to Lakebase: {e}") raise finally: diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 48823715f..ca0a40808 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -217,37 +217,16 @@ class LakebaseConnectionConfig: port: str = LAKEBASE_DEFAULT_PORT @staticmethod - def _parse_connection_string(connection_string: str | None) -> "LakebaseConnectionConfig": - """ - Parse PostgreSQL connection string to extract connection parameters. - - Expected format: postgresql://user:password@instance_name:port/database?params. - - Args: - connection_string: Lakebase (SQLAlchemy) connection string. - - Returns: - Instance of LakebaseConnectionConfig with extracted parameters. - - Raises: - ValueError: If the URL format is invalid or required components are missing. - """ + def parse_connection_string(connection_string: str | None) -> "LakebaseConnectionConfig": if not connection_string: raise ValueError("Connection string cannot be empty or None.") - try: - parsed = urlparse(connection_string) - except Exception as e: - raise ValueError(f"Failed to parse URL '{connection_string}': {e}") from e + parsed = urlparse(connection_string) if parsed.scheme != "postgresql": raise ValueError(f"Invalid URL scheme '{parsed.scheme}'. Expected 'postgresql' for Lakebase connections.") - try: - user = unquote(parsed.username) if parsed.username else None - except Exception as e: - raise ValueError(f"Failed to decode username from URL: {e}") from e - + user = unquote(parsed.username) if parsed.username else None if not user: raise ValueError(f"Missing username in URL: {connection_string}") @@ -255,22 +234,9 @@ def _parse_connection_string(connection_string: str | None) -> "LakebaseConnecti if not instance_name: raise ValueError(f"Missing hostname in URL: {connection_string}") - port = LAKEBASE_DEFAULT_PORT - if parsed.port: - try: - port = str(parsed.port) - except (ValueError, TypeError) as e: - raise ValueError(f"Invalid port '{parsed.port}' in URL: {e}") from e - - database = None - if parsed.path: - try: - database = parsed.path.lstrip("/") - if not database: - raise ValueError("Database name is missing") - except Exception as e: - raise ValueError(f"Failed to extract database name from connection string '{parsed.path}': {e}") from e + port = str(parsed.port) if parsed.port else LAKEBASE_DEFAULT_PORT + database = parsed.path.lstrip("/") if parsed.path else None if not database: raise ValueError(f"Missing required database name in connection string: {connection_string}") @@ -313,7 +279,7 @@ def __post_init__(self): if self.connection_string and self.connection_string.startswith("postgresql://"): try: - LakebaseConnectionConfig._parse_connection_string(self.connection_string) + LakebaseConnectionConfig.parse_connection_string(self.connection_string) except Exception as e: raise ValueError(f"Failed to parse connection string '{self.connection_string}': {e}") from e From 75bb153f9e5aaa0d4e670f6b1fb1b5fd01cfb871 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 30 Sep 2025 11:06:31 +0200 Subject: [PATCH 036/125] refactor: reorder import statements for better organization --- tests/unit/test_save_checks.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_save_checks.py b/tests/unit/test_save_checks.py index 51ce75cf0..bbd8e939b 100644 --- a/tests/unit/test_save_checks.py +++ b/tests/unit/test_save_checks.py @@ -1,13 +1,14 @@ import os import json -import yaml from unittest.mock import create_autospec -from sqlalchemy import create_engine, select -from pyspark.sql import SparkSession -from databricks.sdk import WorkspaceClient +import yaml import pytest import testing.postgresql +from pyspark.sql import SparkSession +from sqlalchemy import create_engine, select + +from databricks.sdk import WorkspaceClient from databricks.labs.dqx.engine import DQEngineCore from databricks.labs.dqx.checks_storage import LakebaseChecksStorageHandler From bd069f89c9e08552a5ab695b3c3959648be345a7 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 30 Sep 2025 11:20:59 +0200 Subject: [PATCH 037/125] refactor: simplify saving and loading checks to Lakebase --- src/databricks/labs/dqx/checks_storage.py | 120 +++++++++++++--------- 1 file changed, 69 insertions(+), 51 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 89265c9b8..87f479753 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -211,6 +211,63 @@ def get_table_definition(self, schema_name: str, table_name: str) -> Table: Column("user_metadata", JSONB), ) + 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 + """ + schema_name, table_name = self.get_schema_and_table_name(config) + + with engine.begin() as conn: + if not conn.dialect.has_schema(conn, schema_name): + conn.execute(CreateSchema(schema_name)) + logger.info(f"Successfully created schema '{schema_name}'.") + + table = self.get_table_definition(schema_name, table_name) + table.metadata.create_all(engine, checkfirst=True) + logger.info(f"Successfully created or verified table '{schema_name}.{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"Successfully deleted {result.rowcount} existing checks for run_config_name '{config.run_config_name}'" + ) + + insert_stmt = insert(table) + conn.execute(insert_stmt, checks) + + 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. + """ + schema_name, table_name = self.get_schema_and_table_name(config) + table = self.get_table_definition(schema_name, table_name) + + stmt = select(table) + if config.run_config_name: + stmt = stmt.where(table.c.run_config_name == config.run_config_name) + + with engine.connect() as conn: + result = conn.execute(stmt) + checks = result.mappings().all() + logger.info(f"Successfully loaded {len(checks)} checks from {schema_name}.{table_name}") + return [dict(check) for check in checks] + @telemetry_logger("load_checks", "lakebase") def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: """ @@ -225,32 +282,17 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: Raises: DatabaseError: If loading checks fails. """ - connection_config = LakebaseConnectionConfig.parse_connection_string(config.connection_string) - + engine = self.engine or None engine_created_internally = False - if not self.engine: + if not engine: + connection_config = LakebaseConnectionConfig.parse_connection_string(config.connection_string) engine = self._get_engine(connection_config) engine_created_internally = True - else: - engine = self.engine try: - logger.info(f"Loading checks from Lakebase instance {connection_config.instance_name}") - - schema_name, table_name = self.get_schema_and_table_name(config) - table = self.get_table_definition(schema_name, table_name) - - stmt = select(table) - if config.run_config_name: - stmt = stmt.where(table.c.run_config_name == config.run_config_name) - - with engine.connect() as conn: - result = conn.execute(stmt) - checks = result.mappings().all() - logger.info(f"Successfully loaded {len(checks)} checks from {schema_name}.{table_name}") - return [dict(check) for check in checks] + return self._load_checks_from_lakebase(config, engine) except DatabaseError as e: - logger.error(f"Failed to load checks from Lakebase instance {connection_config.instance_name}: {e}") + logger.error(f"Failed to load checks from Lakebase: {e}") raise finally: if engine_created_internally: @@ -269,49 +311,25 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: None Raises: + ValueError: If the mode is not 'append' or 'overwrite'. DatabaseError: If saving checks fails. """ if not checks: - logger.warning("No checks provided to save.") - return + raise ValueError("Checks cannot be empty or None.") if config.mode not in ("append", "overwrite"): raise ValueError(f"Invalid mode '{config.mode}'. Must be 'append' or 'overwrite'.") - connection_config = LakebaseConnectionConfig.parse_connection_string(config.connection_string) - + engine = self.engine or None engine_created_internally = False - if not self.engine: + if not engine: + connection_config = LakebaseConnectionConfig.parse_connection_string(config.connection_string) engine = self._get_engine(connection_config) engine_created_internally = True - else: - engine = self.engine - - schema_name, table_name = self.get_schema_and_table_name(config) try: - logger.info(f"Saving {len(checks)} checks to Lakebase instance {connection_config.instance_name}") - - with engine.begin() as conn: - if not conn.dialect.has_schema(conn, schema_name): - conn.execute(CreateSchema(schema_name)) - logger.info( - f"Successfully created schema '{schema_name}' in database '{connection_config.database}'." - ) - - table = self.get_table_definition(schema_name, table_name) - table.metadata.create_all(engine, checkfirst=True) - - 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}'" - ) - - insert_stmt = insert(table) - conn.execute(insert_stmt, checks) - logger.info(f"Inserted {len(checks)} new checks into {schema_name}.{table_name}") + self._save_checks_to_lakebase(checks, config, engine) + logger.info(f"Successfully saved {len(checks)} checks to Lakebase.") except DatabaseError as e: logger.error(f"Failed to save checks to Lakebase: {e}") raise From 29fa2312eb7f55054eceebd70bb174567f89db43 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 30 Sep 2025 12:00:08 +0200 Subject: [PATCH 038/125] refactor: consolidate save and load tests --- src/databricks/labs/dqx/utils.py | 23 ++++++++++++++++++++++- tests/unit/test_load_checks.py | 14 ++------------ tests/unit/test_save_checks.py | 13 ++----------- 3 files changed, 26 insertions(+), 24 deletions(-) diff --git a/src/databricks/labs/dqx/utils.py b/src/databricks/labs/dqx/utils.py index 73d11946a..fbab84fed 100644 --- a/src/databricks/labs/dqx/utils.py +++ b/src/databricks/labs/dqx/utils.py @@ -339,7 +339,8 @@ def safe_json_load(value: str): return json.loads(value) # load as json if possible except json.JSONDecodeError: return value - + + def sort_key(check: dict[str, Any]) -> str: """ Sorts a checks dictionary by the 'name' field. @@ -351,3 +352,23 @@ def sort_key(check: dict[str, Any]) -> str: The name of the check as a string, or an empty string if not found. """ return str(check.get('name', '')) + + +def compare_checks(result, expected): + """ + 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 r, e in zip(sorted_result, sorted_expected): + for key in e: + assert r.get(key) == e[key], f"Mismatch for key '{key}': {r.get(key)} != {e[key]}" diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index 9fce84116..15582b590 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -12,7 +12,7 @@ from databricks.labs.dqx.checks_storage import VolumeFileChecksStorageHandler, LakebaseChecksStorageHandler from databricks.labs.dqx.config import LakebaseChecksStorageConfig, VolumeFileChecksStorageConfig from databricks.labs.dqx.engine import DQEngineCore -from databricks.labs.dqx.utils import sort_key +from databricks.labs.dqx.utils import compare_checks TEST_CHECKS = [ @@ -129,14 +129,4 @@ def test_lakebase_checks_storage_handler_load(): result = handler.load(config) - assert len(result) == len(TEST_CHECKS), f"Expected {len(TEST_CHECKS)} checks, got {len(result)}" - - sorted_result = sorted(result, key=sort_key) - sorted_expected = sorted(TEST_CHECKS, key=sort_key) - - for result, expected in zip(sorted_result, sorted_expected): - for key in expected: - print(result.get(key), expected[key]) - assert ( - result.get(key) == expected[key] - ), f"Mismatch for key '{key}': {result.get(key)} != {expected[key]}" + compare_checks(result, TEST_CHECKS) diff --git a/tests/unit/test_save_checks.py b/tests/unit/test_save_checks.py index bbd8e939b..8df628338 100644 --- a/tests/unit/test_save_checks.py +++ b/tests/unit/test_save_checks.py @@ -14,7 +14,7 @@ from databricks.labs.dqx.checks_storage import LakebaseChecksStorageHandler from databricks.labs.dqx.config import LakebaseChecksStorageConfig from databricks.labs.dqx.config import LakebaseConnectionConfig -from databricks.labs.dqx.utils import sort_key +from databricks.labs.dqx.utils import compare_checks TEST_CHECKS = [ @@ -97,16 +97,7 @@ def test_lakebase_checks_storage_handler_save(): result = conn.execute(select(table)).mappings().all() result = [dict(check) for check in result] - assert len(result) == len(TEST_CHECKS), f"Expected {len(TEST_CHECKS)} checks, got {len(result)}" - - sorted_result = sorted(result, key=sort_key) - sorted_expected = sorted(TEST_CHECKS, key=sort_key) - - for result, expected in zip(sorted_result, sorted_expected): - for key in expected: - assert ( - result.get(key) == expected[key] - ), f"Mismatch for key '{key}': {result.get(key)} != {expected[key]}" + compare_checks(result, TEST_CHECKS) def test_installation_checks_storage_handler_postgresql_parsing(): From f117c5646adcd94874aa0a06d3374cf54350f9f7 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 30 Sep 2025 12:13:05 +0200 Subject: [PATCH 039/125] fix: add type annotations to utils file --- src/databricks/labs/dqx/utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/databricks/labs/dqx/utils.py b/src/databricks/labs/dqx/utils.py index fbab84fed..2ce654067 100644 --- a/src/databricks/labs/dqx/utils.py +++ b/src/databricks/labs/dqx/utils.py @@ -354,7 +354,7 @@ def sort_key(check: dict[str, Any]) -> str: return str(check.get('name', '')) -def compare_checks(result, expected): +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. @@ -369,6 +369,6 @@ def compare_checks(result, expected): 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 r, e in zip(sorted_result, sorted_expected): - for key in e: - assert r.get(key) == e[key], f"Mismatch for key '{key}': {r.get(key)} != {e[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]}" From 5b3368939999dab2fa808cbcfb43e08459064807 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 30 Sep 2025 12:17:31 +0200 Subject: [PATCH 040/125] refactor: remove redundant code --- src/databricks/labs/dqx/checks_storage.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 87f479753..760d66875 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -282,7 +282,7 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: Raises: DatabaseError: If loading checks fails. """ - engine = self.engine or None + engine = self.engine engine_created_internally = False if not engine: connection_config = LakebaseConnectionConfig.parse_connection_string(config.connection_string) @@ -320,7 +320,7 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: if config.mode not in ("append", "overwrite"): raise ValueError(f"Invalid mode '{config.mode}'. Must be 'append' or 'overwrite'.") - engine = self.engine or None + engine = self.engine engine_created_internally = False if not engine: connection_config = LakebaseConnectionConfig.parse_connection_string(config.connection_string) From a977bc46d1f3833b5c00459f620d5f5cd24bec3f Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 30 Sep 2025 12:18:08 +0200 Subject: [PATCH 041/125] refactor: harmonize test checks --- tests/unit/test_load_checks.py | 2 +- tests/unit/test_save_checks.py | 23 +++++++++++------------ 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index 15582b590..e583bf200 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -22,7 +22,7 @@ "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, }, { - "name": "id_is_null", + "name": "name_is_null", "criticality": "error", "check": { "function": "is_not_null", diff --git a/tests/unit/test_save_checks.py b/tests/unit/test_save_checks.py index 8df628338..dfa5be9ed 100644 --- a/tests/unit/test_save_checks.py +++ b/tests/unit/test_save_checks.py @@ -19,20 +19,19 @@ TEST_CHECKS = [ { - 'name': 'id_is_null', - 'criticality': 'error', - 'check': {'function': 'is_not_null', 'arguments': {'column': 'id'}}, - 'filter': None, - 'run_config_name': 'default', - 'user_metadata': None, + "name": "id_is_null", + "criticality": "error", + "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, }, { - 'name': 'name_is_null', - 'criticality': 'warning', - 'check': {'function': 'is_not_null', 'arguments': {'column': 'name'}}, - 'filter': "col1 < 3", - 'run_config_name': 'default', - 'user_metadata': {'team': 'data-engineers'}, + "name": "name_is_null", + "criticality": "error", + "check": { + "function": "is_not_null", + "for_each_column": ["col1", "col2"], + "arguments": {}, + "user_metadata": {"rule_type": "completeness"}, + }, }, ] From ec9cffffc9db33f6348fc8b82a930f005f67c3c8 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 30 Sep 2025 15:27:41 +0200 Subject: [PATCH 042/125] refactor: update import statement for testing.postgresql --- tests/unit/test_load_checks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index e583bf200..b476bde17 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -1,7 +1,7 @@ from unittest.mock import create_autospec import pytest -import testing.postgresql +from testing.postgresql import Postgresql from pyspark.sql import SparkSession from sqlalchemy import create_engine, insert @@ -114,7 +114,7 @@ def test_lakebase_checks_storage_handler_load(): spark = create_autospec(SparkSession) location = "test.public.checks" - with testing.postgresql.Postgresql() as postgresql: + with Postgresql() as postgresql: connection_string = postgresql.url() engine = create_engine(connection_string) handler = LakebaseChecksStorageHandler(ws, spark, engine) From 573126c7f985ffedad3c973adb09db3ac91c3313 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 30 Sep 2025 15:28:17 +0200 Subject: [PATCH 043/125] refactor: improve error handling in Lakebase checks storage --- src/databricks/labs/dqx/checks_storage.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 760d66875..72fecb79f 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -291,9 +291,8 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: try: return self._load_checks_from_lakebase(config, engine) - except DatabaseError as e: - logger.error(f"Failed to load checks from Lakebase: {e}") - raise + except: + raise DatabaseError("Failed to load checks from Lakebase.") finally: if engine_created_internally: engine.dispose() @@ -330,9 +329,8 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: try: self._save_checks_to_lakebase(checks, config, engine) logger.info(f"Successfully saved {len(checks)} checks to Lakebase.") - except DatabaseError as e: - logger.error(f"Failed to save checks to Lakebase: {e}") - raise + except: + raise DatabaseError("Failed to save checks to Lakebase.") finally: if engine_created_internally: engine.dispose() From 3579c400cb1f9700f2c32665c8d028ebcd1bed45 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 30 Sep 2025 16:00:14 +0200 Subject: [PATCH 044/125] fix: update database instance reference in LakebaseChecksStorageHandler --- src/databricks/labs/dqx/checks_storage.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 083c71da4..d63eece26 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -149,9 +149,9 @@ def _get_connection_url(self, connection_config: LakebaseConnectionConfig) -> st Returns: Lakebase connection URL. """ - instance = self.ws.database.get_database_instance(connection_config.instance_name) + instance = self.ws.database.get_database_instance(connection_config.database) cred = self.ws.database.generate_database_credential( - request_id=str(uuid.uuid4()), instance_names=[connection_config.instance_name] + request_id=str(uuid.uuid4()), instance_names=[connection_config.database] ) host = instance.read_write_dns password = cred.token From e08911060ea79821142e4ae8a226c39d16a0ad8e Mon Sep 17 00:00:00 2001 From: Manuel Tilgner <58051108+tlgnr@users.noreply.github.com> Date: Wed, 1 Oct 2025 11:53:03 +0200 Subject: [PATCH 045/125] refactor: use InvalidParameterError instead of ValueError Co-authored-by: Marcin Wojtyczka --- src/databricks/labs/dqx/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 7b4541209..526b5fd6d 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -276,7 +276,7 @@ def __post_init__(self): raise ValueError("The location ('location' field) must not be empty or None.") if not self.connection_string: - raise ValueError("The connection string ('connection_string' field) must not be empty or None.") + raise InvalidParameterError("The connection string ('connection_string' field) must not be empty or None.") if self.connection_string and self.connection_string.startswith("postgresql://"): try: From 60d1db1b2d240b3537ad6ccf90e7cf4ba6659da0 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner <58051108+tlgnr@users.noreply.github.com> Date: Wed, 1 Oct 2025 11:53:43 +0200 Subject: [PATCH 046/125] refactor: use InvalidParameterError instead of ValueError Co-authored-by: Marcin Wojtyczka --- src/databricks/labs/dqx/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 526b5fd6d..718e6f068 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -233,7 +233,7 @@ def parse_connection_string(connection_string: str | None) -> "LakebaseConnectio instance_name = parsed.hostname if not instance_name: - raise ValueError(f"Missing hostname in URL: {connection_string}") + raise InvalidParameterError(f"Missing hostname in URL: {connection_string}") port = str(parsed.port) if parsed.port else LAKEBASE_DEFAULT_PORT From 5b8d964b43a72fea2e3b48c828f8f07eed819f6b Mon Sep 17 00:00:00 2001 From: Manuel Tilgner <58051108+tlgnr@users.noreply.github.com> Date: Wed, 1 Oct 2025 11:54:05 +0200 Subject: [PATCH 047/125] refactor: use InvalidParameterError instead of ValueError Co-authored-by: Marcin Wojtyczka --- src/databricks/labs/dqx/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 718e6f068..45d8cc351 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -282,7 +282,7 @@ def __post_init__(self): try: LakebaseConnectionConfig.parse_connection_string(self.connection_string) except Exception as e: - raise ValueError(f"Failed to parse connection string '{self.connection_string}': {e}") from e + raise InvalidParameterError(f"Failed to parse connection string '{self.connection_string}': {e}") from e @dataclass From 628e4517254433ff24e46b60224088485dd42c8a Mon Sep 17 00:00:00 2001 From: Manuel Tilgner <58051108+tlgnr@users.noreply.github.com> Date: Wed, 1 Oct 2025 11:54:15 +0200 Subject: [PATCH 048/125] refactor: use InvalidParameterError instead of ValueError Co-authored-by: Marcin Wojtyczka --- src/databricks/labs/dqx/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 45d8cc351..a583521c2 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -273,7 +273,7 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig): def __post_init__(self): if not self.location: - raise ValueError("The location ('location' field) must not be empty or None.") + raise InvalidParameterError("The location ('location' field) must not be empty or None.") if not self.connection_string: raise InvalidParameterError("The connection string ('connection_string' field) must not be empty or None.") From 969b2672450ab9c0a1d8983774f4764222880793 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 1 Oct 2025 12:02:50 +0200 Subject: [PATCH 049/125] refactor: move and redefine utility functions for comparing checks --- src/databricks/labs/dqx/utils.py | 33 -------------------------------- tests/conftest.py | 33 ++++++++++++++++++++++++++++++++ tests/unit/test_load_checks.py | 3 ++- tests/unit/test_save_checks.py | 3 ++- 4 files changed, 37 insertions(+), 35 deletions(-) diff --git a/src/databricks/labs/dqx/utils.py b/src/databricks/labs/dqx/utils.py index 985a6168d..4c7f05a20 100644 --- a/src/databricks/labs/dqx/utils.py +++ b/src/databricks/labs/dqx/utils.py @@ -349,36 +349,3 @@ def safe_json_load(value: str): return json.loads(value) # load as json if possible except json.JSONDecodeError: return value - - -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', '')) - - -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]}" diff --git a/tests/conftest.py b/tests/conftest.py index 5766c0846..59477e4d1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,5 @@ import os +from typing import Any from collections.abc import Callable, Generator from dataclasses import replace from functools import cached_property @@ -594,3 +595,35 @@ def delete(volume_file_path: str) -> None: ws.files.delete(volume_file_path) yield from factory("file", create, delete) + +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', '')) + + +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]}" diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index 19827f355..130632486 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -12,9 +12,10 @@ from databricks.labs.dqx.checks_storage import VolumeFileChecksStorageHandler, LakebaseChecksStorageHandler from databricks.labs.dqx.config import LakebaseChecksStorageConfig, VolumeFileChecksStorageConfig from databricks.labs.dqx.engine import DQEngineCore -from databricks.labs.dqx.utils import compare_checks from databricks.labs.dqx.errors import InvalidCheckError, CheckDownloadError, InvalidConfigError +from tests.conftest import compare_checks + TEST_CHECKS = [ { diff --git a/tests/unit/test_save_checks.py b/tests/unit/test_save_checks.py index 125b17cfc..e1f5797b0 100644 --- a/tests/unit/test_save_checks.py +++ b/tests/unit/test_save_checks.py @@ -14,9 +14,10 @@ from databricks.labs.dqx.checks_storage import LakebaseChecksStorageHandler from databricks.labs.dqx.config import LakebaseChecksStorageConfig from databricks.labs.dqx.config import LakebaseConnectionConfig -from databricks.labs.dqx.utils import compare_checks from databricks.labs.dqx.errors import InvalidConfigError +from tests.conftest import compare_checks + TEST_CHECKS = [ { From 4c94611c3409a847f00a6dce2ae3708c132b0a8c Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 1 Oct 2025 12:05:10 +0200 Subject: [PATCH 050/125] chore: import InvalidParameterError --- src/databricks/labs/dqx/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index a583521c2..cb9d5fba7 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -2,7 +2,7 @@ from datetime import datetime, timezone from dataclasses import dataclass, field from urllib.parse import urlparse, unquote -from databricks.labs.dqx.errors import InvalidConfigError +from databricks.labs.dqx.errors import InvalidConfigError, InvalidParameterError __all__ = [ "WorkspaceConfig", From 9a9017b14fb5b9812c60ec99b9ed5e6e8f49aae6 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 1 Oct 2025 16:10:11 +0200 Subject: [PATCH 051/125] feat: add connection_string attribute to RunConfig for Lakebase instances --- src/databricks/labs/dqx/config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index cb9d5fba7..95978e3f4 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -66,6 +66,7 @@ class RunConfig: checks_location: str = ( "checks.yml" # absolute or relative workspace file path or table containing quality rules / checks ) + connection_string: str | None = None # connection string for Lakebase instances 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 From 8f53150fecc80d250427552843d1931bcb4e146f Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 1 Oct 2025 16:10:37 +0200 Subject: [PATCH 052/125] feat: normalize checks and improve error handling in LakebaseChecksStorageHandler - Added _normalize_checks method to standardize the format of checks before saving to Lakebase. - Updated _save_checks_to_lakebase to use normalized checks. - Enhanced error handling in _load_checks_from_lakebase and _save_checks_to_lakebase to provide more specific error messages and raise NotFound exceptions when applicable. --- src/databricks/labs/dqx/checks_storage.py | 49 ++++++++++++++++++----- 1 file changed, 40 insertions(+), 9 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index d63eece26..b6a9fab1d 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -17,6 +17,7 @@ insert, select, delete, + null, ) from sqlalchemy.schema import CreateSchema from sqlalchemy.dialects.postgresql import JSONB @@ -40,6 +41,7 @@ ) from databricks.labs.dqx.errors import InvalidCheckError, InvalidConfigError, CheckDownloadError from databricks.sdk import WorkspaceClient +from databricks.sdk.errors import NotFound from databricks.labs.dqx.checks_serializer import ( serialize_checks_from_dataframe, @@ -213,9 +215,33 @@ def get_table_definition(self, schema_name: str, table_name: str) -> Table: Column("check", JSONB), Column("filter", Text), Column("run_config_name", String(255), server_default="default"), - Column("user_metadata", JSONB), + Column("user_metadata", JSONB, server_default=null()), ) + def _normalize_checks(self, 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: + 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": check.get("user_metadata", null()), + } + 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. @@ -242,12 +268,11 @@ def _save_checks_to_lakebase(self, checks: list[dict], config: LakebaseChecksSto 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"Successfully deleted {result.rowcount} existing checks for run_config_name '{config.run_config_name}'" - ) + 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, checks) + conn.execute(insert_stmt, normalized_checks) def _load_checks_from_lakebase(self, config: LakebaseChecksStorageConfig, engine: Engine) -> list[dict]: """ @@ -296,8 +321,13 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: try: return self._load_checks_from_lakebase(config, engine) - except: - raise DatabaseError("Failed to load checks from Lakebase.") + except DatabaseError as e: + logger.error(f"Failed to load checks from Lakebase: {e}") + if isinstance(e, DatabaseError): + if "does not exist" in str(e).lower() or "relation" in str(e).lower(): + raise NotFound(f"Table '{config.location}' does not exist in the Lakebase instance") from e + else: + raise DatabaseError(f"Database error loading checks from table '{config.location}': {e}") from e finally: if engine_created_internally: engine.dispose() @@ -334,8 +364,9 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: try: self._save_checks_to_lakebase(checks, config, engine) logger.info(f"Successfully saved {len(checks)} checks to Lakebase.") - except: - raise DatabaseError("Failed to save checks to Lakebase.") + except DatabaseError as e: + logger.error(f"Failed to save checks to Lakebase: {e}") + raise DatabaseError(f"Database error saving checks to table '{config.location}': {e}") from e finally: if engine_created_internally: engine.dispose() From 419d688742832909f05d3869cbef821849efa714 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 1 Oct 2025 16:34:28 +0200 Subject: [PATCH 053/125] docs: document Lakebase columns --- src/databricks/labs/dqx/checks_storage.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index b6a9fab1d..6cd363c04 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -204,6 +204,19 @@ def get_table_definition(self, schema_name: str, table_name: str) -> Table: schema_name: The schema where the checks table is located. table_name: The table where the checks are stored. + Note: + The table has the following columns: + * `name` - Optional. Name to use for the check. + * `criticality` - Either "error" (rows go only to "bad" dataset) or "warn" (rows go to both "good" and "bad"). + Defaults to "error". + * `check` - Defines the DQX check to apply, e.g. {"function": "is_not_null", "arguments": {"column": "col1"}}. + * `filter` - Optional. Spark SQL expression to filter rows to which the check is applied, e.g. "business_unit = 'Finance'". + The check function will run only on the rows matching the filter condition. The condition can reference + any column of the validated dataset, not only the one where you apply the check function. + * `run_config_name` - Name of the run config name. Could be any string such as workflow name. Useful for selecting + applicable checks. Defaults to `"default"`. + * `user_metadata` - Optional. Custom metadata to add to any row-level warnings or errors generated by the check. + Returns: SQLAlchemy table definition for the Lakebase instance. """ From dc9e05b4b9c26bde46443f9c6368c5a6cc256e05 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 1 Oct 2025 16:38:27 +0200 Subject: [PATCH 054/125] feat: add pytest fixture for Lakebase connection_string --- tests/conftest.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/conftest.py b/tests/conftest.py index 59477e4d1..9fa677356 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -596,6 +596,13 @@ def delete(volume_file_path: str) -> None: yield from factory("file", create, delete) + +@pytest.fixture +def connection_string(env_or_skip): + """Get Lakebase connection string from environment or skip test if not available.""" + return env_or_skip("DQX_LAKEBASE_CONNECTION_STRING") + + def sort_key(check: dict[str, Any]) -> str: """ Sorts a checks dictionary by the 'name' field. From 839e4b6ded9cb7b2209e72423d5a9fc82341e2d7 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 1 Oct 2025 16:38:37 +0200 Subject: [PATCH 055/125] feat: add integration tests --- ...ave_and_load_checks_from_lakebase_table.py | 144 ++++++++++++++++++ 1 file changed, 144 insertions(+) create mode 100644 tests/integration/test_save_and_load_checks_from_lakebase_table.py 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..f721d20b8 --- /dev/null +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -0,0 +1,144 @@ +import pytest + +from databricks.sdk.errors import NotFound + +from databricks.labs.dqx.config import InstallationChecksStorageConfig, LakebaseChecksStorageConfig +from databricks.labs.dqx.engine import DQEngine + +from tests.conftest import compare_checks, connection_string + + +TEST_CHECKS = [ + { + "name": "col1_is_null", + "criticality": "error", + "check": {"function": "is_not_null", "arguments": {"column": "col1"}}, + "user_metadata": {"check_type": "completeness", "check_owner": "someone@email.com"}, + }, + { + "name": "col2_is_null", + "criticality": "error", + "check": {"function": "is_not_null", "arguments": {"column": "col2"}}, + "user_metadata": {"check_type": "completeness", "check_owner": "someone@email.com"}, + }, + { + "name": "column_not_less_than", + "criticality": "warn", + "check": { + "function": "is_not_less_than", + "arguments": {"column": "col_2", "limit": 1}, + }, + "user_metadata": {"check_type": "standardization", "check_owner": "someone_else@email.com"}, + }, + { + "name": "column_in_list", + "criticality": "warn", + "check": {"function": "is_in_list", "arguments": {"column": "col_2", "allowed": [1, 2]}}, + }, +] + +location = "dqx.config.checks" + + +def test_load_checks_when_checks_lakebase_table_does_not_exist(ws, spark, connection_string): + with pytest.raises(NotFound, match=f"Table '{location}' does not exist in the Lakebase instance"): + dq_engine = DQEngine(ws) + config = LakebaseChecksStorageConfig(location=location, connection_string=connection_string) + dq_engine.load_checks(config=config) + + +def test_save_and_load_checks_from_lakebase_table(ws, make_schema, make_random, spark, connection_string): + dq_engine = DQEngine(ws) + config = LakebaseChecksStorageConfig(location=location, connection_string=connection_string) + 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, connection_string): + dq_engine = DQEngine(ws, spark) + run_config_name = "workflow_001" + config_save = LakebaseChecksStorageConfig( + location=location, connection_string=connection_string, run_config_name=run_config_name + ) + dq_engine.save_checks(TEST_CHECKS[:1], config=config_save) + config_load = LakebaseChecksStorageConfig( + location=location, connection_string=connection_string, run_config_name=run_config_name + ) + checks = dq_engine.load_checks(config=config_load) + compare_checks(checks, TEST_CHECKS[:1]) + + run_config_name2 = "workflow_002" + config_save2 = LakebaseChecksStorageConfig( + location=location, + connection_string=connection_string, + run_config_name=run_config_name2, + mode="overwrite", + ) + dq_engine.save_checks(TEST_CHECKS[1:], config=config_save2) + config_load2 = LakebaseChecksStorageConfig( + location=location, connection_string=connection_string, run_config_name=run_config_name + ) + checks = dq_engine.load_checks(config=config_load2) + compare_checks(checks, TEST_CHECKS[:1]) + + dq_engine.save_checks( + TEST_CHECKS[1:], config=LakebaseChecksStorageConfig(location=location, connection_string=connection_string) + ) + checks = dq_engine.load_checks( + config=LakebaseChecksStorageConfig(location=location, connection_string=connection_string) + ) + compare_checks(checks, TEST_CHECKS[1:]) + + +def test_save_and_load_checks_to_lakebase_table_output_modes(ws, make_schema, make_random, spark, connection_string): + dq_engine = DQEngine(ws, spark) + run_config_name = "workflow_003" + dq_engine.save_checks( + TEST_CHECKS[:1], + config=LakebaseChecksStorageConfig( + location=location, connection_string=connection_string, run_config_name=run_config_name, mode="append" + ), + ) + checks = dq_engine.load_checks( + config=LakebaseChecksStorageConfig( + location=location, connection_string=connection_string, 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=location, connection_string=connection_string, run_config_name=run_config_name, mode="overwrite" + ), + ) + checks = dq_engine.load_checks( + config=LakebaseChecksStorageConfig( + location=location, connection_string=connection_string, run_config_name=run_config_name + ) + ) + compare_checks(checks, TEST_CHECKS[1:]) + + +def test_save_load_checks_from_lakebase_table_in_user_installation(ws, spark, installation_ctx, connection_string): + config = installation_ctx.config + run_config = config.get_run_config() + run_config.checks_location = location + run_config.connection_string = connection_string + installation_ctx.installation.save(installation_ctx.config) + product_name = installation_ctx.product_info.product_name() + + dq_engine = DQEngine(ws, spark) + config = InstallationChecksStorageConfig( + location=location, + connection_string=connection_string, + run_config_name=run_config.name, + assume_user=True, + product_name=product_name, + ) + + dq_engine.save_checks(TEST_CHECKS, config=config) + checks = dq_engine.load_checks(config=config) + compare_checks(checks, TEST_CHECKS) From fe923acb805cd297e4c5cef266d4202e983fdf0f Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 1 Oct 2025 17:09:47 +0200 Subject: [PATCH 056/125] fix: enhance error handling in LakebaseChecksStorageHandler --- src/databricks/labs/dqx/checks_storage.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 6cd363c04..40a499772 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -340,7 +340,12 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: if "does not exist" in str(e).lower() or "relation" in str(e).lower(): raise NotFound(f"Table '{config.location}' does not exist in the Lakebase instance") from e else: - raise DatabaseError(f"Database error loading checks from table '{config.location}': {e}") from e + raise DatabaseError( + f"Database error loading checks from table '{config.location}': {e}", + statement=getattr(e, 'statement', None), + params=getattr(e, 'params', None), + orig=e, + ) from e finally: if engine_created_internally: engine.dispose() @@ -379,7 +384,12 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: logger.info(f"Successfully saved {len(checks)} checks to Lakebase.") except DatabaseError as e: logger.error(f"Failed to save checks to Lakebase: {e}") - raise DatabaseError(f"Database error saving checks to table '{config.location}': {e}") from e + raise DatabaseError( + f"Database error saving checks to table '{config.location}': {e}", + statement=getattr(e, 'statement', None), + params=getattr(e, 'params', None), + orig=e, + ) from e finally: if engine_created_internally: engine.dispose() From a78c2cca5ba7688a0da89a7c4bae554a637182fb Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 1 Oct 2025 17:09:59 +0200 Subject: [PATCH 057/125] refactor: remove redundant import --- .../test_save_and_load_checks_from_lakebase_table.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index f721d20b8..1e0c28f9a 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -5,7 +5,7 @@ from databricks.labs.dqx.config import InstallationChecksStorageConfig, LakebaseChecksStorageConfig from databricks.labs.dqx.engine import DQEngine -from tests.conftest import compare_checks, connection_string +from tests.conftest import compare_checks TEST_CHECKS = [ From bf311b73b96db2c70a4c2486c8f2653b03dea246 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 1 Oct 2025 18:12:46 +0200 Subject: [PATCH 058/125] refactor: standardize LOCATION variable usage in Lakebase checks tests - Updated all instances of the 'location' variable to 'LOCATION' for consistency across test files. - Enhanced error handling in tests to reflect the updated variable name. --- ...ave_and_load_checks_from_lakebase_table.py | 32 +++++++++---------- tests/unit/test_load_checks.py | 5 +-- tests/unit/test_save_checks.py | 11 ++----- 3 files changed, 20 insertions(+), 28 deletions(-) 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 index 1e0c28f9a..8a6d56989 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -37,19 +37,19 @@ }, ] -location = "dqx.config.checks" +LOCATION = "dqx.config.checks" def test_load_checks_when_checks_lakebase_table_does_not_exist(ws, spark, connection_string): - with pytest.raises(NotFound, match=f"Table '{location}' does not exist in the Lakebase instance"): + with pytest.raises(NotFound, match=f"Table '{LOCATION}' does not exist in the Lakebase instance"): dq_engine = DQEngine(ws) - config = LakebaseChecksStorageConfig(location=location, connection_string=connection_string) + config = LakebaseChecksStorageConfig(LOCATION=LOCATION, connection_string=connection_string) dq_engine.load_checks(config=config) def test_save_and_load_checks_from_lakebase_table(ws, make_schema, make_random, spark, connection_string): dq_engine = DQEngine(ws) - config = LakebaseChecksStorageConfig(location=location, connection_string=connection_string) + config = LakebaseChecksStorageConfig(LOCATION=LOCATION, connection_string=connection_string) dq_engine.save_checks(checks=TEST_CHECKS, config=config) checks = dq_engine.load_checks(config=config) compare_checks(checks, TEST_CHECKS) @@ -59,34 +59,34 @@ def test_save_and_load_checks_from_lakebase_table_with_run_config(ws, spark, con dq_engine = DQEngine(ws, spark) run_config_name = "workflow_001" config_save = LakebaseChecksStorageConfig( - location=location, connection_string=connection_string, run_config_name=run_config_name + LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name ) dq_engine.save_checks(TEST_CHECKS[:1], config=config_save) config_load = LakebaseChecksStorageConfig( - location=location, connection_string=connection_string, run_config_name=run_config_name + LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name ) checks = dq_engine.load_checks(config=config_load) compare_checks(checks, TEST_CHECKS[:1]) run_config_name2 = "workflow_002" config_save2 = LakebaseChecksStorageConfig( - location=location, + LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name2, mode="overwrite", ) dq_engine.save_checks(TEST_CHECKS[1:], config=config_save2) config_load2 = LakebaseChecksStorageConfig( - location=location, connection_string=connection_string, run_config_name=run_config_name + LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name ) checks = dq_engine.load_checks(config=config_load2) compare_checks(checks, TEST_CHECKS[:1]) dq_engine.save_checks( - TEST_CHECKS[1:], config=LakebaseChecksStorageConfig(location=location, connection_string=connection_string) + TEST_CHECKS[1:], config=LakebaseChecksStorageConfig(LOCATION=LOCATION, connection_string=connection_string) ) checks = dq_engine.load_checks( - config=LakebaseChecksStorageConfig(location=location, connection_string=connection_string) + config=LakebaseChecksStorageConfig(LOCATION=LOCATION, connection_string=connection_string) ) compare_checks(checks, TEST_CHECKS[1:]) @@ -97,12 +97,12 @@ def test_save_and_load_checks_to_lakebase_table_output_modes(ws, make_schema, ma dq_engine.save_checks( TEST_CHECKS[:1], config=LakebaseChecksStorageConfig( - location=location, connection_string=connection_string, run_config_name=run_config_name, mode="append" + LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name, mode="append" ), ) checks = dq_engine.load_checks( config=LakebaseChecksStorageConfig( - location=location, connection_string=connection_string, run_config_name=run_config_name + LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name ) ) compare_checks(checks, TEST_CHECKS[:1]) @@ -111,12 +111,12 @@ def test_save_and_load_checks_to_lakebase_table_output_modes(ws, make_schema, ma dq_engine.save_checks( TEST_CHECKS[1:], config=LakebaseChecksStorageConfig( - location=location, connection_string=connection_string, run_config_name=run_config_name, mode="overwrite" + LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name, mode="overwrite" ), ) checks = dq_engine.load_checks( config=LakebaseChecksStorageConfig( - location=location, connection_string=connection_string, run_config_name=run_config_name + LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name ) ) compare_checks(checks, TEST_CHECKS[1:]) @@ -125,14 +125,14 @@ def test_save_and_load_checks_to_lakebase_table_output_modes(ws, make_schema, ma def test_save_load_checks_from_lakebase_table_in_user_installation(ws, spark, installation_ctx, connection_string): config = installation_ctx.config run_config = config.get_run_config() - run_config.checks_location = location + run_config.checks_LOCATION = LOCATION run_config.connection_string = connection_string installation_ctx.installation.save(installation_ctx.config) product_name = installation_ctx.product_info.product_name() dq_engine = DQEngine(ws, spark) config = InstallationChecksStorageConfig( - location=location, + LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config.name, assume_user=True, diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index 130632486..374f13371 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -2,7 +2,6 @@ import pytest from testing.postgresql import Postgresql -from pyspark.sql import SparkSession from sqlalchemy import create_engine, insert from databricks.sdk import WorkspaceClient @@ -111,9 +110,7 @@ def test_file_download_contents_read_none(): handler.load(VolumeFileChecksStorageConfig(location="test_path")) -def test_lakebase_checks_storage_handler_load(): - ws = create_autospec(WorkspaceClient) - spark = create_autospec(SparkSession) +def test_lakebase_checks_storage_handler_load(ws, spark): location = "test.public.checks" with Postgresql() as postgresql: diff --git a/tests/unit/test_save_checks.py b/tests/unit/test_save_checks.py index e1f5797b0..4c44c5762 100644 --- a/tests/unit/test_save_checks.py +++ b/tests/unit/test_save_checks.py @@ -1,14 +1,11 @@ import os import json -from unittest.mock import create_autospec import yaml import pytest import testing.postgresql -from pyspark.sql import SparkSession from sqlalchemy import create_engine, select -from databricks.sdk import WorkspaceClient from databricks.labs.dqx.engine import DQEngineCore from databricks.labs.dqx.checks_storage import LakebaseChecksStorageHandler @@ -78,9 +75,7 @@ def _validate_file(file_path: str, file_format: str = "yaml") -> None: yaml.safe_load(file) -def test_lakebase_checks_storage_handler_save(): - ws = create_autospec(WorkspaceClient) - spark = create_autospec(SparkSession) +def test_lakebase_checks_storage_handler_save(ws, spark): location = "test.public.checks" with testing.postgresql.Postgresql() as postgresql: @@ -103,11 +98,11 @@ def test_lakebase_checks_storage_handler_save(): def test_installation_checks_storage_handler_postgresql_parsing(): connection_string = ( - "postgresql://user@databricks.com:password@instance-test.database.azuredatabricks.net:5432/dqx?sslmode=require" + "postgresql://user@domain.com:password@instance-test.database.azuredatabricks.net:5432/dqx?sslmode=require" ) connection_config = LakebaseConnectionConfig.parse_connection_string(connection_string) - assert connection_config.user == "user@databricks.com" + assert connection_config.user == "user@domain.com" assert connection_config.instance_name == "instance-test.database.azuredatabricks.net" assert connection_config.port == "5432" assert connection_config.database == "dqx" From b36a5d34f5ee1147c1fa19e1c056718250bf878d Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 1 Oct 2025 18:13:39 +0200 Subject: [PATCH 059/125] refactor: streamline error handling in LakebaseChecksStorageHandler --- src/databricks/labs/dqx/checks_storage.py | 22 +++++----------------- 1 file changed, 5 insertions(+), 17 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 40a499772..6a01e8a9c 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -41,7 +41,6 @@ ) from databricks.labs.dqx.errors import InvalidCheckError, InvalidConfigError, CheckDownloadError from databricks.sdk import WorkspaceClient -from databricks.sdk.errors import NotFound from databricks.labs.dqx.checks_serializer import ( serialize_checks_from_dataframe, @@ -324,6 +323,7 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: Raises: DatabaseError: If loading checks fails. + NotFound: If the table does not exist in the Lakebase instance. """ engine = self.engine engine_created_internally = False @@ -336,16 +336,9 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: return self._load_checks_from_lakebase(config, engine) except DatabaseError as e: logger.error(f"Failed to load checks from Lakebase: {e}") - if isinstance(e, DatabaseError): - if "does not exist" in str(e).lower() or "relation" in str(e).lower(): - raise NotFound(f"Table '{config.location}' does not exist in the Lakebase instance") from e - else: - raise DatabaseError( - f"Database error loading checks from table '{config.location}': {e}", - statement=getattr(e, 'statement', None), - params=getattr(e, 'params', None), - orig=e, - ) from e + if "does not exist" in str(e).lower() or "relation" in str(e).lower(): + raise NotFound(f"Table '{config.location}' does not exist in the Lakebase instance") from e + raise DatabaseError(getattr(e, 'statement', None), getattr(e, 'params', None), e) from e finally: if engine_created_internally: engine.dispose() @@ -384,12 +377,7 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: logger.info(f"Successfully saved {len(checks)} checks to Lakebase.") except DatabaseError as e: logger.error(f"Failed to save checks to Lakebase: {e}") - raise DatabaseError( - f"Database error saving checks to table '{config.location}': {e}", - statement=getattr(e, 'statement', None), - params=getattr(e, 'params', None), - orig=e, - ) from e + raise DatabaseError(getattr(e, 'statement', None), getattr(e, 'params', None), e) from e finally: if engine_created_internally: engine.dispose() From 2a30ce8080bd1f2968f81c6ca9f8c89139d2e9f9 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 1 Oct 2025 20:42:57 +0200 Subject: [PATCH 060/125] fix: correct substitution error --- ...ave_and_load_checks_from_lakebase_table.py | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) 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 index 8a6d56989..dc52d4fde 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -43,13 +43,13 @@ def test_load_checks_when_checks_lakebase_table_does_not_exist(ws, spark, connection_string): with pytest.raises(NotFound, match=f"Table '{LOCATION}' does not exist in the Lakebase instance"): dq_engine = DQEngine(ws) - config = LakebaseChecksStorageConfig(LOCATION=LOCATION, connection_string=connection_string) + config = LakebaseChecksStorageConfig(location=LOCATION, connection_string=connection_string) dq_engine.load_checks(config=config) def test_save_and_load_checks_from_lakebase_table(ws, make_schema, make_random, spark, connection_string): dq_engine = DQEngine(ws) - config = LakebaseChecksStorageConfig(LOCATION=LOCATION, connection_string=connection_string) + config = LakebaseChecksStorageConfig(location=LOCATION, connection_string=connection_string) dq_engine.save_checks(checks=TEST_CHECKS, config=config) checks = dq_engine.load_checks(config=config) compare_checks(checks, TEST_CHECKS) @@ -59,34 +59,34 @@ def test_save_and_load_checks_from_lakebase_table_with_run_config(ws, spark, con dq_engine = DQEngine(ws, spark) run_config_name = "workflow_001" config_save = LakebaseChecksStorageConfig( - LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name + location=LOCATION, connection_string=connection_string, run_config_name=run_config_name ) dq_engine.save_checks(TEST_CHECKS[:1], config=config_save) config_load = LakebaseChecksStorageConfig( - LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name + location=LOCATION, connection_string=connection_string, run_config_name=run_config_name ) checks = dq_engine.load_checks(config=config_load) compare_checks(checks, TEST_CHECKS[:1]) run_config_name2 = "workflow_002" config_save2 = LakebaseChecksStorageConfig( - LOCATION=LOCATION, + location=LOCATION, connection_string=connection_string, run_config_name=run_config_name2, mode="overwrite", ) dq_engine.save_checks(TEST_CHECKS[1:], config=config_save2) config_load2 = LakebaseChecksStorageConfig( - LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name + location=LOCATION, connection_string=connection_string, run_config_name=run_config_name ) checks = dq_engine.load_checks(config=config_load2) compare_checks(checks, TEST_CHECKS[:1]) dq_engine.save_checks( - TEST_CHECKS[1:], config=LakebaseChecksStorageConfig(LOCATION=LOCATION, connection_string=connection_string) + TEST_CHECKS[1:], config=LakebaseChecksStorageConfig(location=LOCATION, connection_string=connection_string) ) checks = dq_engine.load_checks( - config=LakebaseChecksStorageConfig(LOCATION=LOCATION, connection_string=connection_string) + config=LakebaseChecksStorageConfig(location=LOCATION, connection_string=connection_string) ) compare_checks(checks, TEST_CHECKS[1:]) @@ -97,12 +97,12 @@ def test_save_and_load_checks_to_lakebase_table_output_modes(ws, make_schema, ma dq_engine.save_checks( TEST_CHECKS[:1], config=LakebaseChecksStorageConfig( - LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name, mode="append" + location=LOCATION, connection_string=connection_string, run_config_name=run_config_name, mode="append" ), ) checks = dq_engine.load_checks( config=LakebaseChecksStorageConfig( - LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name + location=LOCATION, connection_string=connection_string, run_config_name=run_config_name ) ) compare_checks(checks, TEST_CHECKS[:1]) @@ -111,12 +111,12 @@ def test_save_and_load_checks_to_lakebase_table_output_modes(ws, make_schema, ma dq_engine.save_checks( TEST_CHECKS[1:], config=LakebaseChecksStorageConfig( - LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name, mode="overwrite" + location=LOCATION, connection_string=connection_string, run_config_name=run_config_name, mode="overwrite" ), ) checks = dq_engine.load_checks( config=LakebaseChecksStorageConfig( - LOCATION=LOCATION, connection_string=connection_string, run_config_name=run_config_name + location=LOCATION, connection_string=connection_string, run_config_name=run_config_name ) ) compare_checks(checks, TEST_CHECKS[1:]) @@ -132,7 +132,7 @@ def test_save_load_checks_from_lakebase_table_in_user_installation(ws, spark, in dq_engine = DQEngine(ws, spark) config = InstallationChecksStorageConfig( - LOCATION=LOCATION, + location=LOCATION, connection_string=connection_string, run_config_name=run_config.name, assume_user=True, From aecc1d12c8f910fb23f91829795efb4226315258 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 1 Oct 2025 20:48:16 +0200 Subject: [PATCH 061/125] refactor: remove code redundancies --- ...ave_and_load_checks_from_lakebase_table.py | 31 +------------------ tests/unit/test_load_checks.py | 20 +----------- 2 files changed, 2 insertions(+), 49 deletions(-) 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 index dc52d4fde..4789bbb82 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -6,36 +6,7 @@ from databricks.labs.dqx.engine import DQEngine from tests.conftest import compare_checks - - -TEST_CHECKS = [ - { - "name": "col1_is_null", - "criticality": "error", - "check": {"function": "is_not_null", "arguments": {"column": "col1"}}, - "user_metadata": {"check_type": "completeness", "check_owner": "someone@email.com"}, - }, - { - "name": "col2_is_null", - "criticality": "error", - "check": {"function": "is_not_null", "arguments": {"column": "col2"}}, - "user_metadata": {"check_type": "completeness", "check_owner": "someone@email.com"}, - }, - { - "name": "column_not_less_than", - "criticality": "warn", - "check": { - "function": "is_not_less_than", - "arguments": {"column": "col_2", "limit": 1}, - }, - "user_metadata": {"check_type": "standardization", "check_owner": "someone_else@email.com"}, - }, - { - "name": "column_in_list", - "criticality": "warn", - "check": {"function": "is_in_list", "arguments": {"column": "col_2", "allowed": [1, 2]}}, - }, -] +from tests.integration.test_save_and_load_checks_from_table import EXPECTED_CHECKS as TEST_CHECKS LOCATION = "dqx.config.checks" diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index 374f13371..151aacceb 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -14,25 +14,7 @@ from databricks.labs.dqx.errors import InvalidCheckError, CheckDownloadError, InvalidConfigError from tests.conftest import compare_checks - - -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", - "for_each_column": ["col1", "col2"], - "arguments": {}, - "user_metadata": {"rule_type": "completeness"}, - }, - }, -] +from tests.unit.test_save_checks import TEST_CHECKS def test_load_checks_from_local_file_json(make_local_check_file_as_json, expected_checks): From 2e5993038d1e6962708904d506feff48e1ea6519 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 2 Oct 2025 10:23:39 +0200 Subject: [PATCH 062/125] feat: add DATABRICKS_HOST secret environment variable --- .github/workflows/push.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index d45a722ed..c70d14ec2 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -40,6 +40,8 @@ jobs: # click 8.3+ introduced bug for hatch pip install "hatch==1.13.0" "click<8.3" make test + env: + DATABRICKS_HOST: ${{ secrets.DATABRICKS_HOST }} fmt: runs-on: ubuntu-latest From c34fe8c5f205d65c6b78331ba9774563d043ed8b Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 2 Oct 2025 13:50:53 +0200 Subject: [PATCH 063/125] refactor: upgrade databricks-sdk version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 7b5b6da3a..c3a331825 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ 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", ] From 42f059ee872c9f6b0b58d71bce6c78f771ab8b9f Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 2 Oct 2025 15:40:14 +0200 Subject: [PATCH 064/125] refactor: move unit tests to integration tests due to spark argument --- ...ave_and_load_checks_from_postgres_table.py | 48 +++++++++++++++++++ tests/unit/test_load_checks.py | 30 +----------- tests/unit/test_save_checks.py | 28 ----------- 3 files changed, 50 insertions(+), 56 deletions(-) create mode 100644 tests/integration/test_save_and_load_checks_from_postgres_table.py diff --git a/tests/integration/test_save_and_load_checks_from_postgres_table.py b/tests/integration/test_save_and_load_checks_from_postgres_table.py new file mode 100644 index 000000000..483c4fbbe --- /dev/null +++ b/tests/integration/test_save_and_load_checks_from_postgres_table.py @@ -0,0 +1,48 @@ +from testing.postgresql import Postgresql + +from sqlalchemy import create_engine, insert, select + +from databricks.labs.dqx.checks_storage import LakebaseChecksStorageHandler +from databricks.labs.dqx.config import LakebaseChecksStorageConfig + +from tests.conftest import compare_checks +from tests.unit.test_save_checks import TEST_CHECKS + +LOCATION = "test.public.checks" + + +def test_lakebase_checks_storage_handler_save(ws, spark): + with Postgresql() as postgresql: + connection_string = postgresql.url() + engine = create_engine(connection_string) + handler = LakebaseChecksStorageHandler(ws, spark, engine) + config = LakebaseChecksStorageConfig(LOCATION, connection_string) + schema_name, table_name = handler.get_schema_and_table_name(config) + table = handler.get_table_definition(schema_name, table_name) + + handler.save(TEST_CHECKS, config) + + with engine.connect() as conn: + result = conn.execute(select(table)).mappings().all() + result = [dict(check) for check in result] + + compare_checks(result, TEST_CHECKS) + + +def test_lakebase_checks_storage_handler_load(ws, spark): + with Postgresql() as postgresql: + connection_string = postgresql.url() + engine = create_engine(connection_string) + handler = LakebaseChecksStorageHandler(ws, spark, engine) + config = LakebaseChecksStorageConfig(LOCATION, connection_string) + schema_name, table_name = handler.get_schema_and_table_name(config) + table = handler.get_table_definition(schema_name, table_name) + + table.metadata.create_all(engine, checkfirst=True) + + with engine.begin() as conn: + conn.execute(insert(table), TEST_CHECKS) + + result = handler.load(config) + + compare_checks(result, TEST_CHECKS) diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index 151aacceb..49019c42b 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -1,21 +1,16 @@ from unittest.mock import create_autospec import pytest -from testing.postgresql import Postgresql -from sqlalchemy import create_engine, insert 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, LakebaseChecksStorageHandler -from databricks.labs.dqx.config import LakebaseChecksStorageConfig, VolumeFileChecksStorageConfig +from databricks.labs.dqx.checks_storage import VolumeFileChecksStorageHandler +from databricks.labs.dqx.config import VolumeFileChecksStorageConfig from databricks.labs.dqx.engine import DQEngineCore from databricks.labs.dqx.errors import InvalidCheckError, CheckDownloadError, InvalidConfigError -from tests.conftest import compare_checks -from tests.unit.test_save_checks import TEST_CHECKS - def test_load_checks_from_local_file_json(make_local_check_file_as_json, expected_checks): file = make_local_check_file_as_json @@ -90,24 +85,3 @@ def test_file_download_contents_read_none(): with pytest.raises(NotFound, match="No contents at Unity Catalog volume path"): handler.load(VolumeFileChecksStorageConfig(location="test_path")) - - -def test_lakebase_checks_storage_handler_load(ws, spark): - location = "test.public.checks" - - with Postgresql() as postgresql: - connection_string = postgresql.url() - engine = create_engine(connection_string) - handler = LakebaseChecksStorageHandler(ws, spark, engine) - config = LakebaseChecksStorageConfig(location, connection_string) - - schema_name, table_name = handler.get_schema_and_table_name(config) - table = handler.get_table_definition(schema_name=schema_name, table_name=table_name) - table.metadata.create_all(engine, checkfirst=True) - - with engine.begin() as conn: - conn.execute(insert(table), TEST_CHECKS) - - result = handler.load(config) - - compare_checks(result, TEST_CHECKS) diff --git a/tests/unit/test_save_checks.py b/tests/unit/test_save_checks.py index 4c44c5762..1b0cc4547 100644 --- a/tests/unit/test_save_checks.py +++ b/tests/unit/test_save_checks.py @@ -3,18 +3,11 @@ import yaml import pytest -import testing.postgresql -from sqlalchemy import create_engine, select - from databricks.labs.dqx.engine import DQEngineCore -from databricks.labs.dqx.checks_storage import LakebaseChecksStorageHandler -from databricks.labs.dqx.config import LakebaseChecksStorageConfig from databricks.labs.dqx.config import LakebaseConnectionConfig from databricks.labs.dqx.errors import InvalidConfigError -from tests.conftest import compare_checks - TEST_CHECKS = [ { @@ -75,27 +68,6 @@ def _validate_file(file_path: str, file_format: str = "yaml") -> None: yaml.safe_load(file) -def test_lakebase_checks_storage_handler_save(ws, spark): - location = "test.public.checks" - - with testing.postgresql.Postgresql() as postgresql: - connection_string = postgresql.url() - - engine = create_engine(connection_string) - handler = LakebaseChecksStorageHandler(ws, spark, engine) - config = LakebaseChecksStorageConfig(location, connection_string) - schema_name, table_name = handler.get_schema_and_table_name(config) - table = handler.get_table_definition(schema_name, table_name) - - handler.save(TEST_CHECKS, config) - - with engine.connect() as conn: - result = conn.execute(select(table)).mappings().all() - result = [dict(check) for check in result] - - compare_checks(result, TEST_CHECKS) - - def test_installation_checks_storage_handler_postgresql_parsing(): connection_string = ( "postgresql://user@domain.com:password@instance-test.database.azuredatabricks.net:5432/dqx?sslmode=require" From caeb5d7009b6e18e1f7ad45e23f5c474ce720927 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 2 Oct 2025 15:45:54 +0200 Subject: [PATCH 065/125] refactor: update integration tests to use pytest fixture --- tests/conftest.py | 33 ++++++++++++++++--- ...ave_and_load_checks_from_lakebase_table.py | 31 +++++++++++------ 2 files changed, 50 insertions(+), 14 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 9fa677356..8ebdba738 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -19,9 +19,10 @@ 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 -@pytest.fixture +@pytest.fixture(scope="session") def debug_env_name(): return "ws" # Specify the name of the debug environment from ~/.databricks/debug-env.json @@ -598,9 +599,33 @@ def delete(volume_file_path: str) -> None: @pytest.fixture -def connection_string(env_or_skip): - """Get Lakebase connection string from environment or skip test if not available.""" - return env_or_skip("DQX_LAKEBASE_CONNECTION_STRING") +def make_lakebase_instance_and_catalog(ws, make_random): + database_instance_name = f"dqxtest{make_random(10)}" + database_name = f"dqx" # does not need to be random + catalog_name = f"dqxtest{make_random(10)}" + capacity = "CU_2" + + def create() -> str: + instance = ws.database.create_database_instance_and_wait( + database_instance=DatabaseInstance(name=database_instance_name, capacity=capacity) + ) + + database = ws.database.create_database_catalog( + DatabaseCatalog( + name=catalog_name, + database_name=database_name, + database_instance_name=database_instance_name, + create_database_if_not_exists=True, + ) + ) + + return f"postgresql://{instance.creator}:password@{instance.read_only_dns}:5432/{database.name}?sslmode=require" + + def delete(_: str) -> None: + ws.database.delete_database_catalog(name=catalog_name) + ws.database.delete_database_instance(name=database_instance_name) + + yield from factory("lakebase", create, delete) def sort_key(check: dict[str, Any]) -> str: 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 index 4789bbb82..b76a24ed4 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -11,23 +11,29 @@ LOCATION = "dqx.config.checks" -def test_load_checks_when_checks_lakebase_table_does_not_exist(ws, spark, connection_string): - with pytest.raises(NotFound, match=f"Table '{LOCATION}' does not exist in the Lakebase instance"): - dq_engine = DQEngine(ws) +def test_load_checks_when_checks_lakebase_table_does_not_exist(ws, make_lakebase_instance_and_catalog, spark): + connection_string = make_lakebase_instance_and_catalog() + + with pytest.raises(NotFound, match=f"Resource not found"): + dq_engine = DQEngine(ws, spark) config = LakebaseChecksStorageConfig(location=LOCATION, connection_string=connection_string) dq_engine.load_checks(config=config) -def test_save_and_load_checks_from_lakebase_table(ws, make_schema, make_random, spark, connection_string): - dq_engine = DQEngine(ws) +def test_save_and_load_checks_from_lakebase_table(ws, make_lakebase_instance_and_catalog, spark): + connection_string = make_lakebase_instance_and_catalog() + dq_engine = DQEngine(ws, spark) config = LakebaseChecksStorageConfig(location=LOCATION, connection_string=connection_string) 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, connection_string): +def test_save_and_load_checks_from_lakebase_table_with_run_config(ws, spark, make_lakebase_instance_and_catalog): + connection_string = make_lakebase_instance_and_catalog() dq_engine = DQEngine(ws, spark) + + # test first run config run_config_name = "workflow_001" config_save = LakebaseChecksStorageConfig( location=LOCATION, connection_string=connection_string, run_config_name=run_config_name @@ -39,12 +45,12 @@ def test_save_and_load_checks_from_lakebase_table_with_run_config(ws, spark, con checks = dq_engine.load_checks(config=config_load) compare_checks(checks, TEST_CHECKS[:1]) + # test second run config run_config_name2 = "workflow_002" config_save2 = LakebaseChecksStorageConfig( location=LOCATION, connection_string=connection_string, run_config_name=run_config_name2, - mode="overwrite", ) dq_engine.save_checks(TEST_CHECKS[1:], config=config_save2) config_load2 = LakebaseChecksStorageConfig( @@ -53,6 +59,7 @@ def test_save_and_load_checks_from_lakebase_table_with_run_config(ws, spark, con checks = dq_engine.load_checks(config=config_load2) compare_checks(checks, TEST_CHECKS[:1]) + # test default config dq_engine.save_checks( TEST_CHECKS[1:], config=LakebaseChecksStorageConfig(location=LOCATION, connection_string=connection_string) ) @@ -62,7 +69,8 @@ def test_save_and_load_checks_from_lakebase_table_with_run_config(ws, spark, con compare_checks(checks, TEST_CHECKS[1:]) -def test_save_and_load_checks_to_lakebase_table_output_modes(ws, make_schema, make_random, spark, connection_string): +def test_save_and_load_checks_to_lakebase_table_output_modes(ws, spark, make_lakebase_instance_and_catalog): + connection_string = make_lakebase_instance_and_catalog() dq_engine = DQEngine(ws, spark) run_config_name = "workflow_003" dq_engine.save_checks( @@ -93,7 +101,10 @@ def test_save_and_load_checks_to_lakebase_table_output_modes(ws, make_schema, ma compare_checks(checks, TEST_CHECKS[1:]) -def test_save_load_checks_from_lakebase_table_in_user_installation(ws, spark, installation_ctx, connection_string): +def test_save_load_checks_from_lakebase_table_in_user_installation( + ws, spark, installation_ctx, make_lakebase_instance_and_catalog +): + connection_string = make_lakebase_instance_and_catalog() config = installation_ctx.config run_config = config.get_run_config() run_config.checks_LOCATION = LOCATION @@ -106,8 +117,8 @@ def test_save_load_checks_from_lakebase_table_in_user_installation(ws, spark, in location=LOCATION, connection_string=connection_string, run_config_name=run_config.name, - assume_user=True, product_name=product_name, + assume_user=True, ) dq_engine.save_checks(TEST_CHECKS, config=config) From 45f5d20d144b664f37a45f702673fd79d3767ea0 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 2 Oct 2025 15:47:00 +0200 Subject: [PATCH 066/125] refactor: update error message --- tests/conftest.py | 2 +- .../test_save_and_load_checks_from_lakebase_table.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 8ebdba738..e1722c180 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -601,7 +601,7 @@ def delete(volume_file_path: str) -> None: @pytest.fixture def make_lakebase_instance_and_catalog(ws, make_random): database_instance_name = f"dqxtest{make_random(10)}" - database_name = f"dqx" # does not need to be random + database_name = "dqx" # does not need to be random catalog_name = f"dqxtest{make_random(10)}" capacity = "CU_2" 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 index b76a24ed4..3ab8630d7 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -14,7 +14,7 @@ def test_load_checks_when_checks_lakebase_table_does_not_exist(ws, make_lakebase_instance_and_catalog, spark): connection_string = make_lakebase_instance_and_catalog() - with pytest.raises(NotFound, match=f"Resource not found"): + with pytest.raises(NotFound, match="Resource not found"): dq_engine = DQEngine(ws, spark) config = LakebaseChecksStorageConfig(location=LOCATION, connection_string=connection_string) dq_engine.load_checks(config=config) From 4318d1dd6b53845d479c0d73d83b79a74955757b Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 2 Oct 2025 17:10:28 +0200 Subject: [PATCH 067/125] refactor: place create schema and write to table in two separate transactions --- src/databricks/labs/dqx/checks_storage.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 6a01e8a9c..24b72b427 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -273,6 +273,7 @@ def _save_checks_to_lakebase(self, checks: list[dict], config: LakebaseChecksSto conn.execute(CreateSchema(schema_name)) logger.info(f"Successfully created schema '{schema_name}'.") + with engine.begin() as conn: table = self.get_table_definition(schema_name, table_name) table.metadata.create_all(engine, checkfirst=True) logger.info(f"Successfully created or verified table '{schema_name}.{table_name}'.") From 554be45eca13eec59b4722c18d0a1ca439f52c92 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 2 Oct 2025 17:40:46 +0200 Subject: [PATCH 068/125] fix: correct argument in connection string --- tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index e1722c180..751287969 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -619,7 +619,7 @@ def create() -> str: ) ) - return f"postgresql://{instance.creator}:password@{instance.read_only_dns}:5432/{database.name}?sslmode=require" + return f"postgresql://{instance.creator}:password@{instance.read_only_dns}:5432/{database_name}?sslmode=require" def delete(_: str) -> None: ws.database.delete_database_catalog(name=catalog_name) From c703f976bbb38c10ac522f118d028fdd34553a42 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Fri, 3 Oct 2025 14:15:35 +0200 Subject: [PATCH 069/125] fmt --- .github/workflows/push.yml | 2 -- tests/conftest.py | 6 +++--- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index c70d14ec2..d45a722ed 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -40,8 +40,6 @@ jobs: # click 8.3+ introduced bug for hatch pip install "hatch==1.13.0" "click<8.3" make test - env: - DATABRICKS_HOST: ${{ secrets.DATABRICKS_HOST }} fmt: runs-on: ubuntu-latest diff --git a/tests/conftest.py b/tests/conftest.py index 83236a23a..811662f36 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -639,9 +639,9 @@ def delete(volume_file_path: str) -> None: @pytest.fixture def make_lakebase_instance_and_catalog(ws, make_random): - database_instance_name = f"dqxtest{make_random(10)}" + database_instance_name = f"dqxtest-{make_random(10).lower()}" database_name = "dqx" # does not need to be random - catalog_name = f"dqxtest{make_random(10)}" + catalog_name = f"dqxtest-{make_random(10).lower()}" capacity = "CU_2" def create() -> str: @@ -649,7 +649,7 @@ def create() -> str: database_instance=DatabaseInstance(name=database_instance_name, capacity=capacity) ) - database = ws.database.create_database_catalog( + ws.database.create_database_catalog( DatabaseCatalog( name=catalog_name, database_name=database_name, From 48b7a5627cfcc43144e585321860615a9bb7df9e Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Fri, 3 Oct 2025 15:06:00 +0200 Subject: [PATCH 070/125] refactor --- src/databricks/labs/dqx/checks_storage.py | 62 ++-- src/databricks/labs/dqx/config.py | 59 ++-- ...ave_and_load_checks_from_postgres_table.py | 6 +- tests/unit/test_lakebase_config.py | 267 ++++++++++++++++++ tests/unit/test_save_checks.py | 2 +- 5 files changed, 330 insertions(+), 66 deletions(-) create mode 100644 tests/unit/test_lakebase_config.py diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 24b72b427..18788f7ce 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -172,30 +172,8 @@ def _get_engine(self, connection_config: LakebaseConnectionConfig) -> Engine: connection_url = self._get_connection_url(connection_config) return create_engine(connection_url) - def get_schema_and_table_name(self, config: LakebaseChecksStorageConfig) -> tuple[str, str]: - """ - Extract the schema and table name from the fully qualified table name. - - Args: - config: Configuration for saving and loading checks to Lakebase. - - Returns: - A tuple containing the schema and table name. - - Raises: - ValueError: If the fully qualified table name is not in the correct format. - """ - location_components = config.location.split(".") - - if len(location_components) != 3: - raise ValueError( - f"Invalid Lakebase table name '{config.location}'. Must be in the format 'database.schema.table'." - ) - - _, schema_name, table_name = location_components - return schema_name, table_name - - def get_table_definition(self, schema_name: str, table_name: str) -> Table: + @staticmethod + def get_table_definition(schema_name: str, table_name: str) -> Table: """ Create a table definition for consistency between the load and save methods. @@ -230,7 +208,8 @@ def get_table_definition(self, schema_name: str, table_name: str) -> Table: Column("user_metadata", JSONB, server_default=null()), ) - def _normalize_checks(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> list[dict]: + @staticmethod + def _normalize_checks(checks: list[dict], config: LakebaseChecksStorageConfig) -> list[dict]: """ Normalize the checks to be compatible with the Lakebase table. @@ -266,17 +245,15 @@ def _save_checks_to_lakebase(self, checks: list[dict], config: LakebaseChecksSto Returns: None """ - schema_name, table_name = self.get_schema_and_table_name(config) - with engine.begin() as conn: - if not conn.dialect.has_schema(conn, schema_name): - conn.execute(CreateSchema(schema_name)) - logger.info(f"Successfully created schema '{schema_name}'.") + 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(schema_name, table_name) + 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 '{schema_name}.{table_name}'.") + logger.info(f"Successfully created or verified table '{config.schema_name}.{config.table_name}'.") if config.mode == "overwrite": delete_stmt = delete(table).where(table.c.run_config_name == config.run_config_name) @@ -298,8 +275,7 @@ def _load_checks_from_lakebase(self, config: LakebaseChecksStorageConfig, engine Returns: List of dq rules. """ - schema_name, table_name = self.get_schema_and_table_name(config) - table = self.get_table_definition(schema_name, table_name) + table = self.get_table_definition(config.schema_name, config.table_name) stmt = select(table) if config.run_config_name: @@ -308,7 +284,7 @@ def _load_checks_from_lakebase(self, config: LakebaseChecksStorageConfig, engine with engine.connect() as conn: result = conn.execute(stmt) checks = result.mappings().all() - logger.info(f"Successfully loaded {len(checks)} checks from {schema_name}.{table_name}") + logger.info(f"Successfully loaded {len(checks)} checks from {config.schema_name}.{config.table_name}") return [dict(check) for check in checks] @telemetry_logger("load_checks", "lakebase") @@ -329,8 +305,7 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: engine = self.engine engine_created_internally = False if not engine: - connection_config = LakebaseConnectionConfig.parse_connection_string(config.connection_string) - engine = self._get_engine(connection_config) + engine = self._get_engine(config.connection_config) engine_created_internally = True try: @@ -357,20 +332,16 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: None Raises: - ValueError: If the mode is not 'append' or 'overwrite'. + InvalidCheckError: If any check is invalid or unsupported. DatabaseError: If saving checks fails. """ if not checks: - raise ValueError("Checks cannot be empty or None.") - - if config.mode not in ("append", "overwrite"): - raise ValueError(f"Invalid mode '{config.mode}'. Must be 'append' or 'overwrite'.") + raise InvalidCheckError("Checks cannot be empty or None.") engine = self.engine engine_created_internally = False if not engine: - connection_config = LakebaseConnectionConfig.parse_connection_string(config.connection_string) - engine = self._get_engine(connection_config) + engine = self._get_engine(config.connection_config) engine_created_internally = True try: @@ -559,13 +530,14 @@ def _get_storage_handler_and_config( matches_table_pattern = TABLE_PATTERN.match(config.location) and not config.location.lower().endswith( tuple(FILE_SERIALIZERS.keys()) ) - matches_lakebase_pattern = config.connection_string and config.connection_string.startswith("postgresql://") + matches_lakebase_pattern = config.connection_string.startswith("postgresql://") if matches_table_pattern and matches_lakebase_pattern: return self.lakebase_handler, config if matches_table_pattern: return self.table_handler, config + if config.location.startswith("/Volumes/"): return self.volume_handler, config diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index a62e80d07..bcdca5f93 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -1,4 +1,5 @@ import abc +from functools import cached_property from datetime import datetime, timezone from dataclasses import dataclass, field from urllib.parse import urlparse, unquote @@ -220,28 +221,28 @@ class LakebaseConnectionConfig: port: str = LAKEBASE_DEFAULT_PORT @staticmethod - def parse_connection_string(connection_string: str | None) -> "LakebaseConnectionConfig": + def create(connection_string: str) -> "LakebaseConnectionConfig": if not connection_string: - raise ValueError("Connection string cannot be empty or None.") + raise InvalidParameterError("Connection string cannot be empty or None.") parsed = urlparse(connection_string) if parsed.scheme != "postgresql": - raise ValueError(f"Invalid URL scheme '{parsed.scheme}'. Expected 'postgresql' for Lakebase connections.") + raise InvalidParameterError(f"Invalid URL scheme '{parsed.scheme}'. Expected 'postgresql'.") user = unquote(parsed.username) if parsed.username else None if not user: - raise ValueError(f"Missing username in URL: {connection_string}") + raise InvalidParameterError(f"Missing username in connection string: {connection_string}") instance_name = parsed.hostname if not instance_name: - raise InvalidParameterError(f"Missing hostname in URL: {connection_string}") + raise InvalidParameterError(f"Missing hostname in connection string: {connection_string}") port = str(parsed.port) if parsed.port else LAKEBASE_DEFAULT_PORT database = parsed.path.lstrip("/") if parsed.path else None if not database: - raise ValueError(f"Missing required database name in connection string: {connection_string}") + raise InvalidParameterError(f"Missing database name in connection string: {connection_string}") return LakebaseConnectionConfig( instance_name=instance_name, @@ -269,22 +270,48 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig): """ location: str - connection_string: str | None = None + connection_string: str run_config_name: str = "default" mode: str = "overwrite" def __post_init__(self): if not self.location: - raise InvalidParameterError("The location ('location' field) must not be empty or None.") + raise InvalidParameterError("The 'location' field must not be empty or None.") - if not self.connection_string: - raise InvalidParameterError("The connection string ('connection_string' field) must not be empty or None.") + # Eager validation of connection string + try: + self._connection_config = LakebaseConnectionConfig.create(self.connection_string) + except Exception as e: + raise InvalidParameterError(f"Failed to parse connection string '{self.connection_string}': {e}") from e - if self.connection_string and self.connection_string.startswith("postgresql://"): - try: - LakebaseConnectionConfig.parse_connection_string(self.connection_string) - except Exception as e: - raise InvalidParameterError(f"Failed to parse connection string '{self.connection_string}': {e}") from e + 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'.") + + @cached_property + def connection_config(self) -> LakebaseConnectionConfig: + """Parses and validates the Lakebase connection string.""" + return LakebaseConnectionConfig.create(self.connection_string) + + @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] + + def _split_location(self) -> tuple[str, ...]: + """Splits 'database.schema.table' into components.""" + return tuple(self.location.split(".")) @dataclass @@ -325,8 +352,8 @@ class InstallationChecksStorageConfig( * Global directory if `DQX_FORCE_INSTALL=global`: "/Applications/dqx" """ - connection_string: str | None = None # retrieved from the installation config location: str = "installation" # retrieved from the installation config + connection_string: str = "installation" # (optional) retrieved from the installation config for database connection run_config_name: str = "default" # to retrieve run config product_name: str = "dqx" assume_user: bool = True diff --git a/tests/integration/test_save_and_load_checks_from_postgres_table.py b/tests/integration/test_save_and_load_checks_from_postgres_table.py index 483c4fbbe..a9e078abc 100644 --- a/tests/integration/test_save_and_load_checks_from_postgres_table.py +++ b/tests/integration/test_save_and_load_checks_from_postgres_table.py @@ -17,8 +17,7 @@ def test_lakebase_checks_storage_handler_save(ws, spark): engine = create_engine(connection_string) handler = LakebaseChecksStorageHandler(ws, spark, engine) config = LakebaseChecksStorageConfig(LOCATION, connection_string) - schema_name, table_name = handler.get_schema_and_table_name(config) - table = handler.get_table_definition(schema_name, table_name) + table = handler.get_table_definition(config.schema_name, config.table_name) handler.save(TEST_CHECKS, config) @@ -35,8 +34,7 @@ def test_lakebase_checks_storage_handler_load(ws, spark): engine = create_engine(connection_string) handler = LakebaseChecksStorageHandler(ws, spark, engine) config = LakebaseChecksStorageConfig(LOCATION, connection_string) - schema_name, table_name = handler.get_schema_and_table_name(config) - table = handler.get_table_definition(schema_name, table_name) + table = handler.get_table_definition(config.schema_name, config.table_name) table.metadata.create_all(engine, checkfirst=True) diff --git a/tests/unit/test_lakebase_config.py b/tests/unit/test_lakebase_config.py new file mode 100644 index 000000000..03ba2671a --- /dev/null +++ b/tests/unit/test_lakebase_config.py @@ -0,0 +1,267 @@ +"""Unit tests for LakebaseConnectionConfig and LakebaseChecksStorageConfig.""" + +import pytest +from databricks.labs.dqx.config import LakebaseConnectionConfig, LakebaseChecksStorageConfig +from databricks.labs.dqx.errors import InvalidParameterError, InvalidConfigError + + +def test_create_valid_connection_string(): + """Test creating LakebaseConnectionConfig with valid connection string.""" + connection_string = "postgresql://testuser:password@testhost:5432/testdb" + config = LakebaseConnectionConfig.create(connection_string) + + assert config.instance_name == "testhost" + assert config.database == "testdb" + assert config.user == "testuser" + assert config.port == "5432" + + +def test_create_valid_connection_string_default_port(): + """Test creating LakebaseConnectionConfig with valid connection string without port.""" + connection_string = "postgresql://testuser:password@testhost/testdb" + config = LakebaseConnectionConfig.create(connection_string) + + assert config.instance_name == "testhost" + assert config.database == "testdb" + assert config.user == "testuser" + assert config.port == "5432" # Default port + + +def test_create_valid_connection_string_with_encoded_username(): + """Test creating LakebaseConnectionConfig with URL-encoded username.""" + connection_string = "postgresql://test%40user:password@testhost:5432/testdb" + config = LakebaseConnectionConfig.create(connection_string) + + assert config.instance_name == "testhost" + assert config.database == "testdb" + assert config.user == "test@user" # URL-decoded + assert config.port == "5432" + + +@pytest.mark.parametrize( + "connection_string,expected_error", + [ + # Empty/None connection string + ("", "Connection string cannot be empty or None."), + (None, "Connection string cannot be empty or None."), + # Invalid URL scheme + ("mysql://user:password@host:5432/db", "Invalid URL scheme 'mysql'. Expected 'postgresql'."), + ("http://user:password@host:5432/db", "Invalid URL scheme 'http'. Expected 'postgresql'."), + ("ftp://user:password@host:5432/db", "Invalid URL scheme 'ftp'. Expected 'postgresql'."), + # Missing username + ( + "postgresql://:password@host:5432/db", + "Missing username in connection string: postgresql://:password@host:5432/db", + ), + ("postgresql://host:5432/db", "Missing username in connection string: postgresql://host:5432/db"), + # Missing hostname + ( + "postgresql://user:password@:5432/db", + "Missing hostname in connection string: postgresql://user:password@:5432/db", + ), + ("postgresql://user:password@/db", "Missing hostname in connection string: postgresql://user:password@/db"), + # Missing database name + ( + "postgresql://user:password@host:5432/", + "Missing database name in connection string: postgresql://user:password@host:5432/", + ), + ( + "postgresql://user:password@host:5432", + "Missing database name in connection string: postgresql://user:password@host:5432", + ), + ], +) +def test_create_invalid_connection_strings(connection_string, expected_error): + """Test LakebaseConnectionConfig.create() with invalid connection strings.""" + with pytest.raises(InvalidParameterError, match=expected_error): + LakebaseConnectionConfig.create(connection_string) + + +def test_create_malformed_url(): + """Test creating LakebaseConnectionConfig with malformed URL.""" + with pytest.raises(InvalidParameterError, match="Invalid URL scheme"): + LakebaseConnectionConfig.create("not-a-valid-url") + + +@pytest.fixture +def valid_connection_string(): + """Valid connection string for testing.""" + return "postgresql://testuser:password@testhost:5432/testdb" + + +def test_valid_config(valid_connection_string): + """Test creating LakebaseChecksStorageConfig with valid parameters.""" + config = LakebaseChecksStorageConfig( + location="testdb.testschema.testtable", + connection_string=valid_connection_string, + run_config_name="test_run", + mode="append", + ) + + assert config.location == "testdb.testschema.testtable" + assert config.connection_string == valid_connection_string + 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_config_defaults(valid_connection_string): + """Test creating LakebaseChecksStorageConfig with default values.""" + config = LakebaseChecksStorageConfig( + location="testdb.testschema.testtable", connection_string=valid_connection_string + ) + + assert config.run_config_name == "default" + assert config.mode == "overwrite" + + +@pytest.mark.parametrize( + "location,expected_error", + [ + # Empty/None location + ("", "The 'location' field must not be empty or None."), + (None, "The 'location' field must not be empty or None."), + ], +) +def test_invalid_location_empty(location, expected_error, valid_connection_string): + """Test LakebaseChecksStorageConfig with empty/None location.""" + with pytest.raises(InvalidParameterError, match=expected_error): + LakebaseChecksStorageConfig(location=location, connection_string=valid_connection_string) + + +@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'.", + ), + ("", "The 'location' field must not be empty or None."), + ], +) +def test_invalid_location_format(location, expected_error, valid_connection_string): + """Test LakebaseChecksStorageConfig with invalid location format.""" + if location == "": + # Empty location raises InvalidParameterError + with pytest.raises(InvalidParameterError, match=expected_error): + LakebaseChecksStorageConfig(location=location, connection_string=valid_connection_string) + else: + # Invalid format raises InvalidConfigError + with pytest.raises(InvalidConfigError, match=expected_error): + LakebaseChecksStorageConfig(location=location, connection_string=valid_connection_string) + + +@pytest.mark.parametrize( + "connection_string,expected_error_pattern", + [ + # Invalid connection strings (delegated to LakebaseConnectionConfig.create) + ("", "Failed to parse connection string '': Connection string cannot be empty or None."), + ( + "mysql://user:password@host:5432/db", + "Failed to parse connection string 'mysql://user:password@host:5432/db': Invalid URL scheme 'mysql'. Expected 'postgresql'.", + ), + ( + "postgresql://host:5432/db", + "Failed to parse connection string 'postgresql://host:5432/db': Missing username in connection string", + ), + ( + "postgresql://user:password@:5432/db", + "Failed to parse connection string 'postgresql://user:password@:5432/db': Missing hostname in connection string", + ), + ( + "postgresql://user:password@host:5432/", + "Failed to parse connection string 'postgresql://user:password@host:5432/': Missing database name in connection string", + ), + ], +) +def test_invalid_connection_string(connection_string, expected_error_pattern): + """Test LakebaseChecksStorageConfig with invalid connection strings.""" + with pytest.raises(InvalidParameterError, match=expected_error_pattern): + LakebaseChecksStorageConfig(location="testdb.testschema.testtable", connection_string=connection_string) + + +@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, valid_connection_string): + """Test LakebaseChecksStorageConfig with invalid mode.""" + with pytest.raises(InvalidConfigError, match=expected_error): + LakebaseChecksStorageConfig( + location="testdb.testschema.testtable", connection_string=valid_connection_string, mode=mode + ) + + +def test_valid_modes(valid_connection_string): + """Test LakebaseChecksStorageConfig with valid modes.""" + # Test 'append' mode + config_append = LakebaseChecksStorageConfig( + location="testdb.testschema.testtable", connection_string=valid_connection_string, mode="append" + ) + assert config_append.mode == "append" + + # Test 'overwrite' mode + config_overwrite = LakebaseChecksStorageConfig( + location="testdb.testschema.testtable", connection_string=valid_connection_string, mode="overwrite" + ) + assert config_overwrite.mode == "overwrite" + + +def test_connection_config_property(valid_connection_string): + """Test that connection_config property returns correct LakebaseConnectionConfig.""" + config = LakebaseChecksStorageConfig( + location="testdb.testschema.testtable", connection_string=valid_connection_string + ) + + conn_config = config.connection_config + assert isinstance(conn_config, LakebaseConnectionConfig) + assert conn_config.instance_name == "testhost" + assert conn_config.database == "testdb" + assert conn_config.user == "testuser" + assert conn_config.port == "5432" + + +def test_location_properties(valid_connection_string): + """Test database_name, schema_name, and table_name properties.""" + config = LakebaseChecksStorageConfig(location="mydb.myschema.mytable", connection_string=valid_connection_string) + + assert config.database_name == "mydb" + assert config.schema_name == "myschema" + assert config.table_name == "mytable" + + +def test_complex_location_with_special_characters(valid_connection_string): + """Test location parsing with special characters in names.""" + config = LakebaseChecksStorageConfig(location="my_db.my_schema.my_table", connection_string=valid_connection_string) + + assert config.database_name == "my_db" + assert config.schema_name == "my_schema" + assert config.table_name == "my_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="Failed to parse connection string"): + LakebaseChecksStorageConfig( + location="testdb.testschema.testtable", connection_string="invalid-connection-string" + ) + + +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 connection string validation + with pytest.raises(InvalidParameterError, match="The 'location' field must not be empty or None"): + LakebaseChecksStorageConfig(location="", connection_string="invalid-connection-string") diff --git a/tests/unit/test_save_checks.py b/tests/unit/test_save_checks.py index 1b0cc4547..5a8fe7c68 100644 --- a/tests/unit/test_save_checks.py +++ b/tests/unit/test_save_checks.py @@ -72,7 +72,7 @@ def test_installation_checks_storage_handler_postgresql_parsing(): connection_string = ( "postgresql://user@domain.com:password@instance-test.database.azuredatabricks.net:5432/dqx?sslmode=require" ) - connection_config = LakebaseConnectionConfig.parse_connection_string(connection_string) + connection_config = LakebaseConnectionConfig.create(connection_string) assert connection_config.user == "user@domain.com" assert connection_config.instance_name == "instance-test.database.azuredatabricks.net" From 93184fa06b0173e5cd6341cbae15eb13cd3e6e21 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Fri, 3 Oct 2025 20:00:01 +0200 Subject: [PATCH 071/125] force delete db instance in fixture --- tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 9e9b195e2..ae41398bc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -669,7 +669,7 @@ def create() -> str: def delete(_: str) -> None: ws.database.delete_database_catalog(name=catalog_name) - ws.database.delete_database_instance(name=database_instance_name) + ws.database.delete_database_instance(name=database_instance_name, force=True) yield from factory("lakebase", create, delete) From 1dd0fbf65352ec9cdafef00695c11db876618467 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 6 Oct 2025 17:47:38 +0200 Subject: [PATCH 072/125] fix: remove errant argument from database delete function --- tests/conftest.py | 50 ++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 41 insertions(+), 9 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index ae41398bc..cec60abb0 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -646,12 +646,18 @@ def delete(volume_file_path: str) -> None: @pytest.fixture def make_lakebase_instance_and_catalog(ws, make_random): - database_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" + import logging + + logger = logging.getLogger(__name__) + + instances = {} def create() -> str: + database_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" + instance = ws.database.create_database_instance_and_wait( database_instance=DatabaseInstance(name=database_instance_name, capacity=capacity) ) @@ -665,11 +671,37 @@ def create() -> str: ) ) - return f"postgresql://{instance.creator}:password@{instance.read_only_dns}:5432/{database_name}?sslmode=require" + connection_string = ( + f"postgresql://{instance.creator}:password@{instance.read_only_dns}:5432/{database_name}?sslmode=require" + ) + + instances[connection_string] = {"database_instance_name": database_instance_name, "catalog_name": catalog_name} + + return connection_string + + def delete(connection_string: str) -> None: + if connection_string not in instances: + logger.warning(f"No resources found for connection string: {connection_string}") + return + + metadata = instances[connection_string] + database_instance_name = metadata["database_instance_name"] + catalog_name = metadata["catalog_name"] + + try: + ws.database.delete_database_catalog(name=catalog_name) + logger.info(f"Successfully deleted database catalog: {catalog_name}") + except Exception as e: + logger.warning(f"Failed to delete database catalog {catalog_name}: {e}") - def delete(_: str) -> None: - ws.database.delete_database_catalog(name=catalog_name) - ws.database.delete_database_instance(name=database_instance_name, force=True) + try: + ws.database.delete_database_instance(name=database_instance_name) + logger.info(f"Successfully deleted database instance: {database_instance_name}") + except Exception as e: + logger.error(f"Failed to delete database instance {database_instance_name}: {e}") + raise + finally: + instances.pop(connection_string, None) yield from factory("lakebase", create, delete) @@ -684,7 +716,7 @@ def sort_key(check: dict[str, Any]) -> str: Returns: The name of the check as a string, or an empty string if not found. """ - return str(check.get('name', '')) + return str(check.get("name", "")) def compare_checks(result: list[dict[str, Any]], expected: list[dict[str, Any]]) -> None: From 3cd2a314119c4f151c34568549323d91cdb43037 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 7 Oct 2025 16:24:41 +0200 Subject: [PATCH 073/125] fix: clean up old test instances --- ...st_save_and_load_checks_from_lakebase_table.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) 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 index 3ab8630d7..1ced6f153 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -124,3 +124,18 @@ def test_save_load_checks_from_lakebase_table_in_user_installation( dq_engine.save_checks(TEST_CHECKS, config=config) checks = dq_engine.load_checks(config=config) compare_checks(checks, TEST_CHECKS) + + +def delete_all_leftover_instances(ws): + import re + + pattern = re.compile(r"^dqxtest-[A-Za-z0-9]{10}$") + + instances = [] + + for instance in ws.database.list_database_instances(): + if pattern.match(instance.name): + instances.append(instance.name) + + for instance in instances: + ws.database.delete_database_instance(name=instance) From 6e77a5bb17eef80c3b23d6778e22c3a023f8bf4d Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 7 Oct 2025 16:31:26 +0200 Subject: [PATCH 074/125] docs: reformat docstring --- src/databricks/labs/dqx/checks_storage.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 49c7a9aa1..64b22a3a2 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -182,16 +182,16 @@ def get_table_definition(schema_name: str, table_name: str) -> Table: Note: The table has the following columns: - * `name` - Optional. Name to use for the check. - * `criticality` - Either "error" (rows go only to "bad" dataset) or "warn" (rows go to both "good" and "bad"). + `name` - Optional. Name to use for the check. + `criticality` - Either "error" (rows go only to "bad" dataset) or "warn" (rows go to both "good" and "bad"). Defaults to "error". - * `check` - Defines the DQX check to apply, e.g. {"function": "is_not_null", "arguments": {"column": "col1"}}. - * `filter` - Optional. Spark SQL expression to filter rows to which the check is applied, e.g. "business_unit = 'Finance'". + `check` - Defines the DQX check to apply, e.g. {"function": "is_not_null", "arguments": {"column": "col1"}}. + `filter` - Optional. Spark SQL expression to filter rows to which the check is applied, e.g. "business_unit = 'Finance'". The check function will run only on the rows matching the filter condition. The condition can reference any column of the validated dataset, not only the one where you apply the check function. - * `run_config_name` - Name of the run config name. Could be any string such as workflow name. Useful for selecting + `run_config_name` - Name of the run config name. Could be any string such as workflow name. Useful for selecting applicable checks. Defaults to `"default"`. - * `user_metadata` - Optional. Custom metadata to add to any row-level warnings or errors generated by the check. + `user_metadata` - Optional. Custom metadata to add to any row-level warnings or errors generated by the check. Returns: SQLAlchemy table definition for the Lakebase instance. From 6e971f75f1c335037aaccb5bb033ffbf79d5d521 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 8 Oct 2025 14:55:24 +0200 Subject: [PATCH 075/125] refactor: remove connection string argument --- src/databricks/labs/dqx/checks_storage.py | 25 +++-- src/databricks/labs/dqx/config.py | 110 ++++++---------------- 2 files changed, 41 insertions(+), 94 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 64b22a3a2..3383c55c8 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -30,7 +30,6 @@ from databricks.sdk.service.workspace import ImportFormat from databricks.labs.dqx.config import ( - LakebaseConnectionConfig, TableChecksStorageConfig, LakebaseChecksStorageConfig, FileChecksStorageConfig, @@ -139,36 +138,36 @@ def __init__(self, ws: WorkspaceClient, spark: SparkSession, engine: Engine | No self.spark = spark self.engine = engine - def _get_connection_url(self, connection_config: LakebaseConnectionConfig) -> str: + def _get_connection_url(self, config: LakebaseChecksStorageConfig) -> str: """ Generate a Lakebase connection URL. Args: - connection_config: Configuration for a Lakebase connection. + config: Configuration for saving and loading checks to Lakebase. Returns: Lakebase connection URL. """ - instance = self.ws.database.get_database_instance(connection_config.database) + instance = self.ws.database.get_database_instance(config.instance_name) cred = self.ws.database.generate_database_credential( - request_id=str(uuid.uuid4()), instance_names=[connection_config.database] + request_id=str(uuid.uuid4()), instance_names=[config.instance_name] ) host = instance.read_write_dns password = cred.token - return f"postgresql://{connection_config.user}:{password}@{host}:{connection_config.port}/{connection_config.database}?sslmode=require" + return f"postgresql://{config.user}:{password}@{host}:{config.port}/{config.database_name}?sslmode=require" - def _get_engine(self, connection_config: LakebaseConnectionConfig) -> Engine: + def _get_engine(self, config: LakebaseChecksStorageConfig) -> Engine: """ Create a SQLAlchemy engine for the Lakebase instance. Args: - connection_config: Configuration for a Lakebase connection. + config: Configuration for saving and loading checks to Lakebase. Returns: SQLAlchemy engine for the Lakebase instance. """ - connection_url = self._get_connection_url(connection_config) + connection_url = self._get_connection_url(config) return create_engine(connection_url) @staticmethod @@ -304,7 +303,7 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: engine = self.engine engine_created_internally = False if not engine: - engine = self._get_engine(config.connection_config) + engine = self._get_engine(config) engine_created_internally = True try: @@ -340,7 +339,7 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: engine = self.engine engine_created_internally = False if not engine: - engine = self._get_engine(config.connection_config) + engine = self._get_engine(config) engine_created_internally = True try: @@ -532,9 +531,9 @@ def _get_storage_handler_and_config( config.location = checks_location matches_table_pattern = is_table_location(config.location) - matches_lakebase_pattern = config.connection_string.startswith("postgresql://") + is_lakebase = hasattr(config, 'instance_name') and config.instance_name - if matches_table_pattern and matches_lakebase_pattern: + if matches_table_pattern and is_lakebase: return self.lakebase_handler, config if matches_table_pattern: diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 8f874da27..3ffa5eac1 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -2,7 +2,6 @@ from functools import cached_property from datetime import datetime, timezone from dataclasses import dataclass, field -from urllib.parse import urlparse, unquote from databricks.labs.dqx.errors import InvalidConfigError, InvalidParameterError __all__ = [ @@ -68,7 +67,10 @@ class RunConfig: checks_location: str = ( "checks.yml" # absolute or relative workspace file path or table containing quality rules / checks ) - connection_string: str | None = None # connection string for Lakebase instances + # Lakebase connection parameters (only used when checks_location points to a Lakebase instance) + lakebase_instance_name: str | None = None # Lakebase instance name + lakebase_user: str | None = None # Lakebase username + lakebase_port: str = LAKEBASE_DEFAULT_PORT # Lakebase port 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 @@ -201,104 +203,49 @@ def __post_init__(self): raise InvalidConfigError("The table name ('location' field) must not be empty or None.") -@dataclass -class LakebaseConnectionConfig: - """ - Configuration class for a Lakebase connection. - - Note: - The connection string format includes a placeholder for a password, - e.g., postgresql://user:password@host:port/database, but the password - is not stored or parsed by this class. - - Args: - instance_name: The Lakebase instance name (hostname). - database: The Lakebase database name. - user: The Lakebase username. - port: The Lakebase port (default is '5432'). - """ - - instance_name: str - database: str - user: str - port: str = LAKEBASE_DEFAULT_PORT - - @staticmethod - def create(connection_string: str) -> "LakebaseConnectionConfig": - if not connection_string: - raise InvalidParameterError("Connection string cannot be empty or None.") - - parsed = urlparse(connection_string) - - if parsed.scheme != "postgresql": - raise InvalidParameterError(f"Invalid URL scheme '{parsed.scheme}'. Expected 'postgresql'.") - - user = unquote(parsed.username) if parsed.username else None - if not user: - raise InvalidParameterError(f"Missing username in connection string: {connection_string}") - - instance_name = parsed.hostname - if not instance_name: - raise InvalidParameterError(f"Missing hostname in connection string: {connection_string}") - - port = str(parsed.port) if parsed.port else LAKEBASE_DEFAULT_PORT - - database = parsed.path.lstrip("/") if parsed.path else None - if not database: - raise InvalidParameterError(f"Missing database name in connection string: {connection_string}") - - return LakebaseConnectionConfig( - instance_name=instance_name, - database=database, - user=user, - port=port, - ) - - @dataclass class LakebaseChecksStorageConfig(BaseChecksStorageConfig): """ Configuration class for storing checks in a Lakebase table. - Note: - The `schema` and `table` fields are not parsed from the connection string, but - must be provided in the `location` field in the format 'database.schema.table'. - Args: - location: Fully qualified Lakebase table name in the format 'database.schema.table'. - connection_string: (SQLAlchemy) Lakebase connection string, e.g., postgresql://user:password@host:port/database. + 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'). + only replaces checks for the specific run config and not all checks in the table (default is 'overwrite'). """ + instance_name: str + user: str location: str - connection_string: str + port: str = LAKEBASE_DEFAULT_PORT run_config_name: str = "default" mode: str = "overwrite" def __post_init__(self): - if not self.location: - raise InvalidParameterError("The 'location' field must not be empty or None.") + 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.") - # Eager validation of connection string - try: - self._connection_config = LakebaseConnectionConfig.create(self.connection_string) - except Exception as e: - raise InvalidParameterError(f"Failed to parse connection string '{self.connection_string}': {e}") from e + 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'." + 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'.") - @cached_property - def connection_config(self) -> LakebaseConnectionConfig: - """Parses and validates the Lakebase connection string.""" - return LakebaseConnectionConfig.create(self.connection_string) + def _split_location(self) -> tuple[str, ...]: + """Splits 'database.schema.table' into components.""" + return tuple(self.location.split(".")) @cached_property def database_name(self) -> str: @@ -312,10 +259,6 @@ def schema_name(self) -> str: def table_name(self) -> str: return self._split_location()[2] - def _split_location(self) -> tuple[str, ...]: - """Splits 'database.schema.table' into components.""" - return tuple(self.location.split(".")) - @dataclass class VolumeFileChecksStorageConfig(BaseChecksStorageConfig): @@ -344,9 +287,12 @@ class InstallationChecksStorageConfig( Configuration class for storing checks in an installation. Args: + instance_name: The name of the Lakebase instance (default is 'installation'). location: The installation path where the checks are stored (e.g., table name, file path). Not used when using installation method, as it is retrieved from the installation config, unless overwrite_location is enabled. + user: The user for the Lakebase connection. + port: The port for the Lakebase connection (default is '5432'). run_config_name: The name of the run configuration to use for checks (default is 'default'). product_name: The product name for retrieving checks from the installation (default is 'dqx'). assume_user: Whether to assume the user is the owner of the checks (default is True). @@ -357,8 +303,10 @@ class InstallationChecksStorageConfig( overwrite_location: Whether to overwrite the location from run config if provided (default is False). """ + instance_name: str | None = None location: str = "installation" # retrieved from the installation config - connection_string: str = "installation" # (optional) retrieved from the installation config for database connection + user: str | None = None + port: str = LAKEBASE_DEFAULT_PORT # default port for Lakebase connection run_config_name: str = "default" # to retrieve run config product_name: str = "dqx" assume_user: bool = True From 02734c5f089bf560cf1f6d30b145aad90bd5199a Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 8 Oct 2025 15:42:35 +0200 Subject: [PATCH 076/125] refactor: removed postgresql tests --- pyproject.toml | 4 +- ...ave_and_load_checks_from_postgres_table.py | 46 ------------------- 2 files changed, 1 insertion(+), 49 deletions(-) delete mode 100644 tests/integration/test_save_and_load_checks_from_postgres_table.py diff --git a/pyproject.toml b/pyproject.toml index 49bd6a9e4..e114047a6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,8 +102,6 @@ dependencies = [ "dbldatagen~=0.4.0", "pyparsing~=3.2.3", "jmespath~=1.0.1", - "testing.postgresql~=1.3.0", - "testing.common.database~=2.0.3", "psycopg2-binary~=2.9.10", ] @@ -167,7 +165,7 @@ profile = "black" exclude = ['venv', '.venv', 'demos/*', 'tests/e2e/*'] [[tool.mypy.overrides]] -module = ["google", "google.*", "testing.*"] # Google and testing libraries are not type annotated +module = ["google", "google.*"] # Google and testing libraries are not type annotated ignore_missing_imports = true [tool.pytest.ini_options] diff --git a/tests/integration/test_save_and_load_checks_from_postgres_table.py b/tests/integration/test_save_and_load_checks_from_postgres_table.py deleted file mode 100644 index a9e078abc..000000000 --- a/tests/integration/test_save_and_load_checks_from_postgres_table.py +++ /dev/null @@ -1,46 +0,0 @@ -from testing.postgresql import Postgresql - -from sqlalchemy import create_engine, insert, select - -from databricks.labs.dqx.checks_storage import LakebaseChecksStorageHandler -from databricks.labs.dqx.config import LakebaseChecksStorageConfig - -from tests.conftest import compare_checks -from tests.unit.test_save_checks import TEST_CHECKS - -LOCATION = "test.public.checks" - - -def test_lakebase_checks_storage_handler_save(ws, spark): - with Postgresql() as postgresql: - connection_string = postgresql.url() - engine = create_engine(connection_string) - handler = LakebaseChecksStorageHandler(ws, spark, engine) - config = LakebaseChecksStorageConfig(LOCATION, connection_string) - table = handler.get_table_definition(config.schema_name, config.table_name) - - handler.save(TEST_CHECKS, config) - - with engine.connect() as conn: - result = conn.execute(select(table)).mappings().all() - result = [dict(check) for check in result] - - compare_checks(result, TEST_CHECKS) - - -def test_lakebase_checks_storage_handler_load(ws, spark): - with Postgresql() as postgresql: - connection_string = postgresql.url() - engine = create_engine(connection_string) - handler = LakebaseChecksStorageHandler(ws, spark, engine) - config = LakebaseChecksStorageConfig(LOCATION, connection_string) - table = handler.get_table_definition(config.schema_name, config.table_name) - - table.metadata.create_all(engine, checkfirst=True) - - with engine.begin() as conn: - conn.execute(insert(table), TEST_CHECKS) - - result = handler.load(config) - - compare_checks(result, TEST_CHECKS) From e6f07a91b45da5212b43337b80112a890405878d Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 8 Oct 2025 15:48:35 +0200 Subject: [PATCH 077/125] chore: update unit tests --- tests/unit/test_lakebase_config.py | 218 ++++++++++------------------- tests/unit/test_save_checks.py | 13 -- 2 files changed, 75 insertions(+), 156 deletions(-) diff --git a/tests/unit/test_lakebase_config.py b/tests/unit/test_lakebase_config.py index 03ba2671a..7787db2c0 100644 --- a/tests/unit/test_lakebase_config.py +++ b/tests/unit/test_lakebase_config.py @@ -1,105 +1,69 @@ """Unit tests for LakebaseConnectionConfig and LakebaseChecksStorageConfig.""" import pytest -from databricks.labs.dqx.config import LakebaseConnectionConfig, LakebaseChecksStorageConfig +from databricks.labs.dqx.config import LakebaseChecksStorageConfig from databricks.labs.dqx.errors import InvalidParameterError, InvalidConfigError -def test_create_valid_connection_string(): - """Test creating LakebaseConnectionConfig with valid connection string.""" - connection_string = "postgresql://testuser:password@testhost:5432/testdb" - config = LakebaseConnectionConfig.create(connection_string) +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 == "testdb" + assert config.database_name == "testdb" assert config.user == "testuser" assert config.port == "5432" -def test_create_valid_connection_string_default_port(): - """Test creating LakebaseConnectionConfig with valid connection string without port.""" - connection_string = "postgresql://testuser:password@testhost/testdb" - config = LakebaseConnectionConfig.create(connection_string) +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 == "testdb" + assert config.database_name == "testdb" assert config.user == "testuser" assert config.port == "5432" # Default port -def test_create_valid_connection_string_with_encoded_username(): - """Test creating LakebaseConnectionConfig with URL-encoded username.""" - connection_string = "postgresql://test%40user:password@testhost:5432/testdb" - config = LakebaseConnectionConfig.create(connection_string) - - assert config.instance_name == "testhost" - assert config.database == "testdb" - assert config.user == "test@user" # URL-decoded - assert config.port == "5432" - - @pytest.mark.parametrize( - "connection_string,expected_error", + "instance_name,location,user,port,expected_error", [ - # Empty/None connection string - ("", "Connection string cannot be empty or None."), - (None, "Connection string cannot be empty or None."), - # Invalid URL scheme - ("mysql://user:password@host:5432/db", "Invalid URL scheme 'mysql'. Expected 'postgresql'."), - ("http://user:password@host:5432/db", "Invalid URL scheme 'http'. Expected 'postgresql'."), - ("ftp://user:password@host:5432/db", "Invalid URL scheme 'ftp'. Expected 'postgresql'."), - # Missing username - ( - "postgresql://:password@host:5432/db", - "Missing username in connection string: postgresql://:password@host:5432/db", - ), - ("postgresql://host:5432/db", "Missing username in connection string: postgresql://host:5432/db"), - # Missing hostname - ( - "postgresql://user:password@:5432/db", - "Missing hostname in connection string: postgresql://user:password@:5432/db", - ), - ("postgresql://user:password@/db", "Missing hostname in connection string: postgresql://user:password@/db"), - # Missing database name - ( - "postgresql://user:password@host:5432/", - "Missing database name in connection string: postgresql://user:password@host:5432/", - ), - ( - "postgresql://user:password@host:5432", - "Missing database name in connection string: postgresql://user:password@host:5432", - ), + # 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_connection_strings(connection_string, expected_error): - """Test LakebaseConnectionConfig.create() with invalid connection strings.""" +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): - LakebaseConnectionConfig.create(connection_string) + LakebaseChecksStorageConfig(instance_name=instance_name, location=location, user=user, port=port) -def test_create_malformed_url(): - """Test creating LakebaseConnectionConfig with malformed URL.""" - with pytest.raises(InvalidParameterError, match="Invalid URL scheme"): - LakebaseConnectionConfig.create("not-a-valid-url") - - -@pytest.fixture -def valid_connection_string(): - """Valid connection string for testing.""" - return "postgresql://testuser:password@testhost:5432/testdb" - - -def test_valid_config(valid_connection_string): +def test_valid_lakebase_checks_storage_config(): """Test creating LakebaseChecksStorageConfig with valid parameters.""" config = LakebaseChecksStorageConfig( + instance_name="testinstance", location="testdb.testschema.testtable", - connection_string=valid_connection_string, + user="testuser", + port="5432", run_config_name="test_run", mode="append", ) + assert config.instance_name == "testinstance" assert config.location == "testdb.testschema.testtable" - assert config.connection_string == valid_connection_string + assert config.user == "testuser" + assert config.port == "5432" assert config.run_config_name == "test_run" assert config.mode == "append" assert config.database_name == "testdb" @@ -107,28 +71,29 @@ def test_valid_config(valid_connection_string): assert config.table_name == "testtable" -def test_valid_config_defaults(valid_connection_string): +def test_valid_lakebase_checks_storage_config_defaults(): """Test creating LakebaseChecksStorageConfig with default values.""" config = LakebaseChecksStorageConfig( - location="testdb.testschema.testtable", connection_string=valid_connection_string + 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 - ("", "The 'location' field must not be empty or None."), - (None, "The 'location' field must not be empty or None."), + ("", "Location must not be empty or None."), + (None, "Location must not be empty or None."), ], ) -def test_invalid_location_empty(location, expected_error, valid_connection_string): +def test_invalid_location_empty(location, expected_error): """Test LakebaseChecksStorageConfig with empty/None location.""" with pytest.raises(InvalidParameterError, match=expected_error): - LakebaseChecksStorageConfig(location=location, connection_string=valid_connection_string) + LakebaseChecksStorageConfig(instance_name="testinstance", location=location, user="testuser") @pytest.mark.parametrize( @@ -141,48 +106,27 @@ def test_invalid_location_empty(location, expected_error, valid_connection_strin "db.schema.table.extra", "Invalid Lakebase table name 'db.schema.table.extra'. Must be in the format 'database.schema.table'.", ), - ("", "The 'location' field must not be empty or None."), ], ) -def test_invalid_location_format(location, expected_error, valid_connection_string): +def test_invalid_location_format(location, expected_error): """Test LakebaseChecksStorageConfig with invalid location format.""" - if location == "": - # Empty location raises InvalidParameterError - with pytest.raises(InvalidParameterError, match=expected_error): - LakebaseChecksStorageConfig(location=location, connection_string=valid_connection_string) - else: - # Invalid format raises InvalidConfigError - with pytest.raises(InvalidConfigError, match=expected_error): - LakebaseChecksStorageConfig(location=location, connection_string=valid_connection_string) + with pytest.raises(InvalidConfigError, match=expected_error): + LakebaseChecksStorageConfig(instance_name="testinstance", location=location, user="testuser") @pytest.mark.parametrize( - "connection_string,expected_error_pattern", + "instance_name,expected_error", [ - # Invalid connection strings (delegated to LakebaseConnectionConfig.create) - ("", "Failed to parse connection string '': Connection string cannot be empty or None."), - ( - "mysql://user:password@host:5432/db", - "Failed to parse connection string 'mysql://user:password@host:5432/db': Invalid URL scheme 'mysql'. Expected 'postgresql'.", - ), - ( - "postgresql://host:5432/db", - "Failed to parse connection string 'postgresql://host:5432/db': Missing username in connection string", - ), - ( - "postgresql://user:password@:5432/db", - "Failed to parse connection string 'postgresql://user:password@:5432/db': Missing hostname in connection string", - ), - ( - "postgresql://user:password@host:5432/", - "Failed to parse connection string 'postgresql://user:password@host:5432/': Missing database name in connection string", - ), + ("", "Instance name must not be empty or None."), + (None, "Instance name must not be empty or None."), ], ) -def test_invalid_connection_string(connection_string, expected_error_pattern): - """Test LakebaseChecksStorageConfig with invalid connection strings.""" - with pytest.raises(InvalidParameterError, match=expected_error_pattern): - LakebaseChecksStorageConfig(location="testdb.testschema.testtable", connection_string=connection_string) +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( @@ -196,72 +140,60 @@ def test_invalid_connection_string(connection_string, expected_error_pattern): ("", "Invalid mode ''. Must be 'append' or 'overwrite'."), ], ) -def test_invalid_mode(mode, expected_error, valid_connection_string): +def test_invalid_mode(mode, expected_error): """Test LakebaseChecksStorageConfig with invalid mode.""" with pytest.raises(InvalidConfigError, match=expected_error): LakebaseChecksStorageConfig( - location="testdb.testschema.testtable", connection_string=valid_connection_string, mode=mode + instance_name="testinstance", location="testdb.testschema.testtable", user="testuser", mode=mode ) -def test_valid_modes(valid_connection_string): +def test_valid_modes(): """Test LakebaseChecksStorageConfig with valid modes.""" # Test 'append' mode config_append = LakebaseChecksStorageConfig( - location="testdb.testschema.testtable", connection_string=valid_connection_string, mode="append" + instance_name="testinstance", location="testdb.testschema.testtable", user="testuser", mode="append" ) assert config_append.mode == "append" # Test 'overwrite' mode config_overwrite = LakebaseChecksStorageConfig( - location="testdb.testschema.testtable", connection_string=valid_connection_string, mode="overwrite" + instance_name="testinstance", location="testdb.testschema.testtable", user="testuser", mode="overwrite" ) assert config_overwrite.mode == "overwrite" -def test_connection_config_property(valid_connection_string): - """Test that connection_config property returns correct LakebaseConnectionConfig.""" +def test_location_properties(): + """Test database_name, schema_name, and table_name properties.""" config = LakebaseChecksStorageConfig( - location="testdb.testschema.testtable", connection_string=valid_connection_string + instance_name="testinstance", location="testdb.testschema.testtable", user="testuser" ) - conn_config = config.connection_config - assert isinstance(conn_config, LakebaseConnectionConfig) - assert conn_config.instance_name == "testhost" - assert conn_config.database == "testdb" - assert conn_config.user == "testuser" - assert conn_config.port == "5432" - - -def test_location_properties(valid_connection_string): - """Test database_name, schema_name, and table_name properties.""" - config = LakebaseChecksStorageConfig(location="mydb.myschema.mytable", connection_string=valid_connection_string) - - assert config.database_name == "mydb" - assert config.schema_name == "myschema" - assert config.table_name == "mytable" + assert config.database_name == "testdb" + assert config.schema_name == "testschema" + assert config.table_name == "testtable" -def test_complex_location_with_special_characters(valid_connection_string): +def test_complex_location_with_special_characters(): """Test location parsing with special characters in names.""" - config = LakebaseChecksStorageConfig(location="my_db.my_schema.my_table", connection_string=valid_connection_string) + config = LakebaseChecksStorageConfig( + instance_name="testinstance", location="test_db.test_schema.test_table", user="testuser" + ) - assert config.database_name == "my_db" - assert config.schema_name == "my_schema" - assert config.table_name == "my_table" + 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="Failed to parse connection string"): - LakebaseChecksStorageConfig( - location="testdb.testschema.testtable", connection_string="invalid-connection-string" - ) + 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 connection string validation - with pytest.raises(InvalidParameterError, match="The 'location' field must not be empty or None"): - LakebaseChecksStorageConfig(location="", connection_string="invalid-connection-string") + # 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_save_checks.py b/tests/unit/test_save_checks.py index 5a8fe7c68..1e773a6b7 100644 --- a/tests/unit/test_save_checks.py +++ b/tests/unit/test_save_checks.py @@ -5,7 +5,6 @@ import pytest from databricks.labs.dqx.engine import DQEngineCore -from databricks.labs.dqx.config import LakebaseConnectionConfig from databricks.labs.dqx.errors import InvalidConfigError @@ -66,15 +65,3 @@ def _validate_file(file_path: str, file_format: str = "yaml") -> None: if file_format == "json": json.load(file) yaml.safe_load(file) - - -def test_installation_checks_storage_handler_postgresql_parsing(): - connection_string = ( - "postgresql://user@domain.com:password@instance-test.database.azuredatabricks.net:5432/dqx?sslmode=require" - ) - connection_config = LakebaseConnectionConfig.create(connection_string) - - assert connection_config.user == "user@domain.com" - assert connection_config.instance_name == "instance-test.database.azuredatabricks.net" - assert connection_config.port == "5432" - assert connection_config.database == "dqx" From 6d50bf15239dbc2ed2e7ad7fdeb66441f38e643a Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 8 Oct 2025 15:49:01 +0200 Subject: [PATCH 078/125] chore: update InstallationChecksStorageConfig --- src/databricks/labs/dqx/config.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 3ffa5eac1..0521d1c40 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -228,7 +228,7 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig): 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.") @@ -303,9 +303,9 @@ class InstallationChecksStorageConfig( overwrite_location: Whether to overwrite the location from run config if provided (default is False). """ - instance_name: str | None = None + instance_name: str = "installation" # retrieved from the installation config location: str = "installation" # retrieved from the installation config - user: str | None = None + user: str = "installation" # retrieved from the installation config port: str = LAKEBASE_DEFAULT_PORT # default port for Lakebase connection run_config_name: str = "default" # to retrieve run config product_name: str = "dqx" From b085f94bc8c96707f7c2a5aba2354446bb97f888 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 8 Oct 2025 15:50:39 +0200 Subject: [PATCH 079/125] chore: update integration tests --- tests/conftest.py | 66 +++++---- ...ave_and_load_checks_from_lakebase_table.py | 126 +++++++++++++----- 2 files changed, 136 insertions(+), 56 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index cec60abb0..413d46ae3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ import os from typing import Any import re +import logging from collections.abc import Callable, Generator from dataclasses import replace from functools import cached_property @@ -22,6 +23,8 @@ from databricks.sdk.service.workspace import ImportFormat from databricks.sdk.service.database import DatabaseInstance, DatabaseCatalog +logger = logging.getLogger(__name__) + @pytest.fixture(scope="session") def debug_env_name(): @@ -644,48 +647,65 @@ def delete(volume_file_path: str) -> None: yield from factory("file", create, delete) -@pytest.fixture -def make_lakebase_instance_and_catalog(ws, make_random): - import logging +@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. - logger = logging.getLogger(__name__) + 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 + + +@pytest.fixture +def make_lakebase_instance_and_catalog(ws, make_random): instances = {} def create() -> str: - database_instance_name = f"dqxtest-{make_random(10).lower()}" + 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" - instance = ws.database.create_database_instance_and_wait( - database_instance=DatabaseInstance(name=database_instance_name, capacity=capacity) + _ = 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=database_instance_name, + database_instance_name=instance_name, create_database_if_not_exists=True, ) ) + logger.info(f"Successfully created database catalog: {catalog_name}") - connection_string = ( - f"postgresql://{instance.creator}:password@{instance.read_only_dns}:5432/{database_name}?sslmode=require" - ) - - instances[connection_string] = {"database_instance_name": database_instance_name, "catalog_name": catalog_name} + instances[instance_name] = { + "instance_name": instance_name, + "catalog_name": catalog_name, + } - return connection_string + return instance_name - def delete(connection_string: str) -> None: - if connection_string not in instances: - logger.warning(f"No resources found for connection string: {connection_string}") + def delete(instance_name: str) -> None: + if instance_name not in instances: + logger.warning(f"No resources found for database instance name: {instance_name}") return - metadata = instances[connection_string] - database_instance_name = metadata["database_instance_name"] + metadata = instances[instance_name] + instance_name = metadata["instance_name"] catalog_name = metadata["catalog_name"] try: @@ -695,13 +715,13 @@ def delete(connection_string: str) -> None: logger.warning(f"Failed to delete database catalog {catalog_name}: {e}") try: - ws.database.delete_database_instance(name=database_instance_name) - logger.info(f"Successfully deleted database instance: {database_instance_name}") + 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 {database_instance_name}: {e}") + logger.error(f"Failed to delete database instance {instance_name}: {e}") raise finally: - instances.pop(connection_string, None) + instances.pop(instance_name, None) yield from factory("lakebase", create, delete) 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 index 1ced6f153..95813d591 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -1,9 +1,12 @@ import pytest -from databricks.sdk.errors import NotFound +from pyspark.sql.types import StructType, StructField, StringType, IntegerType from databricks.labs.dqx.config import InstallationChecksStorageConfig, LakebaseChecksStorageConfig from databricks.labs.dqx.engine import DQEngine +from databricks.labs.dqx.profiler.profiler import DQProfiler +from databricks.labs.dqx.profiler.generator import DQGenerator +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 @@ -11,36 +14,39 @@ LOCATION = "dqx.config.checks" -def test_load_checks_when_checks_lakebase_table_does_not_exist(ws, make_lakebase_instance_and_catalog, spark): - connection_string = make_lakebase_instance_and_catalog() +def test_load_checks_when_lakebase_table_does_not_exist(ws, spark, make_lakebase_instance_and_catalog, lakebase_user): + instance_name = make_lakebase_instance_and_catalog() + config = LakebaseChecksStorageConfig(location=LOCATION, user=lakebase_user, instance_name=instance_name) - with pytest.raises(NotFound, match="Resource not found"): + with pytest.raises(NotFound, match=f"Table '{config.location}' does not exist in the Lakebase instance"): dq_engine = DQEngine(ws, spark) - config = LakebaseChecksStorageConfig(location=LOCATION, connection_string=connection_string) + config = LakebaseChecksStorageConfig(location=LOCATION, user=lakebase_user, instance_name=instance_name) dq_engine.load_checks(config=config) -def test_save_and_load_checks_from_lakebase_table(ws, make_lakebase_instance_and_catalog, spark): - connection_string = make_lakebase_instance_and_catalog() +def test_save_and_load_checks_from_lakebase_table(ws, make_lakebase_instance_and_catalog, spark, lakebase_user): dq_engine = DQEngine(ws, spark) - config = LakebaseChecksStorageConfig(location=LOCATION, connection_string=connection_string) + instance_name = make_lakebase_instance_and_catalog() + config = LakebaseChecksStorageConfig(location=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_and_catalog): - connection_string = make_lakebase_instance_and_catalog() +def test_save_and_load_checks_from_lakebase_table_with_run_config( + ws, spark, make_lakebase_instance_and_catalog, lakebase_user +): + instance_name = make_lakebase_instance_and_catalog() dq_engine = DQEngine(ws, spark) # test first run config run_config_name = "workflow_001" config_save = LakebaseChecksStorageConfig( - location=LOCATION, connection_string=connection_string, run_config_name=run_config_name + location=LOCATION, user=lakebase_user, instance_name=instance_name, run_config_name=run_config_name ) dq_engine.save_checks(TEST_CHECKS[:1], config=config_save) config_load = LakebaseChecksStorageConfig( - location=LOCATION, connection_string=connection_string, run_config_name=run_config_name + location=LOCATION, user=lakebase_user, instance_name=instance_name, run_config_name=run_config_name ) checks = dq_engine.load_checks(config=config_load) compare_checks(checks, TEST_CHECKS[:1]) @@ -48,40 +54,49 @@ def test_save_and_load_checks_from_lakebase_table_with_run_config(ws, spark, mak # test second run config run_config_name2 = "workflow_002" config_save2 = LakebaseChecksStorageConfig( - location=LOCATION, - connection_string=connection_string, - run_config_name=run_config_name2, + location=LOCATION, user=lakebase_user, instance_name=instance_name, run_config_name=run_config_name2 ) dq_engine.save_checks(TEST_CHECKS[1:], config=config_save2) config_load2 = LakebaseChecksStorageConfig( - location=LOCATION, connection_string=connection_string, run_config_name=run_config_name + location=LOCATION, user=lakebase_user, instance_name=instance_name, run_config_name=run_config_name ) checks = dq_engine.load_checks(config=config_load2) compare_checks(checks, TEST_CHECKS[:1]) # test default config dq_engine.save_checks( - TEST_CHECKS[1:], config=LakebaseChecksStorageConfig(location=LOCATION, connection_string=connection_string) + TEST_CHECKS[1:], + config=LakebaseChecksStorageConfig(location=LOCATION, user=lakebase_user, instance_name=instance_name), ) checks = dq_engine.load_checks( - config=LakebaseChecksStorageConfig(location=LOCATION, connection_string=connection_string) + config=LakebaseChecksStorageConfig(location=LOCATION, user=lakebase_user, instance_name=instance_name) ) compare_checks(checks, TEST_CHECKS[1:]) -def test_save_and_load_checks_to_lakebase_table_output_modes(ws, spark, make_lakebase_instance_and_catalog): - connection_string = make_lakebase_instance_and_catalog() +def test_save_and_load_checks_to_lakebase_table_output_modes( + ws, + spark, + make_lakebase_instance_and_catalog, + lakebase_user, +): + instance_name = make_lakebase_instance_and_catalog() dq_engine = DQEngine(ws, spark) + run_config_name = "workflow_003" dq_engine.save_checks( TEST_CHECKS[:1], config=LakebaseChecksStorageConfig( - location=LOCATION, connection_string=connection_string, run_config_name=run_config_name, mode="append" + location=LOCATION, + instance_name=instance_name, + user=lakebase_user, + run_config_name=run_config_name, + mode="append", ), ) checks = dq_engine.load_checks( config=LakebaseChecksStorageConfig( - location=LOCATION, connection_string=connection_string, run_config_name=run_config_name + location=LOCATION, instance_name=instance_name, user=lakebase_user, run_config_name=run_config_name ) ) compare_checks(checks, TEST_CHECKS[:1]) @@ -90,42 +105,87 @@ def test_save_and_load_checks_to_lakebase_table_output_modes(ws, spark, make_lak dq_engine.save_checks( TEST_CHECKS[1:], config=LakebaseChecksStorageConfig( - location=LOCATION, connection_string=connection_string, run_config_name=run_config_name, mode="overwrite" + location=LOCATION, + instance_name=instance_name, + user=lakebase_user, + run_config_name=run_config_name, + mode="overwrite", ), ) checks = dq_engine.load_checks( config=LakebaseChecksStorageConfig( - location=LOCATION, connection_string=connection_string, run_config_name=run_config_name + location=LOCATION, instance_name=instance_name, user=lakebase_user, run_config_name=run_config_name ) ) compare_checks(checks, TEST_CHECKS[1:]) def test_save_load_checks_from_lakebase_table_in_user_installation( - ws, spark, installation_ctx, make_lakebase_instance_and_catalog + ws, spark, installation_ctx, make_lakebase_instance_and_catalog, lakebase_user ): - connection_string = make_lakebase_instance_and_catalog() + instance_name = make_lakebase_instance_and_catalog() + config = installation_ctx.config run_config = config.get_run_config() - run_config.checks_LOCATION = LOCATION - run_config.connection_string = connection_string + run_config.checks_location = LOCATION + installation_ctx.installation.save(installation_ctx.config) - product_name = installation_ctx.product_info.product_name() - dq_engine = DQEngine(ws, spark) config = InstallationChecksStorageConfig( - location=LOCATION, - connection_string=connection_string, + instance_name=instance_name, + user=lakebase_user, run_config_name=run_config.name, - product_name=product_name, + product_name="dqx", 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_save_and_load_checks_from_profiler(ws, spark, make_lakebase_instance_and_catalog, lakebase_user): + df = 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), + ] + ), + ) + profiler = DQProfiler(ws) + generator = DQGenerator(ws) + dq_engine = DQEngine(ws, spark) + _, profiles = profiler.profile(df) + checks = generator.generate_dq_rules(profiles) + instance_name = make_lakebase_instance_and_catalog() + dq_engine.save_checks( + checks, config=LakebaseChecksStorageConfig(location=LOCATION, instance_name=instance_name, user=lakebase_user) + ) + loaded_checks = dq_engine.load_checks( + config=LakebaseChecksStorageConfig(location=LOCATION, instance_name=instance_name, user=lakebase_user) + ) + compare_checks(loaded_checks, checks) + + def delete_all_leftover_instances(ws): import re From 6117b22a81c3db41fcef8d1d441a1ee53a641b1f Mon Sep 17 00:00:00 2001 From: Manuel Tilgner <58051108+tlgnr@users.noreply.github.com> Date: Wed, 8 Oct 2025 16:04:06 +0200 Subject: [PATCH 080/125] chore: update import order Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/unit/test_load_checks.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py index 49019c42b..777528294 100644 --- a/tests/unit/test_load_checks.py +++ b/tests/unit/test_load_checks.py @@ -1,7 +1,6 @@ from unittest.mock import create_autospec import pytest - from databricks.sdk import WorkspaceClient from databricks.sdk.errors import NotFound from databricks.sdk.service.files import DownloadResponse From 5532e30c38465499cca81dcf419161478244acb0 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 8 Oct 2025 16:05:23 +0200 Subject: [PATCH 081/125] chore: update import order --- tests/unit/test_save_checks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_save_checks.py b/tests/unit/test_save_checks.py index 1e773a6b7..4b0fd0593 100644 --- a/tests/unit/test_save_checks.py +++ b/tests/unit/test_save_checks.py @@ -1,7 +1,7 @@ import os import json - import yaml + import pytest from databricks.labs.dqx.engine import DQEngineCore From 7ba287e9e2805d0dc3c78ba28dddd8411b2e7638 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 8 Oct 2025 16:10:01 +0200 Subject: [PATCH 082/125] chore: update lakebase run config name test --- ...est_save_and_load_checks_from_lakebase_table.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) 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 index 95813d591..99d2163b0 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -52,16 +52,16 @@ def test_save_and_load_checks_from_lakebase_table_with_run_config( compare_checks(checks, TEST_CHECKS[:1]) # test second run config - run_config_name2 = "workflow_002" - config_save2 = LakebaseChecksStorageConfig( - location=LOCATION, user=lakebase_user, instance_name=instance_name, run_config_name=run_config_name2 + run_config_name = "workflow_002" + config_save = LakebaseChecksStorageConfig( + location=LOCATION, user=lakebase_user, instance_name=instance_name, run_config_name=run_config_name ) - dq_engine.save_checks(TEST_CHECKS[1:], config=config_save2) - config_load2 = LakebaseChecksStorageConfig( + dq_engine.save_checks(TEST_CHECKS[1:], config=config_save) + config_load = LakebaseChecksStorageConfig( location=LOCATION, user=lakebase_user, instance_name=instance_name, run_config_name=run_config_name ) - checks = dq_engine.load_checks(config=config_load2) - compare_checks(checks, TEST_CHECKS[:1]) + checks = dq_engine.load_checks(config=config_load) + compare_checks(checks, TEST_CHECKS[1:]) # test default config dq_engine.save_checks( From 824409d06376a4ad73c908e7481b2f3fde1846e6 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 8 Oct 2025 16:13:17 +0200 Subject: [PATCH 083/125] refactor: remove sqlalchemy null function --- src/databricks/labs/dqx/checks_storage.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 3383c55c8..254fca680 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -17,7 +17,6 @@ insert, select, delete, - null, ) from sqlalchemy.schema import CreateSchema from sqlalchemy.dialects.postgresql import JSONB @@ -203,7 +202,7 @@ def get_table_definition(schema_name: str, table_name: str) -> Table: Column("check", JSONB), Column("filter", Text), Column("run_config_name", String(255), server_default="default"), - Column("user_metadata", JSONB, server_default=null()), + Column("user_metadata", JSONB), ) @staticmethod @@ -226,7 +225,7 @@ def _normalize_checks(checks: list[dict], config: LakebaseChecksStorageConfig) - "check": check.get("check"), "filter": check.get("filter"), "run_config_name": check.get("run_config_name", config.run_config_name), - "user_metadata": check.get("user_metadata", null()), + "user_metadata": check.get("user_metadata"), } normalized_checks.append(normalized_check) return normalized_checks From 2c06a2ca7cc8aa1591a4d52ef6e64e11b0d9fe0f Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 8 Oct 2025 16:30:53 +0200 Subject: [PATCH 084/125] fix: remove old instances --- .../test_save_and_load_checks_from_lakebase_table.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index 99d2163b0..023abdb08 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -186,7 +186,7 @@ def test_save_and_load_checks_from_profiler(ws, spark, make_lakebase_instance_an compare_checks(loaded_checks, checks) -def delete_all_leftover_instances(ws): +def test_delete_all_leftover_instances(ws): import re pattern = re.compile(r"^dqxtest-[A-Za-z0-9]{10}$") From b17d53d8c9e1916b9682b8dabc88d4f4571c861a Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 8 Oct 2025 17:10:33 +0200 Subject: [PATCH 085/125] chore: remove instance cleanup function --- ...st_save_and_load_checks_from_lakebase_table.py | 15 --------------- 1 file changed, 15 deletions(-) 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 index 023abdb08..2e487a643 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -184,18 +184,3 @@ def test_save_and_load_checks_from_profiler(ws, spark, make_lakebase_instance_an config=LakebaseChecksStorageConfig(location=LOCATION, instance_name=instance_name, user=lakebase_user) ) compare_checks(loaded_checks, checks) - - -def test_delete_all_leftover_instances(ws): - import re - - pattern = re.compile(r"^dqxtest-[A-Za-z0-9]{10}$") - - instances = [] - - for instance in ws.database.list_database_instances(): - if pattern.match(instance.name): - instances.append(instance.name) - - for instance in instances: - ws.database.delete_database_instance(name=instance) From 600117fc3e5c891e960f063b263f2d30c415f8ca Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 9 Oct 2025 10:03:01 +0200 Subject: [PATCH 086/125] fix: make product name dynamic in test --- .../test_save_and_load_checks_from_lakebase_table.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 index 2e487a643..e602a180e 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -128,14 +128,14 @@ def test_save_load_checks_from_lakebase_table_in_user_installation( config = installation_ctx.config run_config = config.get_run_config() run_config.checks_location = 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="dqx", + product_name=product_name, assume_user=True, ) From 780914b9e1b1ef7f1f41b983306de86b5d32dc3f Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 9 Oct 2025 16:29:26 +0200 Subject: [PATCH 087/125] refactor: streamline lakebase integration tests --- tests/conftest.py | 37 ++++- ...ave_and_load_checks_from_lakebase_table.py | 144 +++++++++--------- 2 files changed, 111 insertions(+), 70 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 413d46ae3..8b95b1adb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,6 +7,7 @@ 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 @@ -411,6 +412,40 @@ 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 location(): + return "dqx.config.checks" + + @pytest.fixture def make_local_check_file_as_yaml(checks_yaml_content, make_random): file_path = f"checks_{make_random(10).lower()}.yml" @@ -668,7 +703,7 @@ def lakebase_user(): @pytest.fixture -def make_lakebase_instance_and_catalog(ws, make_random): +def make_lakebase_instance(ws, make_random): instances = {} def create() -> str: 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 index e602a180e..53f36e6ec 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -1,6 +1,6 @@ -import pytest +import yaml -from pyspark.sql.types import StructType, StructField, StringType, IntegerType +import pytest from databricks.labs.dqx.config import InstallationChecksStorageConfig, LakebaseChecksStorageConfig from databricks.labs.dqx.engine import DQEngine @@ -11,42 +11,37 @@ from tests.conftest import compare_checks from tests.integration.test_save_and_load_checks_from_table import EXPECTED_CHECKS as TEST_CHECKS -LOCATION = "dqx.config.checks" - -def test_load_checks_when_lakebase_table_does_not_exist(ws, spark, make_lakebase_instance_and_catalog, lakebase_user): - instance_name = make_lakebase_instance_and_catalog() - config = LakebaseChecksStorageConfig(location=LOCATION, user=lakebase_user, instance_name=instance_name) +def test_load_checks_when_lakebase_table_does_not_exist(ws, spark, make_lakebase_instance, user, location): + instance_name = make_lakebase_instance() + dq_engine = DQEngine(ws, spark) + config = LakebaseChecksStorageConfig(location=location, user=user, instance_name=instance_name) with pytest.raises(NotFound, match=f"Table '{config.location}' does not exist in the Lakebase instance"): - dq_engine = DQEngine(ws, spark) - config = LakebaseChecksStorageConfig(location=LOCATION, user=lakebase_user, instance_name=instance_name) dq_engine.load_checks(config=config) -def test_save_and_load_checks_from_lakebase_table(ws, make_lakebase_instance_and_catalog, spark, lakebase_user): +def test_save_and_load_checks_from_lakebase_table(ws, spark, make_lakebase_instance, user, location): + instance_name = make_lakebase_instance() dq_engine = DQEngine(ws, spark) - instance_name = make_lakebase_instance_and_catalog() - config = LakebaseChecksStorageConfig(location=LOCATION, user=lakebase_user, instance_name=instance_name) + config = LakebaseChecksStorageConfig(location=location, user=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_and_catalog, lakebase_user -): - instance_name = make_lakebase_instance_and_catalog() +def test_save_and_load_checks_from_lakebase_table_with_run_config(ws, spark, make_lakebase_instance, user, location): + instance_name = make_lakebase_instance() dq_engine = DQEngine(ws, spark) # test first run config run_config_name = "workflow_001" config_save = LakebaseChecksStorageConfig( - location=LOCATION, user=lakebase_user, instance_name=instance_name, run_config_name=run_config_name + location=location, user=user, instance_name=instance_name, run_config_name=run_config_name ) dq_engine.save_checks(TEST_CHECKS[:1], config=config_save) config_load = LakebaseChecksStorageConfig( - location=LOCATION, user=lakebase_user, instance_name=instance_name, run_config_name=run_config_name + location=location, user=user, instance_name=instance_name, run_config_name=run_config_name ) checks = dq_engine.load_checks(config=config_load) compare_checks(checks, TEST_CHECKS[:1]) @@ -54,11 +49,11 @@ def test_save_and_load_checks_from_lakebase_table_with_run_config( # test second run config run_config_name = "workflow_002" config_save = LakebaseChecksStorageConfig( - location=LOCATION, user=lakebase_user, instance_name=instance_name, run_config_name=run_config_name + location=location, user=user, instance_name=instance_name, run_config_name=run_config_name ) dq_engine.save_checks(TEST_CHECKS[1:], config=config_save) config_load = LakebaseChecksStorageConfig( - location=LOCATION, user=lakebase_user, instance_name=instance_name, run_config_name=run_config_name + location=location, user=user, instance_name=instance_name, run_config_name=run_config_name ) checks = dq_engine.load_checks(config=config_load) compare_checks(checks, TEST_CHECKS[1:]) @@ -66,37 +61,38 @@ def test_save_and_load_checks_from_lakebase_table_with_run_config( # test default config dq_engine.save_checks( TEST_CHECKS[1:], - config=LakebaseChecksStorageConfig(location=LOCATION, user=lakebase_user, instance_name=instance_name), + config=LakebaseChecksStorageConfig(location=location, user=user, instance_name=instance_name), ) checks = dq_engine.load_checks( - config=LakebaseChecksStorageConfig(location=LOCATION, user=lakebase_user, instance_name=instance_name) + config=LakebaseChecksStorageConfig(location=location, user=user, instance_name=instance_name) ) compare_checks(checks, TEST_CHECKS[1:]) -def test_save_and_load_checks_to_lakebase_table_output_modes( +def test_save_and_load_checks_from_lakebase_table_with_output_modes( ws, spark, - make_lakebase_instance_and_catalog, - lakebase_user, + make_lakebase_instance, + user, + location, ): - instance_name = make_lakebase_instance_and_catalog() + instance_name = make_lakebase_instance() dq_engine = DQEngine(ws, spark) run_config_name = "workflow_003" dq_engine.save_checks( TEST_CHECKS[:1], config=LakebaseChecksStorageConfig( - location=LOCATION, + location=location, instance_name=instance_name, - user=lakebase_user, + user=user, run_config_name=run_config_name, mode="append", ), ) checks = dq_engine.load_checks( config=LakebaseChecksStorageConfig( - location=LOCATION, instance_name=instance_name, user=lakebase_user, run_config_name=run_config_name + location=location, instance_name=instance_name, user=user, run_config_name=run_config_name ) ) compare_checks(checks, TEST_CHECKS[:1]) @@ -105,35 +101,35 @@ def test_save_and_load_checks_to_lakebase_table_output_modes( dq_engine.save_checks( TEST_CHECKS[1:], config=LakebaseChecksStorageConfig( - location=LOCATION, + location=location, instance_name=instance_name, - user=lakebase_user, + user=user, run_config_name=run_config_name, mode="overwrite", ), ) checks = dq_engine.load_checks( config=LakebaseChecksStorageConfig( - location=LOCATION, instance_name=instance_name, user=lakebase_user, run_config_name=run_config_name + location=location, instance_name=instance_name, user=user, run_config_name=run_config_name ) ) compare_checks(checks, TEST_CHECKS[1:]) -def test_save_load_checks_from_lakebase_table_in_user_installation( - ws, spark, installation_ctx, make_lakebase_instance_and_catalog, lakebase_user +def test_save_and_load_checks_from_lakebase_table_with_user_installation( + ws, spark, installation_ctx, make_lakebase_instance, user, location ): - instance_name = make_lakebase_instance_and_catalog() - config = installation_ctx.config run_config = config.get_run_config() - run_config.checks_location = LOCATION + run_config.checks_location = location installation_ctx.installation.save(installation_ctx.config) product_name = installation_ctx.product_info.product_name() + instance_name = make_lakebase_instance() + config = InstallationChecksStorageConfig( instance_name=instance_name, - user=lakebase_user, + user=user, run_config_name=run_config.name, product_name=product_name, assume_user=True, @@ -145,42 +141,52 @@ def test_save_load_checks_from_lakebase_table_in_user_installation( compare_checks(checks, TEST_CHECKS) -def test_save_and_load_checks_from_profiler(ws, spark, make_lakebase_instance_and_catalog, lakebase_user): - df = 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), - ] - ), - ) +def test_save_and_load_checks_from_lakebase_table_with_profiler(ws, spark, make_lakebase_instance, df, user, location): profiler = DQProfiler(ws) - generator = DQGenerator(ws) - dq_engine = DQEngine(ws, spark) _, profiles = profiler.profile(df) + generator = DQGenerator(ws) checks = generator.generate_dq_rules(profiles) - instance_name = make_lakebase_instance_and_catalog() + + instance_name = make_lakebase_instance() + dq_engine = DQEngine(ws, spark) + dq_engine.save_checks( - checks, config=LakebaseChecksStorageConfig(location=LOCATION, instance_name=instance_name, user=lakebase_user) + checks, config=LakebaseChecksStorageConfig(location=location, instance_name=instance_name, user=user) ) loaded_checks = dq_engine.load_checks( - config=LakebaseChecksStorageConfig(location=LOCATION, instance_name=instance_name, user=lakebase_user) + config=LakebaseChecksStorageConfig(location=location, instance_name=instance_name, user=user) ) compare_checks(loaded_checks, checks) + + +def test_save_and_load_checks_from_lakebase_table_with_check_validaton( + ws, spark, make_lakebase_instance, user, location +): + dq_engine = DQEngine(ws, spark) + + checks = yaml.safe_load( + """ + - criticality: error + check: + function: is_not_null_and_not_empty + arguments: + column: name + + - criticality: error + check: + function: is_not_null + for_each_column: + - region + - state + """ + ) + + instance_name = make_lakebase_instance() + + dq_engine.save_checks( + checks, config=LakebaseChecksStorageConfig(location=location, instance_name=instance_name, user=user) + ) + loaded_checks = dq_engine.load_checks( + config=LakebaseChecksStorageConfig(location=location, instance_name=instance_name, user=user) + ) + assert not dq_engine.validate_checks(loaded_checks).has_errors From f679f3a6fbfa5059b1106322f0acddf21e03f0ac Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 9 Oct 2025 16:44:05 +0200 Subject: [PATCH 088/125] feat: extend DQEngine with Lakebase storage handler and config --- src/databricks/labs/dqx/checks_storage.py | 52 +++++++++++++++++++++++ src/databricks/labs/dqx/engine.py | 9 ++-- 2 files changed, 58 insertions(+), 3 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 254fca680..748c87a88 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -36,6 +36,7 @@ InstallationChecksStorageConfig, BaseChecksStorageConfig, VolumeFileChecksStorageConfig, + RunConfig, ) from databricks.labs.dqx.errors import InvalidCheckError, InvalidConfigError, CheckDownloadError from databricks.sdk import WorkspaceClient @@ -699,6 +700,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, + 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/engine.py b/src/databricks/labs/dqx/engine.py index 71dadfb0c..d5473e822 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); @@ -853,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. """ @@ -864,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) From 214e3ea35029e63f0d4b0e3b5ebd617215e52650 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 9 Oct 2025 16:47:14 +0200 Subject: [PATCH 089/125] chore: clean up pyproject.toml --- pyproject.toml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e114047a6..1333a80a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,7 +102,6 @@ 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 @@ -165,7 +164,7 @@ profile = "black" exclude = ['venv', '.venv', 'demos/*', 'tests/e2e/*'] [[tool.mypy.overrides]] -module = ["google", "google.*"] # Google and testing libraries are not type annotated +module = ["google", "google.*"] # Google libraries are not type annotated ignore_missing_imports = true [tool.pytest.ini_options] From d95d190fc4a3c1ffb7e8ab6d284d13143821a81c Mon Sep 17 00:00:00 2001 From: Manuel Tilgner <58051108+tlgnr@users.noreply.github.com> Date: Thu, 9 Oct 2025 16:51:42 +0200 Subject: [PATCH 090/125] docs: make docstring more concise Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/databricks/labs/dqx/checks_storage.py | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 748c87a88..d76b82a0c 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -172,26 +172,12 @@ def _get_engine(self, config: LakebaseChecksStorageConfig) -> Engine: @staticmethod def get_table_definition(schema_name: str, table_name: str) -> Table: - """ - Create a table definition for consistency between the load and save methods. + 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. - Note: - The table has the following columns: - `name` - Optional. Name to use for the check. - `criticality` - Either "error" (rows go only to "bad" dataset) or "warn" (rows go to both "good" and "bad"). - Defaults to "error". - `check` - Defines the DQX check to apply, e.g. {"function": "is_not_null", "arguments": {"column": "col1"}}. - `filter` - Optional. Spark SQL expression to filter rows to which the check is applied, e.g. "business_unit = 'Finance'". - The check function will run only on the rows matching the filter condition. The condition can reference - any column of the validated dataset, not only the one where you apply the check function. - `run_config_name` - Name of the run config name. Could be any string such as workflow name. Useful for selecting - applicable checks. Defaults to `"default"`. - `user_metadata` - Optional. Custom metadata to add to any row-level warnings or errors generated by the check. - Returns: SQLAlchemy table definition for the Lakebase instance. """ From 07bb2c66d11f2b8e398174f77573ed898299dcf5 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 9 Oct 2025 17:00:13 +0200 Subject: [PATCH 091/125] fix: harmonize fixture names --- tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 8b95b1adb..1f6212455 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -683,7 +683,7 @@ def delete(volume_file_path: str) -> None: @pytest.fixture(scope="session") -def lakebase_user(): +def user(): """ Get the Lakebase user (service principal client ID) from environment variable. From 4507513b60067b041675a997701c631abb10b5f2 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 9 Oct 2025 17:03:30 +0200 Subject: [PATCH 092/125] fix: add back quotation marks copilot review deleted --- src/databricks/labs/dqx/checks_storage.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index d76b82a0c..4121a515e 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -172,6 +172,7 @@ def _get_engine(self, config: LakebaseChecksStorageConfig) -> Engine: @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: From 2395102e33f3616fb58ddc39cc6dcc3a46a66bca Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Thu, 9 Oct 2025 17:20:07 +0200 Subject: [PATCH 093/125] refactor: improve error handling in load checks method --- src/databricks/labs/dqx/checks_storage.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 4121a515e..28cd41e93 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -20,7 +20,7 @@ ) from sqlalchemy.schema import CreateSchema from sqlalchemy.dialects.postgresql import JSONB -from sqlalchemy.exc import DatabaseError +from sqlalchemy.exc import DatabaseError, ProgrammingError import yaml @@ -295,11 +295,17 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: try: return self._load_checks_from_lakebase(config, engine) - except DatabaseError as e: + except ProgrammingError as e: + # ProgrammingError typically indicates SQL syntax errors or missing objects + # PostgreSQL error code 42P01 = undefined_table logger.error(f"Failed to load checks from Lakebase: {e}") - if "does not exist" in str(e).lower() or "relation" in str(e).lower(): + orig_error = getattr(e, 'orig', None) + if orig_error and hasattr(orig_error, 'pgcode') and orig_error.pgcode == '42P01': raise NotFound(f"Table '{config.location}' does not exist in the Lakebase instance") from e - raise DatabaseError(getattr(e, 'statement', None), getattr(e, 'params', None), e) from e + raise + except DatabaseError as e: + logger.error(f"Failed to load checks from Lakebase: {e}") + raise finally: if engine_created_internally: engine.dispose() From 5591bbf415330fb241dd69250bcf945dac1ea0a0 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 10 Oct 2025 09:33:04 +0200 Subject: [PATCH 094/125] style: edit whitespace --- src/databricks/labs/dqx/engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py index d5473e822..d13eee819 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -855,7 +855,7 @@ def _apply_checks_for_run_config(self, run_config: RunConfig) -> None: 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 + 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: From ea805bfcd369263e8762e19b0f58ee73469783c2 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 10 Oct 2025 09:33:25 +0200 Subject: [PATCH 095/125] fix: add abstract method to BaseChecksStorageHandlerFactory --- src/databricks/labs/dqx/checks_storage.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 28cd41e93..ea46f73a1 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -638,6 +638,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): From 4d26ee1cbab605a075462ae05d289fca3cabb26f Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 10 Oct 2025 13:40:28 +0200 Subject: [PATCH 096/125] fix: add back psycopg2 library --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 1333a80a8..6e4c6a614 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,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 From 7d6b033ceaf55f24a02bc62445f2da029620c3e8 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 10 Oct 2025 15:30:33 +0200 Subject: [PATCH 097/125] refactor: check if database exists and improve error handling --- src/databricks/labs/dqx/checks_storage.py | 69 +++++++++++++++++++++-- 1 file changed, 64 insertions(+), 5 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index ea46f73a1..c6066542b 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -20,7 +20,7 @@ ) from sqlalchemy.schema import CreateSchema from sqlalchemy.dialects.postgresql import JSONB -from sqlalchemy.exc import DatabaseError, ProgrammingError +from sqlalchemy.exc import DatabaseError, ProgrammingError, OperationalError import yaml @@ -55,6 +55,9 @@ logger = logging.getLogger(__name__) T = TypeVar("T", bound=BaseChecksStorageConfig) +# PostgreSQL error codes +POSTGRES_UNDEFINED_TABLE_ERROR = '42P01' + class ChecksStorageHandler(ABC, Generic[T]): """ @@ -218,6 +221,51 @@ def _normalize_checks(checks: list[dict], config: LakebaseChecksStorageConfig) - normalized_checks.append(normalized_check) return normalized_checks + def _get_postgres_error_code(exception: Exception) -> str | None: + """ + Safely extract PostgreSQL error code from a SQLAlchemy exception. + + This function defensively handles different SQLAlchemy versions and driver exceptions. + + Args: + exception: The exception to extract the error code from. + + Returns: + PostgreSQL error code (e.g., '42P01') or None if not available. + """ + # Check standard SQLAlchemy .orig attribute + try: + orig = getattr(exception, 'orig', None) + if orig and hasattr(orig, 'pgcode'): + return orig.pgcode + except (AttributeError, TypeError): + pass + + # Check if the exception itself has pgcode + try: + if hasattr(exception, 'pgcode'): + return exception.pgcode + except (AttributeError, TypeError): + pass + + # Check __cause__ chain (exception chaining) + try: + cause = exception.__cause__ + if cause and hasattr(cause, 'pgcode'): + return cause.pgcode + except (AttributeError, TypeError): + pass + + # Check __context__ chain (implicit exception chaining) + try: + context = exception.__context__ + if context and hasattr(context, 'pgcode'): + return context.pgcode + except (AttributeError, TypeError): + pass + + return None + def _save_checks_to_lakebase(self, checks: list[dict], config: LakebaseChecksStorageConfig, engine: Engine) -> None: """ Save dq rules (checks) to a Lakebase table. @@ -229,7 +277,20 @@ def _save_checks_to_lakebase(self, checks: list[dict], config: LakebaseChecksSto 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)) @@ -296,11 +357,9 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: try: return self._load_checks_from_lakebase(config, engine) except ProgrammingError as e: - # ProgrammingError typically indicates SQL syntax errors or missing objects - # PostgreSQL error code 42P01 = undefined_table logger.error(f"Failed to load checks from Lakebase: {e}") - orig_error = getattr(e, 'orig', None) - if orig_error and hasattr(orig_error, 'pgcode') and orig_error.pgcode == '42P01': + pgcode = self._get_postgres_error_code(e) + if pgcode == POSTGRES_UNDEFINED_TABLE_ERROR: raise NotFound(f"Table '{config.location}' does not exist in the Lakebase instance") from e raise except DatabaseError as e: From abd7aed8c434f3019c18ec0d2198fc9654d30108 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 10 Oct 2025 15:30:42 +0200 Subject: [PATCH 098/125] fix: correct spelling mistake --- .../test_save_and_load_checks_from_lakebase_table.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index 53f36e6ec..7726e1ab1 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -159,7 +159,7 @@ def test_save_and_load_checks_from_lakebase_table_with_profiler(ws, spark, make_ compare_checks(loaded_checks, checks) -def test_save_and_load_checks_from_lakebase_table_with_check_validaton( +def test_save_and_load_checks_from_lakebase_table_with_check_validation( ws, spark, make_lakebase_instance, user, location ): dq_engine = DQEngine(ws, spark) From 58e85c30deb0cd4862613cd5f83d1d98b44bf974 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 10 Oct 2025 16:08:56 +0200 Subject: [PATCH 099/125] refactor: further refine error handling --- src/databricks/labs/dqx/checks_storage.py | 107 ++++++++++------------ 1 file changed, 49 insertions(+), 58 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index c6066542b..1c66300c5 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -20,7 +20,7 @@ ) from sqlalchemy.schema import CreateSchema from sqlalchemy.dialects.postgresql import JSONB -from sqlalchemy.exc import DatabaseError, ProgrammingError, OperationalError +from sqlalchemy.exc import DatabaseError, ProgrammingError, OperationalError, IntegrityError import yaml @@ -221,51 +221,6 @@ def _normalize_checks(checks: list[dict], config: LakebaseChecksStorageConfig) - normalized_checks.append(normalized_check) return normalized_checks - def _get_postgres_error_code(exception: Exception) -> str | None: - """ - Safely extract PostgreSQL error code from a SQLAlchemy exception. - - This function defensively handles different SQLAlchemy versions and driver exceptions. - - Args: - exception: The exception to extract the error code from. - - Returns: - PostgreSQL error code (e.g., '42P01') or None if not available. - """ - # Check standard SQLAlchemy .orig attribute - try: - orig = getattr(exception, 'orig', None) - if orig and hasattr(orig, 'pgcode'): - return orig.pgcode - except (AttributeError, TypeError): - pass - - # Check if the exception itself has pgcode - try: - if hasattr(exception, 'pgcode'): - return exception.pgcode - except (AttributeError, TypeError): - pass - - # Check __cause__ chain (exception chaining) - try: - cause = exception.__cause__ - if cause and hasattr(cause, 'pgcode'): - return cause.pgcode - except (AttributeError, TypeError): - pass - - # Check __context__ chain (implicit exception chaining) - try: - context = exception.__context__ - if context and hasattr(context, 'pgcode'): - return context.pgcode - except (AttributeError, TypeError): - pass - - return None - def _save_checks_to_lakebase(self, checks: list[dict], config: LakebaseChecksStorageConfig, engine: Engine) -> None: """ Save dq rules (checks) to a Lakebase table. @@ -333,6 +288,27 @@ def _load_checks_from_lakebase(self, config: LakebaseChecksStorageConfig, engine logger.info(f"Successfully loaded {len(checks)} checks from {config.schema_name}.{config.table_name}") return [dict(check) for check in checks] + def _check_for_undefined_table_error(self, e: ProgrammingError, config: LakebaseChecksStorageConfig): + """ + Check if the error is an undefined table error. + + Args: + e: Programming error. + config: Configuration for saving and loading checks to Lakebase. + + Returns: + + Raises: + NotFound: If the table does not exist in the Lakebase instance. + """ + try: + orig_error = getattr(e, 'orig', None) + if orig_error and hasattr(orig_error, 'pgcode') and orig_error.pgcode == POSTGRES_UNDEFINED_TABLE_ERROR: + raise NotFound(f"Table '{config.location}' does not exist in the Lakebase instance") from e + except (AttributeError, TypeError): + pass + raise + @telemetry_logger("load_checks", "lakebase") def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: """ @@ -345,8 +321,10 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: List of dq rules or error if loading checks fails. Raises: - DatabaseError: If loading checks fails. NotFound: If the table does not exist in the Lakebase instance. + ProgrammingError: If SQL syntax errors or missing objects. + OperationalError: If other operational errors occur. + DatabaseError: If other database operations fail. """ engine = self.engine engine_created_internally = False @@ -356,15 +334,15 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: try: return self._load_checks_from_lakebase(config, engine) + except ProgrammingError as e: - logger.error(f"Failed to load checks from Lakebase: {e}") - pgcode = self._get_postgres_error_code(e) - if pgcode == POSTGRES_UNDEFINED_TABLE_ERROR: - raise NotFound(f"Table '{config.location}' does not exist in the Lakebase instance") from e - raise - except DatabaseError as e: - logger.error(f"Failed to load checks from Lakebase: {e}") + logger.error(f"Programming error while loading checks from Lakebase: {e}") + self._check_for_undefined_table_error(e, config) + + except (OperationalError, DatabaseError) as e: + logger.error(f"Database error while loading checks from Lakebase: {e}") raise + finally: if engine_created_internally: engine.dispose() @@ -383,7 +361,10 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: Raises: InvalidCheckError: If any check is invalid or unsupported. - DatabaseError: If saving checks fails. + IntegrityError: If constraint violations occur (e.g., duplicate keys). + ProgrammingError: If SQL syntax errors or missing objects. + OperationalError: If operational errors occur. + DatabaseError: If other database operations fail. """ if not checks: raise InvalidCheckError("Checks cannot be empty or None.") @@ -397,9 +378,19 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: try: self._save_checks_to_lakebase(checks, config, engine) logger.info(f"Successfully saved {len(checks)} checks to Lakebase.") - except DatabaseError as e: - logger.error(f"Failed to save checks to Lakebase: {e}") - raise DatabaseError(getattr(e, 'statement', None), getattr(e, 'params', None), e) from e + + 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 (OperationalError, DatabaseError) as e: + logger.error(f"Database error while saving checks to Lakebase: {e}") + raise + finally: if engine_created_internally: engine.dispose() From 067a4e7a759f1490741b00c70b1f34841d543000 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 10 Oct 2025 16:25:07 +0200 Subject: [PATCH 100/125] refactor: simplify error handling --- src/databricks/labs/dqx/checks_storage.py | 33 +++++++++++------------ 1 file changed, 16 insertions(+), 17 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 1c66300c5..c9ba9851f 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -5,7 +5,7 @@ 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, @@ -288,26 +288,27 @@ def _load_checks_from_lakebase(self, config: LakebaseChecksStorageConfig, engine logger.info(f"Successfully loaded {len(checks)} checks from {config.schema_name}.{config.table_name}") return [dict(check) for check in checks] - def _check_for_undefined_table_error(self, e: ProgrammingError, config: LakebaseChecksStorageConfig): + def _check_for_undefined_table_error(self, e: ProgrammingError, config: LakebaseChecksStorageConfig) -> NoReturn: """ - Check if the error is an undefined table error. + 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. + e: Programming error to check. config: Configuration for saving and loading checks to Lakebase. - Returns: - Raises: - NotFound: If the table does not exist in the Lakebase instance. + 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. """ try: orig_error = getattr(e, 'orig', None) if orig_error and hasattr(orig_error, 'pgcode') and orig_error.pgcode == POSTGRES_UNDEFINED_TABLE_ERROR: raise NotFound(f"Table '{config.location}' does not exist in the Lakebase instance") from e - except (AttributeError, TypeError): - pass - raise + except (AttributeError, TypeError) as exc: + raise e from exc + raise e @telemetry_logger("load_checks", "lakebase") def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: @@ -322,9 +323,8 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: Raises: NotFound: If the table does not exist in the Lakebase instance. - ProgrammingError: If SQL syntax errors or missing objects. - OperationalError: If other operational errors occur. - DatabaseError: If other database operations fail. + 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 @@ -339,7 +339,7 @@ def load(self, config: LakebaseChecksStorageConfig) -> list[dict]: logger.error(f"Programming error while loading checks from Lakebase: {e}") self._check_for_undefined_table_error(e, config) - except (OperationalError, DatabaseError) as e: + except DatabaseError as e: logger.error(f"Database error while loading checks from Lakebase: {e}") raise @@ -363,8 +363,7 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: 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. - OperationalError: If operational errors occur. - DatabaseError: If other database operations fail. + DatabaseError: If other database operations fail (includes OperationalError, DataError, etc.). """ if not checks: raise InvalidCheckError("Checks cannot be empty or None.") @@ -387,7 +386,7 @@ def save(self, checks: list[dict], config: LakebaseChecksStorageConfig) -> None: logger.error(f"Programming error while saving checks to Lakebase: {e}") raise - except (OperationalError, DatabaseError) as e: + except DatabaseError as e: logger.error(f"Database error while saving checks to Lakebase: {e}") raise From 0a9e64e9a0de3d5d704a176152f0dad9804c1752 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Fri, 10 Oct 2025 16:40:52 +0200 Subject: [PATCH 101/125] fix: correct indent --- .../test_save_and_load_checks_from_lakebase_table.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 index 7726e1ab1..a26182c18 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -167,13 +167,13 @@ def test_save_and_load_checks_from_lakebase_table_with_check_validation( checks = yaml.safe_load( """ - criticality: error - check: + check: function: is_not_null_and_not_empty arguments: - column: name + column: name - criticality: error - check: + check: function: is_not_null for_each_column: - region From e39479c1fdc6a6a91269ce49b816d95c4f0b52f3 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 13 Oct 2025 10:07:20 +0200 Subject: [PATCH 102/125] fix: properly handle None values in InstallationChecksStorageHandler and Config --- src/databricks/labs/dqx/checks_storage.py | 11 ++++++++++- src/databricks/labs/dqx/config.py | 18 ++++++++++-------- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index c9ba9851f..8b6998f9d 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -151,6 +151,13 @@ def _get_connection_url(self, config: LakebaseChecksStorageConfig) -> str: 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] @@ -573,9 +580,11 @@ def _get_storage_handler_and_config( config.location = checks_location matches_table_pattern = is_table_location(config.location) - is_lakebase = hasattr(config, 'instance_name') and config.instance_name + is_lakebase = hasattr(config, 'instance_name') and config.instance_name is not None if matches_table_pattern and is_lakebase: + if not config.user: + raise InvalidConfigError("user must be provided when using Lakebase storage") return self.lakebase_handler, config if matches_table_pattern: diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 0521d1c40..6243906f8 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -218,9 +218,9 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig): only replaces checks for the specific run config and not all checks in the table (default is 'overwrite'). """ - instance_name: str - user: str - location: str + instance_name: str | None = None + user: str | None = None + location: str | None = None port: str = LAKEBASE_DEFAULT_PORT run_config_name: str = "default" mode: str = "overwrite" @@ -245,6 +245,8 @@ def __post_init__(self): 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 @@ -287,11 +289,11 @@ class InstallationChecksStorageConfig( Configuration class for storing checks in an installation. Args: - instance_name: The name of the Lakebase instance (default is 'installation'). + instance_name: The name of the Lakebase instance (optional, required only for Lakebase storage). location: The installation path where the checks are stored (e.g., table name, file path). Not used when using installation method, as it is retrieved from the installation config, unless overwrite_location is enabled. - user: The user for the Lakebase connection. + user: The user for the Lakebase connection (optional, required only for Lakebase storage). port: The port for the Lakebase connection (default is '5432'). run_config_name: The name of the run configuration to use for checks (default is 'default'). product_name: The product name for retrieving checks from the installation (default is 'dqx'). @@ -303,10 +305,10 @@ class InstallationChecksStorageConfig( overwrite_location: Whether to overwrite the location from run config if provided (default is False). """ - instance_name: str = "installation" # retrieved from the installation config - location: str = "installation" # retrieved from the installation config - user: str = "installation" # retrieved from the installation config + instance_name: str | None = None + user: str | None = None port: str = LAKEBASE_DEFAULT_PORT # default port for Lakebase connection + location: str = "installation" # retrieved from the installation config run_config_name: str = "default" # to retrieve run config product_name: str = "dqx" assume_user: bool = True From bdcdd2a769a7c39be9ec25932ca1ef10ef07bd2c Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 13 Oct 2025 10:11:48 +0200 Subject: [PATCH 103/125] style: extract yaml string --- .../test_save_and_load_checks_from_lakebase_table.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) 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 index a26182c18..82934ca85 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -164,8 +164,7 @@ def test_save_and_load_checks_from_lakebase_table_with_check_validation( ): dq_engine = DQEngine(ws, spark) - checks = yaml.safe_load( - """ + yaml_str = """ - criticality: error check: function: is_not_null_and_not_empty @@ -179,7 +178,7 @@ def test_save_and_load_checks_from_lakebase_table_with_check_validation( - region - state """ - ) + checks = yaml.safe_load(yaml_str) instance_name = make_lakebase_instance() From b6687ef7d213e84b15a7ab1f1c6a7d537978e862 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 13 Oct 2025 10:30:22 +0200 Subject: [PATCH 104/125] refactor: simplify error handling --- src/databricks/labs/dqx/checks_storage.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 8b6998f9d..f81632077 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -309,12 +309,9 @@ def _check_for_undefined_table_error(self, e: ProgrammingError, config: Lakebase 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. """ - try: - orig_error = getattr(e, 'orig', None) - if orig_error and hasattr(orig_error, 'pgcode') and orig_error.pgcode == POSTGRES_UNDEFINED_TABLE_ERROR: - raise NotFound(f"Table '{config.location}' does not exist in the Lakebase instance") from e - except (AttributeError, TypeError) as exc: - raise e from exc + pgcode = getattr(getattr(e, 'orig', None), 'pgcode', None) + 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") From dfd263607441c34d8d1c4880a539d2917f7e2326 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 13 Oct 2025 11:45:03 +0200 Subject: [PATCH 105/125] docs: extend docs with references to LakebaseChecksStorageConfig --- docs/dqx/docs/guide/data_profiling.mdx | 2 +- docs/dqx/docs/guide/quality_checks_apply.mdx | 4 ++-- .../dqx/docs/guide/quality_checks_storage.mdx | 19 +++++++++++++++++-- 3 files changed, 20 insertions(+), 5 deletions(-) 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 198a39f56..3a4c8640e 100644 --- a/docs/dqx/docs/guide/quality_checks_apply.mdx +++ b/docs/dqx/docs/guide/quality_checks_apply.mdx @@ -573,7 +573,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". @@ -688,7 +688,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 58fc63988..35f8ed338 100644 --- a/docs/dqx/docs/guide/quality_checks_storage.mdx +++ b/docs/dqx/docs/guide/quality_checks_storage.mdx @@ -22,6 +22,13 @@ 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. @@ -49,7 +56,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 +89,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 dq_engine.save_checks(checks, config=InstallationChecksStorageConfig(assume_user=True, run_config_name="default")) @@ -194,7 +205,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 @@ -216,6 +228,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 checks: list[dict] = dq_engine.load_checks(config=InstallationChecksStorageConfig(run_config_name="default")) From ce8bbd891d801ed83837d9cb595c54fea3e37914 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 13 Oct 2025 12:20:58 +0200 Subject: [PATCH 106/125] refactor: transfer attribute values from run config if not already set --- src/databricks/labs/dqx/checks_storage.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index f81632077..7d7850b17 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -37,6 +37,7 @@ BaseChecksStorageConfig, VolumeFileChecksStorageConfig, RunConfig, + LAKEBASE_DEFAULT_PORT, ) from databricks.labs.dqx.errors import InvalidCheckError, InvalidConfigError, CheckDownloadError from databricks.sdk import WorkspaceClient @@ -569,6 +570,14 @@ def _get_storage_handler_and_config( install_folder=config.install_folder, ) checks_location = run_config.checks_location + + # Transfer Lakebase fields from run_config to 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 + if run_config.lakebase_port and config.port == LAKEBASE_DEFAULT_PORT: + config.port = run_config.lakebase_port installation = self._run_config_loader.get_installation( config.assume_user, config.product_name, config.install_folder From 9815f274fa4ccf190f64fdcc83d0f721f84298aa Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 13 Oct 2025 15:14:47 +0200 Subject: [PATCH 107/125] refactor: add logging and improve run config handling --- src/databricks/labs/dqx/checks_storage.py | 60 +++++++++++++++++------ 1 file changed, 44 insertions(+), 16 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 7d7850b17..15bc3533b 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -262,7 +262,9 @@ def _save_checks_to_lakebase(self, checks: list[dict], config: LakebaseChecksSto 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.schema_name}.{config.table_name}'.") + 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) @@ -272,6 +274,10 @@ def _save_checks_to_lakebase(self, checks: list[dict], config: LakebaseChecksSto 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]: """ @@ -288,12 +294,24 @@ def _load_checks_from_lakebase(self, config: LakebaseChecksStorageConfig, engine 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.schema_name}.{config.table_name}") + 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: @@ -560,24 +578,27 @@ def save(self, checks: list[dict], config: InstallationChecksStorageConfig) -> N def _get_storage_handler_and_config( self, config: InstallationChecksStorageConfig ) -> tuple[ChecksStorageHandler, InstallationChecksStorageConfig]: + # Always load run_config to get Lakebase parameters, even when overwriting location + run_config = self._run_config_loader.load_run_config( + run_config_name=config.run_config_name, + assume_user=config.assume_user, + product_name=config.product_name, + install_folder=config.install_folder, + ) + + # Transfer Lakebase fields from run_config to 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 + if run_config.lakebase_port and config.port == LAKEBASE_DEFAULT_PORT: + config.port = run_config.lakebase_port + + # Overwrite location if overwrite_location is set if config.overwrite_location: checks_location = config.location else: - run_config = self._run_config_loader.load_run_config( - run_config_name=config.run_config_name, - assume_user=config.assume_user, - product_name=config.product_name, - install_folder=config.install_folder, - ) checks_location = run_config.checks_location - - # Transfer Lakebase fields from run_config to 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 - if run_config.lakebase_port and config.port == LAKEBASE_DEFAULT_PORT: - config.port = run_config.lakebase_port installation = self._run_config_loader.get_installation( config.assume_user, config.product_name, config.install_folder @@ -588,9 +609,16 @@ def _get_storage_handler_and_config( matches_table_pattern = is_table_location(config.location) is_lakebase = hasattr(config, 'instance_name') and config.instance_name is not None + logger.info( + f"InstallationChecksStorageHandler: location='{config.location}', " + f"matches_table_pattern={matches_table_pattern}, is_lakebase={is_lakebase}, " + f"instance_name={config.instance_name if is_lakebase else 'N/A'}" + ) + if matches_table_pattern and is_lakebase: if not config.user: raise InvalidConfigError("user must be provided when using Lakebase storage") + logger.info(f"Using LakebaseChecksStorageHandler for location '{config.location}'") return self.lakebase_handler, config if matches_table_pattern: From ea7e4cd200985a0410763548860f9e0de32dc195 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Mon, 13 Oct 2025 15:16:06 +0200 Subject: [PATCH 108/125] refactor: reduce parallelism in integration tests --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 6e4c6a614..2638b7462 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -113,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 .", From a8c80033b9aa0cb56fae72f2e7c4be097f1d8c94 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 14 Oct 2025 13:13:19 +0200 Subject: [PATCH 109/125] fix: clean up old test instances --- ...st_save_and_load_checks_from_lakebase_table.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) 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 index 82934ca85..a6870a7b3 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -189,3 +189,18 @@ def test_save_and_load_checks_from_lakebase_table_with_check_validation( config=LakebaseChecksStorageConfig(location=location, instance_name=instance_name, user=user) ) assert not dq_engine.validate_checks(loaded_checks).has_errors + + +def delete_all_leftover_instances(ws): + import re + + pattern = re.compile(r"^dqxtest-[A-Za-z0-9]{10}$") + + instances = [] + + for instance in ws.database.list_database_instances(): + if pattern.match(instance.name): + instances.append(instance.name) + + for instance in instances: + ws.database.delete_database_instance(name=instance) From b3b3da0b93128eee45d3e4756ab17c597dca03dc Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 14 Oct 2025 15:04:17 +0200 Subject: [PATCH 110/125] fix: clean up left over lakebase instances --- .../test_save_and_load_checks_from_lakebase_table.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index a6870a7b3..8ed4c61db 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -191,7 +191,7 @@ def test_save_and_load_checks_from_lakebase_table_with_check_validation( assert not dq_engine.validate_checks(loaded_checks).has_errors -def delete_all_leftover_instances(ws): +def test_delete_all_leftover_instances(ws): import re pattern = re.compile(r"^dqxtest-[A-Za-z0-9]{10}$") From 63b25734eefa1edfdd499e913fcac968f65c4053 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 14 Oct 2025 17:17:23 +0200 Subject: [PATCH 111/125] chore: remove cleanup step --- ...st_save_and_load_checks_from_lakebase_table.py | 15 --------------- 1 file changed, 15 deletions(-) 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 index 8ed4c61db..82934ca85 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -189,18 +189,3 @@ def test_save_and_load_checks_from_lakebase_table_with_check_validation( config=LakebaseChecksStorageConfig(location=location, instance_name=instance_name, user=user) ) assert not dq_engine.validate_checks(loaded_checks).has_errors - - -def test_delete_all_leftover_instances(ws): - import re - - pattern = re.compile(r"^dqxtest-[A-Za-z0-9]{10}$") - - instances = [] - - for instance in ws.database.list_database_instances(): - if pattern.match(instance.name): - instances.append(instance.name) - - for instance in instances: - ws.database.delete_database_instance(name=instance) From 941e7f08b1a8c8ebfe6b1d09b8d953f09a689d8f Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 14 Oct 2025 17:26:24 +0200 Subject: [PATCH 112/125] refactor: set port in RunConfig by default to None --- src/databricks/labs/dqx/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 6243906f8..2ba575e3b 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -70,7 +70,7 @@ class RunConfig: # Lakebase connection parameters (only used when checks_location points to a Lakebase instance) lakebase_instance_name: str | None = None # Lakebase instance name lakebase_user: str | None = None # Lakebase username - lakebase_port: str = LAKEBASE_DEFAULT_PORT # Lakebase port + lakebase_port: str | None = None # Lakebase port 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 From 39943c47a9973ec25aef16b7fef5bddebfd3d193 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 14 Oct 2025 18:03:01 +0200 Subject: [PATCH 113/125] refactor: ensure empty user_metadata is null --- src/databricks/labs/dqx/checks_storage.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 15bc3533b..98e040241 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -17,6 +17,7 @@ insert, select, delete, + null, ) from sqlalchemy.schema import CreateSchema from sqlalchemy.dialects.postgresql import JSONB @@ -218,13 +219,14 @@ def _normalize_checks(checks: list[dict], config: LakebaseChecksStorageConfig) - """ 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": check.get("user_metadata"), + "user_metadata": null() if user_metadata is None else user_metadata, } normalized_checks.append(normalized_check) return normalized_checks From 884a597b987d44d301cbfcb4d40d5c97f3e4b2eb Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Tue, 14 Oct 2025 18:07:26 +0200 Subject: [PATCH 114/125] refactor: harmonize RunConfig and LakebaseChecksStorageConfig port defaults --- src/databricks/labs/dqx/config.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 2ba575e3b..87c2399d2 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -221,7 +221,7 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig): instance_name: str | None = None user: str | None = None location: str | None = None - port: str = LAKEBASE_DEFAULT_PORT + port: str | None = None run_config_name: str = "default" mode: str = "overwrite" @@ -243,6 +243,9 @@ def __post_init__(self): if self.mode not in ("append", "overwrite"): raise InvalidConfigError(f"Invalid mode '{self.mode}'. Must be 'append' or 'overwrite'.") + if self.port is None: + self.port = LAKEBASE_DEFAULT_PORT + def _split_location(self) -> tuple[str, ...]: """Splits 'database.schema.table' into components.""" if not self.location: From 0347223f1cf580c34597799dd51e56af1129c8b3 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 15 Oct 2025 12:25:46 +0200 Subject: [PATCH 115/125] refactor to remove duplicated fields --- src/databricks/labs/dqx/checks_storage.py | 3 +-- src/databricks/labs/dqx/config.py | 10 +--------- 2 files changed, 2 insertions(+), 11 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 98e040241..314845c5f 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -38,7 +38,6 @@ BaseChecksStorageConfig, VolumeFileChecksStorageConfig, RunConfig, - LAKEBASE_DEFAULT_PORT, ) from databricks.labs.dqx.errors import InvalidCheckError, InvalidConfigError, CheckDownloadError from databricks.sdk import WorkspaceClient @@ -593,7 +592,7 @@ def _get_storage_handler_and_config( config.instance_name = run_config.lakebase_instance_name if run_config.lakebase_user and not config.user: config.user = run_config.lakebase_user - if run_config.lakebase_port and config.port == LAKEBASE_DEFAULT_PORT: + if run_config.lakebase_port and config.port != run_config.lakebase_port: config.port = run_config.lakebase_port # Overwrite location if overwrite_location is set diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 87c2399d2..37e00ee1e 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -20,8 +20,6 @@ "VolumeFileChecksStorageConfig", ] -LAKEBASE_DEFAULT_PORT = "5432" - @dataclass class InputConfig: @@ -244,7 +242,7 @@ def __post_init__(self): raise InvalidConfigError(f"Invalid mode '{self.mode}'. Must be 'append' or 'overwrite'.") if self.port is None: - self.port = LAKEBASE_DEFAULT_PORT + self.port = "5432" def _split_location(self) -> tuple[str, ...]: """Splits 'database.schema.table' into components.""" @@ -292,12 +290,9 @@ class InstallationChecksStorageConfig( Configuration class for storing checks in an installation. Args: - instance_name: The name of the Lakebase instance (optional, required only for Lakebase storage). location: The installation path where the checks are stored (e.g., table name, file path). Not used when using installation method, as it is retrieved from the installation config, unless overwrite_location is enabled. - user: The user for the Lakebase connection (optional, required only for Lakebase storage). - port: The port for the Lakebase connection (default is '5432'). run_config_name: The name of the run configuration to use for checks (default is 'default'). product_name: The product name for retrieving checks from the installation (default is 'dqx'). assume_user: Whether to assume the user is the owner of the checks (default is True). @@ -308,9 +303,6 @@ class InstallationChecksStorageConfig( overwrite_location: Whether to overwrite the location from run config if provided (default is False). """ - instance_name: str | None = None - user: str | None = None - port: str = LAKEBASE_DEFAULT_PORT # default port for Lakebase connection location: str = "installation" # retrieved from the installation config run_config_name: str = "default" # to retrieve run config product_name: str = "dqx" From f068741d1bf7091a4ccf6189ce9883a554b359ec Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 15 Oct 2025 13:01:46 +0200 Subject: [PATCH 116/125] refactor --- tests/conftest.py | 82 ++++++++----------- ...ave_and_load_checks_from_lakebase_table.py | 79 +++++++++--------- 2 files changed, 75 insertions(+), 86 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 1f6212455..d52b0e65c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,7 +3,7 @@ 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 @@ -441,11 +441,6 @@ def df(spark): ) -@pytest.fixture -def location(): - return "dqx.config.checks" - - @pytest.fixture def make_local_check_file_as_yaml(checks_yaml_content, make_random): file_path = f"checks_{make_random(10).lower()}.yml" @@ -683,7 +678,7 @@ def delete(volume_file_path: str) -> None: @pytest.fixture(scope="session") -def user(): +def lakebase_user(): """ Get the Lakebase user (service principal client ID) from environment variable. @@ -702,11 +697,17 @@ def user(): return client_id +@dataclass +class LakebaseInstance: + name: str + catalog: str + location: str + + @pytest.fixture def make_lakebase_instance(ws, make_random): - instances = {} - def create() -> str: + 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()}" @@ -727,53 +728,27 @@ def create() -> str: ) logger.info(f"Successfully created database catalog: {catalog_name}") - instances[instance_name] = { - "instance_name": instance_name, - "catalog_name": catalog_name, - } - - return instance_name - - def delete(instance_name: str) -> None: - if instance_name not in instances: - logger.warning(f"No resources found for database instance name: {instance_name}") - return - - metadata = instances[instance_name] - instance_name = metadata["instance_name"] - catalog_name = metadata["catalog_name"] + return LakebaseInstance( + name=instance_name, catalog=catalog_name, location=f"{database_name}.{catalog_name}.checks" + ) + def delete(instance: LakebaseInstance) -> None: try: - ws.database.delete_database_catalog(name=catalog_name) - logger.info(f"Successfully deleted database catalog: {catalog_name}") + 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 {catalog_name}: {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}") + 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}") + logger.error(f"Failed to delete database instance {instance.name}: {e}") raise - finally: - instances.pop(instance_name, None) yield from factory("lakebase", create, delete) -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", "")) - - def compare_checks(result: list[dict[str, Any]], expected: list[dict[str, Any]]) -> None: """ Compares two lists of checks dictionaries for equality, ensuring @@ -787,8 +762,21 @@ def compare_checks(result: list[dict[str, Any]], expected: list[dict[str, Any]]) 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) + 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 index 82934ca85..27f75f28e 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -12,36 +12,36 @@ 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, user, location): - instance_name = make_lakebase_instance() +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=location, user=user, instance_name=instance_name) + 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, user, location): - instance_name = make_lakebase_instance() +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=location, user=user, instance_name=instance_name) + 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, user, location): - instance_name = make_lakebase_instance() +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 run_config_name = "workflow_001" config_save = LakebaseChecksStorageConfig( - location=location, user=user, instance_name=instance_name, run_config_name=run_config_name + location=instance.location, user=lakebase_user, instance_name=instance.name, run_config_name=run_config_name ) dq_engine.save_checks(TEST_CHECKS[:1], config=config_save) config_load = LakebaseChecksStorageConfig( - location=location, user=user, instance_name=instance_name, run_config_name=run_config_name + location=instance.location, user=lakebase_user, instance_name=instance.name, run_config_name=run_config_name ) checks = dq_engine.load_checks(config=config_load) compare_checks(checks, TEST_CHECKS[:1]) @@ -49,11 +49,11 @@ def test_save_and_load_checks_from_lakebase_table_with_run_config(ws, spark, mak # test second run config run_config_name = "workflow_002" config_save = LakebaseChecksStorageConfig( - location=location, user=user, instance_name=instance_name, run_config_name=run_config_name + location=instance.location, user=lakebase_user, instance_name=instance.name, run_config_name=run_config_name ) dq_engine.save_checks(TEST_CHECKS[1:], config=config_save) config_load = LakebaseChecksStorageConfig( - location=location, user=user, instance_name=instance_name, run_config_name=run_config_name + location=instance.location, user=lakebase_user, instance_name=instance.name, run_config_name=run_config_name ) checks = dq_engine.load_checks(config=config_load) compare_checks(checks, TEST_CHECKS[1:]) @@ -61,10 +61,10 @@ def test_save_and_load_checks_from_lakebase_table_with_run_config(ws, spark, mak # test default config dq_engine.save_checks( TEST_CHECKS[1:], - config=LakebaseChecksStorageConfig(location=location, user=user, instance_name=instance_name), + config=LakebaseChecksStorageConfig(location=instance.location, user=lakebase_user, instance_name=instance.name), ) checks = dq_engine.load_checks( - config=LakebaseChecksStorageConfig(location=location, user=user, instance_name=instance_name) + config=LakebaseChecksStorageConfig(location=instance.location, user=lakebase_user, instance_name=instance.name) ) compare_checks(checks, TEST_CHECKS[1:]) @@ -73,26 +73,25 @@ def test_save_and_load_checks_from_lakebase_table_with_output_modes( ws, spark, make_lakebase_instance, - user, - location, + lakebase_user, ): - instance_name = make_lakebase_instance() + instance = make_lakebase_instance() dq_engine = DQEngine(ws, spark) run_config_name = "workflow_003" dq_engine.save_checks( TEST_CHECKS[:1], config=LakebaseChecksStorageConfig( - location=location, - instance_name=instance_name, - user=user, + 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=location, instance_name=instance_name, user=user, run_config_name=run_config_name + location=instance.location, instance_name=instance.name, user=lakebase_user, run_config_name=run_config_name ) ) compare_checks(checks, TEST_CHECKS[:1]) @@ -101,35 +100,35 @@ def test_save_and_load_checks_from_lakebase_table_with_output_modes( dq_engine.save_checks( TEST_CHECKS[1:], config=LakebaseChecksStorageConfig( - location=location, - instance_name=instance_name, - user=user, + 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=location, instance_name=instance_name, user=user, run_config_name=run_config_name + 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, user, location + 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 = location + run_config.checks_location = instance.location installation_ctx.installation.save(installation_ctx.config) product_name = installation_ctx.product_info.product_name() - instance_name = make_lakebase_instance() - config = InstallationChecksStorageConfig( - instance_name=instance_name, - user=user, + instance_name=instance.name, + user=lakebase_user, run_config_name=run_config.name, product_name=product_name, assume_user=True, @@ -141,26 +140,27 @@ def test_save_and_load_checks_from_lakebase_table_with_user_installation( compare_checks(checks, TEST_CHECKS) -def test_save_and_load_checks_from_lakebase_table_with_profiler(ws, spark, make_lakebase_instance, df, user, location): +def test_save_and_load_checks_from_lakebase_table_with_profiler(ws, spark, make_lakebase_instance, df, lakebase_user): profiler = DQProfiler(ws) _, profiles = profiler.profile(df) generator = DQGenerator(ws) checks = generator.generate_dq_rules(profiles) - instance_name = make_lakebase_instance() + instance = make_lakebase_instance() dq_engine = DQEngine(ws, spark) dq_engine.save_checks( - checks, config=LakebaseChecksStorageConfig(location=location, instance_name=instance_name, user=user) + checks, + config=LakebaseChecksStorageConfig(location=instance.location, instance_name=instance.name, user=lakebase_user), ) loaded_checks = dq_engine.load_checks( - config=LakebaseChecksStorageConfig(location=location, instance_name=instance_name, user=user) + config=LakebaseChecksStorageConfig(location=instance.location, instance_name=instance.name, user=lakebase_user) ) compare_checks(loaded_checks, checks) def test_save_and_load_checks_from_lakebase_table_with_check_validation( - ws, spark, make_lakebase_instance, user, location + ws, spark, make_lakebase_instance, lakebase_user ): dq_engine = DQEngine(ws, spark) @@ -180,12 +180,13 @@ def test_save_and_load_checks_from_lakebase_table_with_check_validation( """ checks = yaml.safe_load(yaml_str) - instance_name = make_lakebase_instance() + instance = make_lakebase_instance() dq_engine.save_checks( - checks, config=LakebaseChecksStorageConfig(location=location, instance_name=instance_name, user=user) + checks, + config=LakebaseChecksStorageConfig(location=instance.location, instance_name=instance.name, user=lakebase_user), ) loaded_checks = dq_engine.load_checks( - config=LakebaseChecksStorageConfig(location=location, instance_name=instance_name, user=user) + config=LakebaseChecksStorageConfig(location=instance.location, instance_name=instance.name, user=lakebase_user) ) assert not dq_engine.validate_checks(loaded_checks).has_errors From 17eb866a9f9731ea127f00fd9b268ab1e5e898e2 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 15 Oct 2025 14:55:12 +0200 Subject: [PATCH 117/125] refactor --- docs/dqx/docs/guide/quality_checks_storage.mdx | 1 + tests/conftest.py | 6 ++++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/docs/dqx/docs/guide/quality_checks_storage.mdx b/docs/dqx/docs/guide/quality_checks_storage.mdx index 35f8ed338..9d68edb4c 100644 --- a/docs/dqx/docs/guide/quality_checks_storage.mdx +++ b/docs/dqx/docs/guide/quality_checks_storage.mdx @@ -35,6 +35,7 @@ Saving and loading methods accept a storage backend configuration as input. The * `run_config_name` (optional) - run configuration name to load (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). diff --git a/tests/conftest.py b/tests/conftest.py index d52b0e65c..9f607ae0a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -29,7 +29,7 @@ @pytest.fixture(scope="session") def debug_env_name(): - return "ws" # Specify the name of the debug environment from ~/.databricks/debug-env.json + return "ws2" # Specify the name of the debug environment from ~/.databricks/debug-env.json @pytest.fixture @@ -712,6 +712,8 @@ def create() -> LakebaseInstance: 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) @@ -729,7 +731,7 @@ def create() -> LakebaseInstance: logger.info(f"Successfully created database catalog: {catalog_name}") return LakebaseInstance( - name=instance_name, catalog=catalog_name, location=f"{database_name}.{catalog_name}.checks" + name=instance_name, catalog=catalog_name, location=f"{database_name}.{schema_name}.{table_name}" ) def delete(instance: LakebaseInstance) -> None: From c5b9a60bda29ff625159a780491466b22cfad4de Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 15 Oct 2025 15:57:02 +0200 Subject: [PATCH 118/125] fmt --- tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 9f607ae0a..cc8ba996a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -713,7 +713,7 @@ def create() -> LakebaseInstance: 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 + 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) From 4761283eff11550a00bc54017358b5875216e38a Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 15 Oct 2025 16:01:49 +0200 Subject: [PATCH 119/125] refactor --- src/databricks/labs/dqx/checks_storage.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 8d1677826..ab0faf084 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -56,9 +56,6 @@ logger = logging.getLogger(__name__) T = TypeVar("T", bound=BaseChecksStorageConfig) -# PostgreSQL error codes -POSTGRES_UNDEFINED_TABLE_ERROR = '42P01' - class ChecksStorageHandler(ABC, Generic[T]): """ @@ -330,7 +327,8 @@ def _check_for_undefined_table_error(self, e: ProgrammingError, config: Lakebase ProgrammingError: Re-raises the original error if it's not an undefined table error. """ pgcode = getattr(getattr(e, 'orig', None), 'pgcode', None) - if pgcode == POSTGRES_UNDEFINED_TABLE_ERROR: + 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 From 2c67137e22a865a077ceca7baa32754faa117986 Mon Sep 17 00:00:00 2001 From: Manuel Tilgner Date: Wed, 15 Oct 2025 17:48:40 +0200 Subject: [PATCH 120/125] docs: add example for Lakebase in config.yml --- docs/dqx/docs/installation.mdx | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/dqx/docs/installation.mdx b/docs/dqx/docs/installation.mdx index 8894f868a..e5f6ec1ba 100644 --- a/docs/dqx/docs/installation.mdx +++ b/docs/dqx/docs/installation.mdx @@ -236,7 +236,12 @@ 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 + checks_location: dqx.config.checks # <- fully qualified Lakebase table for storing quality rules (checks) + lakebase_instance_name: my-lakebase-instance # <- the name of a Lakebase instance + lakebase_user: 00000000-0000-0000-0000-000000000000 # <- the user to connect to 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 + input_config: ... ``` From 0b2da5bdb1c0aa0b8340063911a9ec820105bada Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 15 Oct 2025 17:54:03 +0200 Subject: [PATCH 121/125] updated logic --- src/databricks/labs/dqx/checks_storage.py | 49 ++++++++++------------- src/databricks/labs/dqx/config.py | 19 ++++----- 2 files changed, 32 insertions(+), 36 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index ab0faf084..4272a2a64 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -577,28 +577,28 @@ def save(self, checks: list[dict], config: InstallationChecksStorageConfig) -> N def _get_storage_handler_and_config( self, config: InstallationChecksStorageConfig ) -> tuple[ChecksStorageHandler, InstallationChecksStorageConfig]: - # Always load run_config to get Lakebase parameters, even when overwriting location - run_config = self._run_config_loader.load_run_config( - run_config_name=config.run_config_name, - assume_user=config.assume_user, - product_name=config.product_name, - install_folder=config.install_folder, - ) - - # Transfer Lakebase fields from run_config to 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 - if run_config.lakebase_port and config.port != run_config.lakebase_port: - config.port = run_config.lakebase_port - # Overwrite location if overwrite_location is set if config.overwrite_location: checks_location = config.location else: + run_config = self._run_config_loader.load_run_config( + run_config_name=config.run_config_name, + assume_user=config.assume_user, + 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 ) @@ -606,30 +606,25 @@ def _get_storage_handler_and_config( config.location = checks_location matches_table_pattern = is_table_location(config.location) - is_lakebase = hasattr(config, 'instance_name') and config.instance_name is not None - - logger.info( - f"InstallationChecksStorageHandler: location='{config.location}', " - f"matches_table_pattern={matches_table_pattern}, is_lakebase={is_lakebase}, " - f"instance_name={config.instance_name if is_lakebase else 'N/A'}" - ) + is_lakebase_storage = config.instance_name is not None - if matches_table_pattern and is_lakebase: - if not config.user: - raise InvalidConfigError("user must be provided when using Lakebase storage") - logger.info(f"Using LakebaseChecksStorageHandler for location '{config.location}'") + 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 diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index ca0213990..d1df0b93f 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -62,20 +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 ) - # Lakebase connection parameters (only used when checks_location points to a Lakebase instance) - lakebase_instance_name: str | None = None # Lakebase instance name - lakebase_user: str | None = None # Lakebase username - lakebase_port: str | None = None # Lakebase port + 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() @@ -219,7 +223,7 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig): instance_name: str | None = None user: str | None = None location: str | None = None - port: str | None = None + port: str = "5432" run_config_name: str = "default" mode: str = "overwrite" @@ -241,9 +245,6 @@ def __post_init__(self): if self.mode not in ("append", "overwrite"): raise InvalidConfigError(f"Invalid mode '{self.mode}'. Must be 'append' or 'overwrite'.") - if self.port is None: - self.port = "5432" - def _split_location(self) -> tuple[str, ...]: """Splits 'database.schema.table' into components.""" if not self.location: From 805ebae5b842b815be1dd8563ee7a9b5e7694d60 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 15 Oct 2025 17:58:39 +0200 Subject: [PATCH 122/125] refactor --- docs/dqx/docs/installation.mdx | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/docs/dqx/docs/installation.mdx b/docs/dqx/docs/installation.mdx index e5f6ec1ba..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 @@ -237,11 +243,6 @@ 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 - checks_location: dqx.config.checks # <- fully qualified Lakebase table for storing quality rules (checks) - lakebase_instance_name: my-lakebase-instance # <- the name of a Lakebase instance - lakebase_user: 00000000-0000-0000-0000-000000000000 # <- the user to connect to 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 - input_config: ... ``` From 1d56cc355904b9d64866940fb0881f148ac0c2e6 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 15 Oct 2025 18:04:06 +0200 Subject: [PATCH 123/125] fmt --- src/databricks/labs/dqx/checks_storage.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 4272a2a64..5fa940d95 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -840,7 +840,7 @@ def create_for_run_config(self, run_config: RunConfig) -> tuple[ChecksStorageHan location=run_config.checks_location, instance_name=run_config.lakebase_instance_name, user=run_config.lakebase_user, - port=run_config.lakebase_port, + port=run_config.lakebase_port or "5432", run_config_name=run_config.name, ), ) From 35352a7bd67356919f20af7017c0f732bc686e17 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 15 Oct 2025 18:24:27 +0200 Subject: [PATCH 124/125] updated tests --- tests/conftest.py | 2 +- ...ave_and_load_checks_from_lakebase_table.py | 100 +++++------------- 2 files changed, 30 insertions(+), 72 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index cc8ba996a..f3ed47c16 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -29,7 +29,7 @@ @pytest.fixture(scope="session") def debug_env_name(): - return "ws2" # Specify the name of the debug environment from ~/.databricks/debug-env.json + return "ws" # Specify the name of the debug environment from ~/.databricks/debug-env.json @pytest.fixture 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 index 27f75f28e..6740ebef2 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -1,11 +1,8 @@ -import yaml import pytest from databricks.labs.dqx.config import InstallationChecksStorageConfig, LakebaseChecksStorageConfig from databricks.labs.dqx.engine import DQEngine -from databricks.labs.dqx.profiler.profiler import DQProfiler -from databricks.labs.dqx.profiler.generator import DQGenerator from databricks.sdk.errors import NotFound from tests.conftest import compare_checks @@ -25,8 +22,10 @@ def test_save_and_load_checks_from_lakebase_table(ws, spark, make_lakebase_insta 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) @@ -35,37 +34,21 @@ def test_save_and_load_checks_from_lakebase_table_with_run_config(ws, spark, mak dq_engine = DQEngine(ws, spark) # test first run config - run_config_name = "workflow_001" - config_save = LakebaseChecksStorageConfig( - location=instance.location, user=lakebase_user, instance_name=instance.name, run_config_name=run_config_name - ) - dq_engine.save_checks(TEST_CHECKS[:1], config=config_save) - config_load = LakebaseChecksStorageConfig( - location=instance.location, user=lakebase_user, instance_name=instance.name, run_config_name=run_config_name + config = LakebaseChecksStorageConfig( + location=instance.location, user=lakebase_user, instance_name=instance.name, run_config_name="workflow_001" ) - checks = dq_engine.load_checks(config=config_load) + 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 - run_config_name = "workflow_002" - config_save = LakebaseChecksStorageConfig( - location=instance.location, user=lakebase_user, instance_name=instance.name, run_config_name=run_config_name + 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=config_save) - config_load = LakebaseChecksStorageConfig( - location=instance.location, user=lakebase_user, instance_name=instance.name, run_config_name=run_config_name - ) - checks = dq_engine.load_checks(config=config_load) - compare_checks(checks, TEST_CHECKS[1:]) + dq_engine.save_checks(TEST_CHECKS[1:], config=second_config) + checks = dq_engine.load_checks(config=second_config) - # test default config - dq_engine.save_checks( - TEST_CHECKS[1:], - config=LakebaseChecksStorageConfig(location=instance.location, user=lakebase_user, instance_name=instance.name), - ) - checks = dq_engine.load_checks( - config=LakebaseChecksStorageConfig(location=instance.location, user=lakebase_user, instance_name=instance.name) - ) compare_checks(checks, TEST_CHECKS[1:]) @@ -135,58 +118,33 @@ def test_save_and_load_checks_from_lakebase_table_with_user_installation( ) 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_save_and_load_checks_from_lakebase_table_with_profiler(ws, spark, make_lakebase_instance, df, lakebase_user): - profiler = DQProfiler(ws) - _, profiles = profiler.profile(df) - generator = DQGenerator(ws) - checks = generator.generate_dq_rules(profiles) +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() - dq_engine = DQEngine(ws, spark) - dq_engine.save_checks( - checks, - config=LakebaseChecksStorageConfig(location=instance.location, instance_name=instance.name, user=lakebase_user), - ) - loaded_checks = dq_engine.load_checks( - config=LakebaseChecksStorageConfig(location=instance.location, instance_name=instance.name, user=lakebase_user) - ) - compare_checks(loaded_checks, checks) - - -def test_save_and_load_checks_from_lakebase_table_with_check_validation( - ws, spark, make_lakebase_instance, lakebase_user -): - dq_engine = DQEngine(ws, spark) + 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" - yaml_str = """ - - criticality: error - check: - function: is_not_null_and_not_empty - arguments: - column: name - - - criticality: error - check: - function: is_not_null - for_each_column: - - region - - state - """ - checks = yaml.safe_load(yaml_str) + installation_ctx.installation.save(installation_ctx.config) - instance = make_lakebase_instance() + installation_ctx.deployed_workflows.run_workflow("profiler", run_config.name) - dq_engine.save_checks( - checks, - config=LakebaseChecksStorageConfig(location=instance.location, instance_name=instance.name, user=lakebase_user), - ) - loaded_checks = dq_engine.load_checks( - config=LakebaseChecksStorageConfig(location=instance.location, instance_name=instance.name, user=lakebase_user) + 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 not dq_engine.validate_checks(loaded_checks).has_errors + assert checks, "Checks are missing" From 6b9db1a6bef4b293d7bdde1aa5b3a891dfeea8fe Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 15 Oct 2025 18:28:41 +0200 Subject: [PATCH 125/125] fmt --- .../integration/test_save_and_load_checks_from_lakebase_table.py | 1 - 1 file changed, 1 deletion(-) 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 index 6740ebef2..4ce99d908 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/tests/integration/test_save_and_load_checks_from_lakebase_table.py @@ -1,4 +1,3 @@ - import pytest from databricks.labs.dqx.config import InstallationChecksStorageConfig, LakebaseChecksStorageConfig