Adding Lakebase checks storage backend (#550)

## Changes
<!-- Summary of your changes that are easy to understand. Add
screenshots when necessary -->

Adding Lakebase checks storage backend.

### Linked issues
<!-- DOC: Link issue with a keyword: close, closes, closed, fix, fixes,
fixed, resolve, resolves, resolved. See
https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword
-->

Resolves #444

### Tests
<!-- How is this tested? Please see the checklist below and also
describe any other relevant tests -->

- [x] manually tested
- [ ] added unit tests
- [x] added integration tests
- [ ] added end-to-end tests

---------

Co-authored-by: Copilot <[email protected]>
Co-authored-by: Marcin Wojtyczka <[email protected]>

Author: 3 people

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".

docs/dqx/docs/guide/quality_checks_apply.mdx

@@ -575,7 +575,7 @@ When running the quality checker workflow using Databricks API or UI, you have t
   - If the `checks_location` in the run config points to a table, the checks will be directly loaded from that table.
     If the `checks_location` in the run config points to a file, file name is replaced with "<input_table>.yml". In addition, if the location is specified as a relative path, it is prefixed with the workspace installation folder.
     For example:
-    - If "checks_location=catalog.schema.table", the location will be resolved to "catalog.schema.table".
+    - If "checks_location=catalog.schema.table", the location will be resolved to "catalog.schema.table" or "database.schema.table" in case of using Lakebase to store checks.
     - If "checks_location=folder/checks.yml", the  location will be resolved to "install_folder/folder/<input_table>.yml".
     - If "checks_location=/App/checks.yml", the location will be resolved to "/App/<input_table>.yml".
     - If "checks_location=/Volume/catalog/schema/folder/checks.yml", the location will be resolved to "/Volume/catalog/schema/folder/<input_table>.yml".
@@ -690,7 +690,7 @@ When running the e2e workflow using Databricks API or UI, you have the same exec
   - If the `checks_location` in the run config points to a table, the checks will be directly loaded from that table.
     If the `checks_location` in the run config points to a file, file name is replaced with <input_table>.yml. In addition, if the location is specified as a relative path, it is prefixed with the workspace installation folder.
     For example:
-    - If "checks_location=catalog.schema.table", the location will be resolved to "catalog.schema.table".
+    - If "checks_location=catalog.schema.table", the location will be resolved to "catalog.schema.table" or "database.schema.table" in case of using Lakebase to store checks.
     - If "checks_location=folder/checks.yml", the  location will be resolved to "install_folder/folder/<input_table>.yml".
     - If "checks_location=/App/checks.yml", the location will be resolved to "/App/<input_table>.yml".
     - If "checks_location=/Volume/catalog/schema/folder/checks.yml", the location will be resolved to "/Volume/catalog/schema/folder/<input_table>.yml".

docs/dqx/docs/guide/quality_checks_storage.mdx

@@ -22,12 +22,20 @@ Saving and loading methods accept a storage backend configuration as input. The
     * `mode`: (optional) write mode for saving checks (`overwrite` or `append`, default is `overwrite`). The `overwrite` mode will only replace checks for the specific run config and not all checks in the table.
 * `VolumeFileChecksStorageConfig`: Unity Catalog Volume (JSON/YAML file). Containing fields:
     * `location`: Unity Catalog Volume file path (JSON or YAML).
+* `LakebaseChecksStorageConfig`: Lakebase table. Containing fields:
+    * `instance_name`: name of the Lakebase instance, e.g., "my-instance".
+    * `user`: user to connect to the Lakebase instance, e.g., "[email protected]" or Databricks service principal client ID.
+    * `location`: fully-qualified table name in the format "database.schema.table".
+    * `port`: (optional) port on which to connect to the Lakebase instance (use 5432 if not provided).
+    * `run_config_name`: (optional) run configuration name to load (use "default" if not provided).
+    * `mode`: (optional) write mode for saving checks (`overwrite` or `append`, default is `overwrite`). The `overwrite` mode will only replace checks for the specific run config and not all checks in the table.
 * `InstallationChecksStorageConfig`: installation-managed location from the run config, ignores location and infers it from `checks_location` in the run config. Containing fields:
     * `location` (optional): automatically set based on the `checks_location` field from the run configuration.
     * `install_folder`: (optional) installation folder where DQX is installed, only required when custom installation folder is used.
     * `run_config_name` (optional) - run configuration name to load (it can be any string), e.g. input table or job name (use "default" if not provided).
     * `product_name`: (optional) name of the product (use "dqx" if not provided).
     * `assume_user`: (optional) if True, assume user installation, otherwise global installation (skipped if `install_folder` is provided).
+    * the config inherits from the specific configs such as `WorkspaceFileChecksStorageConfig`, `TableChecksStorageConfig`, `VolumeFileChecksStorageConfig`, and `LakebaseChecksStorageConfig` so relevant fields from these specific configs can be provided (e.g. instance_name and user for lakebase).
 
 You can find details on how to define checks [here](/docs/guide/quality_checks_definition).
 
@@ -49,7 +57,8 @@ If you create checks as a list of `DQRule` objects, you can convert them to meta
       WorkspaceFileChecksStorageConfig,
       InstallationChecksStorageConfig,
       TableChecksStorageConfig,
-      VolumeFileChecksStorageConfig
+      VolumeFileChecksStorageConfig,
+      LakebaseChecksStorageConfig,
     )
     from databricks.sdk import WorkspaceClient
 
@@ -81,6 +90,9 @@ If you create checks as a list of `DQRule` objects, you can convert them to meta
     # save checks as a YAML in a Unity Catalog Volume location (overwrite the file)
     dq_engine.save_checks(checks, config=VolumeFileChecksStorageConfig(location="/Volumes/dq/config/checks_volume/App1/checks.yml"))
 
+    # save checks as a Lakebase table using a Databricks service principal
+    dq_engine.save_checks(checks, config=LakebaseChecksStorageConfig(instance_name="my-instance", user="00000000-0000-0000-0000-000000000000", location="dqx.config.checks"))
+    
     # save checks as a YAML file or table defined in 'checks_location' of the run config
     # only works if DQX is installed in the workspace
     # the run config name can be any string, e.g. input table or job name
@@ -195,7 +207,8 @@ If you create checks as a list of DQRule objects, you can convert them using the
       WorkspaceFileChecksStorageConfig,
       InstallationChecksStorageConfig,
       TableChecksStorageConfig,
-      VolumeFileChecksStorageConfig
+      VolumeFileChecksStorageConfig,
+      LakebaseChecksStorageConfig,
     )
     from databricks.sdk import WorkspaceClient
 
@@ -217,6 +230,9 @@ If you create checks as a list of DQRule objects, you can convert them using the
     # load checks from a Unity Catalog Volume
     checks: list[dict] = dq_engine.load_checks(config=VolumeFileChecksStorageConfig(location="/Volumes/dq/config/checks_volume/App1/checks.yml"))
 
+    # load checks from a Lakebase table using a Databricks service principal
+    checks: list[dict] = dq_engine.load_checks(config=LakebaseChecksStorageConfig(instance_name="my-instance", user="00000000-0000-0000-0000-000000000000", location="dqx.config.checks"))
+
     # load checks from a file or table defined in the run config ('checks_location' field)
     # only works if DQX is installed in the workspace
     # the run config name is a string (e.g. input table or job name)

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., [email protected] or a Databricks service principal client ID
+  # lakebase_port: 5432                                 # <- optional port to connect to Lakebase, default is 5432
+
   custom_check_functions:                # <- optional mapping of custom check function name to Python file (module) containing check function definition
     my_func: custom_checks/my_funcs.py                  # relative workspace path (installation folder prefix applied)
     my_other: /Workspace/Shared/MyApp/my_funcs.py       # or absolute workspace path
@@ -236,7 +242,7 @@ run_configs:                        # <- list of run configurations, each run co
 
   warehouse_id: your-warehouse-id         # <- warehouse id for refreshing dashboard
 
-- name: another_run_config                # <- unique name of the run config
+- name: another_run_config                            # <- unique name of the run config
   ...
 ```
 

pyproject.toml

@@ -29,8 +29,9 @@ classifiers = [
     "Topic :: Utilities",
 ]
 dependencies = ["databricks-labs-blueprint>=0.9.1,<0.10",
-                "databricks-sdk~=0.57",
+                "databricks-sdk~=0.67",
                 "databricks-labs-lsql>=0.5,<=0.16",
+                "sqlalchemy>=1.4,<3.0",
 ]
 
 [project.optional-dependencies]
@@ -101,6 +102,7 @@ dependencies = [
     "dbldatagen~=0.4.0",
     "pyparsing~=3.2.3",
     "jmespath~=1.0.1",
+    "psycopg2-binary~=2.9.10",
 ]
 
 python="3.12" # must match the version required by databricks-connect and python version on the test clusters
@@ -111,7 +113,7 @@ path = ".venv"
 [tool.hatch.envs.default.scripts]
 test        = "pytest tests/unit/ -n 10 --cov --cov-report=xml:coverage-unit.xml --timeout 30 --durations 20"
 coverage    = "pytest tests/ -n 10 --cov --cov-report=html --timeout 600 --durations 20"
-integration = "pytest tests/integration/ -n 10 --cov --cov-report=xml --timeout 1200 --durations 20"
+integration = "pytest tests/integration/ -n 5 --cov --cov-report=xml --timeout 1200 --durations 20"
 e2e         = "pytest tests/e2e/ -n 10 --cov --cov-report=xml --timeout 600 --durations 20"
 perf         = "pytest tests/perf/ -n 10 --cov --cov-report=xml --timeout 600 --durations 20"
 fmt         = ["black .",