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 c5f248709..cce1de22c 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -24,7 +24,6 @@ from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.exc import DatabaseError, ProgrammingError, OperationalError, IntegrityError - import yaml from pyspark.sql import SparkSession from databricks.sdk.errors import NotFound @@ -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(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 [] 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) 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 994620017..be01680cb 100644 --- a/tests/integration/test_save_and_load_checks_from_table.py +++ b/tests/integration/test_save_and_load_checks_from_table.py @@ -77,6 +77,18 @@ def test_load_checks_when_checks_table_does_not_exist(ws, make_schema, make_rand 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 schema_name = make_schema(catalog_name=catalog_name).name 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(