datamasque-python 1.1.1__py3-none-any.whl → 1.1.3__py3-none-any.whl

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

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

@@ -1,7 +1,7 @@
-datamasque/client/__init__.py,sha256=GmpVFNWZPu6jyadhY6SOZp4hLR1wUc4cgzwLcqoEA5o,6312
+datamasque/client/__init__.py,sha256=bW7Iy-iNxpaYgc_X035KAL-hmOVgr8YcYh5s398zawo,6378
 datamasque/client/base.py,sha256=vP1LuYbqq6U8NAEJWnwNrdt_Fa2wjQTJXHf0gSFO79c,14115
 datamasque/client/connections.py,sha256=EFinx8fJRme0mTxuWY3d29UnmUFbsQhMaUQT0Ma2PK4,2885
-datamasque/client/discovery.py,sha256=j83yCBlG2xRNQlv2zb6plwL3qcCT1m3nXABG0JRdh_w,21009
+datamasque/client/discovery.py,sha256=MTDiif6YWkv64hHXZ8YJxtRhoeKIvQVfjNYvjCIZCXE,21261
 datamasque/client/discovery_configs.py,sha256=G-W-X_KizjjQLa-Rb6ixLjlkxbkCe4WGCOfYZrGtxIk,6420
 datamasque/client/dmclient.py,sha256=tA7V8GEQmnsalTHUZch8YmIUynYw2ldK-mmUGbofFL4,1650
 datamasque/client/exceptions.py,sha256=Tx7tE6i1X_oemR3EOJC97_sihZiToHqGuOQmWWxs_BE,3131
@@ -15,7 +15,7 @@ datamasque/client/runs.py,sha256=l_JpYgelqsDKW_O6LtsHC60IMXHG9ak_RwVnuvJKE1I,770
 datamasque/client/settings.py,sha256=Ui8AyR2XdoW8MZ9FIrGn2jm8DLzUGWIL8i2DOTvu7hc,2898
 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=TSv-bbOFlYCRBZHJiRgiwNuN0es_GdJd-_L3LUsQ-dI,15773
+datamasque/client/models/connection.py,sha256=Mh7MyTQmF0bllu0rurnFXmbiiq7Htc4bTwK5WQiPXSk,16823
 datamasque/client/models/data_selection.py,sha256=406yyUZ5NmLBSql2lYM1gsTWY5GWmnutmF-DjznAoLc,1904
 datamasque/client/models/discovery.py,sha256=WHSIkFebQ5m47K_BdxY1lvCeeF23q_BbZan7Ryuh5vY,13198
 datamasque/client/models/discovery_config.py,sha256=1CHMgSxMXuarPGfRcz9wX7jeZLn1SDw0UEe-L5naE54,1917
@@ -28,9 +28,9 @@ datamasque/client/models/pagination.py,sha256=egg9aO2cf6KUDwDANPu5RxpJbdRKdwUdCQ
 datamasque/client/models/ruleset.py,sha256=ZYv6SmcrONV5rLK0tKjhlKCwZiCU1dMYtH7olxI4l5A,1678
 datamasque/client/models/ruleset_library.py,sha256=5UbL_M47mgPedo8xt4q5crUh4aeOVoEMAEKbGO8CfiY,950
 datamasque/client/models/runs.py,sha256=UtPMGCJFLFP1f2nyVGAbl3A5_wAONrRNJ63y43Hbpi0,6213
-datamasque/client/models/status.py,sha256=AcKK7lelGeRCKGuWL3K9i7SiTpOIVTglyveu0tzJEVc,2374
+datamasque/client/models/status.py,sha256=UUEDHd3CQZA1Mf-1LyhAiE4COuCWvAz5oJx5ECVIwp4,2868
 datamasque/client/models/user.py,sha256=UGAUzgJkf78m24_zFXXoA99zdut48BXkX_ivV8yq1Vc,2043
-datamasque_python-1.1.1.dist-info/METADATA,sha256=TZeVZqLa38CsvBQS1wmS0xXh3nSWGyn5GKnIz5j2RXg,4187
-datamasque_python-1.1.1.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
-datamasque_python-1.1.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-datamasque_python-1.1.1.dist-info/RECORD,,
+datamasque_python-1.1.3.dist-info/METADATA,sha256=1OFcE21laJ8mFsaUCuJaGFssUJFsZzsDen7tpW59gUA,4187
+datamasque_python-1.1.3.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+datamasque_python-1.1.3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+datamasque_python-1.1.3.dist-info/RECORD,,