diff --git a/services/data/postgres_async_db.py b/services/data/postgres_async_db.py index dbde13e60..141734d6c 100644 --- a/services/data/postgres_async_db.py +++ b/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,7 +721,7 @@ 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) @@ -667,14 +729,14 @@ async def get_steps(self, flow_id: str, run_id: str): 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): diff --git a/services/metadata_service/api/run.py b/services/metadata_service/api/run.py index 3f26eed92..28bd28af2 100644 --- a/services/metadata_service/api/run.py +++ b/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 + ) diff --git a/services/metadata_service/api/task.py b/services/metadata_service/api/task.py index 40f42706a..c78a4b922 100644 --- a/services/metadata_service/api/task.py +++ b/services/metadata_service/api/task.py @@ -34,6 +34,11 @@ def __init__(self, app): "/flows/{flow_id}/runs/{run_number}/steps/{step_name}/task", self.create_task, ) + app.router.add_route( + "PATCH", + "/flows/{flow_id}/runs/{run_number}/steps/{step_name}/tasks/{task_id}/status", + self.update_status, + ) app.router.add_route("POST", "/flows/{flow_id}/runs/{run_number}/steps/{step_name}/tasks/{task_id}/heartbeat", self.tasks_heartbeat) @@ -319,3 +324,111 @@ async def tasks_heartbeat(self, request): return await self._async_table.update_heartbeat(flow_name, run_number, step_name, task_id) + + @format_response + @handle_exceptions + async def update_status(self, request): + """ + --- + description: update task (attempt) status + tags: + - Tasks + 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: "step_name" + in: "path" + description: "step_name" + required: true + type: "string" + - name: "task_id" + in: "path" + description: "task_id" + required: true + type: "string" + - name: "body" + in: "body" + description: "body" + required: true + schema: + type: object + properties: + status: + type: string + attempt: + type: int + produces: + - 'text/plain' + responses: + "200": + description: successful operation. Return newly registered run + "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") + step_name = request.match_info.get("step_name") + task_id = request.match_info.get("task_id") + + body = await read_body(request.content) + new_status = body.get("status", None) + attempt = body.get("attempt", 0) + 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._async_metadata_table.update_row_with_attempt( + filter_dict={ + "flow_id": flow_name, + "run_number": run_number, + "step_name": step_name, + "task_id": task_id, + "field_name": "_status", + }, + update_dict={ + "value": "'%s'" % new_status, + }, + attempt_id=attempt + ) + + if res.response_code == 200: + return res + + # Status Metadata missing, so we need to register the metadata for the task. + + task = await self._async_table.get_task( + flow_id=flow_name, + run_id=run_number, + step_name=step_name, + task_id=task_id, + expanded=True + ) + + if task.response_code != 200: + return task + + task = task.body + return await self._async_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"] + ["attempt_id:%i" % attempt], + system_tags=task["system_tags"], + field_name="_status", + type="_status", + value=new_status + ) diff --git a/services/ui_backend_service/api/notify.py b/services/ui_backend_service/api/notify.py index d77fa7992..912676b4f 100644 --- a/services/ui_backend_service/api/notify.py +++ b/services/ui_backend_service/api/notify.py @@ -124,6 +124,33 @@ async def handle_trigger_msg(self, msg: str): filter_dict={"attempt_id": _attempt_id} ) + # Notify related resources on _status updates. + if table.table_name == self.db.metadata_table_postgres.table_name and data["field_name"] == "_status": + + # Run level status. + if data["step_name"] == "_parameters": + await _broadcast( + event_emitter=self.event_emitter, + operation="UPDATE", + table=self.db.run_table_postgres, + data=data, + ) + # Task level status + else: + try: + attempt_tag = [t for t in data['tags'] if t.startswith('attempt_id')][0] + attempt_id = attempt_tag.split(":")[1] + except Exception: + self.logger.exception("Failed to load attempt_id from _status metadata") + pass + await _broadcast( + event_emitter=self.event_emitter, + operation="UPDATE", + table=self.db.task_table_postgres, + data=data, + filter_dict={"attempt_id": attempt_id} if attempt_id else {} + ) + # Notify related resources once attempt_ok for task has been saved. if operation == "INSERT" and \ table.table_name == self.db.metadata_table_postgres.table_name and \ diff --git a/services/ui_backend_service/data/db/tables/metadata.py b/services/ui_backend_service/data/db/tables/metadata.py index 046a7eb5d..0b4603e2b 100644 --- a/services/ui_backend_service/data/db/tables/metadata.py +++ b/services/ui_backend_service/data/db/tables/metadata.py @@ -13,12 +13,13 @@ class AsyncMetadataTablePostgres(AsyncPostgresTable): keys = MetaserviceMetadataTable.keys primary_keys = MetaserviceMetadataTable.primary_keys trigger_keys = MetaserviceMetadataTable.trigger_keys - trigger_operations = ["INSERT"] + trigger_operations = ["INSERT", "UPDATE"] trigger_conditions = [ "NEW.field_name = 'attempt'", "NEW.field_name = 'attempt_ok'", "NEW.field_name = 'code-package'", "NEW.field_name = 'code-package-url'", + "NEW.field_name = '_status'", ] @property diff --git a/services/ui_backend_service/data/db/tables/run.py b/services/ui_backend_service/data/db/tables/run.py index 51f3cc36d..090be3248 100644 --- a/services/ui_backend_service/data/db/tables/run.py +++ b/services/ui_backend_service/data/db/tables/run.py @@ -58,6 +58,22 @@ class AsyncRunTablePostgres(AsyncPostgresTable): table_name=table_name, metadata_table=metadata_table ), """ + LEFT JOIN LATERAL ( + SELECT + value + FROM {metadata_table} as status + WHERE + {table_name}.flow_id = status.flow_id AND + {table_name}.run_number = status.run_number AND + status.step_name = '_parameters' AND + status.field_name = '_status' + ORDER BY ts_epoch DESC + LIMIT 1 + ) as status_metadata ON true + """.format( + table_name=table_name, metadata_table=metadata_table + ), + """ LEFT JOIN LATERAL ( SELECT ts_epoch FROM {metadata_table} as attempt @@ -128,6 +144,8 @@ def select_columns(self): THEN 'completed' WHEN end_attempt_ok IS NOT NULL AND end_attempt_ok.value IS FALSE THEN 'failed' + WHEN status_metadata IS NOT NULL + THEN status_metadata.value WHEN {table_name}.last_heartbeat_ts IS NOT NULL AND @(extract(epoch from now())-{table_name}.last_heartbeat_ts)<={heartbeat_cutoff} THEN 'running' diff --git a/services/ui_backend_service/data/db/tables/task.py b/services/ui_backend_service/data/db/tables/task.py index 996a177e9..3bdb9dcb0 100644 --- a/services/ui_backend_service/data/db/tables/task.py +++ b/services/ui_backend_service/data/db/tables/task.py @@ -129,6 +129,18 @@ class AsyncTaskTablePostgres(AsyncPostgresTable): (attempt.attempt_id + 1) = (next_attempt_start.value::jsonb->>0)::int LIMIT 1 ) as next_attempt_start ON true + LEFT JOIN LATERAL ( + SELECT value + FROM {metadata_table} as attempt_status + WHERE + {table_name}.flow_id = attempt_status.flow_id AND + {table_name}.run_number = attempt_status.run_number AND + {table_name}.step_name = attempt_status.step_name AND + {table_name}.task_id = attempt_status.task_id AND + attempt_status.field_name = '_status' AND + attempt_status.tags @> CONCAT('"', 'attempt_id:', attempt.attempt_id, '"')::jsonb + LIMIT 1 + ) as status_metadata ON true """.format( table_name=table_name, metadata_table=metadata_table, @@ -174,9 +186,6 @@ def select_columns(self): THEN 'completed' WHEN attempt.attempt_ok IS FALSE THEN 'failed' - WHEN COALESCE(attempt.attempt_finished_at, attempt.task_ok_finished_at) IS NOT NULL - AND attempt_ok IS NULL - THEN 'unknown' WHEN COALESCE(attempt.attempt_finished_at, attempt.task_ok_finished_at) IS NOT NULL THEN 'completed' WHEN next_attempt_start.ts_epoch IS NOT NULL @@ -185,6 +194,8 @@ def select_columns(self): AND @(extract(epoch from now())-{table_name}.last_heartbeat_ts)>{heartbeat_threshold} AND {finished_at_column} IS NULL THEN 'failed' + WHEN status_metadata IS NOT NULL + THEN status_metadata.value WHEN {table_name}.last_heartbeat_ts IS NULL AND @(extract(epoch from now())*1000 - COALESCE(attempt.started_at, {table_name}.ts_epoch))>{cutoff} AND {finished_at_column} IS NULL