[AL-7595] Add data_row identifier validation for running foundry app (#1323)

Co-authored-by: Val Brodsky <vbrodsky@users.noreply.github.com>

Author: whistler

labelbox/schema/foundry/foundry_client.py

@@ -1,8 +1,7 @@
 from typing import Union
 from labelbox import exceptions
 from labelbox.schema.foundry.app import App, APP_FIELD_NAMES
-from labelbox.schema.foundry.model import Model, MODEL_FIELD_NAMES
-from labelbox.schema.identifiables import DataRowIds, GlobalKeys
+from labelbox.schema.identifiables import DataRowIds, GlobalKeys, IdType
 from labelbox.schema.task import Task
 
 
@@ -78,23 +77,20 @@ def run_app(self, model_run_name: str,
                 data_rows: Union[DataRowIds, GlobalKeys], app_id: str) -> Task:
         app = self._get_app(app_id)
 
-        data_rows_query = self.client.build_catalog_query(data_rows)
-
         params = {
             "modelId": str(app.model_id),
             "name": model_run_name,
             "classToSchemaId": app.class_to_schema_id,
             "inferenceParams": app.inference_params,
-            "searchQuery": {
-                "query": [data_rows_query],
-                "scope": None
-            },
             "ontologyId": app.ontology_id
         }
 
+        data_rows_key = "dataRowIds" if data_rows.id_type == IdType.DataRowId else "globalKeys"
+        params[data_rows_key] = list(data_rows)
+
         query = """
-        mutation CreateModelJobPyApi($input: CreateModelJobInput!) {
-            createModelJob(input: $input) {
+        mutation CreateModelJobPyApi($input: CreateModelJobForDataRowsInput!) {
+            createModelJobForDataRows(input: $input) {
                 taskId
                 __typename
             }
@@ -104,5 +100,5 @@ def run_app(self, model_run_name: str,
             response = self.client.execute(query, {"input": params})
         except Exception as e:
             raise exceptions.LabelboxError('Unable to run foundry app', e)
-        task_id = response["createModelJob"]["taskId"]
+        task_id = response["createModelJobForDataRows"]["taskId"]
         return Task.get_task(self.client, task_id)

tests/integration/test_foundry.py

@@ -18,6 +18,23 @@ def foundry_client(client):
     return FoundryClient(client)
 
 
+@pytest.fixture()
+def text_data_row(dataset, random_str):
+    global_key = "https://storage.googleapis.com/lb-artifacts-testing-public/sdk_integration_test/sample-text-1.txt-{random_str}"
+    task = dataset.create_data_rows([{
+        "row_data":
+            "https://storage.googleapis.com/lb-artifacts-testing-public/sdk_integration_test/sample-text-1.txt",
+        "media_type":
+            "TEXT",
+        "global_key":
+            global_key
+    }])
+    task.wait_till_done()
+    dr = dataset.data_rows().get_one()
+    yield dr
+    dr.delete()
+
+
 @pytest.fixture()
 def ontology(client, random_str):
     object_features = [
@@ -108,13 +125,23 @@ def test_run_foundry_app_returns_model_run_id(foundry_client, data_row, app):
     assert model_run.uid == model_run_id
 
 
-def test_run_foundry_app_with_non_existent_data_rows(foundry_client, data_row,
-                                                     app, random_str):
-    data_rows = lb.GlobalKeys([data_row.global_key, "non-existent-global-key"])
-    task = foundry_client.run_app(
-        model_run_name=f"test-app-with-wrong-key-{random_str}",
-        data_rows=data_rows,
-        app_id=app.id)
-    task.wait_till_done()
-    # The incorrect data row is filtered out and the task still completes with the correct data row
-    assert task.status == 'COMPLETE'
+def test_run_foundry_with_invalid_data_row_id(foundry_client, app, random_str):
+    invalid_datarow_id = 'invalid-global-key'
+    data_rows = lb.GlobalKeys([invalid_datarow_id])
+    with pytest.raises(lb.exceptions.LabelboxError) as exception:
+        foundry_client.run_app(
+            model_run_name=f"test-app-with-invalid-datarow-id-{random_str}",
+            data_rows=data_rows,
+            app_id=app.id)
+        assert invalid_datarow_id in exception.value
+
+
+def test_run_foundry_with_invalid_global_key(foundry_client, app, random_str):
+    invalid_global_key = 'invalid-global-key'
+    data_rows = lb.GlobalKeys([invalid_global_key])
+    with pytest.raises(lb.exceptions.LabelboxError) as exception:
+        foundry_client.run_app(
+            model_run_name=f"test-app-with-invalid-global-key-{random_str}",
+            data_rows=data_rows,
+            app_id=app.id)
+        assert invalid_global_key in exception.value