update to have optional include metadata field

Author: jtsodapop

labelbox/schema/batch.py

@@ -77,7 +77,7 @@ def remove_queued_data_rows(self) -> None:
                 },
                             experimental=True)
 
-    def export_data_rows(self, timeout_seconds=120) -> Generator:
+    def export_data_rows(self, timeout_seconds=120, include_metadata: bool=False) -> Generator:
         """ Returns a generator that produces all data rows that are currently
         in this batch.
 
@@ -92,22 +92,23 @@ def export_data_rows(self, timeout_seconds=120) -> Generator:
             LabelboxError: if the export fails or is unable to download within the specified time.
         """
         id_param = "batchId"
-        query_str = """mutation GetBatchDataRowsExportUrlPyApi($%s: ID!)
-            {exportBatchDataRows(data:{batchId: $%s }) {downloadUrl createdAt status}}
-        """ % (id_param, id_param)
+        metadata_param = "includeMetadataInput"
+        query_str = """mutation GetBatchDataRowsExportUrlPyApi($%s: ID!, $%s: Boolean!)
+            {exportBatchDataRows(data:{batchId: $%s , includeMetadataInput: $%s}) {downloadUrl createdAt status}}
+        """ % (id_param, metadata_param, id_param, metadata_param)
         sleep_time = 2
         while True:
-            res = self.client.execute(query_str, {id_param: self.uid})
+            res = self.client.execute(query_str, {id_param: self.uid, metadata_param: include_metadata})
             res = res["exportBatchDataRows"]
             if res["status"] == "COMPLETE":
                 download_url = res["downloadUrl"]
                 response = requests.get(download_url)
                 response.raise_for_status()
                 reader = ndjson.reader(StringIO(response.text))
-                # TODO: Update result to parse metadataFields when resolver returns
                 return (Entity.DataRow(self.client, {
-                    **result, 'metadataFields': [],
-                    'customMetadata': []
+                    **result, 
+                    'customMetadata': result['metadata'],
+                    'metadataFields': result['metadataFields']
                 }) for result in reader)
             elif res["status"] == "FAILED":
                 raise LabelboxError("Data row export failed.")

labelbox/schema/dataset.py

@@ -413,7 +413,7 @@ def data_row_for_external_id(self, external_id) -> "DataRow":
                 external_id)
         return data_rows[0]
 
-    def export_data_rows(self, timeout_seconds=120) -> Generator:
+    def export_data_rows(self, timeout_seconds=120, include_metadata: bool=False) -> Generator:
         """ Returns a generator that produces all data rows that are currently
         attached to this dataset.
 
@@ -428,22 +428,23 @@ def export_data_rows(self, timeout_seconds=120) -> Generator:
             LabelboxError: if the export fails or is unable to download within the specified time.
         """
         id_param = "datasetId"
-        query_str = """mutation GetDatasetDataRowsExportUrlPyApi($%s: ID!)
-            {exportDatasetDataRows(data:{datasetId: $%s }) {downloadUrl createdAt status}}
-        """ % (id_param, id_param)
+        metadata_param = "includeMetadataInput"
+        query_str = """mutation GetDatasetDataRowsExportUrlPyApi($%s: ID!, $%s: Boolean!)
+            {exportDatasetDataRows(data:{datasetId: $%s , includeMetadataInput: $%s}) {downloadUrl createdAt status}}
+        """ % (id_param, metadata_param, id_param, metadata_param)
         sleep_time = 2
         while True:
-            res = self.client.execute(query_str, {id_param: self.uid})
+            res = self.client.execute(query_str, {id_param: self.uid, metadata_param: include_metadata})
             res = res["exportDatasetDataRows"]
             if res["status"] == "COMPLETE":
                 download_url = res["downloadUrl"]
                 response = requests.get(download_url)
                 response.raise_for_status()
                 reader = ndjson.reader(StringIO(response.text))
-                # TODO: Update result to parse metadataFields when resolver returns
                 return (Entity.DataRow(self.client, {
-                    **result, 'metadataFields': [],
-                    'customMetadata': []
+                    **result, 
+                    'customMetadata': result['metadata'],
+                    'metadataFields': result['metadataFields']
                 }) for result in reader)
             elif res["status"] == "FAILED":
                 raise LabelboxError("Data row export failed.")
@@ -457,3 +458,5 @@ def export_data_rows(self, timeout_seconds=120) -> Generator:
             logger.debug("Dataset '%s' data row export, waiting for server...",
                          self.uid)
             time.sleep(sleep_time)
+
+

labelbox/schema/project.py

@@ -186,7 +186,7 @@ def labels(self, datasets=None, order_by=None) -> PaginatedCollection:
                                    ["project", "labels"], Label)
 
     def export_queued_data_rows(self,
-                                timeout_seconds=120) -> List[Dict[str, str]]:
+                                timeout_seconds=120,  include_metadata: bool=False) -> List[Dict[str, str]]:
         """ Returns all data rows that are currently enqueued for this project.
 
         Args:
@@ -197,12 +197,13 @@ def export_queued_data_rows(self,
             LabelboxError: if the export fails or is unable to download within the specified time.
         """
         id_param = "projectId"
-        query_str = """mutation GetQueuedDataRowsExportUrlPyApi($%s: ID!)
-            {exportQueuedDataRows(data:{projectId: $%s }) {downloadUrl createdAt status} }
-        """ % (id_param, id_param)
+        metadata_param = "includeMetadataInput"
+        query_str = """mutation GetQueuedDataRowsExportUrlPyApi($%s: ID!, $%s: Boolean!)
+            {exportQueuedDataRows(data:{projectId: $%s , includeMetadataInput: $%s}) {downloadUrl createdAt status} }
+        """ % (id_param, metadata_param, id_param, metadata_param)
         sleep_time = 2
         while True:
-            res = self.client.execute(query_str, {id_param: self.uid})
+            res = self.client.execute(query_str, {id_param: self.uid, metadata_param: include_metadata})
             res = res["exportQueuedDataRows"]
             if res["status"] == "COMPLETE":
                 download_url = res["downloadUrl"]