From 6c9ebbf40c9ed6a4f3cf35f77d576150849ca6fe Mon Sep 17 00:00:00 2001
From: Val Brodsky <vbrodsky@labelbox.com>
Date: Fri, 9 Aug 2024 14:04:44 -0700
Subject: [PATCH] Add client get_task_by_id

---
 libs/labelbox/src/labelbox/client.py         | 50 +++++++++++++++++++-
 libs/labelbox/src/labelbox/schema/task.py    |  2 +-
 libs/labelbox/tests/integration/test_task.py | 13 +++++
 3 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/libs/labelbox/src/labelbox/client.py b/libs/labelbox/src/labelbox/client.py
index a5c398df1..dd70baaef 100644
--- a/libs/labelbox/src/labelbox/client.py
+++ b/libs/labelbox/src/labelbox/client.py
@@ -51,7 +51,7 @@
 from labelbox.schema.send_to_annotate_params import SendToAnnotateFromCatalogParams, build_destination_task_queue_input, \
     build_predictions_input, build_annotations_input
 from labelbox.schema.slice import CatalogSlice, ModelSlice
-from labelbox.schema.task import Task
+from labelbox.schema.task import Task, DataUpsertTask
 from labelbox.schema.user import User
 from labelbox.schema.label_score import LabelScore
 from labelbox.schema.ontology_kind import (OntologyKind, EditorTaskTypeMapper,
@@ -2424,3 +2424,51 @@ def get_labeling_service_dashboards(
         return LabelingServiceDashboard.get_all(self,
                                                 after,
                                                 search_query=search_query)
+
+    def get_task_by_id(self, task_id: str) -> Union[Task, DataUpsertTask]:
+        """
+        Fetches a task by ID.
+
+        Args:
+            task_id (str): The ID of the task.
+
+        Returns:
+            Task or DataUpsertTask
+        
+        Throws:
+            ResourceNotFoundError: If the task does not exist.
+
+        NOTE: Export task is not supported yet
+        """
+        user = self.get_user()
+        query = """
+            query GetUserCreatedTasksPyApi($userId: ID!, $taskId: ID!) {
+            user(where: {id: $userId}) {
+                createdTasks(where: {id: $taskId} skip: 0 first: 1) {
+                    completionPercentage
+                    createdAt
+                    errors
+                    metadata
+                    name
+                    result
+                    status
+                    type
+                    id
+                    updatedAt
+                }
+            }
+        }
+        """
+        result = self.execute(query, {"userId": user.uid, "taskId": task_id})
+        data = result.get("user", {}).get("createdTasks", [])
+        if not data:
+            raise labelbox.exceptions.ResourceNotFoundError(
+                message=f"The task {task_id} does not exist.")
+        task_data = data[0]
+        if task_data["type"].lower() == 'adv-upsert-data-rows':
+            task = DataUpsertTask(self, task_data)
+        else:
+            task = Task(self, task_data)
+
+        task._user = user
+        return task
diff --git a/libs/labelbox/src/labelbox/schema/task.py b/libs/labelbox/src/labelbox/schema/task.py
index 95febc386..19d27c325 100644
--- a/libs/labelbox/src/labelbox/schema/task.py
+++ b/libs/labelbox/src/labelbox/schema/task.py
@@ -2,7 +2,7 @@
 import logging
 import requests
 import time
-from typing import TYPE_CHECKING, Callable, Optional, Dict, Any, List, Union, Final
+from typing import TYPE_CHECKING, Callable, Optional, Dict, Any, List, Union
 from labelbox import parser
 
 from labelbox.exceptions import ResourceNotFoundError
diff --git a/libs/labelbox/tests/integration/test_task.py b/libs/labelbox/tests/integration/test_task.py
index bec0d0414..8ab4bd4e3 100644
--- a/libs/labelbox/tests/integration/test_task.py
+++ b/libs/labelbox/tests/integration/test_task.py
@@ -1,4 +1,5 @@
 import json
+from labelbox.schema.disconnected_task import DisconnectedTask
 import pytest
 import collections.abc
 from labelbox import DataRow
@@ -31,6 +32,13 @@ def test_task_errors(dataset, image_url, snapshot):
         0]['message']
     assert len(task.failed_data_rows[0]['failedDataRows'][0]['metadata']) == 2
 
+    dt = client.get_task_by_id(task.uid)
+    assert dt.status == "COMPLETE"
+    assert len(dt.errors) == 1
+    assert dt.errors[0]['message'].startswith(
+        "A schemaId can only be specified once per DataRow")
+    assert dt.result is None
+
 
 def test_task_success_json(dataset, image_url, snapshot):
     client = dataset.client
@@ -57,3 +65,8 @@ def test_task_success_json(dataset, image_url, snapshot):
     snapshot.assert_match(json.dumps(task_result),
                           'test_task.test_task_success_json.json')
     assert len(task.result)
+
+    dt = client.get_task_by_id(task.uid)
+    assert dt.status == "COMPLETE"
+    assert len(dt.result) == 1
+    assert dt.errors is None