Refactored methods for loading and saving checks

demos/dqx_demo_asset_bundle/dqx_demo_notebook.py

@@ -581,7 +581,7 @@
 
 # DBTITLE 1,Common Imports
 from databricks.sdk import WorkspaceClient
-from databricks.labs.dqx.config import InputConfig, OutputConfig
+from databricks.labs.dqx.config import InputConfig, OutputConfig, WorkspaceFileChecksStorageConfig
 from databricks.labs.dqx.engine import DQEngine
 
 # COMMAND ----------
@@ -594,7 +594,7 @@
 displayHTML(f'<a href="/#workspace{sensor_rules_file_path}" target="_blank">Quality rules file for sensor dataset</a>')
 
 # Load the checks
-maintenance_quality_checks = dq_engine.load_checks_from_workspace_file(workspace_path=sensor_rules_file_path)
+maintenance_quality_checks = dq_engine.load_checks(config=WorkspaceFileChecksStorageConfig(location=sensor_rules_file_path))
 
 # Apply the checks and write the output data
 dq_engine.apply_checks_by_metadata_and_save_in_table(
@@ -614,7 +614,7 @@
 displayHTML(f'<a href="/#workspace{maintenance_rules_file_path}" target="_blank">Quality rules file for maintenance dataset</a>')
 
 # Load the checks
-maintenance_quality_checks = dq_engine.load_checks_from_workspace_file(workspace_path=maintenance_rules_file_path)
+maintenance_quality_checks = dq_engine.load_checks(config=WorkspaceFileChecksStorageConfig(location=maintenance_rules_file_path))
 
 # Apply the checks and write the output data
 dq_engine.apply_checks_by_metadata_and_save_in_table(

demos/dqx_demo_dbt/models/dummy_model_dq.py

@@ -1,4 +1,6 @@
 import yaml
+
+from databricks.labs.dqx.config import WorkspaceFileChecksStorageConfig
 from databricks.labs.dqx.engine import DQEngine
 from databricks.sdk import WorkspaceClient
 
@@ -26,7 +28,7 @@ def model(dbt, session):
 
     # Checks can also be loaded from a file in the workspace or delta table
     # checks_path = dbt.config.get("checks_file_path")  # get from dbt var
-    # checks = dq_engine.load_checks_from_workspace_file(checks_path)
+    # checks = dq_engine.load_checks(config=WorkspaceFileChecksStorageConfig(location=checks_path))
 
     # apply quality checks with issues reported in _warnings and _errors columns
     df = dq_engine.apply_checks_by_metadata(input_df, checks)

demos/dqx_demo_library.py

@@ -30,6 +30,7 @@
 from databricks.labs.dqx.profiler.profiler import DQProfiler
 from databricks.labs.dqx.profiler.generator import DQGenerator
 from databricks.labs.dqx.profiler.dlt_generator import DQDltGenerator
+from databricks.labs.dqx.config import WorkspaceFileChecksStorageConfig, TableChecksStorageConfig
 from databricks.labs.dqx.engine import DQEngine
 from databricks.sdk import WorkspaceClient
 import os
@@ -81,10 +82,13 @@
 user_name = spark.sql("select current_user() as user").collect()[0]["user"]
 checks_file = f"{demo_file_directory}/dqx_demo_checks.yml"
 dq_engine = DQEngine(ws)
-dq_engine.save_checks_in_workspace_file(checks=checks, workspace_path=checks_file)
+dq_engine.save_checks(checks=checks, config=WorkspaceFileChecksStorageConfig(location=checks_file))
 
 # save generated checks in a Delta table
-dq_engine.save_checks_in_table(checks=checks, table_name=f"{demo_catalog_name}.{demo_schema_name}.dqx_checks_table", mode="overwrite")
+dq_engine.save_checks(
+    checks=checks,
+    config=TableChecksStorageConfig(location=f"{demo_catalog_name}.{demo_schema_name}.dqx_checks_table", mode="overwrite")
+)
 
 # COMMAND ----------
 
@@ -95,12 +99,13 @@
 
 from databricks.labs.dqx.engine import DQEngine
 from databricks.sdk import WorkspaceClient
+from databricks.labs.dqx.config import WorkspaceFileChecksStorageConfig
 
 input_df = spark.createDataFrame([[1, 3, 3, 2], [3, 3, None, 1]], schema)
 
 # load checks from a file
 dq_engine = DQEngine(WorkspaceClient())
-checks = dq_engine.load_checks_from_workspace_file(workspace_path=checks_file)
+checks = dq_engine.load_checks(config=WorkspaceFileChecksStorageConfig(location=checks_file))
 
 # Option 1: apply quality rules and quarantine invalid records
 valid_df, quarantine_df = dq_engine.apply_checks_by_metadata_and_split(input_df, checks)
@@ -120,12 +125,13 @@
 
 from databricks.labs.dqx.engine import DQEngine
 from databricks.sdk import WorkspaceClient
+from databricks.labs.dqx.config import TableChecksStorageConfig
 
 input_df = spark.createDataFrame([[1, 3, 3, 2], [3, 3, None, 1]], schema)
 
 # load checks from a Delta table
 dq_engine = DQEngine(WorkspaceClient())
-checks = dq_engine.load_checks_from_table(table_name=f"{demo_catalog_name}.{demo_schema_name}.dqx_checks_table")
+checks = dq_engine.load_checks(config=TableChecksStorageConfig(location=f"{demo_catalog_name}.{demo_schema_name}.dqx_checks_table"))
 
 # Option 1: apply quality rules and quarantine invalid records
 valid_df, quarantine_df = dq_engine.apply_checks_by_metadata_and_split(input_df, checks)

demos/dqx_demo_tool.py

@@ -94,15 +94,19 @@
 from databricks.labs.dqx.profiler.profiler import DQProfiler
 from databricks.labs.dqx.profiler.generator import DQGenerator
 from databricks.labs.dqx.engine import DQEngine
+from databricks.labs.dqx.config import InstallationChecksStorageConfig, WorkspaceFileChecksStorageConfig
+from databricks.labs.dqx.config_loader import RunConfigLoader
 from databricks.labs.dqx.utils import read_input_data
 from databricks.sdk import WorkspaceClient
 
+
 dqx_product_name = dbutils.widgets.get("dqx_product_name")
 
-# setup the DQEngine
 ws = WorkspaceClient()
 dq_engine = DQEngine(ws)
-run_config = dq_engine.load_run_config(run_config_name="default", assume_user=True, product_name=dqx_product_name)
+
+# load the run configuration
+run_config = RunConfigLoader(ws).load_run_config(run_config_name="default", product_name=dqx_product_name)
 
 # read the input data, limit to 1000 rows for demo purpose
 input_df = read_input_data(spark, run_config.input_config).limit(1000)
@@ -120,9 +124,10 @@
 print(yaml.safe_dump(checks))
 
 # save generated checks to location specified in the default run configuration inside workspace installation folder
-dq_engine.save_checks_in_installation(checks, run_config_name="default", product_name=dqx_product_name)
-# or save it to an arbitrary workspace location
-#dq_engine.save_checks_in_workspace_file(checks, workspace_path="/Shared/App1/checks.yml")
+dq_engine.save_checks(checks, config=InstallationChecksStorageConfig(run_config_name="default", product_name=dqx_product_name))
+
+# or save checks in arbitrary workspace location
+#dq_engine.save_checks(checks, config=WorkspaceFileChecksStorageConfig(location="/Shared/App1/checks.yml"))
 
 # COMMAND ----------
 
@@ -136,6 +141,8 @@
 import yaml
 from databricks.labs.dqx.engine import DQEngine
 from databricks.sdk import WorkspaceClient
+from databricks.labs.dqx.config import InstallationChecksStorageConfig, WorkspaceFileChecksStorageConfig
+
 
 checks = yaml.safe_load("""
 - check:
@@ -180,10 +187,12 @@
 assert not status.has_errors
 
 dq_engine = DQEngine(WorkspaceClient())
+
 # save checks to location specified in the default run configuration inside workspace installation folder
-dq_engine.save_checks_in_installation(checks, run_config_name="default", product_name=dqx_product_name)
-# or save it to an arbitrary workspace location
-#dq_engine.save_checks_in_workspace_file(checks, workspace_path="/Shared/App1/checks.yml")
+dq_engine.save_checks(checks, config=InstallationChecksStorageConfig(run_config_name="default", product_name=dqx_product_name))
+
+# or save checks in arbitrary workspace location
+#dq_engine.save_checks(checks, config=WorkspaceFileChecksStorageConfig(location="/Shared/App1/checks.yml"))
 
 # COMMAND ----------
 
@@ -195,21 +204,27 @@
 from databricks.labs.dqx.engine import DQEngine
 from databricks.labs.dqx.utils import read_input_data
 from databricks.sdk import WorkspaceClient
+from databricks.labs.dqx.config import InstallationChecksStorageConfig, WorkspaceFileChecksStorageConfig
+from databricks.labs.dqx.config_loader import RunConfigLoader
 
-run_config = dq_engine.load_run_config(run_config_name="default", assume_user=True, product_name=dqx_product_name)
+
+dq_engine = DQEngine(WorkspaceClient())
+
+# load the run configuration
+run_config = RunConfigLoader(ws).load_run_config(run_config_name="default", assume_user=True, product_name=dqx_product_name)
 
 # read the data, limit to 1000 rows for demo purpose
 bronze_df = read_input_data(spark, run_config.input_config).limit(1000)
 
 # apply your business logic here
 bronze_transformed_df = bronze_df.filter("vendor_id in (1, 2)")
 
-dq_engine = DQEngine(WorkspaceClient())
-
 # load checks from location defined in the run configuration
-checks = dq_engine.load_checks_from_installation(assume_user=True, run_config_name="default", product_name=dqx_product_name)
+
+checks = dq_engine.load_checks(config=InstallationChecksStorageConfig(assume_user=True, run_config_name="default", product_name=dqx_product_name))
+
 # or load checks from arbitrary workspace file
-# checks = dq_engine.load_checks_from_workspace_file(workspace_path="/Shared/App1/checks.yml")
+#checks = dq_engine.load_checks(config=WorkspaceFileChecksStorageConfig(location="/Shared/App1/checks.yml"))
 print(checks)
 
 # Option 1: apply quality rules and quarantine invalid records

demos/dqx_manufacturing_demo.py

@@ -108,7 +108,7 @@
 from pyspark.sql.types import *
 from pyspark.sql.functions import *
 from datetime import datetime
-import delta
+
 
 if spark.catalog.tableExists(sensor_table) and spark.table(sensor_table).count() > 0:
     print(
@@ -634,14 +634,14 @@
 # COMMAND ----------
 
 # DBTITLE 1,Common Imports
-import os
 import yaml
 from pprint import pprint
 
 from databricks.sdk import WorkspaceClient
 from databricks.labs.dqx.profiler.profiler import DQProfiler
 from databricks.labs.dqx.profiler.generator import DQGenerator
 from databricks.labs.dqx.engine import DQEngine
+from databricks.labs.dqx.config import WorkspaceFileChecksStorageConfig, TableChecksStorageConfig
 
 # COMMAND ----------
 
@@ -697,7 +697,7 @@
 maintenance_dq_rules_yaml = f"{quality_rules_path}/maintenance_dq_rules.yml"
 
 # Save file in a workspace path
-dq_engine.save_checks_in_workspace_file(maintenance_checks, workspace_path=maintenance_dq_rules_yaml)
+dq_engine.save_checks(maintenance_checks, config=WorkspaceFileChecksStorageConfig(location=maintenance_dq_rules_yaml))
 
 # display the link to the saved checks
 displayHTML(f'<a href="/#workspace{maintenance_dq_rules_yaml}" target="_blank">Maintenance Data Quality Rules YAML</a>')
@@ -707,7 +707,7 @@
 # DBTITLE 1,Save the quality rules in delta table
 # or save in delta table
 maintenance_quality_rules_table = f"{database}.{schema}.maintenance_inferred_quality_rules"
-dq_engine.save_checks_in_table(table_name=maintenance_quality_rules_table, checks=maintenance_checks, run_config_name="maintenance")
+dq_engine.save_checks(maintenance_checks, config=TableChecksStorageConfig(location=maintenance_quality_rules_table, run_config_name="maintenance"))
 
 # COMMAND ----------
 
@@ -718,10 +718,9 @@
 # COMMAND ----------
 
 # Load checks from workspace file
-quality_checks = dq_engine.load_checks_from_workspace_file(workspace_path=maintenance_dq_rules_yaml)
-
+quality_checks = dq_engine.load_checks(config=WorkspaceFileChecksStorageConfig(location=maintenance_dq_rules_yaml))
 # or Load checks from a table
-# quality_checks = dq_engine.load_checks_from_table(table_name=fq_tbl, run_config_name="maintenance")
+#quality_checks = dq_engine.load_checks(config=TableChecksStorageConfig(location=maintenance_quality_rules_table, run_config_name="maintenance"))
 
 # Apply checks on input data
 valid_df, quarantined_df = dq_engine.apply_checks_by_metadata_and_split(mntnc_bronze_df, quality_checks)
@@ -810,7 +809,7 @@
 
 # save checks in a workspace location
 sensor_dq_rules_yaml = f"{quality_rules_path}/sensor_dq_rules.yml"
-dq_engine.save_checks_in_workspace_file(sensor_dq_checks, workspace_path=sensor_dq_rules_yaml)
+dq_engine.save_checks(sensor_dq_checks, config=WorkspaceFileChecksStorageConfig(location=sensor_dq_rules_yaml))
 
 # display the link to the saved checks
 displayHTML(f'<a href="/#workspace{sensor_dq_rules_yaml}" target="_blank">Sensor Data Quality Rules YAML</a>')
@@ -827,7 +826,7 @@
 sensor_bronze_df = spark.read.table(sensor_table)
 
 # Load quality rules from YAML file
-sensor_dq_checks = dq_engine.load_checks_from_workspace_file(workspace_path=sensor_dq_rules_yaml)
+sensor_dq_checks = dq_engine.load_checks(config=WorkspaceFileChecksStorageConfig(location=sensor_dq_rules_yaml))
 
 # Apply checks on input data
 valid_df, quarantined_df = dq_engine.apply_checks_by_metadata_and_split(sensor_bronze_df, sensor_dq_checks)
@@ -859,6 +858,7 @@
 from pyspark.sql import Column as col
 from databricks.labs.dqx.check_funcs import make_condition
 
+
 def firmware_version_start_with_v(column: str) -> col:
     column_expr = F.expr(column)
 
@@ -886,7 +886,7 @@ def firmware_version_start_with_v(column: str) -> col:
 
 # Save the YAML file with the new custom DQ rule
 sensor_custom_dq_rules_yaml = f"{quality_rules_path}/sensor_custom_dq_rules.yml"
-dq_engine.save_checks_in_workspace_file(sensor_dq_checks, workspace_path=sensor_custom_dq_rules_yaml)
+dq_engine.save_checks(sensor_dq_checks, config=WorkspaceFileChecksStorageConfig(location=sensor_custom_dq_rules_yaml))
 
 # display the link to the saved checks
 displayHTML(f'<a href="/#workspace{sensor_custom_dq_rules_yaml}" target="_blank">Sensor Custom Data Quality Rules YAML</a>')
@@ -897,8 +897,7 @@ def firmware_version_start_with_v(column: str) -> col:
 # DBTITLE 1,Apply the DQ Rules on Input Data
 dq_engine = DQEngine(WorkspaceClient())
 
-sensor_quality_checks = dq_engine.load_checks_from_workspace_file(
-    workspace_path=sensor_custom_dq_rules_yaml)
+sensor_quality_checks = dq_engine.load_checks(config=WorkspaceFileChecksStorageConfig(location=sensor_custom_dq_rules_yaml))
 
 # Define the custom check 
 custom_check_functions = {"firmware_version_start_with_v": firmware_version_start_with_v}  # list of custom check functions

docs/dqx/docs/guide/data_profiling.mdx

@@ -30,6 +30,7 @@ Data loaded as a DataFrame can be profiled to generate summary statistics and ca
     from databricks.labs.dqx.profiler.profiler import DQProfiler
     from databricks.labs.dqx.profiler.generator import DQGenerator
     from databricks.labs.dqx.profiler.dlt_generator import DQDltGenerator
+    from databricks.labs.dqx.config import WorkspaceFileChecksStorageConfig
     from databricks.labs.dqx.engine import DQEngine
     from databricks.sdk import WorkspaceClient
 
@@ -47,7 +48,7 @@ Data loaded as a DataFrame can be profiled to generate summary statistics and ca
     dq_engine = DQEngine(ws)
 
     # save checks in arbitrary workspace location
-    dq_engine.save_checks_in_workspace_file(checks, workspace_path="/Shared/App1/checks.yml")
+    dq_engine.save_checks(checks, config=WorkspaceFileChecksStorageConfig(location="/Shared/App1/checks.yml"))
 
     # generate Lakeflow Pipeline (DLT) expectations
     dlt_generator = DQDltGenerator(ws)
@@ -353,6 +354,12 @@ You can save checks defined in code or generated by the profiler to a table or f
   <TabItem value="Python" label="Python" default>
     ```python
     from databricks.labs.dqx.engine import DQEngine
+    from databricks.labs.dqx.config import (
+      FileChecksStorageConfig,
+      WorkspaceFileChecksStorageConfig,
+      InstallationChecksStorageConfig,
+      TableChecksStorageConfig
+    )
     from databricks.sdk import WorkspaceClient
 
     dq_engine = DQEngine(WorkspaceClient())
@@ -370,29 +377,29 @@ You can save checks defined in code or generated by the profiler to a table or f
 
     # save checks in a local path
     # always overwrite the file
-    dq_engine.save_checks_in_local_file(checks, path="checks.yml")
+    dq_engine.save_checks(checks, config=FileChecksStorageConfig(location="checks.yml"))
 
     # save checks in arbitrary workspace location
     # always overwrite the file
-    dq_engine.save_checks_in_workspace_file(checks, workspace_path="/Shared/App1/checks.yml")
+    dq_engine.save_checks(checks, config=WorkspaceFileChecksStorageConfig(location="/Shared/App1/checks.yml"))
 
     # save checks in file defined in 'checks_file' in the run config
     # always overwrite the file
     # only works if DQX is installed in the workspace
-    dq_engine.save_checks_in_installation(checks, method="file", assume_user=True, run_config_name="default")
+    dq_engine.save_checks(checks, config=InstallationChecksStorageConfig(assume_user=True, run_config_name="default"))
 
     # save checks in a Delta table with default run config for filtering
     # append checks in the table
-    dq_engine.save_checks_in_table(checks, table_name="dq.config.checks_table", mode="append")
+    dq_engine.save_checks(checks, config=TableChecksStorageConfig(location="dq.config.checks_table", mode="append"))
 
     # save checks in a Delta table with specific run config for filtering
     # overwrite checks in the table for the given run config
-    dq_engine.save_checks_in_table(checks, table_name="dq.config.checks_table", run_config_name="workflow_001", mode="overwrite")
+    dq_engine.save_checks(checks, config=TableChecksStorageConfig(location="dq.config.checks_table", run_config_name="workflow_001", mode="overwrite"))
 
     # save checks in table defined in 'checks_table' in the run config
     # always overwrite checks in the table for the given run config
     # only works if DQX is installed in the workspace
-    dq_engine.save_checks_in_installation(checks, method="table", assume_user=True, run_config_name="default")
+    dq_engine.save_checks(checks, config=InstallationChecksStorageConfig(assume_user=True, run_config_name="default"))
     ```
   </TabItem>
 </Tabs>