[AL-5332] allow user to update datarow metadata by global keys (#1060)

Author: yamini-labelbox

labelbox/schema/data_row_metadata.py

@@ -48,7 +48,8 @@ class DataRowMetadataField(_CamelCaseMixin):
 
 
 class DataRowMetadata(_CamelCaseMixin):
-    data_row_id: str
+    global_key: Optional[str]
+    data_row_id: Optional[str]
     fields: List[DataRowMetadataField]
 
 
@@ -58,7 +59,8 @@ class DeleteDataRowMetadata(_CamelCaseMixin):
 
 
 class DataRowMetadataBatchResponse(_CamelCaseMixin):
-    data_row_id: str
+    global_key: Optional[str]
+    data_row_id: Optional[str]
     error: Optional[str] = None
     fields: List[Union[DataRowMetadataField, SchemaId]]
 
@@ -75,7 +77,8 @@ class _UpsertDataRowMetadataInput(_CamelCaseMixin):
 
 # Batch of upsert values for a datarow
 class _UpsertBatchDataRowMetadata(_CamelCaseMixin):
-    data_row_id: str
+    global_key: Optional[str]
+    data_row_id: Optional[str]
     fields: List[_UpsertDataRowMetadataInput]
 
 
@@ -476,11 +479,12 @@ def parse_metadata_fields(
     def bulk_upsert(
             self, metadata: List[DataRowMetadata]
     ) -> List[DataRowMetadataBatchResponse]:
-        """Upsert datarow metadata
-
+        """Upsert metadata to a list of data rows
+        
+        You may specify data row by either data_row_id or global_key
 
         >>> metadata = DataRowMetadata(
-        >>>                 data_row_id="datarow-id",
+        >>>                 data_row_id="datarow-id", # Alternatively, set global_key="global-key"
         >>>                 fields=[
         >>>                        DataRowMetadataField(schema_id="schema-id", value="my-message"),
         >>>                        ...
@@ -504,6 +508,7 @@ def _batch_upsert(
         ) -> List[DataRowMetadataBatchResponse]:
             query = """mutation UpsertDataRowMetadataBetaPyApi($metadata: [DataRowCustomMetadataBatchUpsertInput!]!) {
                 upsertDataRowCustomMetadata(data: $metadata){
+                    globalKey
                     dataRowId
                     error
                     fields {
@@ -515,7 +520,8 @@ def _batch_upsert(
             res = self._client.execute(
                 query, {"metadata": upserts})['upsertDataRowCustomMetadata']
             return [
-                DataRowMetadataBatchResponse(data_row_id=r['dataRowId'],
+                DataRowMetadataBatchResponse(global_key=r['globalKey'],
+                                             data_row_id=r['dataRowId'],
                                              error=r['error'],
                                              fields=self.parse_metadata(
                                                  [r])[0].fields) for r in res
@@ -525,6 +531,7 @@ def _batch_upsert(
         for m in metadata:
             items.append(
                 _UpsertBatchDataRowMetadata(
+                    global_key=m.global_key,
                     data_row_id=m.data_row_id,
                     fields=list(
                         chain.from_iterable(

tests/integration/test_data_row_metadata.py

@@ -53,11 +53,12 @@ def big_dataset(dataset: Dataset, image_url):
     yield dataset
 
 
-def make_metadata(dr_id) -> DataRowMetadata:
+def make_metadata(dr_id: str = None, gk: str = None) -> DataRowMetadata:
     msg = "A message"
     time = datetime.utcnow()
 
     metadata = DataRowMetadata(
+        global_key=gk,
         data_row_id=dr_id,
         fields=[
             DataRowMetadataField(schema_id=SPLIT_SCHEMA_ID,
@@ -122,6 +123,14 @@ def test_bulk_upsert_datarow_metadata(data_row, mdo: DataRowMetadataOntology):
     assert len([field for field in exported[0].fields]) == 3
 
 
+def test_bulk_upsert_datarow_metadata_by_globalkey(
+        data_rows, mdo: DataRowMetadataOntology):
+    global_keys = [data_row.global_key for data_row in data_rows]
+    metadata = [make_metadata(gk=global_key) for global_key in global_keys]
+    errors = mdo.bulk_upsert(metadata)
+    assert len(errors) == 0
+
+
 @pytest.mark.slow
 def test_large_bulk_upsert_datarow_metadata(big_dataset, mdo):
     metadata = []