Merge pull request #130 from Labelbox/ms/bulk-export-status

Easier fetch bulk import ndjson

Author: msokoloff1

labelbox/client.py

@@ -253,13 +253,15 @@ def upload_file(self, path: str) -> str:
     def upload_data(self,
                     content: bytes,
                     filename: str = None,
-                    content_type: str = None) -> str:
+                    content_type: str = None,
+                    sign: bool = False) -> str:
         """ Uploads the given data (bytes) to Labelbox.
 
         Args:
             content: bytestring to upload
             filename: name of the upload
             content_type: content type of data uploaded
+            sign: whether or not to sign the url
 
         Returns:
             str, the URL of uploaded data.
@@ -274,7 +276,7 @@ def upload_data(self,
                     "variables": {
                         "file": None,
                         "contentLength": len(content),
-                        "sign": False
+                        "sign": sign
                     },
                     "query":
                         """mutation UploadFile($file: Upload!, $contentLength: Int!,

labelbox/schema/bulk_import_request.py

@@ -1,6 +1,7 @@
 import json
 import time
 from uuid import UUID, uuid4
+import functools
 
 import logging
 from pathlib import Path
@@ -9,6 +10,7 @@
 import ndjson
 import requests
 from pydantic import BaseModel, validator
+from requests.api import request
 from typing_extensions import Literal
 from typing import (Any, List, Optional, BinaryIO, Dict, Iterable, Tuple, Union,
                     Type, Set)
@@ -113,6 +115,76 @@ class BulkImportRequest(DbObject):
     project = Relationship.ToOne("Project")
     created_by = Relationship.ToOne("User", False, "created_by")
 
+    @property
+    def inputs(self) -> List[Dict[str, Any]]:
+        """
+        Inputs for each individual annotation uploaded.
+        This should match the ndjson annotations that you have uploaded. 
+        
+        Returns:
+            Uploaded ndjson.
+
+        * This information will expire after 24 hours.    
+        """
+        return self._fetch_remote_ndjson(self.input_file_url)
+
+    @property
+    def errors(self) -> List[Dict[str, Any]]:
+        """
+        Errors for each individual annotation uploaded. This is a subset of statuses
+
+        Returns:
+            List of dicts containing error messages. Empty list means there were no errors
+            See `BulkImportRequest.statuses` for more details.
+
+        * This information will expire after 24 hours.        
+        """
+        self.wait_until_done()
+        return self._fetch_remote_ndjson(self.error_file_url)
+
+    @property
+    def statuses(self) -> List[Dict[str, Any]]:
+        """
+        Status for each individual annotation uploaded.
+
+        Returns:
+            A status for each annotation if the upload is done running.
+            See below table for more details
+            
+        .. list-table::
+           :widths: 15 150
+           :header-rows: 1 
+
+           * - Field
+             - Description
+           * - uuid 
+             - Specifies the annotation for the status row.
+           * - dataRow
+             - JSON object containing the Labelbox data row ID for the annotation.
+           * - status
+             - Indicates SUCCESS or FAILURE.
+           * - errors
+             - An array of error messages included when status is FAILURE. Each error has a name, message and optional (key might not exist) additional_info.
+
+        * This information will expire after 24 hours.        
+        """
+        self.wait_until_done()
+        return self._fetch_remote_ndjson(self.status_file_url)
+
+    @functools.lru_cache()
+    def _fetch_remote_ndjson(self, url: str) -> List[Dict[str, Any]]:
+        """
+        Fetches the remote ndjson file and caches the results.
+
+        Args:
+            url (str): Can be any url pointing to an ndjson file.
+        Returns:
+            ndjson as a list of dicts.
+        """
+        response = requests.get(url)
+        response.raise_for_status()
+        return ndjson.loads(response.text)
+
     def refresh(self) -> None:
         """Synchronizes values of all fields with the database.
         """

labelbox/schema/enums.py

@@ -3,6 +3,19 @@
 
 class BulkImportRequestState(Enum):
     """ State of the import job when importing annotations (RUNNING, FAILED, or FINISHED).
+
+    .. list-table::
+       :widths: 15 150
+       :header-rows: 1 
+
+       * - State
+         - Description
+       * - RUNNING
+         - Indicates that the import job is not done yet.
+       * - FAILED
+         - Indicates the import job failed. Check `BulkImportRequest.errors` for more information
+       * - FINISHED
+         - Indicates the import job is no longer running. Check `BulkImportRequest.statuses` for more information
     """
     RUNNING = "RUNNING"
     FAILED = "FAILED"

tests/integration/bulk_import/test_bulk_import_request.py

@@ -121,17 +121,27 @@ def test_validate_ndjson_uuid(tmp_path, configured_project, predictions):
 
 
 @pytest.mark.slow
-def test_wait_till_done(configured_project):
+def test_wait_till_done(rectangle_inference, configured_project):
     name = str(uuid.uuid4())
-    url = "https://storage.googleapis.com/labelbox-public-bucket/predictions_test_v2.ndjson"
+    url = configured_project.client.upload_data(content=ndjson.dumps(
+        [rectangle_inference]),
+                                                sign=True)
     bulk_import_request = configured_project.upload_annotations(name=name,
                                                                 annotations=url,
                                                                 validate=False)
 
+    assert len(bulk_import_request.inputs) == 1
     bulk_import_request.wait_until_done()
-
-    assert (bulk_import_request.state == BulkImportRequestState.FINISHED or
-            bulk_import_request.state == BulkImportRequestState.FAILED)
+    assert bulk_import_request.state == BulkImportRequestState.FINISHED
+
+    # Check that the status files are being returned as expected
+    assert len(bulk_import_request.errors) == 0
+    assert len(bulk_import_request.inputs) == 1
+    assert bulk_import_request.inputs[0]['uuid'] == rectangle_inference['uuid']
+    assert len(bulk_import_request.statuses) == 1
+    assert bulk_import_request.statuses[0]['status'] == 'SUCCESS'
+    assert bulk_import_request.statuses[0]['uuid'] == rectangle_inference[
+        'uuid']
 
 
 def assert_file_content(url: str, predictions):