From 6cd9f02c78aa99ebbf3b0715f21a396fda211d6e Mon Sep 17 00:00:00 2001 From: Eniwoke Cornelius Date: Thu, 17 Jul 2025 18:41:32 +0100 Subject: [PATCH 1/7] move criticiality of rule inro _validate_attributes --- src/databricks/labs/dqx/rule.py | 26 +++++++++----------------- 1 file changed, 9 insertions(+), 17 deletions(-) diff --git a/src/databricks/labs/dqx/rule.py b/src/databricks/labs/dqx/rule.py index 7d1e23a91..cb2f6c327 100644 --- a/src/databricks/labs/dqx/rule.py +++ b/src/databricks/labs/dqx/rule.py @@ -195,6 +195,15 @@ def _initialize_name_if_missing(self, check_condition: Column): object.__setattr__(self, "name", normalized_name) def _validate_attributes(self) -> None: + """Verify Criticality of rule""" + criticality = self.criticality + if criticality not in {Criticality.WARN.value, Criticality.ERROR.value}: + raise ValueError( + f"Invalid 'criticality' value: '{criticality}'. " + f"Expected '{Criticality.WARN.value}' or '{Criticality.ERROR.value}'. " + f"Check details: {self.name}" + ) + """Validate input attributes.""" if self.column is not None and self.columns is not None: raise ValueError("Both 'column' and 'columns' cannot be provided at the same time.") @@ -206,23 +215,6 @@ def get_check_condition(self) -> Column: :return: The Spark Column representing the check condition. """ - - @ft.cached_property - def check_criticality(self) -> str: - """Criticality of the check. - - :return: string describing criticality - `warn` or `error`. - :raises ValueError: if criticality is invalid. - """ - criticality = self.criticality - if criticality not in {Criticality.WARN.value, Criticality.ERROR.value}: - raise ValueError( - f"Invalid 'criticality' value: '{criticality}'. " - f"Expected '{Criticality.WARN.value}' or '{Criticality.ERROR.value}'. " - f"Check details: {self.name}" - ) - return criticality - @ft.cached_property def columns_as_string_expr(self) -> Column: """Spark Column expression representing the column(s) as a string (not normalized). From 98371b705b18d263252252a469932d8344bc1a5d Mon Sep 17 00:00:00 2001 From: Eniwoke Cornelius Date: Thu, 17 Jul 2025 18:42:07 +0100 Subject: [PATCH 2/7] since criticality is validated after creation, filter by criticality attribute --- 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 28d9f8a4f..72d3074d6 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -401,7 +401,7 @@ def _get_check_columns(checks: list[DQRule], criticality: str) -> list[DQRule]: :param criticality: criticality :return: list of check columns """ - return [check for check in checks if check.check_criticality == criticality] + return [check for check in checks if check.criticality == criticality] def _append_empty_checks(self, df: DataFrame) -> DataFrame: """Append empty checks at the end of dataframe. From d5eb5792bd6735499a77114bf5a65943ab26ffad Mon Sep 17 00:00:00 2001 From: Eniwoke Cornelius Date: Sat, 10 Jan 2026 18:33:51 +0000 Subject: [PATCH 3/7] refactor: use special characters --- src/databricks/labs/dqx/checks_storage.py | 5 +++-- .../test_save_and_load_checks_from_table.py | 11 +++++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index c5f248709..3c39e51bc 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -23,7 +23,7 @@ from sqlalchemy.schema import CreateSchema from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.exc import DatabaseError, ProgrammingError, OperationalError, IntegrityError - +from urllib.parse import quote import yaml from pyspark.sql import SparkSession @@ -107,7 +107,8 @@ def load(self, config: TableChecksStorageConfig) -> list[dict]: NotFound: if the table does not exist in the workspace """ logger.info(f"Loading quality rules (checks) from table '{config.location}'") - if not self.ws.tables.exists(config.location).table_exists: + api_ready_location = quote(config.location.replace("`", "")) + if not self.ws.tables.exists(api_ready_location).table_exists: raise NotFound(f"Checks table {config.location} does not exist in the workspace") rules_df = self.spark.read.table(config.location) return serialize_checks_from_dataframe(rules_df, run_config_name=config.run_config_name) or [] diff --git a/tests/integration/test_save_and_load_checks_from_table.py b/tests/integration/test_save_and_load_checks_from_table.py index 994620017..5124e3487 100644 --- a/tests/integration/test_save_and_load_checks_from_table.py +++ b/tests/integration/test_save_and_load_checks_from_table.py @@ -76,6 +76,17 @@ def test_load_checks_when_checks_table_does_not_exist(ws, make_schema, make_rand config = TableChecksStorageConfig(location=table_name) engine.load_checks(config=config) +def test_load_checks_with_special_characters_in_table_name(ws, make_schema, make_random, spark): + catalog_name = TEST_CATALOG + schema_name = make_schema(catalog_name=catalog_name).name + table_name = f"`{catalog_name}`.`{schema_name}`.`table-with-special_chars$#{make_random(5)}`" + + engine = DQEngine(ws, spark) + config = TableChecksStorageConfig(location=table_name, run_config_name="default") + engine.save_checks(INPUT_CHECKS, config=config) + checks = engine.load_checks(config=config) + assert checks == EXPECTED_CHECKS, "Checks were not loaded correctly." + def test_save_and_load_checks_from_table(ws, make_schema, make_random, spark): catalog_name = TEST_CATALOG From 257cd60ace2ef03b98f4a9ac2527b9f0898a500e Mon Sep 17 00:00:00 2001 From: Eniwoke Cornelius Date: Sat, 10 Jan 2026 23:01:10 +0000 Subject: [PATCH 4/7] chore: create one function to handle SDK table location --- src/databricks/labs/dqx/checks_storage.py | 23 ++++++++++++++++--- .../test_save_and_load_checks_from_table.py | 1 + 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 3c39e51bc..bf965c98c 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -20,10 +20,10 @@ null, event, ) +from urllib.parse import quote from sqlalchemy.schema import CreateSchema from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.exc import DatabaseError, ProgrammingError, OperationalError, IntegrityError -from urllib.parse import quote import yaml from pyspark.sql import SparkSession @@ -107,8 +107,7 @@ def load(self, config: TableChecksStorageConfig) -> list[dict]: NotFound: if the table does not exist in the workspace """ logger.info(f"Loading quality rules (checks) from table '{config.location}'") - api_ready_location = quote(config.location.replace("`", "")) - if not self.ws.tables.exists(api_ready_location).table_exists: + if not self.ws.tables.exists(to_api_identifier(config.location)).table_exists: raise NotFound(f"Checks table {config.location} does not exist in the workspace") rules_df = self.spark.read.table(config.location) return serialize_checks_from_dataframe(rules_df, run_config_name=config.run_config_name) or [] @@ -898,3 +897,21 @@ def is_table_location(location: str) -> bool: bool: True if the location is a valid table name and not a file path, False otherwise. """ return bool(TABLE_PATTERN.match(location)) and not location.lower().endswith(tuple(FILE_SERIALIZERS.keys())) + + +def to_api_identifier(location: str) -> str: + """ + Transform a Spark SQL table identifier into a Databricks SDK-compatible string. + + This function bridges the gap between Spark's SQL requirements (backticks for + special characters) and the Databricks Workspace Client's REST requirements + (literal strings in URL paths). It removes SQL-style backticks and applies + URL encoding to ensure safety for characters like '$' or '#'. + + Args: + location (str): The fully qualified table location (e.g., "`cat`.`schema`.`tab`"). + + Returns: + str: A URL-encoded string compatible with Databricks SDK methods like tables.exists(). + """ + return quote(location.replace("`", "")) diff --git a/tests/integration/test_save_and_load_checks_from_table.py b/tests/integration/test_save_and_load_checks_from_table.py index 5124e3487..697020e93 100644 --- a/tests/integration/test_save_and_load_checks_from_table.py +++ b/tests/integration/test_save_and_load_checks_from_table.py @@ -76,6 +76,7 @@ def test_load_checks_when_checks_table_does_not_exist(ws, make_schema, make_rand config = TableChecksStorageConfig(location=table_name) engine.load_checks(config=config) + def test_load_checks_with_special_characters_in_table_name(ws, make_schema, make_random, spark): catalog_name = TEST_CATALOG schema_name = make_schema(catalog_name=catalog_name).name From 4ab71699211e33098064b21bdeed2ab3e4ba0f69 Mon Sep 17 00:00:00 2001 From: Eniwoke Cornelius Date: Sat, 10 Jan 2026 23:02:41 +0000 Subject: [PATCH 5/7] chore: formatting --- 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 bf965c98c..b5b0a5259 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -6,6 +6,7 @@ from pathlib import Path from abc import ABC, abstractmethod from typing import Generic, TypeVar, NoReturn +from urllib.parse import quote from sqlalchemy import ( Engine, create_engine, @@ -20,7 +21,6 @@ null, event, ) -from urllib.parse import quote from sqlalchemy.schema import CreateSchema from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.exc import DatabaseError, ProgrammingError, OperationalError, IntegrityError From 7a02a390cf368ea9d0b7be0b70bbdc72bad41be5 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 14 Jan 2026 18:04:23 +0100 Subject: [PATCH 6/7] use spark.catalog.tableExists --- pyproject.toml | 2 ++ src/databricks/labs/dqx/checks_storage.py | 21 +------------------ tests/integration/test_config_serializer.py | 15 +++++++++++++ .../test_save_and_load_checks_from_table.py | 2 +- tests/integration/test_summary_metrics.py | 2 +- 5 files changed, 20 insertions(+), 22 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 33bde88ef..9926a6641 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -649,6 +649,8 @@ disable = [ "broad-exception-caught", "rewrite-as-for-loop", "redundant-returns-doc", + # dbx pylint is not up-to-date, e.g. spark.catalog.tableExists is compatible with UC since DBR 14+ + "incompatible-with-uc" ] # Enable the message, report, category or checker with the given id(s). You can diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index b5b0a5259..cce1de22c 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, NoReturn -from urllib.parse import quote from sqlalchemy import ( Engine, create_engine, @@ -107,7 +106,7 @@ def load(self, config: TableChecksStorageConfig) -> list[dict]: NotFound: if the table does not exist in the workspace """ logger.info(f"Loading quality rules (checks) from table '{config.location}'") - if not self.ws.tables.exists(to_api_identifier(config.location)).table_exists: + if not self.spark.catalog.tableExists(config.location): raise NotFound(f"Checks table {config.location} does not exist in the workspace") rules_df = self.spark.read.table(config.location) return serialize_checks_from_dataframe(rules_df, run_config_name=config.run_config_name) or [] @@ -897,21 +896,3 @@ def is_table_location(location: str) -> bool: bool: True if the location is a valid table name and not a file path, False otherwise. """ return bool(TABLE_PATTERN.match(location)) and not location.lower().endswith(tuple(FILE_SERIALIZERS.keys())) - - -def to_api_identifier(location: str) -> str: - """ - Transform a Spark SQL table identifier into a Databricks SDK-compatible string. - - This function bridges the gap between Spark's SQL requirements (backticks for - special characters) and the Databricks Workspace Client's REST requirements - (literal strings in URL paths). It removes SQL-style backticks and applies - URL encoding to ensure safety for characters like '$' or '#'. - - Args: - location (str): The fully qualified table location (e.g., "`cat`.`schema`.`tab`"). - - Returns: - str: A URL-encoded string compatible with Databricks SDK methods like tables.exists(). - """ - return quote(location.replace("`", "")) diff --git a/tests/integration/test_config_serializer.py b/tests/integration/test_config_serializer.py index 4f64cb590..b975810e8 100644 --- a/tests/integration/test_config_serializer.py +++ b/tests/integration/test_config_serializer.py @@ -185,3 +185,18 @@ def test_save_run_config_in_custom_folder_installation(ws, installation_ctx_cust run_config_name="default", install_folder=installation_ctx_custom_install_folder.install_folder ) assert actual_default_run_config == installation_ctx_custom_install_folder.config.get_run_config("default") + + +def test_save_workspace_config_with_extra_fields_in_user_installation(ws, installation_ctx, spark): + installation_ctx.installation.save(installation_ctx.config) + + run_config = RunConfig(name="fake") + config = WorkspaceConfig(run_configs=[run_config]) + # extra configs are skipped + config.extra_field = "extra_value" + + product_name = installation_ctx.product_info.product_name() + ConfigSerializer(ws).save_config(config=config, assume_user=True, product_name=product_name) + + actual_config = ConfigSerializer(ws).load_config(assume_user=True, product_name=product_name) + assert actual_config == config diff --git a/tests/integration/test_save_and_load_checks_from_table.py b/tests/integration/test_save_and_load_checks_from_table.py index 697020e93..be01680cb 100644 --- a/tests/integration/test_save_and_load_checks_from_table.py +++ b/tests/integration/test_save_and_load_checks_from_table.py @@ -80,7 +80,7 @@ def test_load_checks_when_checks_table_does_not_exist(ws, make_schema, make_rand def test_load_checks_with_special_characters_in_table_name(ws, make_schema, make_random, spark): catalog_name = TEST_CATALOG schema_name = make_schema(catalog_name=catalog_name).name - table_name = f"`{catalog_name}`.`{schema_name}`.`table-with-special_chars$#{make_random(5)}`" + table_name = f"`{catalog_name}`.`{schema_name}`.`table-with-special_chars$#@{make_random(5)}`" engine = DQEngine(ws, spark) config = TableChecksStorageConfig(location=table_name, run_config_name="default") diff --git a/tests/integration/test_summary_metrics.py b/tests/integration/test_summary_metrics.py index 0dcf6b157..e85c11746 100644 --- a/tests/integration/test_summary_metrics.py +++ b/tests/integration/test_summary_metrics.py @@ -232,7 +232,7 @@ def test_engine_without_observer_no_metrics_saved(ws, spark, make_schema, make_r else: raise ValueError("Invalid 'apply_checks_method' used for testing observable metrics.") - assert not ws.tables.exists(metrics_table_name).table_exists + assert not spark.catalog.tableExists(metrics_table_name) @pytest.mark.parametrize( From 1850fda2af468ff440ce61b34de38c2aa58d254f Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 14 Jan 2026 18:20:36 +0100 Subject: [PATCH 7/7] improve existing tests --- tests/integration/test_ai_rules_generator.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/integration/test_ai_rules_generator.py b/tests/integration/test_ai_rules_generator.py index fe41b61ec..badc8dab3 100644 --- a/tests/integration/test_ai_rules_generator.py +++ b/tests/integration/test_ai_rules_generator.py @@ -94,7 +94,10 @@ def not_ends_with_suffix(column: str, suffix: str): custom_check_functions = {"not_ends_with_suffix": not_ends_with_suffix} - user_input = USER_INPUT + "\nEmail address must not end with '@gmail.com'. Use not_ends_with_suffix function." + user_input = USER_INPUT + ( + "\nEmail address must not end with '@gmail.com'. " + "Use not_ends_with_suffix function and don't add any extra quotes for suffix." + ) generator = DQGenerator(ws, spark, custom_check_functions=custom_check_functions) actual_checks = generator.generate_dq_rules_ai_assisted(user_input=user_input)