datamasque-python 1.1.1__tar.gz → 1.1.3__tar.gz

{datamasque_python-1.1.1 → datamasque_python-1.1.3}/HISTORY.rst

@@ -2,6 +2,21 @@
 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)
+------------------
+
+* Completed wiring of the ``finished_with_warnings`` status for ``AsyncRulesetGenerationTaskStatus``.
+* Added the ``cancelled`` status to ``AsyncRulesetGenerationTaskStatus``.
+
 1.1.1 (2026-06-25)
 ------------------
 

{datamasque_python-1.1.1 → datamasque_python-1.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datamasque-python
-Version: 1.1.1
+Version: 1.1.3
 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.1.1 → datamasque_python-1.1.3}/datamasque/client/__init__.py

@@ -33,6 +33,7 @@ from datamasque.client.models.connection import (
     DatabaseConnectionConfig,
     DatabaseType,
     DatabricksConnectionConfig,
+    DocumentDbConnectionConfig,
     DynamoConnectionConfig,
     FileConnectionConfig,
     MongoConnectionConfig,
@@ -150,6 +151,7 @@ __all__ = [
     "DiscoveryConfigNotFoundError",
     "DiscoveryConfigType",
     "DiscoveryMatch",
+    "DocumentDbConnectionConfig",
     "DynamoConnectionConfig",
     "FailedToStartError",
     "FileConnectionConfig",

{datamasque_python-1.1.1 → datamasque_python-1.1.3}/datamasque/client/discovery.py

@@ -166,7 +166,7 @@ class DiscoveryClient(BaseClient):
             returns a list containing one ruleset
 
         Raises `AsyncRulesetGenerationInProgressError` if the task hasn't finished yet,
-        and `DataMasqueException` if it failed.
+        and `DataMasqueException` if it failed or was cancelled.
 
         Note that the ruleset(s) have autogenerated names, which you may want to customize before uploading.
         """
@@ -176,7 +176,11 @@ class DiscoveryClient(BaseClient):
             logger.error("Ruleset generation failed for connection: %s", connection_id)
             raise DataMasqueException(f"Ruleset generation failed for connection: {connection_id}")
 
-        if status is not AsyncRulesetGenerationTaskStatus.finished:
+        if status is AsyncRulesetGenerationTaskStatus.cancelled:
+            logger.error("Ruleset generation was cancelled for connection: %s", connection_id)
+            raise DataMasqueException(f"Ruleset generation was cancelled for connection: {connection_id}")
+
+        if not status.is_finished:
             logger.error(
                 "Ruleset generation is still in progress for connection: %s. Status: `%s`",
                 connection_id,

{datamasque_python-1.1.1 → datamasque_python-1.1.3}/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 @@ class MongoConnectionConfig(ConnectionConfig):
             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 @@ class MongoConnectionConfig(ConnectionConfig):
         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 @@ FILE_TYPE_MAP: dict[str, type[FileConnectionConfig]] = {
 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,

{datamasque_python-1.1.1 → datamasque_python-1.1.3}/datamasque/client/models/status.py

@@ -64,15 +64,28 @@ class AsyncRulesetGenerationTaskStatus(enum.Enum):
     failed = "failed"
     running = "running"
     queued = "queued"
+    cancelled = "cancelled"
 
     @classmethod
     def get_final_states(cls) -> set["AsyncRulesetGenerationTaskStatus"]:
         """Returns the list of final statuses, i.e. the ruleset generation has completed, successfully or otherwise."""
 
-        return {cls.finished, cls.failed}
+        return {cls.finished, cls.finished_with_warnings, cls.failed, cls.cancelled}
+
+    @classmethod
+    def get_finished_states(cls) -> set["AsyncRulesetGenerationTaskStatus"]:
+        """Returns the list of statuses that indicate the ruleset generation completed successfully."""
+
+        return {cls.finished, cls.finished_with_warnings}
 
     @property
     def is_in_final_state(self) -> bool:
         """Returns True if this status is a final status."""
 
         return self in self.get_final_states()
+
+    @property
+    def is_finished(self) -> bool:
+        """Returns True if this status is a finished status."""
+
+        return self in self.get_finished_states()

{datamasque_python-1.1.1 → datamasque_python-1.1.3}/pyproject.toml

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

{datamasque_python-1.1.1 → datamasque_python-1.1.3}/setup.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 1.1.1
+current_version = 1.1.3
 commit = True
 tag = True
 

{datamasque_python-1.1.1 → datamasque_python-1.1.3}/tests/test_connections.py

@@ -10,6 +10,7 @@ from datamasque.client.models.connection import (
     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"):

{datamasque_python-1.1.1 → datamasque_python-1.1.3}/tests/test_discovery.py

@@ -204,6 +204,47 @@ def test_get_generated_rulesets_success(client):
         assert rulesets[1].yaml == yaml_content_2.decode("utf-8")
 
 
+def test_get_generated_rulesets_finished_with_warnings_success(client):
+    """A task that finishes with warnings is still successful: its rulesets are returned, not treated as in-progress."""
+    connection_id = ConnectionId("1")
+    yaml_content = b"""
+    version: "1.0"
+    tasks:
+    - type: mask_table
+      table: table1
+      key: id
+      rules:
+      - column: col1
+        masks:
+        - type: do_nothing
+    """
+
+    with requests_mock.Mocker() as m:
+        m.get(
+            f"http://test-server/api/async-generate-ruleset/{connection_id}/",
+            json={"status": "finished_with_warnings"},
+            status_code=200,
+        )
+
+        zip_buffer = BytesIO()
+        with zipfile.ZipFile(zip_buffer, "w") as zip_file:
+            zip_file.writestr("ruleset1.yml", yaml_content.decode("utf-8"))
+        zip_buffer.seek(0)
+
+        m.get(
+            f"http://test-server/api/async-generate-ruleset/{connection_id}/download-rulesets/",
+            content=zip_buffer.getvalue(),
+            headers={"Content-Disposition": 'attachment; filename="rulesets.zip"'},
+            status_code=200,
+        )
+
+        rulesets = client.get_generated_rulesets(connection_id)
+
+        assert len(rulesets) == 1
+        assert rulesets[0].name == "ruleset1"
+        assert rulesets[0].yaml == yaml_content.decode("utf-8")
+
+
 def test_get_generated_rulesets_empty_archive_raises(client):
     """A finished task whose download archive contains no ruleset files raises a clear error."""
     connection_id = ConnectionId("1")
@@ -290,6 +331,21 @@ def test_get_generated_rulesets_failed(client):
             client.get_generated_rulesets(connection_id)
 
 
+def test_get_generated_rulesets_cancelled(client):
+    """A cancelled task is terminal, so it raises a non-retryable DataMasqueException, not the in-progress error."""
+    connection_id = ConnectionId("1")
+
+    with requests_mock.Mocker() as m:
+        m.get(
+            f"http://test-server/api/async-generate-ruleset/{connection_id}/",
+            json={"status": "cancelled"},
+            status_code=200,
+        )
+
+        with pytest.raises(DataMasqueException, match="Ruleset generation was cancelled for connection"):
+            client.get_generated_rulesets(connection_id)
+
+
 def test_get_generated_rulesets_in_progress(client):
     connection_id = ConnectionId("1")
 

{datamasque_python-1.1.1 → datamasque_python-1.1.3}/uv.lock

@@ -428,7 +428,7 @@ toml = [
 
 [[package]]
 name = "datamasque-python"
-version = "1.1.1"
+version = "1.1.3"
 source = { editable = "." }
 dependencies = [
     { name = "pydantic" },