datamasque-python 1.2.4__tar.gz → 1.3.0__tar.gz

{datamasque_python-1.2.4 → datamasque_python-1.3.0}/HISTORY.rst

@@ -2,6 +2,32 @@
 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 inserts.
+
+* ``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)
+------------------
+
+* Added ``iam_role_arn`` to ``DatabaseConnectionConfig``: the IAM role DataMasque assumes to tag
+  the connection's AWS resource, for resources in another AWS account. Sent only for the engines
+  that can be an Amazon RDS instance, Aurora cluster or Redshift cluster.
+
+Requires server version 3.26.16
+
 1.2.4 (2026-08-12)
 ------------------
 

{datamasque_python-1.2.4 → datamasque_python-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: datamasque-python
-Version: 1.2.4
+Version: 1.3.0
 Summary: Official Python client for the DataMasque data-masking API.
 Project-URL: Homepage, https://datamasque.com/
 Project-URL: Documentation, https://datamasque-python.readthedocs.io/

{datamasque_python-1.2.4 → datamasque_python-1.3.0}/datamasque/client/__init__.py

@@ -31,6 +31,7 @@ from datamasque.client.models.connection import (
     AzureConnectionConfig,
     ConnectionConfig,
     ConnectionId,
+    CosmosDbConnectionConfig,
     DatabaseConnectionConfig,
     DatabaseType,
     DatabricksConnectionConfig,
@@ -203,6 +204,7 @@ __all__ = [
     "DiscoveryConfigNotFoundError",
     "DiscoveryConfigType",
     "DiscoveryMatch",
+    "CosmosDbConnectionConfig",
     "DocumentDbConnectionConfig",
     "DynamoConnectionConfig",
     "FailedToStartError",

{datamasque_python-1.2.4 → datamasque_python-1.3.0}/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 @@ class MongoConnectionConfig(ConnectionConfig):
         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 @@ class MongoConnectionConfig(ConnectionConfig):
                 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,29 @@ class DocumentDbConnectionConfig(MongoConnectionConfig):
         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 inserts, 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: with
+    retryable writes on, masking still runs (it only updates) but the run-history insert is
+    rejected, so the run is not recorded.
+    """
+
+    # 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.
@@ -289,6 +318,22 @@ class SnowflakeConnectionConfig(ConnectionConfig):
         return data
 
 
+# Engines that can sit behind an Amazon RDS, Aurora or Redshift endpoint, and so can carry a role
+# for DataMasque to assume when tagging the resource. Aurora needs no entry of its own: it uses
+# these same engines with a cluster endpoint.
+_AWS_TAGGABLE_DATABASE_TYPES = frozenset(
+    {
+        DatabaseType.postgres,
+        DatabaseType.mysql,
+        DatabaseType.mariadb,
+        DatabaseType.oracle,
+        DatabaseType.mssql,
+        DatabaseType.db2_luw,
+        DatabaseType.redshift,
+    }
+)
+
+
 class DatabaseConnectionConfig(ConnectionConfig):
     """
     Connection configuration for a SQL database.
@@ -309,6 +354,11 @@ class DatabaseConnectionConfig(ConnectionConfig):
     is_read_only: bool = False
     s3_bucket_name: Optional[str] = None
     s3_redshift_iam_role: Optional[str] = None
+    # An IAM role for DataMasque to assume when tagging this database's AWS resource after a run.
+    # Only the engines in `_AWS_TAGGABLE_DATABASE_TYPES` can be an RDS instance, Aurora cluster or
+    # Redshift cluster, so it is pruned for the rest. Distinct from `s3_redshift_iam_role`, which is
+    # the role the Redshift cluster itself uses to reach S3.
+    iam_role_arn: Optional[str] = None
 
     @model_validator(mode="after")
     def _reject_special_engines(self) -> "DatabaseConnectionConfig":
@@ -318,6 +368,10 @@ class DatabaseConnectionConfig(ConnectionConfig):
             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
@@ -348,6 +402,8 @@ class DatabaseConnectionConfig(ConnectionConfig):
         if db_type is not DatabaseType.redshift:
             d.pop("s3_bucket_name", None)
             d.pop("s3_redshift_iam_role", None)
+        if db_type not in _AWS_TAGGABLE_DATABASE_TYPES:
+            d.pop("iam_role_arn", None)
         if not d.get("engine_options"):
             d.pop("engine_options", None)
         return d
@@ -467,6 +523,7 @@ DB_TYPE_MAP: dict[str, type[ConnectionConfig]] = {
     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,

{datamasque_python-1.2.4 → datamasque_python-1.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "datamasque-python"
-version = "1.2.4"
+version = "1.3.0"
 description = "Official Python client for the DataMasque data-masking API."
 authors = [
     { name = "DataMasque Ltd" },

{datamasque_python-1.2.4 → datamasque_python-1.3.0}/tests/test_connections.py

@@ -7,6 +7,7 @@ from datamasque.client.exceptions import DataMasqueApiError, DataMasqueException
 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,
         )
 
 
@@ -1400,3 +1469,86 @@ def test_create_or_update_mongo_connection(client):
     assert sent["db_type"] == "mongodb"
     assert sent["dbpassword"] == "hunter2"
     assert sent["replica_set"] == "rs0"
+
+
+def _database_connection(database_type: DatabaseType, **extra) -> DatabaseConnectionConfig:
+    return DatabaseConnectionConfig(
+        name="conn",
+        host="mydb.abc123.ap-southeast-2.rds.amazonaws.com",
+        port=5432,
+        database="db",
+        user="user",
+        password="secret",
+        database_type=database_type,
+        **extra,
+    )
+
+
+@pytest.mark.parametrize(
+    "database_type",
+    [
+        DatabaseType.postgres,
+        DatabaseType.mysql,
+        DatabaseType.mariadb,
+        DatabaseType.oracle,
+        DatabaseType.mssql,
+        DatabaseType.db2_luw,
+        DatabaseType.redshift,
+    ],
+)
+def test_database_connection_sends_tagging_iam_role_for_aws_hosted_engines(database_type):
+    """The engines that can sit behind an RDS, Aurora or Redshift endpoint can carry a tagging role."""
+    extra = (
+        {"s3_bucket_name": "bucket", "s3_redshift_iam_role": "arn:aws:iam::123456789012:role/redshift-s3"}
+        if database_type is DatabaseType.redshift
+        else {}
+    )
+    conn = _database_connection(database_type, iam_role_arn="arn:aws:iam::119836602066:role/tagger", **extra)
+
+    api_dict = conn.model_dump(exclude_none=True, by_alias=True, mode="json")
+
+    assert api_dict["iam_role_arn"] == "arn:aws:iam::119836602066:role/tagger"
+
+
+@pytest.mark.parametrize(
+    "database_type",
+    [
+        DatabaseType.db2i,
+        DatabaseType.informix,
+        DatabaseType.saphana,
+        # The two worth documenting rather than merely covering. `mssql_linked` inherits the field
+        # through `MssqlLinkedServerConnectionConfig` and is excluded only by its own enum value;
+        # `databricks_lakebase` is AWS-hosted, so it reads like an omission from the taggable set
+        # until you know it presents no endpoint tagging can act on.
+        DatabaseType.mssql_linked,
+        DatabaseType.databricks_lakebase,
+    ],
+)
+def test_database_connection_omits_tagging_iam_role_for_engines_that_cannot_be_aws_resources(database_type):
+    """Mirrors how `s3_redshift_iam_role` is pruned: the server has no such field on these engines."""
+    conn = _database_connection(database_type, iam_role_arn="arn:aws:iam::119836602066:role/tagger")
+
+    api_dict = conn.model_dump(exclude_none=True, by_alias=True, mode="json")
+
+    assert "iam_role_arn" not in api_dict
+
+
+def test_database_connection_omits_tagging_iam_role_when_unset():
+    conn = _database_connection(DatabaseType.postgres)
+
+    assert "iam_role_arn" not in conn.model_dump(exclude_none=True, by_alias=True, mode="json")
+
+
+def test_redshift_connection_keeps_its_cluster_s3_role_separate_from_its_tagging_role():
+    """Two ARNs with different meanings: the cluster's own S3 access, and the role DataMasque assumes to tag."""
+    conn = _database_connection(
+        DatabaseType.redshift,
+        s3_bucket_name="bucket",
+        s3_redshift_iam_role="arn:aws:iam::123456789012:role/redshift-s3",
+        iam_role_arn="arn:aws:iam::119836602066:role/tagger",
+    )
+
+    api_dict = conn.model_dump(exclude_none=True, by_alias=True, mode="json")
+
+    assert api_dict["s3_redshift_iam_role"] == "arn:aws:iam::123456789012:role/redshift-s3"
+    assert api_dict["iam_role_arn"] == "arn:aws:iam::119836602066:role/tagger"

{datamasque_python-1.2.4 → datamasque_python-1.3.0}/uv.lock

@@ -419,7 +419,7 @@ toml = [
 
 [[package]]
 name = "datamasque-python"
-version = "1.2.4"
+version = "1.3.0"
 source = { editable = "." }
 dependencies = [
     { name = "pydantic" },