feat: Use DBResponse for get_run_ids and get_task_ids

services/data/postgres_async_db.py

@@ -109,15 +109,20 @@ def get_table_by_name(self, table_name: str):
     async def get_run_ids(self, flow_id: str, run_id: str):
         run = await self.run_table_postgres.get_run(flow_id, run_id,
                                                     expanded=True)
-        return run.body['run_number'], run.body['run_id']
+        if run.response_code != 200:
+            return run
+        body = {'run_number': run.body['run_number'], 'run_id': run.body['run_id']}
+        return DBResponse(response_code=run.response_code, body=body)
 
     async def get_task_ids(self, flow_id: str, run_id: str,
                            step_name: str, task_name: str):
-
         task = await self.task_table_postgres.get_task(flow_id, run_id,
                                                        step_name, task_name,
                                                        expanded=True)
-        return task.body['task_id'], task.body['task_name']
+        if task.response_code != 200:
+            return task
+        body = {'task_id': task.body['task_id'], 'task_name': task.body['task_name']}
+        return DBResponse(response_code=task.response_code, body=body)
 
 
 class AsyncPostgresDB(object):

services/metadata_service/api/artifact.py

@@ -478,18 +478,18 @@ async def create_artifacts(self, request):
         body = await read_body(request.content)
         count = 0
 
-        try:
-            run_number, run_id = await self._db.get_run_ids(flow_name, run_number)
-            task_id, task_name = await self._db.get_task_ids(
-                flow_name, run_number, step_name, task_id
-            )
-        except Exception:
-            return web.Response(
-                status=400,
-                body=json.dumps(
-                    {"message": "need to register run_id and task_id first"}
-                ),
-            )
+        db_response = await self._db.get_run_ids(flow_name, run_number)
+        if db_response.response_code != 200:
+            return web.Response(status=db_response.response_code,
+                                body=json.dumps(http_500(db_response.body)))
+        run_number, run_id = db_response.body["run_number"], db_response.body["run_id"]
+
+        db_response = await self._db.get_task_ids(flow_name, run_number,
+                                                  step_name, task_id)
+        if db_response.response_code != 200:
+            return web.Response(status=db_response.response_code,
+                                body=json.dumps(http_500(db_response.body)))
+        task_id, task_name = db_response.body["task_id"], db_response.body["task_name"]
 
         # todo change to bulk insert
         for artifact in body:

services/metadata_service/api/metadata.py

@@ -2,7 +2,7 @@
 import json
 from services.utils import read_body
 from services.metadata_service.api.utils import format_response, \
-    handle_exceptions
+    handle_exceptions, http_500
 import asyncio
 from services.data.postgres_async_db import AsyncPostgresDB
 
@@ -173,13 +173,19 @@ async def create_metadata(self, request):
 
         body = await read_body(request.content)
         count = 0
-        try:
-            run_number, run_id = await self._db.get_run_ids(flow_name, run_number)
-            task_id, task_name = await self._db.get_task_ids(flow_name, run_number,
-                                                             step_name, task_id)
-        except Exception:
-            return web.Response(status=400, body=json.dumps(
-                {"message": "need to register run_id and task_id first"}))
+
+        db_response = await self._db.get_run_ids(flow_name, run_number)
+        if db_response.response_code != 200:
+            return web.Response(status=db_response.response_code,
+                                body=json.dumps(http_500(db_response.body)))
+        run_number, run_id = db_response.body["run_number"], db_response.body["run_id"]
+
+        db_response = await self._db.get_task_ids(flow_name, run_number,
+                                                  step_name, task_id)
+        if db_response.response_code != 200:
+            return web.Response(status=db_response.response_code,
+                                body=json.dumps(http_500(db_response.body)))
+        task_id, task_name = db_response.body["task_id"], db_response.body["task_name"]
 
         for datum in body:
             values = {

services/metadata_service/api/step.py

@@ -5,7 +5,6 @@
     handle_exceptions
 from services.data.postgres_async_db import AsyncPostgresDB
 
-
 class StepApi(object):
     _step_table = None
 
@@ -157,7 +156,10 @@ async def create_step(self, request):
         tags = body.get("tags")
         system_tags = body.get("system_tags")
 
-        run_number, run_id = await self._db.get_run_ids(flow_id, run_number)
+        db_response = await self._db.get_run_ids(flow_id, run_number)
+        if db_response.response_code != 200:
+            return db_response
+        run_number, run_id = db_response.body["run_number"], db_response.body["run_id"]
 
         step_row = StepRow(
             flow_id, run_number, run_id, user, step_name, tags=tags,

services/metadata_service/api/task.py

@@ -191,7 +191,10 @@ async def create_task(self, request):
             return web.Response(status=400, body=json.dumps(
                 {"message": "provided task_name may not be a numeric"}))
 
-        run_number, run_id = await self._db.get_run_ids(flow_id, run_number)
+        db_response = await self._db.get_run_ids(flow_id, run_number)
+        if db_response.response_code != 200:
+            return db_response
+        run_number, run_id = db_response.body["run_number"], db_response.body["run_id"]
 
         task = TaskRow(
             flow_id=flow_id,

services/metadata_service/tests/integration_tests/artifact_test.py

@@ -88,35 +88,31 @@ async def test_artifact_post(cli, db):
         cli,
         path="/flows/NonExistentFlow/runs/{run_number}/steps/{step_name}/tasks/{task_id}/artifact".format(**_task),
         payload=payload,
-        status=400,
-        expected_body={"message": "need to register run_id and task_id first"}
+        status=404,
     )
 
     # posting on a non-existent run number should result in an error
     await assert_api_post_response(
         cli,
         path="/flows/{flow_id}/runs/1234/steps/{step_name}/tasks/{task_id}/artifact".format(**_task),
         payload=payload,
-        status=400,
-        expected_body={"message": "need to register run_id and task_id first"}
+        status=404,
     )
 
     # posting on a non-existent step_name should result in an error
     await assert_api_post_response(
         cli,
         path="/flows/{flow_id}/runs/{run_number}/steps/nonexistent/tasks/{task_id}/artifact".format(**_task),
         payload=payload,
-        status=400,
-        expected_body={"message": "need to register run_id and task_id first"}
+        status=404,
     )
 
     # posting on a non-existent task_id should result in an error
     await assert_api_post_response(
         cli,
         path="/flows/{flow_id}/runs/{run_number}/steps/{step_name}/tasks/1234/artifact".format(**_task),
         payload=payload,
-        status=400,
-        expected_body={"message": "need to register run_id and task_id first"}
+        status=404,
     )
 
 

services/metadata_service/tests/integration_tests/metadata_test.py

@@ -68,35 +68,31 @@ async def test_metadata_post(cli, db):
         cli,
         path="/flows/NonExistentFlow/runs/{run_number}/steps/{step_name}/tasks/{task_id}/metadata".format(**_task),
         payload=payload,
-        status=400,
-        expected_body={"message": "need to register run_id and task_id first"}
+        status=404,
     )
 
     # posting on a non-existent run number should result in an error
     await assert_api_post_response(
         cli,
         path="/flows/{flow_id}/runs/1234/steps/{step_name}/tasks/{task_id}/metadata".format(**_task),
         payload=payload,
-        status=400,
-        expected_body={"message": "need to register run_id and task_id first"}
+        status=404,
     )
 
     # posting on a non-existent step_name should result in an error
     await assert_api_post_response(
         cli,
         path="/flows/{flow_id}/runs/{run_number}/steps/nonexistent/tasks/{task_id}/metadata".format(**_task),
         payload=payload,
-        status=400,
-        expected_body={"message": "need to register run_id and task_id first"}
+        status=404,
     )
 
     # posting on a non-existent task_id should result in an error
     await assert_api_post_response(
         cli,
         path="/flows/{flow_id}/runs/{run_number}/steps/{step_name}/tasks/1234/metadata".format(**_task),
         payload=payload,
-        status=400,
-        expected_body={"message": "need to register run_id and task_id first"}
+        status=404,
     )
 
 

services/metadata_service/tests/integration_tests/step_test.py

@@ -55,15 +55,15 @@ def _check_response_body(body):
         cli,
         path="/flows/NonExistentFlow/runs/{run_number}/steps/test_step/step".format(**_run),
         payload=payload,
-        status=500
+        status=404
     )
 
     # posting on a non-existent run number should result in an error
     await assert_api_post_response(
         cli,
         path="/flows/{flow_id}/runs/1234/steps/test_step/step".format(**_run),
         payload=payload,
-        status=500
+        status=404
     )
 
 

services/metadata_service/tests/integration_tests/task_test.py

@@ -48,15 +48,15 @@ def _check_response_body(body):
         cli,
         path="/flows/NonExistentFlow/runs/{run_number}/steps/{step_name}/task".format(**_task),
         payload=payload,
-        status=500
+        status=404
     )
 
     # posting on a non-existent run number should result in an error
     await assert_api_post_response(
         cli,
         path="/flows/{flow_id}/runs/1234/steps/{step_name}/task".format(**_task),
         payload=payload,
-        status=500
+        status=404
     )
 
     # posting on a non-existent step_name should result in a 404 due to foreign key constraint