Make tasks cancellable (#1957)

Author: mnoszczak

libs/labelbox/src/labelbox/client.py

@@ -2389,3 +2389,26 @@ def get_task_by_id(self, task_id: str) -> Union[Task, DataUpsertTask]:
 
         task._user = user
         return task
+
+    def cancel_task(self, task_id: str) -> bool:
+        """
+        Cancels a task with the given ID.
+
+        Args:
+            task_id (str): The ID of the task to cancel.
+
+        Returns:
+            bool: True if the task was successfully cancelled.
+
+        Raises:
+            LabelboxError: If the task could not be cancelled.
+        """
+        mutation_str = """
+        mutation CancelTaskPyApi($id: ID!) {
+            cancelBulkOperationJob(id: $id) {
+                success
+            }
+        }
+        """
+        res = self.execute(mutation_str, {"id": task_id})
+        return res["cancelBulkOperationJob"]["success"]

libs/labelbox/tests/data/export/streamable/test_export_data_rows_streamable.py

@@ -1,6 +1,5 @@
 import time
 
-
 from labelbox import DataRow, ExportTask, StreamType
 
 
@@ -117,3 +116,22 @@ def test_with_invalid_id(self, client):
         assert (
             export_task.get_total_lines(stream_type=StreamType.RESULT) is None
         )
+
+    def test_cancel_export_task(
+        self, client, data_row, wait_for_data_row_processing
+    ):
+        data_row = wait_for_data_row_processing(client, data_row)
+        time.sleep(7)  # temp fix for ES indexing delay
+        export_task = DataRow.export(
+            client=client,
+            data_rows=[data_row],
+            task_name="TestExportDataRow:test_cancel_export_task",
+        )
+
+        # Cancel the task before it completes
+        success = client.cancel_task(export_task.uid)
+        assert success is True
+
+        # Verify the task was cancelled
+        cancelled_task = client.get_task_by_id(export_task.uid)
+        assert cancelled_task.status in ["CANCELING", "CANCELED"]