Add sdk method to get data row by global key

Author: Val Brodsky

labelbox/client.py

@@ -24,6 +24,7 @@
 from labelbox.pagination import PaginatedCollection
 from labelbox.schema.data_row_metadata import DataRowMetadataOntology
 from labelbox.schema.dataset import Dataset
+from labelbox.schema.data_row import DataRow
 from labelbox.schema.enums import CollectionJobStatus
 from labelbox.schema.iam_integration import IAMIntegration
 from labelbox.schema import role
@@ -430,6 +431,7 @@ def _get_single(self, db_object_type, uid):
                 of the given type for the given ID.
         """
         query_str, params = query.get_single(db_object_type, uid)
+
         res = self.execute(query_str, params)
         res = res and res.get(utils.camel_case(db_object_type.type_name()))
         if res is None:
@@ -727,6 +729,21 @@ def get_data_row(self, data_row_id):
 
         return self._get_single(Entity.DataRow, data_row_id)
 
+    def get_data_row_by_global_key(self, global_key: str) -> DataRow:
+        """
+            Returns: DataRow: returns a single data row given the global key
+        """
+
+        res = self.get_data_row_ids_for_global_keys([global_key])
+        if res['status'] != "SUCCESS":
+            raise labelbox.exceptions.MalformedQueryException(res['errors'][0])
+        if len(res['results']) == 0:
+            raise labelbox.exceptions.ResourceNotFoundError(
+                Entity.DataRow, {global_key: global_key})
+        data_row_id = res['results'][0]
+
+        return self.get_data_row(data_row_id)
+
     def get_data_row_metadata_ontology(self) -> DataRowMetadataOntology:
         """
 

tests/integration/conftest.py

@@ -195,11 +195,12 @@ def small_dataset(dataset: Dataset):
 
 @pytest.fixture
 def data_row(dataset, image_url, rand_gen):
+    global_key = f"global-key-{rand_gen(str)}"
     task = dataset.create_data_rows([
         {
             "row_data": image_url,
             "external_id": "my-image",
-            "global_key": f"global-key-{rand_gen(str)}"
+            "global_key": global_key
         },
     ])
     task.wait_till_done()
@@ -208,6 +209,22 @@ def data_row(dataset, image_url, rand_gen):
     dr.delete()
 
 
+@pytest.fixture
+def data_row_and_global_key(dataset, image_url, rand_gen):
+    global_key = f"global-key-{rand_gen(str)}"
+    task = dataset.create_data_rows([
+        {
+            "row_data": image_url,
+            "external_id": "my-image",
+            "global_key": global_key
+        },
+    ])
+    task.wait_till_done()
+    dr = dataset.data_rows().get_one()
+    yield dr, global_key
+    dr.delete()
+
+
 # can be used with
 # @pytest.mark.parametrize('data_rows', [<count of data rows>], indirect=True)
 # if omitted, count defaults to 1

tests/integration/test_data_rows.py

@@ -118,6 +118,13 @@ def make_metadata_fields_dict():
     return fields
 
 
+def test_get_data_row_by_global_key(data_row_and_global_key, client, rand_gen):
+    _, global_key = data_row_and_global_key
+    data_row = client.get_data_row_by_global_key(global_key)
+    assert type(data_row) == DataRow
+    assert data_row.global_key == global_key
+
+
 def test_get_data_row(data_row, client):
     assert client.get_data_row(data_row.uid)