[AL-5662] Add batch export v2 (#1116)

Author: mnoszczak

labelbox/schema/batch.py

@@ -1,4 +1,5 @@
 from typing import Generator, TYPE_CHECKING
+
 from labelbox.orm.db_object import DbObject, experimental
 from labelbox.orm import query
 from labelbox.orm.model import Entity, Field, Relationship

labelbox/schema/export_filters.py

@@ -37,7 +37,11 @@ class SharedExportFilters(TypedDict):
 
 
 class ProjectExportFilters(SharedExportFilters):
-    pass
+    batch_id: Optional[str]
+    """ Batch id to export
+    Example:
+    >>> "clgo3lyax0000veeezdbu3ws4"
+    """
 
 
 class DatasetExportFilters(SharedExportFilters):
@@ -178,4 +182,12 @@ def _get_timezone() -> str:
             "operator": "is",
             "type": "data_row_id"
         })
+
+    batch_id = filters.get("batch_id")
+    if batch_id:
+        search_query.append({
+            "ids": [batch_id],
+            "operator": "is",
+            "type": "batch"
+        })
     return search_query

labelbox/schema/project.py

@@ -442,6 +442,7 @@ def export_v2(self,
             "last_activity_at": None,
             "label_created_at": None,
             "data_row_ids": None,
+            "batch_id": None,
         })
 
         mutation_name = "exportDataRowsInProject"

tests/integration/test_batch.py

@@ -204,4 +204,4 @@ def test_delete_labels_with_templates(batch_project: Project,
     exported_data_rows = list(batch.export_data_rows())
     res = batch.delete_labels(labels_as_template=True)
     exported_data_rows = list(batch.export_data_rows())
-    assert len(exported_data_rows) == 5
+    assert len(exported_data_rows) == 5

tests/integration/test_project.py

@@ -1,12 +1,15 @@
 import time
 import os
+from typing import Tuple
 import uuid
 
 import pytest
 import requests
 
 from labelbox import Project, LabelingFrontend, Dataset
 from labelbox.exceptions import InvalidQueryError
+from labelbox.schema.data_row import DataRow
+from labelbox.schema.label import Label
 from labelbox.schema.media_type import MediaType
 from labelbox.schema.queue_mode import QueueMode
 
@@ -42,6 +45,43 @@ def test_project(client, rand_gen):
     assert project not in projects
 
 
+def test_batch_project_export_v2(
+        configured_batch_project_with_label: Tuple[Project, Dataset, DataRow,
+                                                   Label],
+        export_v2_test_helpers, dataset: Dataset, image_url: str):
+    project, dataset, *_ = configured_batch_project_with_label
+
+    batch = list(project.batches())[0]
+    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"],
+        "batch_id": batch.uid,
+    }
+    params = {
+        "include_performance_details": True,
+        "include_labels": True,
+        "media_type_override": MediaType.Image
+    }
+    task_name = "test_batch_export_v2"
+    task = dataset.create_data_rows([
+        {
+            "row_data": image_url,
+            "external_id": "my-image"
+        },
+    ] * 2)
+    task.wait_till_done()
+    data_rows = [dr.uid for dr in list(dataset.export_data_rows())]
+    batch_one = f'batch one {uuid.uuid4()}'
+
+    # This test creates two batches, only one batch should be exporter
+    # Creatin second batch that will not be used in the export due to the filter: batch_id
+    project.create_batch(batch_one, data_rows)
+
+    task_results = export_v2_test_helpers.run_project_export_v2_task(
+        project, task_name=task_name, filters=filters, params=params)
+    assert (batch.size == len(task_results))
+
+
 def test_project_export_v2(client, export_v2_test_helpers,
                            configured_project_with_label,
                            wait_for_data_row_processing):