Update Lakebase authentication method

docs/dqx/docs/installation.mdx

@@ -280,10 +280,10 @@ 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
+  # 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_client_id: 00000000-0000-0000-0000-000000000000  # <- optional service principal client ID to use to connect to Lakebase, if not provided, the caller identity, i.e., user@domain.com, is used
+  # 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)

src/databricks/labs/dqx/checks_storage.py

@@ -18,6 +18,7 @@
     select,
     delete,
     null,
+    event,
 )
 from sqlalchemy.schema import CreateSchema
 from sqlalchemy.dialects.postgresql import JSONB
@@ -135,7 +136,12 @@ class LakebaseChecksStorageHandler(ChecksStorageHandler[LakebaseChecksStorageCon
     Handler for storing dq rules (checks) in a Lakebase table.
     """
 
-    def __init__(self, ws: WorkspaceClient, spark: SparkSession, engine: Engine | None = None):
+    def __init__(
+        self,
+        ws: WorkspaceClient,
+        spark: SparkSession,
+        engine: Engine | None = None,
+    ):
         self.ws = ws
         self.spark = spark
         self.engine = engine
@@ -152,19 +158,43 @@ def _get_connection_url(self, config: LakebaseChecksStorageConfig) -> str:
         """
         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]
-        )
         host = instance.read_write_dns
-        password = cred.token
+        prefix = "postgresql+psycopg2"
+        port = config.port
+        database = config.database_name
+        user = config.client_id if config.client_id else self.ws.current_user.me().user_name
+
+        return f"{prefix}://{user}@{host}:{port}/{database}?sslmode=require"
+
+    def _prepare_before_connect(self, config: LakebaseChecksStorageConfig):
+        """
+        Prepare the before_connect event listener with the instance_name captured in closure.
+
+        Args:
+            config: Configuration for saving and loading checks to Lakebase.
+
+        Returns:
+            The _before_connect event listener function.
 
-        return f"postgresql://{config.user}:{password}@{host}:{config.port}/{config.database_name}?sslmode=require"
+        Raises:
+            InvalidConfigError: If instance_name is not provided.
+        """
+        if not config.instance_name:
+            raise InvalidConfigError("instance_name must be provided for Lakebase storage")
+
+        instance_name = config.instance_name
+
+        def _before_connect(_dialect, _conn_rec, _cargs, cparams) -> None:
+            cred = self.ws.database.generate_database_credential(
+                request_id=str(uuid.uuid4()), instance_names=[instance_name]
+            )
+            cparams["password"] = cred.token
+
+        return _before_connect
 
     def _get_engine(self, config: LakebaseChecksStorageConfig) -> Engine:
         """
@@ -176,8 +206,15 @@ def _get_engine(self, config: LakebaseChecksStorageConfig) -> Engine:
         Returns:
             SQLAlchemy engine for the Lakebase instance.
         """
-        connection_url = self._get_connection_url(config)
-        return create_engine(connection_url)
+        engine_url = self._get_connection_url(config)
+        engine = create_engine(
+            engine_url,
+            pool_recycle=45 * 60,  # recycle connections every 45 minutes
+            connect_args={'sslmode': 'require'},
+            pool_size=4,
+        )
+        event.listen(engine, "do_connect", self._prepare_before_connect(config))
+        return engine
 
     @staticmethod
     def get_table_definition(schema_name: str, table_name: str) -> Table:
@@ -600,8 +637,8 @@ def _get_storage_handler_and_config(
             # 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
+            if run_config.lakebase_client_id and not config.client_id:
+                config.client_id = run_config.lakebase_client_id
             # 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
@@ -822,12 +859,6 @@ def create_for_run_config(self, run_config: RunConfig) -> tuple[ChecksStorageHan
             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. "
@@ -846,7 +877,6 @@ def create_for_run_config(self, run_config: RunConfig) -> tuple[ChecksStorageHan
                 LakebaseChecksStorageConfig(
                     location=run_config.checks_location,
                     instance_name=run_config.lakebase_instance_name,
-                    user=run_config.lakebase_user,
                     port=run_config.lakebase_port or "5432",
                     run_config_name=run_config.name,
                 ),

src/databricks/labs/dqx/config.py

@@ -99,7 +99,7 @@ class RunConfig:
 
     # Lakebase connection parameters, if wanting to store checks in lakebase database
     lakebase_instance_name: str | None = None
-    lakebase_user: str | None = None
+    lakebase_client_id: str | None = None
     lakebase_port: str | None = None
 
 
@@ -272,18 +272,18 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig):
     Configuration class for storing checks in a Lakebase table.
 
     Args:
-        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'.
+        instance_name: Name of the Lakebase instance.
+        client_id: ID of the Databricks service principal to use for the Lakebase connection.
         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').
     """
 
-    instance_name: str | None = None
-    user: str | None = None
     location: str
+    instance_name: str | None = None
+    client_id: str | None = None
     port: str = "5432"
     run_config_name: str = "default"
     mode: str = "overwrite"
@@ -295,9 +295,6 @@ def __post_init__(self):
         if not self.instance_name or self.instance_name == "":
             raise InvalidParameterError("Instance name must not be empty or None.")
 
-        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'."

tests/conftest.py

@@ -685,15 +685,15 @@ def delete(volume_file_path: str) -> None:
 
 
 @pytest.fixture
-def lakebase_user(ws):
+def lakebase_client_id(ws):
     """
-    Get a Lakebase user.
+    Get a Lakebase client ID.
 
     This fixture reads ARM_CLIENT_ID which is set in the CI/CD pipeline.
     For local development where the ARM_CLIENT_ID is not set, the user is fetched from the workspace client.
 
     Returns:
-        The client ID or use name to use as the Lakebase user.
+        The client ID to use for Lakebase connections.
     """
     client_id = os.environ.get("ARM_CLIENT_ID")
     if not client_id:

tests/integration/test_save_and_load_checks_from_lakebase_table.py

@@ -8,24 +8,47 @@
 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, lakebase_user, make_random):
+def test_load_checks_when_lakebase_table_does_not_exist(
+    ws, spark, make_lakebase_instance, lakebase_client_id, make_random
+):
     dq_engine = DQEngine(ws, spark)
 
     instance = make_lakebase_instance()
     lakebase_location = _create_lakebase_location(instance.database_name, make_random)
-    config = LakebaseChecksStorageConfig(location=lakebase_location, user=lakebase_user, instance_name=instance.name)
+
+    config = LakebaseChecksStorageConfig(
+        location=lakebase_location, client_id=lakebase_client_id, 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, lakebase_user, make_random):
+def test_save_and_load_checks_from_lakebase_table_with_client_id(
+    ws, spark, make_lakebase_instance, lakebase_client_id, make_random
+):
     dq_engine = DQEngine(ws, spark)
 
     instance = make_lakebase_instance()
     lakebase_location = _create_lakebase_location(instance.database_name, make_random)
 
-    config = LakebaseChecksStorageConfig(location=lakebase_location, user=lakebase_user, instance_name=instance.name)
+    config = LakebaseChecksStorageConfig(
+        location=lakebase_location, client_id=lakebase_client_id, 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(ws, spark, make_lakebase_instance, make_random):
+    dq_engine = DQEngine(ws, spark)
+
+    instance = make_lakebase_instance()
+    lakebase_location = _create_lakebase_location(instance.database_name, make_random)
+
+    config = LakebaseChecksStorageConfig(location=lakebase_location, instance_name=instance.name)
 
     dq_engine.save_checks(checks=TEST_CHECKS, config=config)
     checks = dq_engine.load_checks(config=config)
@@ -34,25 +57,31 @@ def test_save_and_load_checks_from_lakebase_table(ws, spark, make_lakebase_insta
 
 
 def test_save_and_load_checks_from_lakebase_table_with_run_config(
-    ws, spark, make_lakebase_instance, lakebase_user, make_random
+    ws, spark, make_lakebase_instance, lakebase_client_id, make_random
 ):
     dq_engine = DQEngine(ws, spark)
 
     instance = make_lakebase_instance()
     lakebase_location = _create_lakebase_location(instance.database_name, make_random)
 
-    # test first run config
+    # test the first run config
     config = LakebaseChecksStorageConfig(
-        location=lakebase_location, user=lakebase_user, instance_name=instance.name, run_config_name="workflow_001"
+        location=lakebase_location,
+        client_id=lakebase_client_id,
+        instance_name=instance.name,
+        run_config_name="workflow_001",
     )
     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
+    # test the second run config
     second_config = LakebaseChecksStorageConfig(
-        location=lakebase_location, user=lakebase_user, instance_name=instance.name, run_config_name="workflow_002"
+        location=lakebase_location,
+        client_id=lakebase_client_id,
+        instance_name=instance.name,
+        run_config_name="workflow_002",
     )
     dq_engine.save_checks(TEST_CHECKS[1:], config=second_config)
     checks = dq_engine.load_checks(config=second_config)
@@ -61,7 +90,7 @@ def test_save_and_load_checks_from_lakebase_table_with_run_config(
 
 
 def test_save_and_load_checks_from_lakebase_table_with_output_modes(
-    ws, spark, make_lakebase_instance, lakebase_user, make_random
+    ws, spark, make_lakebase_instance, lakebase_client_id, make_random
 ):
     dq_engine = DQEngine(ws, spark)
 
@@ -73,7 +102,7 @@ def test_save_and_load_checks_from_lakebase_table_with_output_modes(
         TEST_CHECKS[:1],
         config=LakebaseChecksStorageConfig(
             location=lakebase_location,
-            user=lakebase_user,
+            client_id=lakebase_client_id,
             instance_name=instance.name,
             run_config_name=run_config_name,
             mode="append",
@@ -83,7 +112,7 @@ def test_save_and_load_checks_from_lakebase_table_with_output_modes(
     checks = dq_engine.load_checks(
         config=LakebaseChecksStorageConfig(
             location=lakebase_location,
-            user=lakebase_user,
+            client_id=lakebase_client_id,
             instance_name=instance.name,
             run_config_name=run_config_name,
         )
@@ -95,7 +124,7 @@ def test_save_and_load_checks_from_lakebase_table_with_output_modes(
         TEST_CHECKS[1:],
         config=LakebaseChecksStorageConfig(
             location=lakebase_location,
-            user=lakebase_user,
+            client_id=lakebase_client_id,
             instance_name=instance.name,
             run_config_name=run_config_name,
             mode="overwrite",
@@ -104,7 +133,7 @@ def test_save_and_load_checks_from_lakebase_table_with_output_modes(
     checks = dq_engine.load_checks(
         config=LakebaseChecksStorageConfig(
             location=lakebase_location,
-            user=lakebase_user,
+            client_id=lakebase_client_id,
             instance_name=instance.name,
             run_config_name=run_config_name,
         )
@@ -113,7 +142,7 @@ def test_save_and_load_checks_from_lakebase_table_with_output_modes(
 
 
 def test_save_and_load_checks_from_lakebase_table_with_user_installation(
-    ws, spark, installation_ctx, make_lakebase_instance, lakebase_user, make_random
+    ws, spark, installation_ctx, make_lakebase_instance, lakebase_client_id, make_random
 ):
     instance = make_lakebase_instance()
 
@@ -125,7 +154,7 @@ def test_save_and_load_checks_from_lakebase_table_with_user_installation(
 
     config = InstallationChecksStorageConfig(
         instance_name=instance.name,
-        user=lakebase_user,
+        client_id=lakebase_client_id,
         run_config_name=run_config.name,
         product_name=product_name,
         assume_user=True,
@@ -140,7 +169,7 @@ def test_save_and_load_checks_from_lakebase_table_with_user_installation(
 
 
 def test_profiler_workflow_save_to_lakebase(
-    ws, spark, setup_workflows, make_lakebase_instance, lakebase_user, make_random
+    ws, spark, setup_workflows, make_lakebase_instance, lakebase_client_id, make_random
 ):
     installation_ctx, run_config = setup_workflows()
 
@@ -150,7 +179,7 @@ def test_profiler_workflow_save_to_lakebase(
     config = installation_ctx.config
     run_config = config.get_run_config()
     run_config.checks_location = lakebase_location
-    run_config.lakebase_user = lakebase_user
+    run_config.lakebase_client_id = lakebase_client_id
     run_config.lakebase_instance_name = instance.name
     run_config.lakebase_port = "5432"
 
@@ -161,7 +190,10 @@ def test_profiler_workflow_save_to_lakebase(
     dq_engine = DQEngine(ws, spark)
     checks = dq_engine.load_checks(
         config=LakebaseChecksStorageConfig(
-            location=lakebase_location, instance_name=instance.name, user=lakebase_user, run_config_name=run_config.name
+            location=lakebase_location,
+            instance_name=instance.name,
+            client_id=lakebase_client_id,
+            run_config_name=run_config.name,
         )
     )
     assert checks, "Checks are missing"