Merge pull request #29 from datamasque/feat/add-mongodb-tls-and-documentdb-connection

feat: MongoDB TLS CA fields + DocumentDB connection type

Author: jiatolentino

HISTORY.rst

@@ -2,6 +2,15 @@
 History
 =======
 
+1.1.3 (2026-06-26)
+------------------
+
+* Added ``tls_ca_file`` and ``tls_allow_invalid_certificates`` to ``MongoConnectionConfig``,
+  serialized only when TLS is enabled (``tls_ca_file`` only when set,
+  ``tls_allow_invalid_certificates`` only when ``True``).
+* Added ``DatabaseType.documentdb`` and ``DocumentDbConnectionConfig`` for AWS DocumentDB,
+  which is MongoDB wire-compatible and reuses ``MongoConnectionConfig`` (differing only by ``db_type``).
+
 1.1.2 (2026-06-26)
 ------------------
 

datamasque/client/__init__.py

@@ -33,6 +33,7 @@
     DatabaseConnectionConfig,
     DatabaseType,
     DatabricksConnectionConfig,
+    DocumentDbConnectionConfig,
     DynamoConnectionConfig,
     FileConnectionConfig,
     MongoConnectionConfig,
@@ -150,6 +151,7 @@
     "DiscoveryConfigNotFoundError",
     "DiscoveryConfigType",
     "DiscoveryMatch",
+    "DocumentDbConnectionConfig",
     "DynamoConnectionConfig",
     "FailedToStartError",
     "FileConnectionConfig",

datamasque/client/models/connection.py

@@ -44,6 +44,7 @@ class DatabaseType(Enum):
     mssql_linked = "mssql_linked"
     snowflake = "snowflake"
     mongodb = "mongodb"
+    documentdb = "documentdb"
     databricks_lakebase = "databricks_lakebase"
     databricks = "databricks"
     informix = "informix"
@@ -160,6 +161,8 @@ class MongoConnectionConfig(ConnectionConfig):
     password: Optional[str] = None
     auth_source: str = "admin"
     tls: bool = False
+    tls_ca_file: str = ""
+    tls_allow_invalid_certificates: bool = False
     direct_connection: bool = False
     replica_set: str = ""
     is_read_only: bool = False
@@ -180,6 +183,13 @@ def _serialize(self, handler: Callable) -> dict:
             d["dbpassword"] = password
         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 not d.get("direct_connection"):
             d.pop("direct_connection", None)
         if not d.get("replica_set"):
@@ -197,6 +207,23 @@ def _strip_encrypted_password(cls, data: dict) -> dict:
         return data
 
 
+class DocumentDbConnectionConfig(MongoConnectionConfig):
+    """
+    Connection configuration for an AWS DocumentDB cluster.
+
+    DocumentDB is MongoDB wire-compatible,
+    so it reuses `MongoConnectionConfig` (including the TLS handling)
+    and only differs by `db_type`/`database_type`.
+    """
+
+    # Narrowing the inherited Literal is a deliberate Pydantic discriminator override.
+    db_type: Literal["documentdb"] = "documentdb"  # type: ignore[assignment]
+
+    @property
+    def database_type(self) -> DatabaseType:
+        return DatabaseType.documentdb
+
+
 class SnowflakeConnectionConfig(ConnectionConfig):
     """
     Connection configuration for a Snowflake database.
@@ -430,6 +457,7 @@ def _strip_encrypted_token(cls, data: dict) -> dict:
 DB_TYPE_MAP: dict[str, type[ConnectionConfig]] = {
     DatabaseType.dynamodb.value: DynamoConnectionConfig,
     DatabaseType.mongodb.value: MongoConnectionConfig,
+    DatabaseType.documentdb.value: DocumentDbConnectionConfig,
     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.1.2"
+version = "1.1.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.1.2
+current_version = 1.1.3
 commit = True
 tag = True
 

tests/test_connections.py

@@ -10,6 +10,7 @@
     DatabaseConnectionConfig,
     DatabaseType,
     DatabricksConnectionConfig,
+    DocumentDbConnectionConfig,
     DynamoConnectionConfig,
     MongoConnectionConfig,
     MountedShareConnectionConfig,
@@ -1201,6 +1202,111 @@ def test_connection_config_dispatch_picks_mongo_subclass():
     assert isinstance(validate_connection(payload), MongoConnectionConfig)
 
 
+def test_mongo_tls_ca_fields_emitted_when_tls_enabled():
+    conn = MongoConnectionConfig(
+        name="mongo",
+        host="mongo.example",
+        database="people",
+        tls=True,
+        tls_ca_file="/certs/rds-ca.pem",
+        tls_allow_invalid_certificates=True,
+    )
+    d = conn.model_dump(exclude_none=True, by_alias=True, mode="json")
+    assert d["tls"] is True
+    assert d["tls_ca_file"] == "/certs/rds-ca.pem"
+    assert d["tls_allow_invalid_certificates"] is True
+
+
+def test_mongo_tls_ca_fields_omitted_when_unset_but_tls_enabled():
+    conn = MongoConnectionConfig(name="mongo", host="mongo.example", database="people", tls=True)
+    d = conn.model_dump(exclude_none=True, by_alias=True, mode="json")
+    assert d["tls"] is True
+    assert "tls_ca_file" not in d
+    assert "tls_allow_invalid_certificates" not in d
+
+
+def test_mongo_tls_ca_fields_omitted_when_tls_disabled():
+    """Even when set, the TLS-dependent fields are not sent while TLS is disabled."""
+    conn = MongoConnectionConfig(
+        name="mongo",
+        host="mongo.example",
+        database="people",
+        tls=False,
+        tls_ca_file="/certs/rds-ca.pem",
+        tls_allow_invalid_certificates=True,
+    )
+    d = conn.model_dump(exclude_none=True, by_alias=True, mode="json")
+    for absent in ("tls", "tls_ca_file", "tls_allow_invalid_certificates"):
+        assert absent not in d
+
+
+def test_mongo_tls_ca_fields_roundtrip():
+    payload = {
+        "id": "mongo-tls-1",
+        "name": "mongo",
+        "mask_type": "database",
+        "db_type": "mongodb",
+        "host": "mongo.example",
+        "database": "people",
+        "tls": True,
+        "tls_ca_file": "/certs/rds-ca.pem",
+        "tls_allow_invalid_certificates": True,
+    }
+    conn = MongoConnectionConfig.model_validate(payload)
+    assert conn.tls is True
+    assert conn.tls_ca_file == "/certs/rds-ca.pem"
+    assert conn.tls_allow_invalid_certificates is True
+
+
+def test_mongo_tls_ca_fields_default_when_missing():
+    payload = {
+        "id": "mongo-tls-2",
+        "name": "mongo",
+        "mask_type": "database",
+        "db_type": "mongodb",
+        "host": "mongo.example",
+        "database": "people",
+    }
+    conn = MongoConnectionConfig.model_validate(payload)
+    assert conn.tls_ca_file == ""
+    assert conn.tls_allow_invalid_certificates is False
+
+
+def test_documentdb_connection_model_dump():
+    conn = DocumentDbConnectionConfig(
+        name="docdb",
+        host="dtq-documentdb.cluster.example",
+        database="people",
+        user="dmadmin",
+        password="hunter2",
+        tls=True,
+        tls_ca_file="/certs/rds-ca.pem",
+        retry_writes=False,
+    )
+    d = conn.model_dump(exclude_none=True, by_alias=True, mode="json")
+    assert d["db_type"] == "documentdb"
+    assert d["mask_type"] == "database"
+    assert d["dbpassword"] == "hunter2"
+    assert d["tls"] is True
+    assert d["tls_ca_file"] == "/certs/rds-ca.pem"
+    assert d["retry_writes"] is False
+    assert conn.database_type is DatabaseType.documentdb
+
+
+def test_connection_config_dispatch_picks_documentdb_subclass():
+    payload = {
+        "id": "docdb-id-1",
+        "name": "docdb",
+        "mask_type": "database",
+        "db_type": "documentdb",
+        "host": "dtq-documentdb.cluster.example",
+        "database": "people",
+    }
+    conn = validate_connection(payload)
+    assert isinstance(conn, DocumentDbConnectionConfig)
+    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"):