From aa18a55cb83e02f960a6881ce14230b7df465893 Mon Sep 17 00:00:00 2001 From: Aaron Barnes Date: Fri, 4 Sep 2026 15:24:01 +1200 Subject: [PATCH 1/3] feat: DM-4708: add CosmosDbConnectionConfig for Azure Cosmos DB Cosmos DB's MongoDB API is wire-compatible, so the config reuses MongoConnectionConfig and differs only by db_type and two defaults: TLS on, because Cosmos only accepts TLS, and retryable writes off, because Cosmos rejects them. Needed by the AIT framework, which builds its DataMasque connection payloads through this client. --- HISTORY.rst | 10 +++++++ datamasque/client/__init__.py | 2 ++ datamasque/client/models/connection.py | 22 +++++++++++++++ pyproject.toml | 2 +- tests/test_connections.py | 37 ++++++++++++++++++++++++++ uv.lock | 2 +- 6 files changed, 73 insertions(+), 2 deletions(-) diff --git a/HISTORY.rst b/HISTORY.rst index 063183f..bb20a53 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -2,6 +2,16 @@ History ======= +1.3.0 (2026-09-04) +------------------ + +* Added ``CosmosDbConnectionConfig`` and the ``cosmosdb`` ``DatabaseType``, for Azure Cosmos DB for + MongoDB. Cosmos DB's MongoDB API is wire-compatible, so the config reuses ``MongoConnectionConfig`` + and differs only by ``db_type`` and two defaults: ``tls`` on, because Cosmos only accepts TLS + connections, and ``retry_writes`` off, because Cosmos rejects retryable writes. + +Requires server version 3.26.18 + 1.2.5 (2026-08-17) ------------------ diff --git a/datamasque/client/__init__.py b/datamasque/client/__init__.py index 5c150fb..cd9fea7 100644 --- a/datamasque/client/__init__.py +++ b/datamasque/client/__init__.py @@ -31,6 +31,7 @@ AzureConnectionConfig, ConnectionConfig, ConnectionId, + CosmosDbConnectionConfig, DatabaseConnectionConfig, DatabaseType, DatabricksConnectionConfig, @@ -203,6 +204,7 @@ "DiscoveryConfigNotFoundError", "DiscoveryConfigType", "DiscoveryMatch", + "CosmosDbConnectionConfig", "DocumentDbConnectionConfig", "DynamoConnectionConfig", "FailedToStartError", diff --git a/datamasque/client/models/connection.py b/datamasque/client/models/connection.py index 6cf40d6..c874803 100644 --- a/datamasque/client/models/connection.py +++ b/datamasque/client/models/connection.py @@ -45,6 +45,7 @@ class DatabaseType(Enum): snowflake = "snowflake" mongodb = "mongodb" documentdb = "documentdb" + cosmosdb = "cosmosdb" databricks_lakebase = "databricks_lakebase" databricks = "databricks" informix = "informix" @@ -229,6 +230,26 @@ def database_type(self) -> DatabaseType: return DatabaseType.documentdb +class CosmosDbConnectionConfig(MongoConnectionConfig): + """ + Connection configuration for an Azure Cosmos DB for MongoDB account. + + Cosmos DB's MongoDB API is wire-compatible, + so it reuses `MongoConnectionConfig` (including the TLS handling) + and differs only by `db_type`/`database_type` and two defaults: + Cosmos mandates TLS, and rejects retryable writes on every write path. + """ + + # Narrowing the inherited Literal is a deliberate Pydantic discriminator override. + db_type: Literal["cosmosdb"] = "cosmosdb" # type: ignore[assignment] + tls: bool = True + retry_writes: bool = False + + @property + def database_type(self) -> DatabaseType: + return DatabaseType.cosmosdb + + class SnowflakeConnectionConfig(ConnectionConfig): """ Connection configuration for a Snowflake database. @@ -490,6 +511,7 @@ def _strip_encrypted_token(cls, data: dict) -> dict: DatabaseType.dynamodb.value: DynamoConnectionConfig, DatabaseType.mongodb.value: MongoConnectionConfig, DatabaseType.documentdb.value: DocumentDbConnectionConfig, + DatabaseType.cosmosdb.value: CosmosDbConnectionConfig, DatabaseType.snowflake.value: SnowflakeConnectionConfig, DatabaseType.mssql_linked.value: MssqlLinkedServerConnectionConfig, DatabaseType.databricks.value: DatabricksConnectionConfig, diff --git a/pyproject.toml b/pyproject.toml index e42524d..6a68cd1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "datamasque-python" -version = "1.2.5" +version = "1.3.0" description = "Official Python client for the DataMasque data-masking API." authors = [ { name = "DataMasque Ltd" }, diff --git a/tests/test_connections.py b/tests/test_connections.py index 8ae90cc..111a2b4 100644 --- a/tests/test_connections.py +++ b/tests/test_connections.py @@ -7,6 +7,7 @@ from datamasque.client.models.connection import ( AzureConnectionConfig, ConnectionId, + CosmosDbConnectionConfig, DatabaseConnectionConfig, DatabaseType, DatabricksConnectionConfig, @@ -1360,6 +1361,42 @@ def test_connection_config_dispatch_picks_documentdb_subclass(): assert conn.database_type is DatabaseType.documentdb +def test_cosmosdb_connection_defaults_tls_on_and_retry_writes_off(): + """ + Cosmos DB mandates TLS and rejects retryable writes, so the config carries both as defaults. + + Getting `retry_writes` wrong is silent on reads, so it fails only once a masking run writes. + """ + conn = CosmosDbConnectionConfig( + name="cosmos", + host="dtq-cosmos.mongo.cosmos.azure.com", + port=10255, + database="people", + user="dtq-cosmos", + password="hunter2", + ) + d = conn.model_dump(exclude_none=True, by_alias=True, mode="json") + assert d["db_type"] == "cosmosdb" + assert d["mask_type"] == "database" + assert d["tls"] is True + assert d["retry_writes"] is False + assert conn.database_type is DatabaseType.cosmosdb + + +def test_connection_config_dispatch_picks_cosmosdb_subclass(): + payload = { + "id": "cosmos-id-1", + "name": "cosmos", + "mask_type": "database", + "db_type": "cosmosdb", + "host": "dtq-cosmos.mongo.cosmos.azure.com", + "database": "people", + } + conn = validate_connection(payload) + assert isinstance(conn, CosmosDbConnectionConfig) + assert conn.database_type is DatabaseType.cosmosdb + + def test_database_connection_config_rejects_mongodb_database_type(): """`DatabaseConnectionConfig` is for SQL engines; MongoDB users must use `MongoConnectionConfig`.""" with pytest.raises(ValueError, match="For MongoDB"): diff --git a/uv.lock b/uv.lock index 64389ca..0d45760 100644 --- a/uv.lock +++ b/uv.lock @@ -419,7 +419,7 @@ toml = [ [[package]] name = "datamasque-python" -version = "1.2.5" +version = "1.3.0" source = { editable = "." } dependencies = [ { name = "pydantic" }, From d9d86552b146ad12812d5c2130198a4fb8b33084 Mon Sep 17 00:00:00 2001 From: Aaron Barnes Date: Sun, 6 Sep 2026 10:04:59 +1200 Subject: [PATCH 2/3] fix: DM-4708: send Cosmos DB values that differ from its own defaults The Mongo serializer pruned `tls` when falsy and `retry_writes` when truthy, which encodes MongoDB's server-side defaults. Cosmos DB inverts both, so a caller asking for `tls=False` or `retry_writes=True` had the key dropped and the server applied the opposite. Pruning now compares against the concrete class's default, which is the default the server applies for that connection type. `DatabaseConnectionConfig` also now steers documentdb and cosmosdb at their own classes, as it already did for the other special engines, and the release requires 3.26.17 rather than 3.26.18. --- HISTORY.rst | 13 ++++-- datamasque/client/models/connection.py | 23 +++++++--- tests/test_connections.py | 60 +++++++++++++++++++------- 3 files changed, 72 insertions(+), 24 deletions(-) diff --git a/HISTORY.rst b/HISTORY.rst index bb20a53..ff73ea6 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -7,10 +7,17 @@ History * Added ``CosmosDbConnectionConfig`` and the ``cosmosdb`` ``DatabaseType``, for Azure Cosmos DB for MongoDB. Cosmos DB's MongoDB API is wire-compatible, so the config reuses ``MongoConnectionConfig`` - and differs only by ``db_type`` and two defaults: ``tls`` on, because Cosmos only accepts TLS - connections, and ``retry_writes`` off, because Cosmos rejects retryable writes. + and differs only by ``db_type`` and two defaults: ``tls`` on, because Cosmos DB only accepts TLS + connections, and ``retry_writes`` off, because Cosmos DB rejects retryable writes. -Requires server version 3.26.18 +* ``DatabaseConnectionConfig`` now rejects the ``documentdb`` and ``cosmosdb`` database types with a + message naming the class to use, as it already did for MongoDB, DynamoDB, Snowflake and Databricks. + +* The MongoDB serializer prunes ``tls`` and ``retry_writes`` against the concrete class's default + rather than MongoDB's, so a Cosmos DB config that sets either the other way now sends it instead + of leaving the server to apply its own default. + +Requires server version 3.26.17 1.2.5 (2026-08-17) ------------------ diff --git a/datamasque/client/models/connection.py b/datamasque/client/models/connection.py index c874803..0efd125 100644 --- a/datamasque/client/models/connection.py +++ b/datamasque/client/models/connection.py @@ -185,8 +185,11 @@ def _serialize(self, handler: Callable) -> dict: password = d.pop("password", None) if password: d["dbpassword"] = password + # `tls` and `retry_writes` are pruned against this class's own default, not MongoDB's: + # the subclasses invert them, so a hardcoded comparison would drop the very value the + # caller set and let the server apply its default instead. + defaults = type(self).model_fields if not d.get("tls"): - d.pop("tls", None) d.pop("tls_ca_file", None) d.pop("tls_allow_invalid_certificates", None) else: @@ -194,9 +197,11 @@ def _serialize(self, handler: Callable) -> dict: d.pop("tls_ca_file", None) if not d.get("tls_allow_invalid_certificates"): d.pop("tls_allow_invalid_certificates", None) + if d.get("tls") == defaults["tls"].default: + d.pop("tls", None) if not d.get("direct_connection"): d.pop("direct_connection", None) - if d.get("retry_writes", True): + if d.get("retry_writes") == defaults["retry_writes"].default: d.pop("retry_writes", None) if not d.get("replica_set"): d.pop("replica_set", None) @@ -234,10 +239,12 @@ class CosmosDbConnectionConfig(MongoConnectionConfig): """ Connection configuration for an Azure Cosmos DB for MongoDB account. - Cosmos DB's MongoDB API is wire-compatible, - so it reuses `MongoConnectionConfig` (including the TLS handling) - and differs only by `db_type`/`database_type` and two defaults: - Cosmos mandates TLS, and rejects retryable writes on every write path. + Cosmos DB's MongoDB API is wire-compatible, so it reuses `MongoConnectionConfig` and differs + only by `db_type`/`database_type` and two defaults: Cosmos DB only accepts TLS connections, + and rejects retryable writes. + + Both are defaults, not constraints. Setting either the other way sends it to the server, which + is what makes a masking run fail at write time with retryable writes on. """ # Narrowing the inherited Literal is a deliberate Pydantic discriminator override. @@ -360,6 +367,10 @@ 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.documentdb: + raise ValueError("For AWS DocumentDB, use the DocumentDbConnectionConfig class instead") + if self.database_type is DatabaseType.cosmosdb: + raise ValueError("For Azure Cosmos DB, use the CosmosDbConnectionConfig class instead") if self.database_type is DatabaseType.databricks: raise ValueError("For Databricks SQL Warehouse, use the DatabricksConnectionConfig class instead") return self diff --git a/tests/test_connections.py b/tests/test_connections.py index 111a2b4..a14d5d6 100644 --- a/tests/test_connections.py +++ b/tests/test_connections.py @@ -1361,28 +1361,48 @@ def test_connection_config_dispatch_picks_documentdb_subclass(): assert conn.database_type is DatabaseType.documentdb -def test_cosmosdb_connection_defaults_tls_on_and_retry_writes_off(): - """ - Cosmos DB mandates TLS and rejects retryable writes, so the config carries both as defaults. - - Getting `retry_writes` wrong is silent on reads, so it fails only once a masking run writes. - """ - conn = CosmosDbConnectionConfig( +def _cosmosdb_connection(**overrides) -> CosmosDbConnectionConfig: + return CosmosDbConnectionConfig( name="cosmos", host="dtq-cosmos.mongo.cosmos.azure.com", port=10255, database="people", user="dtq-cosmos", password="hunter2", + **overrides, ) + + +def test_cosmosdb_connection_defaults_tls_on_and_retry_writes_off(): + """ + Cosmos DB only accepts TLS connections and rejects retryable writes. + + The config carries both as defaults, omitted from the payload because the server defaults + the same way for this connection type. + """ + conn = _cosmosdb_connection() d = conn.model_dump(exclude_none=True, by_alias=True, mode="json") assert d["db_type"] == "cosmosdb" assert d["mask_type"] == "database" - assert d["tls"] is True - assert d["retry_writes"] is False + assert conn.tls is True + assert conn.retry_writes is False + assert "tls" not in d + assert "retry_writes" not in d assert conn.database_type is DatabaseType.cosmosdb +def test_cosmosdb_connection_sends_values_that_differ_from_its_own_defaults(): + """ + The Mongo serializer prunes a field whose value matches the default the server would apply. + + Cosmos DB inverts two of MongoDB's defaults, so pruning against MongoDB's would drop exactly + the value the caller set and leave the server applying the opposite. + """ + d = _cosmosdb_connection(tls=False, retry_writes=True).model_dump(exclude_none=True, by_alias=True, mode="json") + assert d["tls"] is False + assert d["retry_writes"] is True + + def test_connection_config_dispatch_picks_cosmosdb_subclass(): payload = { "id": "cosmos-id-1", @@ -1397,17 +1417,27 @@ def test_connection_config_dispatch_picks_cosmosdb_subclass(): assert conn.database_type is DatabaseType.cosmosdb -def test_database_connection_config_rejects_mongodb_database_type(): - """`DatabaseConnectionConfig` is for SQL engines; MongoDB users must use `MongoConnectionConfig`.""" - with pytest.raises(ValueError, match="For MongoDB"): +@pytest.mark.parametrize( + ("database_type", "message"), + [ + (DatabaseType.mongodb, "For MongoDB"), + (DatabaseType.documentdb, "For AWS DocumentDB"), + (DatabaseType.cosmosdb, "For Azure Cosmos DB"), + ], +) +def test_database_connection_config_rejects_document_store_database_types( + database_type: DatabaseType, message: str +) -> None: + """`DatabaseConnectionConfig` is for SQL engines; each document store has its own class.""" + with pytest.raises(ValueError, match=message): DatabaseConnectionConfig( - name="mongo", - host="mongo.example", + name="doc-store", + host="doc-store.example", port=27017, database="people", user="alice", password="hunter2", - database_type=DatabaseType.mongodb, + database_type=database_type, ) From 5e1448c412deead57bba3fb61bbd7137625b9ab4 Mon Sep 17 00:00:00 2001 From: Aaron Barnes Date: Mon, 7 Sep 2026 10:07:23 +1200 Subject: [PATCH 3/3] refactor: DM-4708: name the config Azure Cosmos DB and default its port The class described itself as "for MongoDB", which names the API rather than the product. The port now defaults to 10255, so a caller no longer has to know it. --- HISTORY.rst | 8 ++++---- datamasque/client/models/connection.py | 12 ++++++------ tests/test_connections.py | 12 +++++++----- 3 files changed, 17 insertions(+), 15 deletions(-) diff --git a/HISTORY.rst b/HISTORY.rst index ff73ea6..940363d 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -5,10 +5,10 @@ History 1.3.0 (2026-09-04) ------------------ -* Added ``CosmosDbConnectionConfig`` and the ``cosmosdb`` ``DatabaseType``, for Azure Cosmos DB for - MongoDB. Cosmos DB's MongoDB API is wire-compatible, so the config reuses ``MongoConnectionConfig`` - and differs only by ``db_type`` and two defaults: ``tls`` on, because Cosmos DB only accepts TLS - connections, and ``retry_writes`` off, because Cosmos DB rejects retryable writes. +* Added ``CosmosDbConnectionConfig`` and the ``cosmosdb`` ``DatabaseType``, for Azure Cosmos DB. + It differs from ``MongoConnectionConfig`` by ``db_type`` and three defaults: ``port`` 10255, + ``tls`` on, because Cosmos DB only accepts TLS connections, and ``retry_writes`` off, because + Cosmos DB rejects retryable writes. * ``DatabaseConnectionConfig`` now rejects the ``documentdb`` and ``cosmosdb`` database types with a message naming the class to use, as it already did for MongoDB, DynamoDB, Snowflake and Databricks. diff --git a/datamasque/client/models/connection.py b/datamasque/client/models/connection.py index 0efd125..8fd5218 100644 --- a/datamasque/client/models/connection.py +++ b/datamasque/client/models/connection.py @@ -237,18 +237,18 @@ def database_type(self) -> DatabaseType: class CosmosDbConnectionConfig(MongoConnectionConfig): """ - Connection configuration for an Azure Cosmos DB for MongoDB account. + Connection configuration for an Azure Cosmos DB account. - Cosmos DB's MongoDB API is wire-compatible, so it reuses `MongoConnectionConfig` and differs - only by `db_type`/`database_type` and two defaults: Cosmos DB only accepts TLS connections, - and rejects retryable writes. + Cosmos DB listens on 10255, only accepts TLS connections and rejects retryable writes, so it + differs from `MongoConnectionConfig` by `db_type`/`database_type` and those three defaults. - Both are defaults, not constraints. Setting either the other way sends it to the server, which - is what makes a masking run fail at write time with retryable writes on. + They are defaults, not constraints. Setting one the other way sends it to the server, which is + what makes a masking run fail at write time with retryable writes on. """ # Narrowing the inherited Literal is a deliberate Pydantic discriminator override. db_type: Literal["cosmosdb"] = "cosmosdb" # type: ignore[assignment] + port: int = 10255 tls: bool = True retry_writes: bool = False diff --git a/tests/test_connections.py b/tests/test_connections.py index a14d5d6..ed894f9 100644 --- a/tests/test_connections.py +++ b/tests/test_connections.py @@ -1362,10 +1362,10 @@ def test_connection_config_dispatch_picks_documentdb_subclass(): def _cosmosdb_connection(**overrides) -> CosmosDbConnectionConfig: + """Build a Cosmos DB config, leaving the port to the connection type's own default.""" return CosmosDbConnectionConfig( name="cosmos", host="dtq-cosmos.mongo.cosmos.azure.com", - port=10255, database="people", user="dtq-cosmos", password="hunter2", @@ -1373,17 +1373,19 @@ def _cosmosdb_connection(**overrides) -> CosmosDbConnectionConfig: ) -def test_cosmosdb_connection_defaults_tls_on_and_retry_writes_off(): +def test_cosmosdb_connection_defaults_port_tls_and_retry_writes(): """ - Cosmos DB only accepts TLS connections and rejects retryable writes. + Cosmos DB listens on 10255, only accepts TLS connections and rejects retryable writes. - The config carries both as defaults, omitted from the payload because the server defaults - the same way for this connection type. + The config carries all three as defaults. `tls` and `retry_writes` are omitted from the payload + because the server defaults the same way for this connection type; the port is always sent. """ conn = _cosmosdb_connection() d = conn.model_dump(exclude_none=True, by_alias=True, mode="json") assert d["db_type"] == "cosmosdb" assert d["mask_type"] == "database" + assert conn.port == 10255 + assert d["port"] == 10255 assert conn.tls is True assert conn.retry_writes is False assert "tls" not in d