Add static method to export data rows + add tests

Author: mnoszczak

labelbox/schema/data_row.py

@@ -1,14 +1,18 @@
 import logging
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Collection, Dict, List, Optional
 import json
+from labelbox.exceptions import ResourceNotFoundError
 
 from labelbox.orm import query
 from labelbox.orm.db_object import DbObject, Updateable, BulkDeletable
 from labelbox.orm.model import Entity, Field, Relationship
 from labelbox.schema.data_row_metadata import DataRowMetadataField  # type: ignore
+from labelbox.schema.export_params import CatalogExportParams
+from labelbox.schema.task import Task
+from labelbox.schema.user import User  # type: ignore
 
 if TYPE_CHECKING:
-    from labelbox import AssetAttachment
+    from labelbox import AssetAttachment, Client
 
 logger = logging.getLogger(__name__)
 
@@ -150,3 +154,108 @@ def create_attachment(self,
             })
         return Entity.AssetAttachment(self.client,
                                       res["createDataRowAttachment"])
+
+    @staticmethod
+    def export_v2(client: 'Client',
+                  data_rows: List['DataRow'],
+                  task_name: Optional[str] = None,
+                  params: Optional[CatalogExportParams] = None) -> Task:
+        """
+        Creates a data rows export task with the given list, params and returns the task.
+        
+        >>>     dataset = client.get_dataset(DATASET_ID)
+        >>>     task = DataRow.export_v2(
+        >>>         data_rows_ids=[data_row.uid for data_row in dataset.data_rows.list()],
+        >>>         filters={
+        >>>             "last_activity_at": ["2000-01-01 00:00:00", "2050-01-01 00:00:00"],
+        >>>             "label_created_at": ["2000-01-01 00:00:00", "2050-01-01 00:00:00"]
+        >>>         },
+        >>>         params={
+        >>>             "performance_details": False,
+        >>>             "label_details": True
+        >>>         })
+        >>>     task.wait_till_done()
+        >>>     task.result
+        """
+        print('export start')
+
+        _params = params or CatalogExportParams({
+            "attachments": False,
+            "metadata_fields": False,
+            "data_row_details": False,
+            "project_details": False,
+            "performance_details": False,
+            "label_details": False,
+            "media_type_override": None,
+            "model_runs_ids": None,
+            "projects_ids": None,
+        })
+
+        mutation_name = "exportDataRowsInCatalog"
+        create_task_query_str = """mutation exportDataRowsInCatalogPyApi($input: ExportDataRowsInCatalogInput!){
+            %s(input: $input) {taskId} }
+            """ % (mutation_name)
+
+        data_rows_ids = [data_row.uid for data_row in data_rows]
+        search_query: List[Dict[str, Collection[str]]] = []
+        search_query.append({
+            "ids": data_rows_ids,
+            "operator": "is",
+            "type": "data_row_id"
+        })
+
+        print(search_query)
+        media_type_override = _params.get('media_type_override', None)
+
+        if task_name is None:
+            task_name = f"Export v2: data rows (%s)" % len(data_rows_ids)
+        query_params = {
+            "input": {
+                "taskName": task_name,
+                "filters": {
+                    "searchQuery": {
+                        "scope": None,
+                        "query": search_query
+                    }
+                },
+                "params": {
+                    "mediaTypeOverride":
+                        media_type_override.value
+                        if media_type_override is not None else None,
+                    "includeAttachments":
+                        _params.get('attachments', False),
+                    "includeMetadata":
+                        _params.get('metadata_fields', False),
+                    "includeDataRowDetails":
+                        _params.get('data_row_details', False),
+                    "includeProjectDetails":
+                        _params.get('project_details', False),
+                    "includePerformanceDetails":
+                        _params.get('performance_details', False),
+                    "includeLabelDetails":
+                        _params.get('label_details', False)
+                },
+            }
+        }
+
+        print('export execution')
+        print(client)
+
+        res = client.execute(
+            create_task_query_str,
+            query_params,
+        )
+        print(res)
+        res = res[mutation_name]
+        task_id = res["taskId"]
+        user: User = 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

labelbox/schema/dataset.py

@@ -568,7 +568,7 @@ def export_v2(self,
             "label_details": False,
             "media_type_override": None,
             "model_runs_ids": None,
-            "project_ids": None,
+            "projects_ids": None,
         })
 
         _filters = filters or DatasetExportFilters({

tests/integration/test_data_rows.py

@@ -1,4 +1,5 @@
 from tempfile import NamedTemporaryFile
+import time
 import uuid
 from datetime import datetime
 import json
@@ -962,3 +963,13 @@ def test_create_data_row_with_media_type(dataset, image_url):
     assert "Found invalid contents for media type: \'IMAGE\'" in str(exc.value)
 
     dataset.create_data_row(row_data=image_url, media_type="IMAGE")
+
+
+def test_export_data_rows(client, datarow):
+    # Ensure created data rows are indexed
+    time.sleep(10)
+    task = DataRow.export_v2(client=client, data_rows=[datarow])
+    task.wait_till_done()
+    assert task.status == "COMPLETE"
+    assert task.errors is None
+    assert len(task.result) == 1

tests/integration/test_dataset.py

@@ -1,4 +1,5 @@
 import json
+import time
 import pytest
 import requests
 from labelbox import Dataset
@@ -153,10 +154,13 @@ def test_dataset_export_v2(dataset, image_url):
     ids = set()
     for _ in range(n_data_rows):
         ids.add(dataset.create_data_row(row_data=image_url))
+
+    time.sleep(10)
     task = dataset.export_v2(params={
         "performance_details": False,
         "label_details": True
     })
+    task.wait_till_done()
     assert task.status == "COMPLETE"
     assert task.errors is None
     assert len(task.result) == n_data_rows