fetch data row ids by external ids

Author: Matt Sokoloff

labelbox/client.py

@@ -1,6 +1,8 @@
 # type: ignore
 from datetime import datetime, timezone
 import json
+from typing import List, Dict
+from collections import defaultdict
 
 import logging
 import mimetypes
@@ -658,3 +660,29 @@ def create_model(self, name, ontology_id):
             "ontologyId": ontology_id
         })
         return Model(self, result['createModel'])
+
+    def get_data_row_ids_for_external_ids(
+            self, external_ids: List[str]) -> Dict[str, List[str]]:
+        """
+        Returns a list of data row ids for a list of external ids.
+        There is a max of 1500 items returned at a time.
+
+        Args:
+            external_ids: List of external ids to fetch data row ids for
+
+        Returns:
+            A dict of external ids as keys and values as a list of data row ids that correspond to that external id.
+        """
+        query_str = """query externalIdsToDataRowIdsPyApi($externalId_in: [String!]!){
+            externalIdsToDataRowIds(externalId_in: $externalId_in) { dataRowId externalId }
+        }
+        """
+        max_n_per_request = 100
+        result = defaultdict(list)
+        for i in range(0, len(external_ids), max_n_per_request):
+            for row in self.execute(
+                    query_str,
+                {'externalId_in': external_ids[i:i + max_n_per_request]
+                })['externalIdsToDataRowIds']:
+                result[row['externalId']].append(row['dataRowId'])
+        return result

tests/integration/test_data_rows.py

@@ -1,4 +1,6 @@
 from tempfile import NamedTemporaryFile
+import uuid
+import time
 
 import pytest
 import requests
@@ -11,6 +13,36 @@ def test_get_data_row(datarow, client):
     assert client.get_data_row(datarow.uid)
 
 
+def test_lookup_data_rows(client, dataset):
+    uid = str(uuid.uuid4())
+    # 1 external id : 1 uid
+    dr = dataset.create_data_row(row_data="123", external_id=uid)
+    lookup = client.get_data_row_ids_for_external_ids([uid])
+    assert len(lookup) == 1
+    assert lookup[uid][0] == dr.uid
+    # 2 external ids : 1 uid
+    uid2 = str(uuid.uuid4())
+    dr2 = dataset.create_data_row(row_data="123", external_id=uid2)
+    lookup = client.get_data_row_ids_for_external_ids([uid, uid2])
+    assert len(lookup) == 2
+    assert all([len(x) == 1 for x in lookup.values()])
+    assert lookup[uid][0] == dr.uid
+    assert lookup[uid2][0] == dr2.uid
+    #1 external id : 2 uid
+    dr3 = dataset.create_data_row(row_data="123", external_id=uid2)
+    lookup = client.get_data_row_ids_for_external_ids([uid2])
+    assert len(lookup) == 1
+    assert len(lookup[uid2]) == 2
+    assert lookup[uid2][0] == dr2.uid
+    assert lookup[uid2][1] == dr3.uid
+    # Empty args
+    lookup = client.get_data_row_ids_for_external_ids([])
+    assert len(lookup) == 0
+    # Non matching
+    lookup = client.get_data_row_ids_for_external_ids([str(uuid.uuid4())])
+    assert len(lookup) == 0
+
+
 def test_data_row_bulk_creation(dataset, rand_gen, image_url):
     client = dataset.client
     assert len(list(dataset.data_rows())) == 0