Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions HISTORY.rst
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,12 @@
History
=======

1.0.3 (2026-05-27)
------------------

* Added ``databricks`` to ``DatabaseType`` enum.
* Removed ``DatabricksDeltaS3ConnectionConfig``.

1.0.2 (2026-05-14)
------------------

Expand Down
4 changes: 2 additions & 2 deletions datamasque/client/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@
ConnectionId,
DatabaseConnectionConfig,
DatabaseType,
DatabricksDeltaS3ConnectionConfig,
DatabricksConnectionConfig,
DynamoConnectionConfig,
FileConnectionConfig,
MongoConnectionConfig,
Expand Down Expand Up @@ -129,7 +129,7 @@
"DataMasqueUserError",
"DatabaseConnectionConfig",
"DatabaseType",
"DatabricksDeltaS3ConnectionConfig",
"DatabricksConnectionConfig",
"DiscoveryMatch",
"DynamoConnectionConfig",
"FailedToStartError",
Expand Down
33 changes: 27 additions & 6 deletions datamasque/client/models/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ class DatabaseType(Enum):
snowflake = "snowflake"
mongodb = "mongodb"
databricks_lakebase = "databricks_lakebase"
databricks = "databricks"


class SnowflakeStageLocation(str, Enum):
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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`
}

Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -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" },
Expand Down
2 changes: 1 addition & 1 deletion setup.cfg
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
[bumpversion]
current_version = 1.0.2
current_version = 1.0.3
commit = True
tag = True

Expand Down
88 changes: 46 additions & 42 deletions tests/test_connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
ConnectionId,
DatabaseConnectionConfig,
DatabaseType,
DatabricksDeltaS3ConnectionConfig,
DatabricksConnectionConfig,
DynamoConnectionConfig,
MongoConnectionConfig,
MountedShareConnectionConfig,
Expand Down Expand Up @@ -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():
Expand Down
2 changes: 1 addition & 1 deletion uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading