feature: metadata driven run and task status

services/data/postgres_async_db.py

@@ -39,7 +39,7 @@
 operator_match = re.compile('([^:]*):([=><]+)$')
 
 # use a ddmmyyy timestamp as the version for triggers
-TRIGGER_VERSION = "05092024"
+TRIGGER_VERSION = "05022026"
 TRIGGER_NAME_PREFIX = "notify_ui"
 
 
@@ -293,7 +293,7 @@ async def _execute_on_cursor(_cur):
             self.db.logger.exception("Exception occurred")
             return aiopg_exception_handling(error), None
 
-    async def create_record(self, record_dict):
+    async def create_record(self, record_dict, expanded=False):
         # note: need to maintain order
         cols = []
         values = []
@@ -333,7 +333,7 @@ async def create_record(self, record_dict):
                 for key, value in record.items():
                     if key in self.keys:
                         filtered_record[key] = value
-                response_body = self._row_type(**filtered_record).serialize()  # pylint: disable=not-callable
+                response_body = self._row_type(**filtered_record).serialize(expanded)  # pylint: disable=not-callable
                 # todo make sure connection is closed even with error
                 cur.close()
             return DBResponse(response_code=200, body=response_body)
@@ -418,6 +418,68 @@ async def _execute_update_on_cursor(_cur):
             self.db.logger.exception("Exception occurred")
             return aiopg_exception_handling(error)
 
+    # TODO: Cleanup. Placeholder for experimentation
+    async def update_row_with_attempt(self, filter_dict={}, update_dict={}, attempt_id: int = 0, cur: aiopg.Cursor = None):
+        # generate where clause
+        filters = []
+        for col_name, col_val in filter_dict.items():
+            operator = '='
+            v = str(col_val).strip("'")
+            if not v.isnumeric():
+                v = "'" + v + "'"
+            find_operator = operator_match.match(col_name)
+            if find_operator:
+                col_name = find_operator.group(1)
+                operator = find_operator.group(2)
+                filters.append('(%s IS NULL or %s %s %s)' %
+                               (col_name, col_name, operator, str(v)))
+            else:
+                filters.append(col_name + operator + str(v))
+
+        seperator = " and "
+        # always include attempt matching. Assume that attempt number is recorded in tags.
+        filters.append("tags @> '\"attempt_id:%i\"'" % attempt_id)
+
+        where_clause = ""
+        if bool(filter_dict):
+            where_clause = seperator.join(filters)
+
+        sets = []
+        for col_name, col_val in update_dict.items():
+            sets.append(col_name + " = " + str(col_val))
+
+        set_seperator = ", "
+        set_clause = ""
+        if bool(filter_dict):
+            set_clause = set_seperator.join(sets)
+        update_sql = """
+                UPDATE {0} SET {1} WHERE {2};
+        """.format(self.table_name, set_clause, where_clause)
+
+        async def _execute_update_on_cursor(_cur):
+            await _cur.execute(update_sql)
+            if _cur.rowcount < 1:
+                return DBResponse(response_code=404,
+                                  body={"msg": "could not find row"})
+            if _cur.rowcount > 1:
+                return DBResponse(response_code=500,
+                                  body={"msg": "duplicate rows"})
+            return DBResponse(response_code=200, body={"rowcount": _cur.rowcount})
+        if cur:
+            return await _execute_update_on_cursor(cur)
+        try:
+            with (
+                await self.db.pool.cursor(
+                    cursor_factory=psycopg2.extras.DictCursor
+                )
+            ) as cur:
+                db_response = await _execute_update_on_cursor(cur)
+                cur.close()
+                return db_response
+        except (Exception, psycopg2.DatabaseError) as error:
+            self.db.logger.exception("Exception occurred")
+            return aiopg_exception_handling(error)
+
 
 class PostgresUtils(object):
     @staticmethod
@@ -649,7 +711,7 @@ class AsyncStepTablePostgres(AsyncPostgresTable):
     select_columns = keys
     run_table_name = AsyncRunTablePostgres.table_name
 
-    async def add_step(self, step_object: StepRow):
+    async def add_step(self, step_object: StepRow, expanded=False):
         dict = {
             "flow_id": step_object.flow_id,
             "run_number": str(step_object.run_number),
@@ -659,22 +721,22 @@ async def add_step(self, step_object: StepRow):
             "tags": json.dumps(step_object.tags),
             "system_tags": json.dumps(step_object.system_tags),
         }
-        return await self.create_record(dict)
+        return await self.create_record(dict, expanded)
 
     async def get_steps(self, flow_id: str, run_id: str):
         run_id_key, run_id_value = translate_run_key(run_id)
         filter_dict = {"flow_id": flow_id,
                        run_id_key: run_id_value}
         return await self.get_records(filter_dict=filter_dict)
 
-    async def get_step(self, flow_id: str, run_id: str, step_name: str):
+    async def get_step(self, flow_id: str, run_id: str, step_name: str, expanded: bool = False):
         run_id_key, run_id_value = translate_run_key(run_id)
         filter_dict = {
             "flow_id": flow_id,
             run_id_key: run_id_value,
             "step_name": step_name,
         }
-        return await self.get_records(filter_dict=filter_dict, fetch_single=True)
+        return await self.get_records(filter_dict=filter_dict, fetch_single=True, expanded=expanded)
 
 
 class AsyncTaskTablePostgres(AsyncPostgresTable):
@@ -690,7 +752,7 @@ class AsyncTaskTablePostgres(AsyncPostgresTable):
     select_columns = keys
     step_table_name = AsyncStepTablePostgres.table_name
 
-    async def add_task(self, task: TaskRow, fill_heartbeat=False):
+    async def add_task(self, task: TaskRow, fill_heartbeat=False, expanded=False):
         # todo backfill run_number if missing?
         dict = {
             "flow_id": task.flow_id,
@@ -703,16 +765,16 @@ async def add_task(self, task: TaskRow, fill_heartbeat=False):
             "system_tags": json.dumps(task.system_tags),
             "last_heartbeat_ts": str(new_heartbeat_ts()) if fill_heartbeat else None
         }
-        return await self.create_record(dict)
+        return await self.create_record(dict, expanded)
 
-    async def get_tasks(self, flow_id: str, run_id: str, step_name: str):
+    async def get_tasks(self, flow_id: str, run_id: str, step_name: str, expanded: bool = False):
         run_id_key, run_id_value = translate_run_key(run_id)
         filter_dict = {
             "flow_id": flow_id,
             run_id_key: run_id_value,
             "step_name": step_name,
         }
-        return await self.get_records(filter_dict=filter_dict)
+        return await self.get_records(filter_dict=filter_dict, expanded=expanded)
 
     async def get_task(self, flow_id: str, run_id: str, step_name: str,
                        task_id: str, expanded: bool = False):

services/metadata_service/api/run.py

@@ -2,7 +2,7 @@
 from itertools import chain
 
 from services.data.db_utils import DBResponse
-from services.data.models import RunRow
+from services.data.models import RunRow, StepRow, TaskRow
 from services.utils import has_heartbeat_capable_version_tag, read_body
 from services.metadata_service.api.utils import format_response, \
     handle_exceptions
@@ -24,7 +24,13 @@ def __init__(self, app):
         app.router.add_route("PATCH",
                              "/flows/{flow_id}/runs/{run_number}/tag/mutate",
                              self.mutate_user_tags)
+        app.router.add_route("PATCH",
+                             "/flows/{flow_id}/runs/{run_number}/status",
+                             self.update_status)
         self._async_table = AsyncPostgresDB.get_instance().run_table_postgres
+        self._step_table = AsyncPostgresDB.get_instance().step_table_postgres
+        self._task_table = AsyncPostgresDB.get_instance().task_table_postgres
+        self._metadata_table = AsyncPostgresDB.get_instance().metadata_table_postgres
 
     @format_response
     @handle_exceptions
@@ -278,3 +284,133 @@ async def runs_heartbeat(self, request):
         flow_name = request.match_info.get("flow_id")
         run_number = request.match_info.get("run_number")
         return await self._async_table.update_heartbeat(flow_name, run_number)
+
+    @format_response
+    @handle_exceptions
+    async def update_status(self, request):
+        """
+        ---
+        description: update run status
+        tags:
+        - Run
+        parameters:
+        - name: "flow_id"
+          in: "path"
+          description: "flow_id"
+          required: true
+          type: "string"
+        - name: "run_number"
+          in: "path"
+          description: "run_number"
+          required: true
+          type: "string"
+        - name: "body"
+          in: "body"
+          description: "body"
+          required: true
+          schema:
+            type: object
+            properties:
+                status:
+                    type: string
+        produces:
+        - 'text/plain'
+        responses:
+            "200":
+                description: successful operation.
+            "400":
+                description: invalid HTTP Request
+            "405":
+                description: invalid HTTP Method
+        """
+        flow_name = request.match_info.get("flow_id")
+        run_number = request.match_info.get("run_number")
+
+        body = await read_body(request.content)
+        new_status = body.get("status", None)
+        if new_status is None:
+            raise Exception("Can not update without a valid 'status'")
+
+        # TODO: This might not be performant without providing a task_id as part of the query.
+        res = await self._metadata_table.update_row(
+            filter_dict={
+                "flow_id": flow_name,
+                "run_number": run_number,
+                "step_name": "_parameters",
+                "field_name": "_status",
+            },
+            update_dict={
+                "value": "'%s'" % new_status
+            }
+        )
+        if res.response_code == 200:
+            return res
+
+        # Slow path requiring registering additional resources.
+        # metadata record missing, create it. Using the parent elements ensures that all relevant info carries over (e.g. tags / system tags / username).
+        run_response = await self._async_table.get_run(flow_name, run_number, expanded=True)
+
+        if run_response.response_code != 200:
+            return run_response
+
+        run_row = run_response.body
+
+        step = await self._step_table.get_step(
+            flow_id=run_row["flow_id"],
+            run_id=run_row["run_number"],
+            step_name="_parameters",
+            expanded=True
+        )
+        if step.response_code != 200:
+            step = await self._step_table.add_step(
+                StepRow(
+                    flow_id=run_row["flow_id"],
+                    run_number=run_row["run_number"],
+                    run_id=run_row["run_id"],
+                    user_name=run_row["user_name"],
+                    step_name="_parameters",
+                    tags=run_row["tags"],
+                    system_tags=run_row["system_tags"]
+                ),
+                expanded=True
+            )
+
+        # Might be inefficient as we're not providing a task id.
+        step = step.body
+        tasks = await self._task_table.get_tasks(
+            flow_id=step["flow_id"],
+            run_id=step["run_number"],
+            step_name=step["step_name"],
+            expanded=True
+        )
+        if tasks.response_code != 200 or not tasks.body:
+            task = await self._task_table.add_task(
+                TaskRow(
+                    flow_id=step["flow_id"],
+                    run_number=step["run_number"],
+                    run_id=step["run_id"],
+                    step_name=step["step_name"],
+                    user_name=step["user_name"],
+                    tags=step["tags"],
+                    system_tags=step["system_tags"]
+                ),
+                expanded=True
+            )
+            task = task.body
+        else:
+            task = tasks.body[0]
+
+        return await self._metadata_table.add_metadata(
+            flow_id=task["flow_id"],
+            run_id=task["run_id"],
+            run_number=task["run_number"],
+            step_name=task["step_name"],
+            task_id=task["task_id"],
+            task_name=task["task_name"],
+            user_name=task["user_name"],
+            tags=task["tags"],
+            system_tags=task["system_tags"],
+            field_name="_status",
+            type="_status",
+            value=new_status
+        )