datamasque-python 1.2.4__py3-none-any.whl → 1.3.0__py3-none-any.whl

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/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.dist-info → datamasque_python-1.3.0.dist-info}/METADATA

@@ -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.dist-info → datamasque_python-1.3.0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-datamasque/client/__init__.py,sha256=bw_5NEF3vLuOLKrT-z7UCkWaR_nFVkhE_RmY3Q9Y7sw,8707
+datamasque/client/__init__.py,sha256=A_S4sw3m9hNFpW5RvAvK-qpIlLHIuUtH241Dip2VgBk,8769
 datamasque/client/base.py,sha256=cXU4dluL9e1N1F2pV48iIvebv6OHt7Kh4Ka2y9vkOPo,14296
 datamasque/client/connections.py,sha256=EFinx8fJRme0mTxuWY3d29UnmUFbsQhMaUQT0Ma2PK4,2885
 datamasque/client/discovery.py,sha256=xidiAj5LdVSGQy6pjta_2Z_aJmuPPohblyNDfua_Mxk,28859
@@ -19,7 +19,7 @@ datamasque/client/spcs.py,sha256=geHrBQ6WLWsMGpSMV8hmK5SFAwftrOU0DXkrR2orOmo,753
 datamasque/client/table_references.py,sha256=dPyB59Sz4iLuF1lc7xffgf9qeNTFA5srT55E8TnFLvk,6092
 datamasque/client/users.py,sha256=VCUo2CJyOw4-aO_3mp_w_0BcSoa6-h1HcReMcAMyumw,3701
 datamasque/client/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datamasque/client/models/connection.py,sha256=loLy4RCiIfY9kvBFspSqJSRhHqhcJxm_mQzXwMRS4VU,17529
+datamasque/client/models/connection.py,sha256=W5WtzrrEfrZi7rUBMGcIZ7P6HRQ8PAC3btlD-07n4cY,20201
 datamasque/client/models/data_selection.py,sha256=9hRpkQQIQlKGFYDYQCrreggsv96dh6DelppCpCE4S2U,2725
 datamasque/client/models/discovery.py,sha256=J_cAXuIpoQTzHh_-1zvBwXsFnbKI2hRywALb5Yfc5vI,17955
 datamasque/client/models/discovery_config.py,sha256=8sqs9x_X0H_hR1fJDEFmOdJyGRDwgbdb3eddEAe98WE,2420
@@ -38,7 +38,7 @@ datamasque/client/models/safe_data_preview.py,sha256=BiI9d13KdQrTin2JfbjiRiODjns
 datamasque/client/models/status.py,sha256=DXcucL_jRsOQ0vaV0t8CJt6crvCIOLBn07KuwpwW-4A,3247
 datamasque/client/models/table_reference.py,sha256=55fKdHjC2TQnUmbp0XiWqfeFGXqcq4wXjbWgMaN3p3c,3835
 datamasque/client/models/user.py,sha256=UGAUzgJkf78m24_zFXXoA99zdut48BXkX_ivV8yq1Vc,2043
-datamasque_python-1.2.4.dist-info/METADATA,sha256=JfLuGmshJx1zHMoPocIbxfuO3H3jMKAJEa5hPTnKPNM,4497
-datamasque_python-1.2.4.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
-datamasque_python-1.2.4.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-datamasque_python-1.2.4.dist-info/RECORD,,
+datamasque_python-1.3.0.dist-info/METADATA,sha256=cDwoJ6gG2rHLCOv4zhi0OaNIG3VFpfOQkUVgCen5eaE,4497
+datamasque_python-1.3.0.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
+datamasque_python-1.3.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+datamasque_python-1.3.0.dist-info/RECORD,,