Merge pull request #8 from jiatolentino/DM-3293-databricks-support

feat: Add databricks to DatabaseType enum

Author: jiatolentino

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)
 ------------------
 

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

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,26 +394,44 @@ 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]] = {
     DatabaseType.dynamodb.value: DynamoConnectionConfig,
     DatabaseType.mongodb.value: MongoConnectionConfig,
     DatabaseType.snowflake.value: SnowflakeConnectionConfig,
     DatabaseType.mssql_linked.value: MssqlLinkedServerConnectionConfig,
+    DatabaseType.databricks.value: DatabricksConnectionConfig,
     # others use the default `DatabaseConnectionConfig`
 }
 

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" },

setup.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 1.0.2
+current_version = 1.0.3
 commit = True
 tag = True
 

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():