diff --git a/HISTORY.rst b/HISTORY.rst index 503a2a9..f91908d 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -2,6 +2,12 @@ History ======= +1.0.3 (2026-05-27) +------------------ + +* Added ``databricks`` to ``DatabaseType`` enum. +* Removed ``DatabricksDeltaS3ConnectionConfig``. + 1.0.2 (2026-05-14) ------------------ diff --git a/datamasque/client/__init__.py b/datamasque/client/__init__.py index 6493e5a..c4f007d 100644 --- a/datamasque/client/__init__.py +++ b/datamasque/client/__init__.py @@ -30,7 +30,7 @@ ConnectionId, DatabaseConnectionConfig, DatabaseType, - DatabricksDeltaS3ConnectionConfig, + DatabricksConnectionConfig, DynamoConnectionConfig, FileConnectionConfig, MongoConnectionConfig, @@ -129,7 +129,7 @@ "DataMasqueUserError", "DatabaseConnectionConfig", "DatabaseType", - "DatabricksDeltaS3ConnectionConfig", + "DatabricksConnectionConfig", "DiscoveryMatch", "DynamoConnectionConfig", "FailedToStartError", diff --git a/datamasque/client/models/connection.py b/datamasque/client/models/connection.py index 5b3721a..815f1da 100644 --- a/datamasque/client/models/connection.py +++ b/datamasque/client/models/connection.py @@ -45,6 +45,7 @@ class DatabaseType(Enum): snowflake = "snowflake" mongodb = "mongodb" databricks_lakebase = "databricks_lakebase" + databricks = "databricks" class SnowflakeStageLocation(str, Enum): @@ -280,6 +281,8 @@ def _reject_special_engines(self) -> "DatabaseConnectionConfig": raise ValueError("For Snowflake, use the SnowflakeConnectionConfig class instead") if self.database_type is DatabaseType.mongodb: raise ValueError("For MongoDB, use the MongoConnectionConfig class instead") + if self.database_type is DatabaseType.databricks: + raise ValueError("For Databricks SQL Warehouse, use the DatabricksConnectionConfig class instead") return self mask_type: Literal["database"] = "database" @@ -391,19 +394,36 @@ class MountedShareConnectionConfig(FileConnectionConfig): type: Literal["mounted_share_connection"] = "mounted_share_connection" -class DatabricksDeltaS3ConnectionConfig(FileConnectionConfig): - """Connection configuration for Databricks Delta tables stored in S3.""" +class DatabricksConnectionConfig(ConnectionConfig): + """Connection configuration for a Databricks SQL Warehouse.""" - type: Literal["databricks_delta_s3_connection"] = "databricks_delta_s3_connection" - bucket: str = "" - iam_role_arn: Optional[str] = None + server_hostname: str + http_path: str + access_token: Optional[str] = None + catalog: str + db_schema: Optional[str] = Field(default=None, alias="schema") + is_read_only: bool = False + version: str = "1.0" + + mask_type: Literal["database"] = "database" + db_type: Literal["databricks"] = "databricks" + + @property + def database_type(self) -> DatabaseType: + return DatabaseType.databricks + + @model_validator(mode="before") + @classmethod + def _strip_encrypted_token(cls, data: dict) -> dict: + if isinstance(data, dict): + data.pop("access_token_encrypted", None) + return data FILE_TYPE_MAP: dict[str, type[FileConnectionConfig]] = { "s3_connection": S3ConnectionConfig, "azure_blob_connection": AzureConnectionConfig, "mounted_share_connection": MountedShareConnectionConfig, - "databricks_delta_s3_connection": DatabricksDeltaS3ConnectionConfig, } DB_TYPE_MAP: dict[str, type[ConnectionConfig]] = { @@ -411,6 +431,7 @@ class DatabricksDeltaS3ConnectionConfig(FileConnectionConfig): DatabaseType.mongodb.value: MongoConnectionConfig, DatabaseType.snowflake.value: SnowflakeConnectionConfig, DatabaseType.mssql_linked.value: MssqlLinkedServerConnectionConfig, + DatabaseType.databricks.value: DatabricksConnectionConfig, # others use the default `DatabaseConnectionConfig` } diff --git a/pyproject.toml b/pyproject.toml index d95affa..421475a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "datamasque-python" -version = "1.0.2" +version = "1.0.3" description = "Official Python client for the DataMasque data-masking API." authors = [ { name = "DataMasque Ltd" }, diff --git a/setup.cfg b/setup.cfg index e24c7f9..8a4ed42 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,5 +1,5 @@ [bumpversion] -current_version = 1.0.2 +current_version = 1.0.3 commit = True tag = True diff --git a/tests/test_connections.py b/tests/test_connections.py index 4546228..367aed5 100644 --- a/tests/test_connections.py +++ b/tests/test_connections.py @@ -9,7 +9,7 @@ ConnectionId, DatabaseConnectionConfig, DatabaseType, - DatabricksDeltaS3ConnectionConfig, + DatabricksConnectionConfig, DynamoConnectionConfig, MongoConnectionConfig, MountedShareConnectionConfig, @@ -696,63 +696,67 @@ def test_s3_connection_model_validate_no_iam_role(): assert conn.iam_role_arn is None -def test_databricks_delta_s3_connection_model_validate(): +def test_databricks_connection_model_validate(): payload = { - "id": "11223344-5566-7788-99aa-bbccddeeff00", - "name": "delta_s3", - "mask_type": "file", - "type": "databricks_delta_s3_connection", - "base_directory": "delta/", - "is_file_mask_source": True, - "is_file_mask_destination": False, - "bucket": "my-delta-bucket", - "iam_role_arn": "arn:aws:iam::111122223333:role/delta-role", + "id": "db-id-1", + "name": "databricks", + "mask_type": "database", + "db_type": "databricks", + "server_hostname": "adb-1234.azuredatabricks.net", + "http_path": "/sql/1.0/warehouses/abcd1234", + "access_token": "dapi1234", + "catalog": "main", + "schema": "default", + "is_read_only": False, } - conn = DatabricksDeltaS3ConnectionConfig.model_validate(payload) + conn = DatabricksConnectionConfig.model_validate(payload) - assert isinstance(conn, DatabricksDeltaS3ConnectionConfig) - assert conn.id == "11223344-5566-7788-99aa-bbccddeeff00" - assert conn.name == "delta_s3" - assert conn.bucket == "my-delta-bucket" - assert conn.base_directory == "delta/" - assert conn.is_file_mask_source is True - assert conn.is_file_mask_destination is False - assert conn.iam_role_arn == "arn:aws:iam::111122223333:role/delta-role" + assert isinstance(conn, DatabricksConnectionConfig) + assert conn.id == "db-id-1" + assert conn.server_hostname == "adb-1234.azuredatabricks.net" + assert conn.http_path == "/sql/1.0/warehouses/abcd1234" + assert conn.access_token == "dapi1234" + assert conn.catalog == "main" + assert conn.db_schema == "default" + assert conn.database_type is DatabaseType.databricks -def test_databricks_delta_s3_connection_model_validate_no_iam_role(): +def test_databricks_connection_model_validate_blanks_encrypted_token(): payload = { - "id": "id-delta", - "name": "delta_s3", - "mask_type": "file", - "type": "databricks_delta_s3_connection", - "base_directory": "", - "is_file_mask_source": True, - "is_file_mask_destination": False, - "bucket": "my-delta-bucket", + "id": "db-id-2", + "name": "databricks", + "mask_type": "database", + "db_type": "databricks", + "server_hostname": "adb-1234.azuredatabricks.net", + "http_path": "/sql/1.0/warehouses/abcd1234", + "access_token_encrypted": "some_base64_here", + "catalog": "main", + "is_read_only": False, } - conn = DatabricksDeltaS3ConnectionConfig.model_validate(payload) - assert conn.iam_role_arn is None + conn = DatabricksConnectionConfig.model_validate(payload) + + assert isinstance(conn, DatabricksConnectionConfig) + assert conn.access_token is None -def test_validate_connection_dispatches_databricks_delta_s3(): +def test_validate_connection_dispatches_databricks(): payload = { - "id": "aabb-ccdd", - "name": "delta", - "mask_type": "file", - "type": "databricks_delta_s3_connection", - "base_directory": "", - "is_file_mask_source": False, - "is_file_mask_destination": True, - "bucket": "delta-bucket", + "id": "db-id-3", + "name": "databricks", + "mask_type": "database", + "db_type": "databricks", + "server_hostname": "adb-1234.azuredatabricks.net", + "http_path": "/sql/1.0/warehouses/abcd1234", + "catalog": "main", + "is_read_only": False, } conn = validate_connection(payload) - assert isinstance(conn, DatabricksDeltaS3ConnectionConfig) - assert conn.bucket == "delta-bucket" + assert isinstance(conn, DatabricksConnectionConfig) + assert conn.catalog == "main" def test_azure_connection_model_validate_blanks_encrypted_connection_string(): diff --git a/uv.lock b/uv.lock index ab60d30..6fe9d57 100644 --- a/uv.lock +++ b/uv.lock @@ -428,7 +428,7 @@ toml = [ [[package]] name = "datamasque-python" -version = "1.0.2" +version = "1.0.3" source = { editable = "." } dependencies = [ { name = "pydantic" },