diff --git a/HISTORY.rst b/HISTORY.rst index cbfc4d4..51bc067 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -2,6 +2,12 @@ History ======= +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) ------------------ diff --git a/datamasque/client/discovery.py b/datamasque/client/discovery.py index 69862eb..b9fe598 100644 --- a/datamasque/client/discovery.py +++ b/datamasque/client/discovery.py @@ -166,7 +166,7 @@ def get_generated_rulesets(self, connection_id: ConnectionId) -> list[Ruleset]: 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 @@ def get_generated_rulesets(self, connection_id: ConnectionId) -> list[Ruleset]: 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, diff --git a/datamasque/client/models/status.py b/datamasque/client/models/status.py index 93871e8..dc7d0ed 100644 --- a/datamasque/client/models/status.py +++ b/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() diff --git a/tests/test_discovery.py b/tests/test_discovery.py index d954045..6121430 100644 --- a/tests/test_discovery.py +++ b/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")