diff --git a/docs/dqx/docs/installation.mdx b/docs/dqx/docs/installation.mdx index cabf2138b..020fd6601 100644 --- a/docs/dqx/docs/installation.mdx +++ b/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) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index f39e0528f..c5f248709 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/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, ), diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index db727a24a..ff58f72e8 100644 --- a/src/databricks/labs/dqx/config.py +++ b/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'." diff --git a/tests/conftest.py b/tests/conftest.py index 254017117..8accb5731 100644 --- a/tests/conftest.py +++ b/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: diff --git a/tests/integration/test_save_and_load_checks_from_lakebase_table.py b/tests/integration/test_save_and_load_checks_from_lakebase_table.py index f0a44eee9..0eb91b672 100644 --- a/tests/integration/test_save_and_load_checks_from_lakebase_table.py +++ b/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" diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 770ca22fb..ea8719836 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -299,29 +299,31 @@ def test_workspace_config_as_dict(): # Test LakebaseChecksStorageConfig validation and properties def test_lakebase_config_missing_instance_name(): with pytest.raises(InvalidParameterError, match="Instance name must not be empty or None"): - LakebaseChecksStorageConfig(location="db.schema.table", instance_name=None, user="test_user") - - -def test_lakebase_config_missing_user(): - with pytest.raises(InvalidParameterError, match="User must not be empty or None"): - LakebaseChecksStorageConfig(location="db.schema.table", instance_name="instance", user=None) + LakebaseChecksStorageConfig( + location="db.schema.table", instance_name=None, client_id="00000000-0000-0000-0000-000000000000" + ) def test_lakebase_config_invalid_location_format(): with pytest.raises(InvalidConfigError, match="Invalid Lakebase table name.*Must be in the format"): - LakebaseChecksStorageConfig(location="invalid_table", instance_name="instance", user="test_user") + LakebaseChecksStorageConfig( + location="invalid_table", instance_name="instance", client_id="00000000-0000-0000-0000-000000000000" + ) def test_lakebase_config_invalid_mode(): with pytest.raises(InvalidConfigError, match="Invalid mode.*Must be 'append' or 'overwrite'"): LakebaseChecksStorageConfig( - location="db.schema.table", instance_name="instance", user="test_user", mode="invalid" + location="db.schema.table", + instance_name="instance", + client_id="00000000-0000-0000-0000-000000000000", + mode="invalid", ) def test_lakebase_config_properties(): config = LakebaseChecksStorageConfig( - location="my_db.my_schema.my_table", instance_name="instance", user="test_user" + location="my_db.my_schema.my_table", instance_name="instance", client_id="00000000-0000-0000-0000-000000000000" ) assert config.database_name == "my_db" assert config.schema_name == "my_schema" @@ -329,7 +331,9 @@ def test_lakebase_config_properties(): def test_lakebase_config_properties_cached(): - config = LakebaseChecksStorageConfig(location="db1.sch1.tbl1", instance_name="instance", user="test_user") + config = LakebaseChecksStorageConfig( + location="db1.sch1.tbl1", instance_name="instance", client_id="00000000-0000-0000-0000-000000000000" + ) # Access properties multiple times to ensure caching works assert config.database_name == "db1" assert config.database_name == "db1" @@ -341,12 +345,18 @@ def test_lakebase_config_properties_cached(): def test_lakebase_config_valid_modes(): config_append = LakebaseChecksStorageConfig( - location="db.schema.table", instance_name="instance", user="test_user", mode="append" + location="db.schema.table", + instance_name="instance", + client_id="00000000-0000-0000-0000-000000000000", + mode="append", ) assert config_append.mode == "append" config_overwrite = LakebaseChecksStorageConfig( - location="db.schema.table", instance_name="instance", user="test_user", mode="overwrite" + location="db.schema.table", + instance_name="instance", + client_id="00000000-0000-0000-0000-000000000000", + mode="overwrite", ) assert config_overwrite.mode == "overwrite" @@ -437,7 +447,7 @@ def test_run_config_defaults(): assert not config.reference_tables assert not config.custom_check_functions assert config.lakebase_instance_name is None - assert config.lakebase_user is None + assert config.lakebase_client_id is None assert config.lakebase_port is None diff --git a/tests/unit/test_lakebase_config.py b/tests/unit/test_lakebase_config.py index 7787db2c0..ca23fc56b 100644 --- a/tests/unit/test_lakebase_config.py +++ b/tests/unit/test_lakebase_config.py @@ -8,45 +8,47 @@ def test_create_valid_lakebase_check_storage_config(): """Test creating LakebaseConnectionConfig with valid parameters.""" config = LakebaseChecksStorageConfig( - instance_name="testhost", location="testdb.testschema.testtable", user="testuser", port="5432" + instance_name="testhost", + location="testdb.testschema.testtable", + client_id="00000000-0000-0000-0000-000000000000", + port="5432", ) assert config.instance_name == "testhost" assert config.database_name == "testdb" - assert config.user == "testuser" + assert config.client_id == "00000000-0000-0000-0000-000000000000" assert config.port == "5432" def test_create_lakebase_check_storage_config_default_port(): """Test creating LakebaseConnectionConfig with default port.""" config = LakebaseChecksStorageConfig( - instance_name="testhost", location="testdb.testschema.testtable", user="testuser" + instance_name="testhost", + location="testdb.testschema.testtable", + client_id="00000000-0000-0000-0000-000000000000", ) assert config.instance_name == "testhost" assert config.database_name == "testdb" - assert config.user == "testuser" + assert config.client_id == "00000000-0000-0000-0000-000000000000" assert config.port == "5432" # Default port @pytest.mark.parametrize( - "instance_name,location,user,port,expected_error", + "instance_name,location,client_id,port,expected_error", [ # Empty/None instance_name - ("", "testdb.testschema.testtable", "user", "5432", "Instance name must not be empty or None."), - (None, "testdb.testschema.testtable", "user", "5432", "Instance name must not be empty or None."), + ("", "testdb.testschema.testtable", "client_id", "5432", "Instance name must not be empty or None."), + (None, "testdb.testschema.testtable", "client_id", "5432", "Instance name must not be empty or None."), # Empty/None database - ("instance", "", "user", "5432", "Location must not be empty or None."), - ("instance", None, "user", "5432", "Location must not be empty or None."), - # Empty/None user - ("instance", "testdb.testschema.testtable", "", "5432", "User must not be empty or None."), - ("instance", "testdb.testschema.testtable", None, "5432", "User must not be empty or None."), + ("instance", "", "client_id", "5432", "Location must not be empty or None."), + ("instance", None, "client_id", "5432", "Location must not be empty or None."), ], ) -def test_create_invalid_lakebase_check_storage_config(instance_name, location, user, port, expected_error): +def test_create_invalid_lakebase_check_storage_config(instance_name, location, client_id, port, expected_error): """Test LakebaseChecksStorageConfig validation with invalid parameters.""" with pytest.raises(InvalidParameterError, match=expected_error): - LakebaseChecksStorageConfig(instance_name=instance_name, location=location, user=user, port=port) + LakebaseChecksStorageConfig(instance_name=instance_name, location=location, client_id=client_id, port=port) def test_valid_lakebase_checks_storage_config(): @@ -54,7 +56,7 @@ def test_valid_lakebase_checks_storage_config(): config = LakebaseChecksStorageConfig( instance_name="testinstance", location="testdb.testschema.testtable", - user="testuser", + client_id="00000000-0000-0000-0000-000000000000", port="5432", run_config_name="test_run", mode="append", @@ -62,7 +64,7 @@ def test_valid_lakebase_checks_storage_config(): assert config.instance_name == "testinstance" assert config.location == "testdb.testschema.testtable" - assert config.user == "testuser" + assert config.client_id == "00000000-0000-0000-0000-000000000000" assert config.port == "5432" assert config.run_config_name == "test_run" assert config.mode == "append" @@ -74,7 +76,9 @@ def test_valid_lakebase_checks_storage_config(): def test_valid_lakebase_checks_storage_config_defaults(): """Test creating LakebaseChecksStorageConfig with default values.""" config = LakebaseChecksStorageConfig( - instance_name="testinstance", location="testdb.testschema.testtable", user="testuser" + instance_name="testinstance", + location="testdb.testschema.testtable", + client_id="00000000-0000-0000-0000-000000000000", ) assert config.run_config_name == "default" @@ -93,7 +97,9 @@ def test_valid_lakebase_checks_storage_config_defaults(): def test_invalid_location_empty(location, expected_error): """Test LakebaseChecksStorageConfig with empty/None location.""" with pytest.raises(InvalidParameterError, match=expected_error): - LakebaseChecksStorageConfig(instance_name="testinstance", location=location, user="testuser") + LakebaseChecksStorageConfig( + instance_name="testinstance", location=location, client_id="00000000-0000-0000-0000-000000000000" + ) @pytest.mark.parametrize( @@ -111,7 +117,9 @@ def test_invalid_location_empty(location, expected_error): def test_invalid_location_format(location, expected_error): """Test LakebaseChecksStorageConfig with invalid location format.""" with pytest.raises(InvalidConfigError, match=expected_error): - LakebaseChecksStorageConfig(instance_name="testinstance", location=location, user="testuser") + LakebaseChecksStorageConfig( + instance_name="testinstance", location=location, client_id="00000000-0000-0000-0000-000000000000" + ) @pytest.mark.parametrize( @@ -125,7 +133,9 @@ def test_invalid_instance_name(instance_name, expected_error): """Test LakebaseChecksStorageConfig with invalid instance name.""" with pytest.raises(InvalidParameterError, match=expected_error): LakebaseChecksStorageConfig( - instance_name=instance_name, location="testdb.testschema.testtable", user="testuser" + instance_name=instance_name, + location="testdb.testschema.testtable", + client_id="00000000-0000-0000-0000-000000000000", ) @@ -144,7 +154,10 @@ def test_invalid_mode(mode, expected_error): """Test LakebaseChecksStorageConfig with invalid mode.""" with pytest.raises(InvalidConfigError, match=expected_error): LakebaseChecksStorageConfig( - instance_name="testinstance", location="testdb.testschema.testtable", user="testuser", mode=mode + instance_name="testinstance", + location="testdb.testschema.testtable", + client_id="00000000-0000-0000-0000-000000000000", + mode=mode, ) @@ -152,13 +165,19 @@ def test_valid_modes(): """Test LakebaseChecksStorageConfig with valid modes.""" # Test 'append' mode config_append = LakebaseChecksStorageConfig( - instance_name="testinstance", location="testdb.testschema.testtable", user="testuser", mode="append" + instance_name="testinstance", + location="testdb.testschema.testtable", + client_id="00000000-0000-0000-0000-000000000000", + mode="append", ) assert config_append.mode == "append" # Test 'overwrite' mode config_overwrite = LakebaseChecksStorageConfig( - instance_name="testinstance", location="testdb.testschema.testtable", user="testuser", mode="overwrite" + instance_name="testinstance", + location="testdb.testschema.testtable", + client_id="00000000-0000-0000-0000-000000000000", + mode="overwrite", ) assert config_overwrite.mode == "overwrite" @@ -166,7 +185,9 @@ def test_valid_modes(): def test_location_properties(): """Test database_name, schema_name, and table_name properties.""" config = LakebaseChecksStorageConfig( - instance_name="testinstance", location="testdb.testschema.testtable", user="testuser" + instance_name="testinstance", + location="testdb.testschema.testtable", + client_id="00000000-0000-0000-0000-000000000000", ) assert config.database_name == "testdb" @@ -177,7 +198,9 @@ def test_location_properties(): def test_complex_location_with_special_characters(): """Test location parsing with special characters in names.""" config = LakebaseChecksStorageConfig( - instance_name="testinstance", location="test_db.test_schema.test_table", user="testuser" + instance_name="testinstance", + location="test_db.test_schema.test_table", + client_id="00000000-0000-0000-0000-000000000000", ) assert config.database_name == "test_db" @@ -189,11 +212,13 @@ def test_eager_validation_on_init(): """Test that validation happens eagerly during __post_init__.""" # This should fail immediately during construction, not later with pytest.raises(InvalidParameterError, match="Instance name must not be empty or None"): - LakebaseChecksStorageConfig(instance_name="", location="testdb.testschema.testtable", user="testuser") + LakebaseChecksStorageConfig( + instance_name="", location="testdb.testschema.testtable", client_id="00000000-0000-0000-0000-000000000000" + ) def test_multiple_validation_errors_precedence(): """Test that the first validation error is raised when multiple issues exist.""" # Empty location should be caught first, before instance name validation with pytest.raises(InvalidParameterError, match="Location must not be empty or None"): - LakebaseChecksStorageConfig(instance_name="", location="", user="testuser") + LakebaseChecksStorageConfig(instance_name="", location="", client_id="00000000-0000-0000-0000-000000000000")