Add SDK method & test

Author: mnoszczak

labelbox/schema/project.py

@@ -12,16 +12,20 @@
 
 from labelbox import utils
 from labelbox.exceptions import (InvalidQueryError, LabelboxError,
-                                 ProcessingWaitTimeout, ResourceConflict)
+                                 ProcessingWaitTimeout, ResourceConflict,
+                                 ResourceNotFoundError)
 from labelbox.orm import query
 from labelbox.orm.db_object import DbObject, Deletable, Updateable
 from labelbox.orm.model import Entity, Field, Relationship
 from labelbox.pagination import PaginatedCollection
 from labelbox.schema.consensus_settings import ConsensusSettings
 from labelbox.schema.data_row import DataRow
+from labelbox.schema.export_params import ModelRunExportParams
 from labelbox.schema.media_type import MediaType
 from labelbox.schema.queue_mode import QueueMode
 from labelbox.schema.resource_tag import ResourceTag
+from labelbox.schema.task import Task
+from labelbox.schema.user import User
 
 if TYPE_CHECKING:
     from labelbox import BulkImportRequest
@@ -283,10 +287,11 @@ def label_generator(self, timeout_seconds=600, **kwargs):
             return LBV1Converter.deserialize_video(json_data, self.client)
         return LBV1Converter.deserialize(json_data)
 
-    def export_labels(self,
-                      download=False,
-                      timeout_seconds=1800,
-                      **kwargs) -> Optional[Union[str, List[Dict[Any, Any]]]]:
+    def export_labels_v2(
+            self,
+            download=False,
+            timeout_seconds=1800,
+            **kwargs) -> Optional[Union[str, List[Dict[Any, Any]]]]:
         """ Calls the server-side Label exporting that generates a JSON
         payload, and returns the URL to that payload.
 
@@ -371,6 +376,66 @@ def _validate_datetime(string_date: str) -> bool:
                          self.uid)
             time.sleep(sleep_time)
 
+    """
+    Creates a project run export task with the given params and returns the task.
+    
+    >>>    export_task = export_labels_v2("my_export_task", filter={"media_attributes": True})
+    
+    """
+
+    def export_labels_v2(self, task_name: str,
+                         params: Optional[ModelRunExportParams]) -> Task:
+        _params = params or {}
+        mutation_name = "exportDataRows"
+        create_task_query_str = """mutation exportDataRowsPyApi($input: ExportDataRowsInput!){
+          %s(input: $input) {taskId} }
+          """ % (mutation_name)
+        params = {
+            "input": {
+                "taskName": task_name,
+                "filters": {
+                    "modelRunIds": [],
+                    "projectIds": [self.uid]
+                },
+                "params": {
+                    "includeAttachments":
+                        _params.get('include_attachments', False),
+                    "includeMediaAttributes":
+                        _params.get('include_media_attributes', False),
+                    "includeMetadata":
+                        _params.get('include_metadata_fields', False),
+                    "includeDataRowDetails":
+                        _params.get('include_data_row_details', False),
+                    "includeProjectDetails":
+                        _params.get('include_project_details', False),
+                    "includeLabels":
+                        _params.get('include_labels', False),
+                    "includePerformanceDetails":
+                        _params.get('include_performance_details', False),
+                    # Arguments locked based on exectuion context
+                    "includeModelRuns":
+                        False,
+                },
+            }
+        }
+        res = self.client.execute(
+            create_task_query_str,
+            params,
+        )
+        res = res[mutation_name]
+        task_id = res["taskId"]
+        user: User = self.client.get_user()
+        tasks: List[Task] = list(
+            user.created_tasks(where=Entity.Task.uid == task_id))
+        # Cache user in a private variable as the relationship can't be
+        # resolved due to server-side limitations (see Task.created_by)
+        # for more info.
+        if len(tasks) != 1:
+            raise ResourceNotFoundError(Entity.Task, task_id)
+        task: Task = tasks[0]
+        task._user = user
+        return task
+
     def export_issues(self, status=None) -> str:
         """ Calls the server-side Issues exporting that
         returns the URL to that payload.

tests/integration/test_label.py

@@ -43,6 +43,20 @@ def test_label_export(configured_project_with_label):
     # The new exporter doesn't work with the create_label mutation
 
 
+def test_label_export_v2(configured_project_with_label):
+    project, _, _, label = configured_project_with_label
+    label_id = label.uid
+    # Wait for exporter to retrieve latest labels
+    time.sleep(10)
+    task_name = "test_label_export_v2"
+
+    task = project.export_labels_v2(task_name, filter={"project_details": True})
+    assert task.name == task_name
+    task.wait_till_done()
+    assert task.status == "COMPLETE"
+    # TODO: Download result and check it
+
+
 # TODO: Skipping this test in staging due to label not updating
 @pytest.mark.skipif(condition=os.environ['LABELBOX_TEST_ENVIRON'] == "onprem" or
                     os.environ['LABELBOX_TEST_ENVIRON'] == "staging" or