diff --git a/libs/labelbox/src/labelbox/client.py b/libs/labelbox/src/labelbox/client.py index af3b5b3fc..498840c0f 100644 --- a/libs/labelbox/src/labelbox/client.py +++ b/libs/labelbox/src/labelbox/client.py @@ -2339,3 +2339,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"] diff --git a/libs/labelbox/tests/data/export/streamable/test_export_data_rows_streamable.py b/libs/labelbox/tests/data/export/streamable/test_export_data_rows_streamable.py index 8dfa1df72..5ab4d30ec 100644 --- a/libs/labelbox/tests/data/export/streamable/test_export_data_rows_streamable.py +++ b/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"]