Add cancel_jobs option to automation client delete (#6)

Author: lukasbindreiter

.pre-commit-config.yaml

@@ -6,7 +6,7 @@ repos:
       - id: end-of-file-fixer
   - repo: https://github.com/charliermarsh/ruff-pre-commit
     # keep the version here in sync with the version in uv.lock
-    rev: "v0.12.0"
+    rev: "v0.12.2"
     hooks:
       - id: ruff-check
         args: [--fix, --exit-non-zero-on-fix]

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- `tilebox-workflows`: Added `cancel_jobs` argument to `AutomationClient.delete` to optionally cancel all currently queued or running jobs of the automation.
+
 ## [0.38.0] - 2025-06-24
 
 ### Added

tilebox-datasets/tilebox/datasets/datasetsv1/datasets_pb2.py

@@ -65,6 +65,18 @@ class UpdateDatasetDescriptionRequest(_message.Message):
     description: str
     def __init__(self, id: _Optional[_Union[_core_pb2.ID, _Mapping]] = ..., description: _Optional[str] = ...) -> None: ...
 
+class DeleteDatasetRequest(_message.Message):
+    __slots__ = ("id",)
+    ID_FIELD_NUMBER: _ClassVar[int]
+    id: _core_pb2.ID
+    def __init__(self, id: _Optional[_Union[_core_pb2.ID, _Mapping]] = ...) -> None: ...
+
+class DeleteDatasetResponse(_message.Message):
+    __slots__ = ("trashed",)
+    TRASHED_FIELD_NUMBER: _ClassVar[int]
+    trashed: bool
+    def __init__(self, trashed: bool = ...) -> None: ...
+
 class ListDatasetsRequest(_message.Message):
     __slots__ = ("client_info",)
     CLIENT_INFO_FIELD_NUMBER: _ClassVar[int]

tilebox-datasets/tilebox/datasets/datasetsv1/datasets_pb2.pyi

@@ -36,6 +36,11 @@ def __init__(self, channel):
                 request_serializer=datasets_dot_v1_dot_datasets__pb2.UpdateDatasetDescriptionRequest.SerializeToString,
                 response_deserializer=datasets_dot_v1_dot_core__pb2.Dataset.FromString,
                 _registered_method=True)
+        self.DeleteDataset = channel.unary_unary(
+                '/datasets.v1.DatasetService/DeleteDataset',
+                request_serializer=datasets_dot_v1_dot_datasets__pb2.DeleteDatasetRequest.SerializeToString,
+                response_deserializer=datasets_dot_v1_dot_datasets__pb2.DeleteDatasetResponse.FromString,
+                _registered_method=True)
         self.ListDatasets = channel.unary_unary(
                 '/datasets.v1.DatasetService/ListDatasets',
                 request_serializer=datasets_dot_v1_dot_datasets__pb2.ListDatasetsRequest.SerializeToString,
@@ -71,6 +76,12 @@ def UpdateDatasetDescription(self, request, context):
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
+    def DeleteDataset(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
     def ListDatasets(self, request, context):
         """Missing associated documentation comment in .proto file."""
         context.set_code(grpc.StatusCode.UNIMPLEMENTED)
@@ -100,6 +111,11 @@ def add_DatasetServiceServicer_to_server(servicer, server):
                     request_deserializer=datasets_dot_v1_dot_datasets__pb2.UpdateDatasetDescriptionRequest.FromString,
                     response_serializer=datasets_dot_v1_dot_core__pb2.Dataset.SerializeToString,
             ),
+            'DeleteDataset': grpc.unary_unary_rpc_method_handler(
+                    servicer.DeleteDataset,
+                    request_deserializer=datasets_dot_v1_dot_datasets__pb2.DeleteDatasetRequest.FromString,
+                    response_serializer=datasets_dot_v1_dot_datasets__pb2.DeleteDatasetResponse.SerializeToString,
+            ),
             'ListDatasets': grpc.unary_unary_rpc_method_handler(
                     servicer.ListDatasets,
                     request_deserializer=datasets_dot_v1_dot_datasets__pb2.ListDatasetsRequest.FromString,
@@ -225,6 +241,33 @@ def UpdateDatasetDescription(request,
             metadata,
             _registered_method=True)
 
+    @staticmethod
+    def DeleteDataset(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/datasets.v1.DatasetService/DeleteDataset',
+            datasets_dot_v1_dot_datasets__pb2.DeleteDatasetRequest.SerializeToString,
+            datasets_dot_v1_dot_datasets__pb2.DeleteDatasetResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
     @staticmethod
     def ListDatasets(request,
             target,

tilebox-datasets/tilebox/datasets/datasetsv1/datasets_pb2_grpc.py

@@ -18,7 +18,7 @@
     uuid_to_uuid_message,
 )
 from tilebox.workflows.workflowsv1.automation_pb2 import AutomationPrototype as AutomationPrototypeMessage
-from tilebox.workflows.workflowsv1.automation_pb2 import Automations
+from tilebox.workflows.workflowsv1.automation_pb2 import Automations, DeleteAutomationRequest
 from tilebox.workflows.workflowsv1.automation_pb2 import CronTrigger as CronTriggerMessage
 from tilebox.workflows.workflowsv1.automation_pb2 import StorageEventTrigger as StorageEventTriggerMessage
 from tilebox.workflows.workflowsv1.automation_pb2_grpc import AutomationServiceStub
@@ -65,8 +65,8 @@ def GetAutomation(self, req: UUIDMessage) -> AutomationPrototypeMessage:  # noqa
             return self.automations[automation_id]
         raise NotFoundError(f"Automation {automation_id} not found")
 
-    def DeleteAutomation(self, req: UUIDMessage) -> None:  # noqa: N802
-        automation_id = uuid_message_to_uuid(req)
+    def DeleteAutomation(self, req: DeleteAutomationRequest) -> None:  # noqa: N802
+        automation_id = uuid_message_to_uuid(req.automation_id)
         if automation_id in self.automations:
             del self.automations[automation_id]
         else:
@@ -114,7 +114,7 @@ def create_cron_automation(
         cron_triggers: list[CronTrigger],
     ) -> AutomationPrototype:
         self.count_automations += 1
-        triggers = [t.schedule for t in cron_triggers]
+        schedules = [t.schedule for t in cron_triggers]
 
         class TestCronTask(CronTask):
             some_arg: str
@@ -125,7 +125,7 @@ def identifier() -> tuple[str, str]:
 
         task = TestCronTask(task_name)  # task_name reused to serialize the task input
 
-        return self.client.create_cron_automation(task_name, task, triggers, cluster_slug)
+        return self.client.create_cron_automation(task_name, task, schedules, cluster_slug)
 
     @rule(
         target=inserted_automations,

tilebox-workflows/tests/automations/test_client.py

@@ -133,11 +133,12 @@ def create_storage_event_automation(
         )
         return self._service.create(automation)
 
-    def delete(self, automation_or_id: AutomationPrototype | UUID | str) -> None:
+    def delete(self, automation_or_id: AutomationPrototype | UUID | str, cancel_jobs: bool = False) -> None:
         """Delete an automation by id.
 
         Args:
             automation_or_id: The id of the automation to delete or the automation object itself.
+            cancel_jobs: Whether to cancel all currently queued or running jobs of the automation. Defaults to False.
         """
         if isinstance(automation_or_id, str):
             automation_id = UUID(automation_or_id)
@@ -146,4 +147,4 @@ def delete(self, automation_or_id: AutomationPrototype | UUID | str) -> None:
         else:
             automation_id = automation_or_id
 
-        self._service.delete(automation_id)
+        self._service.delete(automation_id, cancel_jobs)

tilebox-workflows/tilebox/workflows/automations/client.py

@@ -10,7 +10,7 @@
     uuid_to_uuid_message,
 )
 from tilebox.workflows.workflowsv1.automation_pb2 import AutomationPrototype as AutomationPrototypeMessage
-from tilebox.workflows.workflowsv1.automation_pb2 import Automations, StorageLocations
+from tilebox.workflows.workflowsv1.automation_pb2 import Automations, DeleteAutomationRequest, StorageLocations
 from tilebox.workflows.workflowsv1.automation_pb2_grpc import AutomationServiceStub
 
 
@@ -45,5 +45,7 @@ def update(self, automation: AutomationPrototype) -> AutomationPrototype:
         response: AutomationPrototypeMessage = self.service.UpdateAutomation(automation.to_message())
         return AutomationPrototype.from_message(response)
 
-    def delete(self, automation_id: UUID) -> None:
-        self.service.DeleteAutomation(uuid_to_uuid_message(automation_id))
+    def delete(self, automation_id: UUID, cancel_jobs: bool = False) -> None:
+        self.service.DeleteAutomation(
+            DeleteAutomationRequest(automation_id=uuid_to_uuid_message(automation_id), cancel_jobs=cancel_jobs)
+        )