DM-4708: add CosmosDbConnectionConfig for Azure Cosmos DB

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

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

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,18 +185,23 @@ 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:
             if not d.get("tls_ca_file"):
                 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,

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

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