Merge pull request #150 from Labelbox/ms/bulk-export-datarows

[DIAG-101] bulk export queued data rows

Author: msokoloff1

labelbox/schema/bulk_import_request.py

@@ -119,12 +119,12 @@ class BulkImportRequest(DbObject):
     def inputs(self) -> List[Dict[str, Any]]:
         """
         Inputs for each individual annotation uploaded.
-        This should match the ndjson annotations that you have uploaded. 
-        
+        This should match the ndjson annotations that you have uploaded.
+
         Returns:
             Uploaded ndjson.
 
-        * This information will expire after 24 hours.    
+        * This information will expire after 24 hours.
         """
         return self._fetch_remote_ndjson(self.input_file_url)
 
@@ -137,7 +137,7 @@ def errors(self) -> List[Dict[str, Any]]:
             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.        
+        * This information will expire after 24 hours.
         """
         self.wait_until_done()
         return self._fetch_remote_ndjson(self.error_file_url)
@@ -150,14 +150,14 @@ def statuses(self) -> List[Dict[str, Any]]:
         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 
+           :header-rows: 1
 
            * - Field
              - Description
-           * - uuid 
+           * - uuid
              - Specifies the annotation for the status row.
            * - dataRow
              - JSON object containing the Labelbox data row ID for the annotation.
@@ -166,7 +166,7 @@ def statuses(self) -> List[Dict[str, Any]]:
            * - 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.        
+        * This information will expire after 24 hours.
         """
         self.wait_until_done()
         return self._fetch_remote_ndjson(self.status_file_url)
@@ -390,27 +390,26 @@ def create_from_local_file(cls,
 
 def _validate_ndjson(lines: Iterable[Dict[str, Any]],
                      project: "labelbox.Project") -> None:
-    """   
+    """
     Client side validation of an ndjson object.
 
     Does not guarentee that an upload will succeed for the following reasons:
         * We are not checking the data row types which will cause the following errors to slip through
             * Missing frame indices will not causes an error for videos
-        * Uploaded annotations for the wrong data type will pass (Eg. entity on images) 
+        * Uploaded annotations for the wrong data type will pass (Eg. entity on images)
         * We are not checking bounds of an asset (Eg. frame index, image height, text location)
- 
+
     Args:
         lines (Iterable[Dict[str,Any]]): An iterable of ndjson lines
         project (Project): id of project for which predictions will be imported
-    
+
     Raises:
         MALValidationError: Raise for invalid NDJson
         UuidError: Duplicate UUID in upload
     """
-    data_row_ids = {
-        data_row.uid: data_row for dataset in project.datasets()
-        for data_row in dataset.data_rows()
-    }
+    export_url = project.export_queued_data_rows()
+    data_row_json = ndjson.loads(requests.get(export_url).text)
+    data_row_ids = {row['id'] for row in data_row_json}
     feature_schemas = get_mal_schemas(project.ontology())
     uids: Set[str] = set()
     for idx, line in enumerate(lines):
@@ -539,7 +538,7 @@ def build(cls: Any, data: Union[dict, BaseModel]) -> "NDBase":
             raises:
                 KeyError: data does not contain the determinant fields for any of the types supported by this SpecialUnion
                 ValidationError: Error while trying to construct a specific object in the union
-            
+
         """
         if isinstance(data, BaseModel):
             data = data.dict()

labelbox/schema/project.py

@@ -11,7 +11,7 @@
 from labelbox.schema.data_row import DataRow
 from labelbox.orm import query
 from labelbox.schema.bulk_import_request import BulkImportRequest
-from labelbox.exceptions import InvalidQueryError
+from labelbox.exceptions import InvalidQueryError, LabelboxError
 from labelbox.orm.db_object import DbObject, Updateable, Deletable
 from labelbox.orm.model import Entity, Field, Relationship
 from labelbox.pagination import PaginatedCollection
@@ -160,6 +160,39 @@ def labels(self, datasets=None, order_by=None):
         return PaginatedCollection(self.client, query_str, {id_param: self.uid},
                                    ["project", "labels"], Label)
 
+    def export_queued_data_rows(self, timeout_seconds=120):
+        """ Returns all data rows that are currently enqueued for this project.
+
+        Args:
+            timeout_seconds (float): Max waiting time, in seconds.
+        Returns:
+            URL of the data file with this DataRow information. If the server didn't
+            generate during the `timeout_seconds` period, None is returned.
+        """
+        id_param = "projectId"
+        query_str = """mutation GetQueuedDataRowsExportUrlPyApi($%s: ID!)
+            {exportQueuedDataRows(data:{projectId: $%s }) {downloadUrl createdAt status} }
+        """ % (id_param, id_param)
+        sleep_time = 2
+        while True:
+            res = self.client.execute(query_str, {id_param: self.uid})
+            res = res["exportQueuedDataRows"]
+            if res["status"] == "COMPLETE":
+                return res["downloadUrl"]
+            elif res["status"] == "FAILED":
+                raise LabelboxError("Data row export failed.")
+
+            timeout_seconds -= sleep_time
+            if timeout_seconds <= 0:
+                raise LabelboxError(
+                    f"Unable to export data rows within {timeout_seconds} seconds."
+                )
+
+            logger.debug(
+                "Project '%s' queued data row export, waiting for server...",
+                self.uid)
+            time.sleep(sleep_time)
+
     def export_labels(self, timeout_seconds=60):
         """ Calls the server-side Label exporting that generates a JSON
         payload, and returns the URL to that payload.
@@ -193,13 +226,13 @@ def export_labels(self, timeout_seconds=60):
             time.sleep(sleep_time)
 
     def export_issues(self, status=None):
-        """ Calls the server-side Issues exporting that 
+        """ Calls the server-side Issues exporting that
         returns the URL to that payload.
 
         Args:
             status (string): valid values: Open, Resolved
         Returns:
-            URL of the data file with this Project's issues. 
+            URL of the data file with this Project's issues.
         """
         id_param = "projectId"
         status_param = "status"
@@ -229,14 +262,14 @@ def export_issues(self, status=None):
     def upsert_instructions(self, instructions_file: str):
         """
         * Uploads instructions to the UI. Running more than once will replace the instructions
-            
+
         Args:
             instructions_file (str): Path to a local file.
                 * Must be either a pdf, text, or html file.
 
         Raises:
             ValueError:
-                * project must be setup 
+                * project must be setup
                 * instructions file must end with one of ".text", ".txt", ".pdf", ".html"
         """
 
@@ -267,18 +300,18 @@ def upsert_instructions(self, instructions_file: str):
 
         self.client.execute(
             """mutation UpdateFrontendWithExistingOptionsPyApi (
-                    $frontendId: ID!, 
-                    $optionsId: ID!, 
-                    $name: String!, 
-                    $description: String!, 
+                    $frontendId: ID!,
+                    $optionsId: ID!,
+                    $name: String!,
+                    $description: String!,
                     $customizationOptions: String!
                 ) {
                     updateLabelingFrontend(
-                        where: {id: $frontendId}, 
+                        where: {id: $frontendId},
                         data: {name: $name, description: $description}
                     ) {id}
                     updateLabelingFrontendOptions(
-                        where: {id: $optionsId}, 
+                        where: {id: $optionsId},
                         data: {customizationOptions: $customizationOptions}
                     ) {id}
                 }""", {
@@ -390,10 +423,10 @@ def validate_labeling_parameter_overrides(self, data):
 
     def set_labeling_parameter_overrides(self, data):
         """ Adds labeling parameter overrides to this project.
-                
+
         See information on priority here:
             https://docs.labelbox.com/en/configure-editor/queue-system#reservation-system
-    
+
             >>> project.set_labeling_parameter_overrides([
             >>>     (data_row_1, 2, 3), (data_row_2, 1, 4)])
 
@@ -407,11 +440,11 @@ def set_labeling_parameter_overrides(self, data):
                         - Minimum priority is 1.
                     * Priority is not the queue position.
                         - The position is determined by the relative priority.
-                        - E.g. [(data_row_1, 5,1), (data_row_2, 2,1), (data_row_3, 10,1)] 
+                        - E.g. [(data_row_1, 5,1), (data_row_2, 2,1), (data_row_3, 10,1)]
                             will be assigned in the following order: [data_row_2, data_row_1, data_row_3]
                     * Datarows with parameter overrides will appear before datarows without overrides.
                     * The priority only effects items in the queue.
-                        - Assigning a priority will not automatically add the item back into the queue.  
+                        - Assigning a priority will not automatically add the item back into the queue.
                 Number of labels:
                     * The number of times a data row should be labeled.
                         - Creates duplicate data rows in a project (one for each number of labels).
@@ -458,7 +491,7 @@ def unset_labeling_parameter_overrides(self, data_rows):
     def upsert_review_queue(self, quota_factor):
         """ Sets the the proportion of total assets in a project to review.
 
-        More information can be found here: 
+        More information can be found here:
             https://docs.labelbox.com/en/quality-assurance/review-labels#configure-review-percentage
 
         Args:

tests/integration/conftest.py

@@ -1,7 +1,6 @@
 from collections import namedtuple
 from enum import Enum
 from datetime import datetime
-from labelbox.orm.db_object import experimental
 from random import randint
 from string import ascii_letters
 from types import SimpleNamespace
@@ -14,6 +13,7 @@
 from labelbox.schema.invite import Invite
 from labelbox.pagination import PaginatedCollection
 from labelbox.schema.user import User
+from labelbox import LabelingFrontend
 from labelbox import Client
 
 IMG_URL = "https://picsum.photos/200/300"
@@ -83,7 +83,7 @@ def get_invites(client):
     Do not use. Only for testing.
     """
     query_str = """query GetOrgInvitationsPyApi($from: ID, $first: PageSize) {
-            organization { id invites(from: $from, first: $first) { 
+            organization { id invites(from: $from, first: $first) {
                 nodes { id createdAt organizationRoleName inviteeEmail } nextCursor }}}"""
     invites = PaginatedCollection(
         client,
@@ -199,13 +199,13 @@ def organization(client):
 def project_based_user(client, rand_gen):
     email = rand_gen(str)
     # Use old mutation because it doesn't require users to accept email invites
-    query_str = """mutation MakeNewUserPyApi { 
-        addMembersToOrganization( 
-            data: { 
-                emails: ["%[email protected]"], 
+    query_str = """mutation MakeNewUserPyApi {
+        addMembersToOrganization(
+            data: {
+                emails: ["%[email protected]"],
                 orgRoleId: "%s",
                 projectRoles: []
-            } 
+            }
         ) {
         newUserId
         }
@@ -227,3 +227,17 @@ def project_pack(client):
     yield projects
     for proj in projects:
         proj.delete()
+
+
+@pytest.fixture
+def configured_project(project, client, rand_gen):
+    dataset = client.create_dataset(name=rand_gen(str), projects=project)
+    dataset.create_data_row(row_data=IMG_URL)
+    editor = list(
+        project.client.get_labeling_frontends(
+            where=LabelingFrontend.name == "editor"))[0]
+    empty_ontology = {"tools": [], "classifications": []}
+    project.setup(editor, empty_ontology)
+    yield project
+    dataset.delete()
+    project.delete()

tests/integration/test_project.py

@@ -1,6 +1,8 @@
-import pytest
 import json
+
 import requests
+import ndjson
+import pytest
 
 from labelbox import Project, LabelingFrontend
 from labelbox.exceptions import InvalidQueryError
@@ -100,3 +102,10 @@ def test_attach_instructions(client, project):
     with pytest.raises(ValueError) as exc_info:
         project.upsert_instructions('/tmp/file.invalid_file_extension')
     assert "instructions_file must end with one of" in str(exc_info.value)
+
+
+def test_queued_data_row_export(configured_project):
+    url = configured_project.export_queued_data_rows()
+    assert url
+    result = ndjson.loads(requests.get(url).text)
+    assert len(result) == 1