diff --git a/HISTORY.rst b/HISTORY.rst index 063183f..940363d 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -2,6 +2,23 @@ History ======= +1.3.0 (2026-09-04) +------------------ + +* 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. + +* 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/__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..8fd5218 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" @@ -184,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: @@ -193,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) @@ -229,6 +235,28 @@ def database_type(self) -> DatabaseType: return DatabaseType.documentdb +class CosmosDbConnectionConfig(MongoConnectionConfig): + """ + Connection configuration for an Azure Cosmos DB account. + + 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. + + 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 + + @property + def database_type(self) -> DatabaseType: + return DatabaseType.cosmosdb + + class SnowflakeConnectionConfig(ConnectionConfig): """ Connection configuration for a Snowflake database. @@ -339,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 @@ -490,6 +522,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..ed894f9 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,17 +1361,85 @@ def test_connection_config_dispatch_picks_documentdb_subclass(): assert conn.database_type is DatabaseType.documentdb -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"): +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", + database="people", + user="dtq-cosmos", + password="hunter2", + **overrides, + ) + + +def test_cosmosdb_connection_defaults_port_tls_and_retry_writes(): + """ + Cosmos DB listens on 10255, only accepts TLS connections and rejects retryable writes. + + 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 + 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", + "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 + + +@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, ) 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" },