Add support to assign multiple datarows to a user

.github/actions/lbox-matrix/index.js

@@ -26814,27 +26814,27 @@ try {
     // To be updated with the new API keys
     {
       "python-version": "3.9",
-      "api-key": "STAGING_LABELBOX_API_KEY_ORG_CMH2FKAOZ032H0735C32V1U63",
+      "api-key": "STAGING_API_KEY_ORG_CMOI3PQ7801GM070D4TPG7FMH",
       "da-test-key": "DA_GCP_LABELBOX_API_KEY"
     },
     {
       "python-version": "3.10",
-      "api-key": "STAGING_LABELBOX_API_KEY_ORG_CMH2G1557037K0726H50N3JQK",
+      "api-key": "STAGING_API_KEY_ORG_CMOI3PQ7801GM070D4TPG7FMH",
       "da-test-key": "DA_GCP_LABELBOX_API_KEY"
     },
     {
       "python-version": "3.11",
-      "api-key": "STAGING_LABELBOX_API_KEY_ORG_CMH2G7STM04WC071F73LG8RSD",
+      "api-key": "STAGING_API_KEY_ORG_CMOI3PQ7801GM070D4TPG7FMH",
       "da-test-key": "DA_GCP_LABELBOX_API_KEY"
     },
     {
       "python-version": "3.12",
-      "api-key": "STAGING_LABELBOX_API_KEY_ORG_CMH2GEJW9033B07299RKLAOFM",
+      "api-key": "STAGING_API_KEY_ORG_CMOI3PQ7801GM070D4TPG7FMH",
       "da-test-key": "DA_GCP_LABELBOX_API_KEY"
     },
     {
       "python-version": "3.13",
-      "api-key": "STAGING_LABELBOX_API_KEY_ORG_CMH2GGV3X04QK071X1EJCH8W0",
+      "api-key": "STAGING_API_KEY_ORG_CMOI3PQ7801GM070D4TPG7FMH",
       "da-test-key": "DA_GCP_LABELBOX_API_KEY"
     },    
   ];

.github/workflows/python-package-develop.yml

@@ -59,19 +59,19 @@ jobs:
       matrix:
         include:
           - python-version: "3.9"
-            api-key: STAGING_LABELBOX_API_KEY_ORG_CMH2FKAOZ032H0735C32V1U63
+            api-key: STAGING_API_KEY_ORG_CMOI3PQ7801GM070D4TPG7FMH 
             da-test-key: DA_GCP_LABELBOX_API_KEY
           - python-version: "3.10"
-            api-key: STAGING_LABELBOX_API_KEY_ORG_CMH2G1557037K0726H50N3JQK
+            api-key: STAGING_API_KEY_ORG_CMOI3PQ7801GM070D4TPG7FMH
             da-test-key: DA_GCP_LABELBOX_API_KEY
           - python-version: "3.11"
-            api-key: STAGING_LABELBOX_API_KEY_ORG_CMH2G7STM04WC071F73LG8RSD
+            api-key: STAGING_API_KEY_ORG_CMOI3PQ7801GM070D4TPG7FMH
             da-test-key: DA_GCP_LABELBOX_API_KEY
           - python-version: "3.12"
-            api-key: STAGING_LABELBOX_API_KEY_ORG_CMH2GEJW9033B07299RKLAOFM
+            api-key: STAGING_API_KEY_ORG_CMOI3PQ7801GM070D4TPG7FMH
             da-test-key: DA_GCP_LABELBOX_API_KEY       
           - python-version: "3.13"
-            api-key: STAGING_LABELBOX_API_KEY_ORG_CMH2GGV3X04QK071X1EJCH8W0
+            api-key: STAGING_API_KEY_ORG_CMOI3PQ7801GM070D4TPG7FMH
             da-test-key: DA_GCP_LABELBOX_API_KEY
     uses: ./.github/workflows/python-package-shared.yml
     with:

libs/labelbox/src/labelbox/__init__.py

@@ -93,6 +93,7 @@
 from labelbox.schema.project_resource_tag import ProjectResourceTag
 from labelbox.schema.media_type import MediaType
 from labelbox.schema.slice import Slice, CatalogSlice, ModelSlice
+from labelbox.schema.task_assignment_status import TaskAssignmentStatus
 from labelbox.schema.task_queue import TaskQueue
 from labelbox.schema.label_score import LabelScore
 from labelbox.schema.identifiables import UniqueIds, GlobalKeys, DataRowIds

libs/labelbox/src/labelbox/schema/api_key.py

@@ -322,27 +322,37 @@ def create_api_key(
         if not user_email or not isinstance(user_email, str):
             raise ValueError("user must be a User object or a valid email")
 
+        role_name = role.name if hasattr(role, "name") else role
+        if not role_name or not isinstance(role_name, str):
+            raise ValueError("role must be a Role object or a valid role name")
+
+        if not isinstance(time_unit, TimeUnit):
+            raise ValueError("time_unit must be a valid TimeUnit enum value")
+
+        if validity < 0:
+            raise ValueError("validity must be a positive integer")
+
+        validity_seconds = validity * time_unit.value
+
+        if validity_seconds < TimeUnit.MINUTE.value:
+            raise ValueError("Minimum validity period is 1 minute")
+
+        max_seconds = 25 * TimeUnit.WEEK.value
+        if validity_seconds > max_seconds:
+            raise ValueError(
+                "Maximum validity period is 6 months (or 25 weeks)"
+            )
+
         # Check if the user exists in the organization
         user_id = ApiKey._get_user(client, user_email)
         if not user_id:
             raise ValueError(
                 f"User with email '{user_email}' does not exist in the organization"
             )
 
-        role_name = role.name if hasattr(role, "name") else role
-        if not role_name or not isinstance(role_name, str):
-            raise ValueError("role must be a Role object or a valid role name")
-
         allowed_roles = ApiKey._get_available_api_key_roles(client)
-        # Determine the exact server role name to pass through.
-        #
-        # - If caller provides a string, require exact match (case-sensitive).
-        # - If caller provides a Role object (which may be normalized by the SDK),
-        #   map it back to the server role name.
         server_role_name: Optional[str] = None
         if hasattr(role, "name"):
-            # Role objects in the SDK are often normalized (e.g. "TENANT_ADMIN").
-            # Map normalized name back to the server-provided role display name.
             normalized_to_server = {format_role(r): r for r in allowed_roles}
             server_role_name = (
                 role_name
@@ -357,24 +367,6 @@ def create_api_key(
                 f"Invalid role specified. Allowed roles are: {allowed_roles}"
             )
 
-        validity_seconds = 0
-        if validity < 0:
-            raise ValueError("validity must be a positive integer")
-
-        if not isinstance(time_unit, TimeUnit):
-            raise ValueError("time_unit must be a valid TimeUnit enum value")
-
-        validity_seconds = validity * time_unit.value
-
-        if validity_seconds < TimeUnit.MINUTE.value:
-            raise ValueError("Minimum validity period is 1 minute")
-
-        max_seconds = 25 * TimeUnit.WEEK.value
-        if validity_seconds > max_seconds:
-            raise ValueError(
-                "Maximum validity period is 6 months (or 25 weeks)"
-            )
-
         query_str = """
          mutation CreateUserApiKeyPyApi($name: String!, $userEmail: String!, $role: String, $validitySeconds: Int) {
              createApiKey(

libs/labelbox/src/labelbox/schema/internal/data_row_upsert_item.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 
-from typing import List, Tuple, Optional
+from typing import List, Tuple, Optional, Union
 
 from labelbox.schema.identifiable import UniqueId, GlobalKey
 from labelbox.schema.data_row import DataRow
@@ -34,7 +34,8 @@ def build(
             if not key:
                 key = {"type": "AUTO", "value": ""}
             elif isinstance(key, key_types):  # type: ignore
-                key = {"type": key.id_type.value, "value": key.key}
+                typed_key: Union[UniqueId, GlobalKey] = key  # type: ignore[assignment]
+                key = {"type": typed_key.id_type.value, "value": typed_key.key}
             else:
                 if not key_types:
                     raise ValueError(

libs/labelbox/src/labelbox/schema/project.py

@@ -48,6 +48,7 @@
 from labelbox.schema.identifiables import (
     DataRowIdentifiers,
 )
+from labelbox.schema.task_assignment_status import TaskAssignmentStatus
 from labelbox.schema.labeling_service import (
     LabelingService,
     LabelingServiceStatus,
@@ -1466,6 +1467,51 @@ def extend_reservations(self, queue_type) -> int:
         res = self.client.execute(query_str, {id_param: self.uid})
         return res["extendReservations"]
 
+    def bulk_assign_data_rows(
+        self,
+        user_id: str,
+        data_row_ids: List[str],
+        allowed_statuses: Optional[List[TaskAssignmentStatus]] = None,
+    ) -> bool:
+        """Assigns multiple data rows to a user in bulk.
+
+        Reserves the specified data rows in the project's initial labeling
+        queue for the given user. Only data rows whose current assignment
+        status matches ``allowed_statuses`` will be assigned.
+
+        Args:
+            user_id: The ID of the user to assign the data rows to.
+            data_row_ids: List of data row IDs to assign.
+            allowed_statuses: Optional list of statuses that a data row must
+                currently have in order to be assigned. Defaults to ``[FREE]``
+                on the server (i.e. only unassigned rows). Pass
+                ``[TaskAssignmentStatus.FREE, TaskAssignmentStatus.RESERVED]``
+                to allow reassignment of already-reserved rows.
+        Returns:
+            True if the bulk assignment succeeded.
+        Raises:
+            LabelboxError: If the GraphQL mutation fails.
+        """
+        if not data_row_ids:
+            return True
+
+        query_str = """mutation BulkAssignDataRowsPyApi($input: BulkAssignDataRowsInput!) {
+            bulkAssignDataRows(input: $input) {
+                success
+            }
+        }"""
+
+        input_dict: Dict[str, Any] = {
+            "projectId": self.uid,
+            "userId": user_id,
+            "dataRowIds": data_row_ids,
+        }
+        if allowed_statuses is not None:
+            input_dict["allowedStatuses"] = [s.value for s in allowed_statuses]
+
+        result = self.client.execute(query_str, {"input": input_dict})
+        return result["bulkAssignDataRows"]["success"]
+
     def enable_model_assisted_labeling(self, toggle: bool = True) -> bool:
         """Turns model assisted labeling either on or off based on input
 

libs/labelbox/src/labelbox/schema/task_assignment_status.py

@@ -0,0 +1,12 @@
+from enum import Enum
+
+
+class TaskAssignmentStatus(str, Enum):
+    """Status filter for bulk data row assignment.
+
+    FREE - only assign data rows that are currently unassigned.
+    RESERVED - only assign data rows that are currently reserved by another user.
+    """
+
+    FREE = "FREE"
+    RESERVED = "RESERVED"

libs/labelbox/tests/integration/test_api_keys.py

@@ -151,8 +151,12 @@ def test_create_api_key_invalid_email_formats(client):
 
 
 def test_create_api_key_invalid_validity_values(client):
-    """Test that providing invalid validity values causes failure."""
-    user_email = client.get_user().email
+    """Test that providing invalid validity values causes failure.
+
+    Validity checks are pure input validation and run before any API calls,
+    so a dummy email is sufficient here.
+    """
+    user_email = "placeholder@labelbox.com"
 
     # Test with negative validity
     with pytest.raises(ValueError) as excinfo:
@@ -200,8 +204,12 @@ def test_create_api_key_invalid_validity_values(client):
 
 
 def test_create_api_key_invalid_time_unit(client):
-    """Test that providing invalid time unit causes failure."""
-    user_email = client.get_user().email
+    """Test that providing invalid time unit causes failure.
+
+    time_unit checks are pure input validation and run before any API calls,
+    so a dummy email is sufficient here.
+    """
+    user_email = "placeholder@labelbox.com"
 
     # Test with None time unit
     with pytest.raises(ValueError) as excinfo:

libs/labelbox/tests/integration/test_bulk_assign.py

@@ -0,0 +1,41 @@
+from labelbox.schema.task_assignment_status import TaskAssignmentStatus
+
+
+def test_bulk_assign_data_rows(
+    configured_batch_project_with_label, project_based_user
+):
+    project, _, data_row, _ = configured_batch_project_with_label
+    user = project_based_user
+
+    result = project.bulk_assign_data_rows(
+        user_id=user.uid,
+        data_row_ids=[data_row.uid],
+    )
+    assert result is True
+
+
+def test_bulk_assign_data_rows_with_allowed_statuses(
+    configured_batch_project_with_label, project_based_user
+):
+    project, _, data_row, _ = configured_batch_project_with_label
+    user = project_based_user
+
+    result = project.bulk_assign_data_rows(
+        user_id=user.uid,
+        data_row_ids=[data_row.uid],
+        allowed_statuses=[
+            TaskAssignmentStatus.FREE,
+            TaskAssignmentStatus.RESERVED,
+        ],
+    )
+    assert result is True
+
+
+def test_bulk_assign_empty_list(configured_batch_project_with_label):
+    project, _, _, _ = configured_batch_project_with_label
+
+    result = project.bulk_assign_data_rows(
+        user_id="any_user_id",
+        data_row_ids=[],
+    )
+    assert result is True