Merge pull request #565 from Labelbox/kkim/AL-2218

[AL-2218] Adding optional metadata_fields to create_data_row()

Author: kkim-labelbox

labelbox/client.py

@@ -37,6 +37,7 @@
 logger = logging.getLogger(__name__)
 
 _LABELBOX_API_KEY = "LABELBOX_API_KEY"
+_DATAROW_METADATA_CREATE_ERROR = "Failed to validate the metadata"
 
 
 class Client:
@@ -90,6 +91,7 @@ def __init__(self,
             'Authorization': 'Bearer %s' % api_key,
             'X-User-Agent': f'python-sdk {SDK_VERSION}'
         }
+        self._data_row_metadata_ontology = None
 
     @retry.Retry(predicate=retry.if_exception_type(
         labelbox.exceptions.InternalServerError))
@@ -269,6 +271,8 @@ def get_error_status_code(error):
 
             if get_error_status_code(internal_server_error) == 400:
                 raise labelbox.exceptions.InvalidQueryError(message)
+            elif _DATAROW_METADATA_CREATE_ERROR in message:
+                raise labelbox.exceptions.ResourceCreationError(message)
             else:
                 raise labelbox.exceptions.InternalServerError(message)
 
@@ -648,7 +652,9 @@ def get_data_row_metadata_ontology(self) -> DataRowMetadataOntology:
             DataRowMetadataOntology: The ontology for Data Row Metadata for an organization
 
         """
-        return DataRowMetadataOntology(self)
+        if self._data_row_metadata_ontology is None:
+            self._data_row_metadata_ontology = DataRowMetadataOntology(self)
+        return self._data_row_metadata_ontology
 
     def get_model(self, model_id) -> Model:
         """ Gets a single Model with the given ID.

labelbox/exceptions.py

@@ -71,6 +71,12 @@ class InvalidQueryError(LabelboxError):
     pass
 
 
+class ResourceCreationError(LabelboxError):
+    """ Indicates that a resource could not be created in the server side
+    due to a validation or transaction error"""
+    pass
+
+
 class NetworkError(LabelboxError):
     """Raised when an HTTPError occurs."""
 
@@ -122,4 +128,4 @@ class MALValidationError(LabelboxError):
 
 class OperationNotAllowedException(Exception):
     """Raised when user does not have permissions to a resource or has exceeded usage limit"""
-    pass
+    pass

labelbox/orm/model.py

@@ -32,8 +32,9 @@ class Field:
     Attributes:
         field_type (Field.Type): The type of the field.
         name (str): name that the attribute has in client-side Python objects
-        grapgql_name (str): name that the attribute has in queries (and in
+        graphql_name (str): name that the attribute has in queries (and in
             server-side database definition).
+        result_subquery (str): graphql query result payload for a field.
     """
 
     class Type(Enum):
@@ -54,6 +55,25 @@ def __init__(self, enum_cls: type):
         def name(self):
             return self.enum_cls.__name__
 
+    class ListType:
+        """ Represents Field that is a list of some object.
+        Args:
+            list_cls (type): Type of object that list is made of.
+            graphql_type (str): Inner object's graphql type.
+                By default, the list_cls's name is used as the graphql type.
+        """
+
+        def __init__(self, list_cls: type, graphql_type=None):
+            self.list_cls = list_cls
+            if graphql_type is None:
+                self.graphql_type = self.list_cls.__name__
+            else:
+                self.graphql_type = graphql_type
+
+        @property
+        def name(self):
+            return f"[{self.graphql_type}]"
+
     class Order(Enum):
         """ Type of sort ordering. """
         Asc = auto()
@@ -91,10 +111,15 @@ def Enum(enum_cls: type, *args):
     def Json(*args):
         return Field(Field.Type.Json, *args)
 
+    @staticmethod
+    def List(list_cls: type, graphql_type=None, **kwargs):
+        return Field(Field.ListType(list_cls, graphql_type), **kwargs)
+
     def __init__(self,
-                 field_type: Union[Type, EnumType],
+                 field_type: Union[Type, EnumType, ListType],
                  name,
-                 graphql_name=None):
+                 graphql_name=None,
+                 result_subquery=None):
         """ Field init.
         Args:
             field_type (Field.Type): The type of the field.
@@ -103,12 +128,14 @@ def __init__(self,
             graphql_name (str): query and server-side name of a database object.
                 If None, it is constructed from the client-side name by converting
                 snake_case (Python convention) into camelCase (GraphQL convention).
+            result_subquery (str): graphql query result payload for a field.
         """
         self.field_type = field_type
         self.name = name
         if graphql_name is None:
             graphql_name = utils.camel_case(name)
         self.graphql_name = graphql_name
+        self.result_subquery = result_subquery
 
     @property
     def asc(self):

labelbox/orm/query.py

@@ -41,7 +41,11 @@ def results_query_part(entity):
         entity (type): The entity which needs fetching.
     """
     # Query for fields
-    fields = [field.graphql_name for field in entity.fields()]
+    fields = [
+        field.result_subquery
+        if field.result_subquery is not None else field.graphql_name
+        for field in entity.fields()
+    ]
 
     # Query for cached relationships
     fields.extend([

labelbox/schema/batch.py

@@ -104,8 +104,10 @@ def export_data_rows(self, timeout_seconds=120) -> Generator:
                 response = requests.get(download_url)
                 response.raise_for_status()
                 reader = ndjson.reader(StringIO(response.text))
-                return (
-                    Entity.DataRow(self.client, result) for result in reader)
+                # TODO: Update result to parse customMetadata when resolver returns
+                return (Entity.DataRow(self.client, {
+                    **result, 'customMetadata': []
+                }) for result in reader)
             elif res["status"] == "FAILED":
                 raise LabelboxError("Data row export failed.")
 

labelbox/schema/data_row.py

@@ -1,10 +1,10 @@
 import logging
-from datetime import datetime
-from typing import List, Dict, Union, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 from labelbox.orm import query
 from labelbox.orm.db_object import DbObject, Updateable, BulkDeletable
 from labelbox.orm.model import Entity, Field, Relationship
+from labelbox.schema.data_row_metadata import DataRowMetadataField  # type: ignore
 
 if TYPE_CHECKING:
     from labelbox import AssetAttachment
@@ -22,6 +22,7 @@ class DataRow(DbObject, Updateable, BulkDeletable):
         updated_at (datetime)
         created_at (datetime)
         media_attributes (dict): generated media attributes for the datarow
+        custom_metadata (list): metadata associated with the datarow
 
         dataset (Relationship): `ToOne` relationship to Dataset
         created_by (Relationship): `ToOne` relationship to User
@@ -34,6 +35,11 @@ class DataRow(DbObject, Updateable, BulkDeletable):
     updated_at = Field.DateTime("updated_at")
     created_at = Field.DateTime("created_at")
     media_attributes = Field.Json("media_attributes")
+    custom_metadata = Field.List(
+        DataRowMetadataField,
+        graphql_type="DataRowCustomMetadataUpsertInput!",
+        name="custom_metadata",
+        result_subquery="customMetadata { value schemaId }")
 
     # Relationships
     dataset = Relationship.ToOne("Dataset")
@@ -95,4 +101,4 @@ def create_attachment(self, attachment_type,
                 data_row_id_param: self.uid
             })
         return Entity.AssetAttachment(self.client,
-                                      res["createDataRowAttachment"])
+                                      res["createDataRowAttachment"])

labelbox/schema/data_row_metadata.py

@@ -200,6 +200,10 @@ def _parse_ontology(raw_ontology) -> List[DataRowMetadataSchema]:
 
         return fields
 
+    def refresh_ontology(self):
+        self._raw_ontology = self._get_ontology()
+        self._build_ontology()
+
     def parse_metadata(
         self, unparsed: List[Dict[str,
                                   List[Union[str,
@@ -221,6 +225,14 @@ def parse_metadata(
         for dr in unparsed:
             fields = []
             for f in dr["fields"]:
+                if f["schemaId"] not in self.fields_by_id:
+                    # Update metadata ontology if field can't be found
+                    self.refresh_ontology()
+                    if f["schemaId"] not in self.fields_by_id:
+                        raise ValueError(
+                            f"Schema Id `{f['schemaId']}` not found in ontology"
+                        )
+
                 schema = self.fields_by_id[f["schemaId"]]
                 if schema.kind == DataRowMetadataKind.enum:
                     continue
@@ -295,9 +307,8 @@ def _batch_upsert(
                     data_row_id=m.data_row_id,
                     fields=list(
                         chain.from_iterable(
-                            self._parse_upsert(m) for m in m.fields))).dict(
+                            self.parse_upsert(m) for m in m.fields))).dict(
                                 by_alias=True))
-
         res = _batch_operations(_batch_upsert, items, self._batch_size)
         return res
 
@@ -393,14 +404,17 @@ def _bulk_export(_data_row_ids: List[str]) -> List[DataRowMetadata]:
                                  data_row_ids,
                                  batch_size=self._batch_size)
 
-    def _parse_upsert(
+    def parse_upsert(
             self, metadatum: DataRowMetadataField
     ) -> List[_UpsertDataRowMetadataInput]:
         """Format for metadata upserts to GQL"""
 
         if metadatum.schema_id not in self.fields_by_id:
-            raise ValueError(
-                f"Schema Id `{metadatum.schema_id}` not found in ontology")
+            # Update metadata ontology if field can't be found
+            self.refresh_ontology()
+            if metadatum.schema_id not in self.fields_by_id:
+                raise ValueError(
+                    f"Schema Id `{metadatum.schema_id}` not found in ontology")
 
         schema = self.fields_by_id[metadatum.schema_id]
 
@@ -428,8 +442,11 @@ def _validate_delete(self, delete: DeleteDataRowMetadata):
         deletes = set()
         for schema_id in delete.fields:
             if schema_id not in self.fields_by_id:
-                raise ValueError(
-                    f"Schema Id `{schema_id}` not found in ontology")
+                # Update metadata ontology if field can't be found
+                self.refresh_ontology()
+                if schema_id not in self.fields_by_id:
+                    raise ValueError(
+                        f"Schema Id `{schema_id}` not found in ontology")
 
             schema = self.fields_by_id[schema_id]
             # handle users specifying enums by adding all option enums

labelbox/schema/dataset.py

@@ -5,7 +5,7 @@
 from collections.abc import Iterable
 import time
 import ndjson
-from itertools import islice
+from itertools import islice, chain
 
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from io import StringIO
@@ -78,6 +78,18 @@ def create_data_row(self, **kwargs) -> "DataRow":
         if os.path.exists(row_data):
             kwargs[DataRow.row_data.name] = self.client.upload_file(row_data)
         kwargs[DataRow.dataset.name] = self
+
+        # Parse metadata fields, if they are provided
+        if DataRow.custom_metadata.name in kwargs:
+            mdo = self.client.get_data_row_metadata_ontology()
+            metadata_fields = kwargs[DataRow.custom_metadata.name]
+            metadata = list(
+                chain.from_iterable(
+                    mdo.parse_upsert(m) for m in metadata_fields))
+            kwargs[DataRow.custom_metadata.name] = [
+                md.dict(by_alias=True) for md in metadata
+            ]
+
         return self.client._create(DataRow, kwargs)
 
     def create_data_rows_sync(self, items) -> None:
@@ -397,8 +409,10 @@ def export_data_rows(self, timeout_seconds=120) -> Generator:
                 response = requests.get(download_url)
                 response.raise_for_status()
                 reader = ndjson.reader(StringIO(response.text))
-                return (
-                    Entity.DataRow(self.client, result) for result in reader)
+                # TODO: Update result to parse customMetadata when resolver returns
+                return (Entity.DataRow(self.client, {
+                    **result, 'customMetadata': []
+                }) for result in reader)
             elif res["status"] == "FAILED":
                 raise LabelboxError("Data row export failed.")
 

tests/integration/test_client_errors.py

@@ -2,11 +2,21 @@
 import os
 import time
 import pytest
+from datetime import datetime
 
 from labelbox import Project, Dataset, User
+from labelbox.schema.data_row_metadata import DataRowMetadataField
 import labelbox.client
 import labelbox.exceptions
 
+SPLIT_SCHEMA_ID = "cko8sbczn0002h2dkdaxb5kal"
+TEST_SPLIT_ID = "cko8scbz70005h2dkastwhgqt"
+EMBEDDING_SCHEMA_ID = "ckpyije740000yxdk81pbgjdc"
+TEXT_SCHEMA_ID = "cko8s9r5v0001h2dk9elqdidh"
+CAPTURE_DT_SCHEMA_ID = "cko8sdzv70006h2dk8jg64zvb"
+IMAGE_EMBEDDING_SCHEMA_ID = "ckrzang79000008l6hb5s6za1"
+TEXT_EMBEDDING_SCHEMA_ID = "ckrzao09x000108l67vrcdnh3"
+
 
 def test_missing_api_key():
     key = os.environ.get(labelbox.client._LABELBOX_API_KEY, None)
@@ -132,3 +142,28 @@ def get(arg):
 
     # Sleep at the end of this test to allow other tests to execute.
     time.sleep(60)
+
+
+@pytest.mark.skip("Staging environment not returning correct exception")
+def test_resource_creation_error(dataset, image_url):
+
+    def make_metadata_fields():
+        embeddings = [0.0] * 128
+        msg = "A message"
+        time = datetime.utcnow()
+
+        fields = [
+            DataRowMetadataField(schema_id=SPLIT_SCHEMA_ID,
+                                 value=TEST_SPLIT_ID),
+            DataRowMetadataField(schema_id=CAPTURE_DT_SCHEMA_ID, value=time),
+            DataRowMetadataField(schema_id=TEXT_SCHEMA_ID, value=msg),
+            DataRowMetadataField(schema_id=EMBEDDING_SCHEMA_ID,
+                                 value=embeddings),
+            DataRowMetadataField(schema_id=EMBEDDING_SCHEMA_ID,
+                                 value=embeddings),
+        ]
+        return fields
+
+    with pytest.raises(labelbox.exceptions.ResourceCreationError) as excinfo:
+        dataset.create_data_row(row_data=image_url,
+                                custom_metadata=make_metadata_fields())

tests/integration/test_data_row_metadata.py

@@ -274,4 +274,4 @@ def test_parse_raw_metadata(mdo):
 
     for row in parsed:
         for field in row.fields:
-            assert mdo._parse_upsert(field)
+            assert mdo.parse_upsert(field)