import psycopg2 import psycopg2.extras from psycopg2.extensions import QuotedString import os import aiopg import json import math import re import time from services.utils import logging, DBType from typing import List, Tuple from .db_utils import DBResponse, DBPagination, aiopg_exception_handling, \ get_db_ts_epoch_str, translate_run_key, translate_task_key, new_heartbeat_ts from .models import FlowRow, RunRow, StepRow, TaskRow, MetadataRow, ArtifactRow from services.utils import DBConfiguration, USE_SEPARATE_READER_POOL from services.data.service_configs import max_connection_retires, \ connection_retry_wait_time_seconds AIOPG_ECHO = os.environ.get("AIOPG_ECHO", 0) == "1" WAIT_TIME = 10 # Create database triggers automatically, disabled by default # Enable with env variable `DB_TRIGGER_CREATE=1` DB_TRIGGER_CREATE = os.environ.get("DB_TRIGGER_CREATE", 0) == "1" # Configure DB Table names. Custom names can be supplied through environment variables, # in case the deployment differs from the default naming scheme from the supplied migrations. FLOW_TABLE_NAME = os.environ.get("DB_TABLE_NAME_FLOWS", "flows_v3") RUN_TABLE_NAME = os.environ.get("DB_TABLE_NAME_RUNS", "runs_v3") STEP_TABLE_NAME = os.environ.get("DB_TABLE_NAME_STEPS", "steps_v3") TASK_TABLE_NAME = os.environ.get("DB_TABLE_NAME_TASKS", "tasks_v3") METADATA_TABLE_NAME = os.environ.get("DB_TABLE_NAME_METADATA", "metadata_v3") ARTIFACT_TABLE_NAME = os.environ.get("DB_TABLE_NAME_ARTIFACT", "artifact_v3") DB_SCHEMA_NAME = os.environ.get("DB_SCHEMA_NAME", "public") operator_match = re.compile('([^:]*):([=><]+)$') # use a ddmmyyy timestamp as the version for triggers TRIGGER_VERSION = "05092024" TRIGGER_NAME_PREFIX = "notify_ui" class _AsyncPostgresDB(object): connection = None flow_table_postgres = None run_table_postgres = None step_table_postgres = None task_table_postgres = None artifact_table_postgres = None metadata_table_postgres = None pool = None reader_pool = None db_conf: DBConfiguration = None def __init__(self, name='global'): self.name = name self.logger = logging.getLogger("AsyncPostgresDB:{name}".format(name=self.name)) tables = [] self.flow_table_postgres = AsyncFlowTablePostgres(self) self.run_table_postgres = AsyncRunTablePostgres(self) self.step_table_postgres = AsyncStepTablePostgres(self) self.task_table_postgres = AsyncTaskTablePostgres(self) self.artifact_table_postgres = AsyncArtifactTablePostgres(self) self.metadata_table_postgres = AsyncMetadataTablePostgres(self) tables.append(self.flow_table_postgres) tables.append(self.run_table_postgres) tables.append(self.step_table_postgres) tables.append(self.task_table_postgres) tables.append(self.artifact_table_postgres) tables.append(self.metadata_table_postgres) self.tables = tables async def _init(self, db_conf: DBConfiguration, create_triggers=DB_TRIGGER_CREATE): # todo make poolsize min and max configurable as well as timeout # todo add retry and better error message retries = max_connection_retires for i in range(retries): try: self.pool = await aiopg.create_pool( db_conf.get_dsn(), minsize=db_conf.pool_min, maxsize=db_conf.pool_max, timeout=db_conf.timeout, pool_recycle=10 * db_conf.timeout, echo=AIOPG_ECHO) self.reader_pool = await aiopg.create_pool( db_conf.get_dsn(type=DBType.READER), minsize=db_conf.pool_min, maxsize=db_conf.pool_max, timeout=db_conf.timeout, pool_recycle=10 * db_conf.timeout, echo=AIOPG_ECHO) if USE_SEPARATE_READER_POOL else self.pool for table in self.tables: await table._init(create_triggers=create_triggers) if USE_SEPARATE_READER_POOL: self.logger.info( "Writer Connection established.\n" " Pool min: {pool_min} max: {pool_max}\n".format( pool_min=self.pool.minsize, pool_max=self.pool.maxsize)) self.logger.info( "Reader Connection established.\n" " Pool min: {pool_min} max: {pool_max}\n".format( pool_min=self.reader_pool.minsize, pool_max=self.reader_pool.maxsize)) else: self.logger.info( "Connection established.\n" " Pool min: {pool_min} max: {pool_max}\n".format( pool_min=self.pool.minsize, pool_max=self.pool.maxsize)) break # Break the retry loop except Exception as e: self.logger.exception("Exception occurred") if retries - i <= 1: raise e time.sleep(connection_retry_wait_time_seconds) def get_table_by_name(self, table_name: str): for table in self.tables: if table.table_name == table_name: return table return None 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'] 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'] class AsyncPostgresDB(object): __instance = None @staticmethod def get_instance(): return AsyncPostgresDB() def __init__(self): if not AsyncPostgresDB.__instance: AsyncPostgresDB.__instance = _AsyncPostgresDB() def __getattribute__(self, name): return getattr(AsyncPostgresDB.__instance, name) class AsyncPostgresTable(object): db = None table_name = None schema_version = 1 keys: List[str] = [] primary_keys: List[str] = None trigger_keys: List[str] = None trigger_operations: List[str] = ["INSERT", "UPDATE", "DELETE"] trigger_conditions: List[str] = None ordering: List[str] = None joins: List[str] = None select_columns: List[str] = keys join_columns: List[str] = None _insert_command = None _filters = None _base_query = "SELECT {0} from" _row_type = None def __init__(self, db: _AsyncPostgresDB = None): self.db = db if self.table_name is None: raise NotImplementedError( "need to specify table name") async def _init(self, create_triggers: bool): if create_triggers: self.db.logger.info( "Setting up notify trigger for {table_name}\n Keys: {keys}".format( table_name=self.table_name, keys=self.trigger_keys)) await PostgresUtils.cleanup_triggers(db=self.db, table_name=self.table_name) if self.trigger_keys and self.trigger_operations: await PostgresUtils.setup_trigger_notify( db=self.db, table_name=self.table_name, keys=self.trigger_keys, operations=self.trigger_operations, conditions=self.trigger_conditions ) async def get_records(self, filter_dict={}, fetch_single=False, ordering: List[str] = None, limit: int = 0, expanded=False, cur: aiopg.Cursor = None) -> DBResponse: conditions = [] values = [] for col_name, col_val in filter_dict.items(): conditions.append("{} = %s".format(col_name)) values.append(col_val) response, _ = await self.find_records( conditions=conditions, values=values, fetch_single=fetch_single, order=ordering, limit=limit, expanded=expanded, cur=cur ) return response async def find_records(self, conditions: List[str] = None, values=[], fetch_single=False, limit: int = 0, offset: int = 0, order: List[str] = None, expanded=False, enable_joins=False, cur: aiopg.Cursor = None) -> Tuple[DBResponse, DBPagination]: sql_template = """ SELECT * FROM ( SELECT {keys} FROM {table_name} {joins} ) T {where} {order_by} {limit} {offset} """ select_sql = sql_template.format( keys=",".join( self.select_columns + (self.join_columns if enable_joins and self.join_columns else [])), table_name=self.table_name, joins=" ".join(self.joins) if enable_joins and self.joins is not None else "", where="WHERE {}".format(" AND ".join(conditions)) if conditions else "", order_by="ORDER BY {}".format(", ".join(order)) if order else "", limit="LIMIT {}".format(limit) if limit else "", offset="OFFSET {}".format(offset) if offset else "" ).strip() return await self.execute_sql(select_sql=select_sql, values=values, fetch_single=fetch_single, expanded=expanded, limit=limit, offset=offset, cur=cur) async def execute_sql(self, select_sql: str, values=[], fetch_single=False, expanded=False, limit: int = 0, offset: int = 0, cur: aiopg.Cursor = None, serialize: bool = True) -> Tuple[DBResponse, DBPagination]: async def _execute_on_cursor(_cur): await _cur.execute(select_sql, values) rows = [] records = await _cur.fetchall() if serialize: for record in records: # pylint-initial-ignore: Lack of __init__ makes this too hard for pylint # pylint: disable=not-callable row = self._row_type(**record) rows.append(row.serialize(expanded)) else: rows = records count = len(rows) # Will raise IndexError in case fetch_single=True and there's no results body = rows[0] if fetch_single else rows pagination = DBPagination( limit=limit, offset=offset, count=count, page=math.floor(int(offset) / max(int(limit), 1)) + 1, ) return body, pagination try: if cur: # if we are using the passed in cursor, we allow any errors to be managed by cursor owner body, pagination = await _execute_on_cursor(cur) return DBResponse(response_code=200, body=body), pagination else: db_pool = self.db.reader_pool if USE_SEPARATE_READER_POOL else self.db.pool with (await db_pool.cursor( cursor_factory=psycopg2.extras.DictCursor )) as cur: body, pagination = await _execute_on_cursor(cur) cur.close() return DBResponse(response_code=200, body=body), pagination except IndexError as error: return aiopg_exception_handling(error), None except (Exception, psycopg2.DatabaseError) as error: self.db.logger.exception("Exception occurred") return aiopg_exception_handling(error), None async def create_record(self, record_dict): # note: need to maintain order cols = [] values = [] for col_name, col_val in record_dict.items(): cols.append(col_name) values.append(col_val) # add create ts cols.append("ts_epoch") values.append(get_db_ts_epoch_str()) str_format = [] for _ in cols: str_format.append("%s") seperator = ", " insert_sql = """ INSERT INTO {0}({1}) VALUES({2}) RETURNING * """.format( self.table_name, seperator.join(cols), seperator.join(str_format) ) try: response_body = {} with ( await self.db.pool.cursor( cursor_factory=psycopg2.extras.DictCursor ) ) as cur: await cur.execute(insert_sql, tuple(values)) records = await cur.fetchall() record = records[0] filtered_record = {} 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 # todo make sure connection is closed even with error cur.close() return DBResponse(response_code=200, body=response_body) except (Exception, psycopg2.DatabaseError) as error: self.db.logger.exception("Exception occurred") return aiopg_exception_handling(error) async def run_in_transaction_with_serializable_isolation_level(self, fun): try: with ( await self.db.pool.cursor( cursor_factory=psycopg2.extras.DictCursor, ) ) as cur: async with cur.begin(): await cur.execute('SET TRANSACTION ISOLATION LEVEL SERIALIZABLE') res = await fun(cur) cur.close() # is this really needed? TODO return res except psycopg2.errors.SerializationFailure: # See https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/409 return DBResponse(response_code=409, body="Conflicting concurrent tag mutation, please retry") except (Exception, psycopg2.DatabaseError) as error: self.db.logger.exception("Exception occurred") return aiopg_exception_handling(error) async def update_row(self, filter_dict={}, update_dict={}, 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 " 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 async def create_trigger_if_missing(db: _AsyncPostgresDB, table_name, trigger_name, commands=[]): "executes the commands only if a trigger with the given name does not already exist on the table" with (await db.pool.cursor()) as cur: try: await cur.execute( """ SELECT * FROM information_schema.triggers WHERE event_object_table = %s AND trigger_name = %s """, (table_name, trigger_name), ) trigger_exist = bool(cur.rowcount) if not trigger_exist: for command in commands: await cur.execute(command) finally: cur.close() @staticmethod async def cleanup_triggers(db: _AsyncPostgresDB, table_name, schema=DB_SCHEMA_NAME): "Cleans up old versions of table triggers" with (await db.pool.cursor()) as cur: try: await cur.execute( """ SELECT DISTINCT trigger_name FROM information_schema.triggers WHERE event_object_table = %s AND trigger_schema = %s """, [table_name, schema] ) results = await cur.fetchall() triggers_to_cleanup = [ res[0] for res in results if res[0].startswith(TRIGGER_NAME_PREFIX) and TRIGGER_VERSION not in res[0] ] if triggers_to_cleanup: logging.getLogger("TriggerSetup").info("Cleaning up old triggers: %s" % triggers_to_cleanup) commands = [] for trigger_name in triggers_to_cleanup: commands += [ (f"DROP TRIGGER IF EXISTS {trigger_name} ON {schema}.{table_name}"), (f"DROP FUNCTION IF EXISTS {schema}.{trigger_name}") ] for command in commands: await cur.execute(command) finally: cur.close() @staticmethod async def setup_trigger_notify( db: _AsyncPostgresDB, table_name, keys: List[str] = None, schema=DB_SCHEMA_NAME, operations: List[str] = None, conditions: List[str] = None ): if not keys: pass name_prefix = "%s_%s" % (TRIGGER_NAME_PREFIX, TRIGGER_VERSION) operations = operations _commands = [""" CREATE OR REPLACE FUNCTION {schema}.{prefix}_{table}() RETURNS trigger LANGUAGE plpgsql AS $$ DECLARE rec RECORD; BEGIN CASE TG_OP WHEN 'INSERT', 'UPDATE' THEN rec := NEW; WHEN 'DELETE' THEN rec := OLD; ELSE RAISE EXCEPTION 'Unknown TG_OP: "%"', TG_OP; END CASE; PERFORM pg_notify('notify', json_build_object( 'table', TG_TABLE_NAME, 'schema', TG_TABLE_SCHEMA, 'operation', TG_OP, 'data', json_build_object({keys}) )::text); RETURN rec; END; $$; """.format( schema=schema, prefix=name_prefix, table=table_name, keys=", ".join(map(lambda k: "'{0}', rec.{0}".format(k), keys)), )] _commands += [""" CREATE TRIGGER {prefix}_{table} AFTER {events} ON {schema}.{table} FOR EACH ROW {conditions} EXECUTE PROCEDURE {schema}.{prefix}_{table}(); """.format( schema=schema, prefix=name_prefix, table=table_name, events=" OR ".join(operations), conditions="WHEN (%s)" % " OR ".join(conditions) if conditions else "" )] # This enables trigger on both replica and non-replica mode _commands += ["ALTER TABLE {schema}.{table} ENABLE ALWAYS TRIGGER {prefix}_{table};".format( schema=schema, prefix=name_prefix, table=table_name )] # NOTE: Only try to setup triggers if they do not already exist. # This will require a table level lock so it should be performed during initial setup at off-peak hours. await PostgresUtils.create_trigger_if_missing( db=db, table_name=table_name, trigger_name="{}_{}".format(name_prefix, table_name), commands=_commands ) class AsyncFlowTablePostgres(AsyncPostgresTable): flow_dict = {} table_name = FLOW_TABLE_NAME keys = ["flow_id", "user_name", "ts_epoch", "tags", "system_tags"] primary_keys = ["flow_id"] trigger_keys = primary_keys select_columns = keys _row_type = FlowRow async def add_flow(self, flow: FlowRow): dict = { "flow_id": flow.flow_id, "user_name": flow.user_name, "tags": json.dumps(flow.tags), "system_tags": json.dumps(flow.system_tags), } return await self.create_record(dict) async def get_flow(self, flow_id: str): filter_dict = {"flow_id": flow_id} return await self.get_records(filter_dict=filter_dict, fetch_single=True) async def get_all_flows(self): return await self.get_records() class AsyncRunTablePostgres(AsyncPostgresTable): run_dict = {} run_by_flow_dict = {} _current_count = 0 _row_type = RunRow table_name = RUN_TABLE_NAME keys = ["flow_id", "run_number", "run_id", "user_name", "ts_epoch", "last_heartbeat_ts", "tags", "system_tags"] primary_keys = ["flow_id", "run_number"] trigger_keys = primary_keys + ["last_heartbeat_ts"] select_columns = keys flow_table_name = AsyncFlowTablePostgres.table_name async def add_run(self, run: RunRow, fill_heartbeat: bool = False): dict = { "flow_id": run.flow_id, "user_name": run.user_name, "tags": json.dumps(run.tags), "system_tags": json.dumps(run.system_tags), "run_id": run.run_id, "last_heartbeat_ts": str(new_heartbeat_ts()) if fill_heartbeat else None } return await self.create_record(dict) async def get_run(self, flow_id: str, run_id: str, expanded: bool = False, cur: aiopg.Cursor = None): key, value = translate_run_key(run_id) filter_dict = {"flow_id": flow_id, key: str(value)} return await self.get_records(filter_dict=filter_dict, fetch_single=True, expanded=expanded, cur=cur) async def get_all_runs(self, flow_id: str): filter_dict = {"flow_id": flow_id} return await self.get_records(filter_dict=filter_dict) async def update_heartbeat(self, flow_id: str, run_id: str): run_key, run_value = translate_run_key(run_id) new_hb = new_heartbeat_ts() filter_dict = {"flow_id": flow_id, run_key: str(run_value), "last_heartbeat_ts:<=": new_hb - WAIT_TIME} set_dict = { "last_heartbeat_ts": new_hb } result = await self.update_row(filter_dict=filter_dict, update_dict=set_dict) body = {"wait_time_in_seconds": WAIT_TIME} return DBResponse(response_code=result.response_code, body=json.dumps(body)) async def update_run_tags(self, flow_id: str, run_id: str, run_tags: list, cur: aiopg.Cursor = None): run_key, run_value = translate_run_key(run_id) filter_dict = {"flow_id": flow_id, run_key: str(run_value)} set_dict = {"tags": QuotedString(json.dumps(run_tags)).getquoted().decode()} return await self.update_row(filter_dict=filter_dict, update_dict=set_dict, cur=cur) class AsyncStepTablePostgres(AsyncPostgresTable): step_dict = {} run_to_step_dict = {} _row_type = StepRow table_name = STEP_TABLE_NAME keys = ["flow_id", "run_number", "run_id", "step_name", "user_name", "ts_epoch", "tags", "system_tags"] primary_keys = ["flow_id", "run_number", "step_name"] trigger_keys = primary_keys select_columns = keys run_table_name = AsyncRunTablePostgres.table_name async def add_step(self, step_object: StepRow): dict = { "flow_id": step_object.flow_id, "run_number": str(step_object.run_number), "run_id": step_object.run_id, "step_name": step_object.step_name, "user_name": step_object.user_name, "tags": json.dumps(step_object.tags), "system_tags": json.dumps(step_object.system_tags), } return await self.create_record(dict) 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): 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) class AsyncTaskTablePostgres(AsyncPostgresTable): task_dict = {} step_to_task_dict = {} _current_count = 0 _row_type = TaskRow table_name = TASK_TABLE_NAME keys = ["flow_id", "run_number", "run_id", "step_name", "task_id", "task_name", "user_name", "ts_epoch", "last_heartbeat_ts", "tags", "system_tags"] primary_keys = ["flow_id", "run_number", "step_name", "task_id"] trigger_keys = primary_keys select_columns = keys step_table_name = AsyncStepTablePostgres.table_name async def add_task(self, task: TaskRow, fill_heartbeat=False): # todo backfill run_number if missing? dict = { "flow_id": task.flow_id, "run_number": str(task.run_number), "run_id": task.run_id, "step_name": task.step_name, "task_name": task.task_name, "user_name": task.user_name, "tags": json.dumps(task.tags), "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) async def get_tasks(self, flow_id: str, run_id: str, step_name: str): 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) async def get_task(self, flow_id: str, run_id: str, step_name: str, task_id: str, expanded: bool = False): run_id_key, run_id_value = translate_run_key(run_id) task_id_key, task_id_value = translate_task_key(task_id) filter_dict = { "flow_id": flow_id, run_id_key: run_id_value, "step_name": step_name, task_id_key: task_id_value, } return await self.get_records(filter_dict=filter_dict, fetch_single=True, expanded=expanded) async def update_heartbeat(self, flow_id: str, run_id: str, step_name: str, task_id: str): run_key, run_value = translate_run_key(run_id) task_key, task_value = translate_task_key(task_id) new_hb = new_heartbeat_ts() filter_dict = {"flow_id": flow_id, run_key: str(run_value), "step_name": step_name, task_key: str(task_value), "last_heartbeat_ts:<=": new_hb - WAIT_TIME} set_dict = { "last_heartbeat_ts": new_hb } result = await self.update_row(filter_dict=filter_dict, update_dict=set_dict) body = {"wait_time_in_seconds": WAIT_TIME} return DBResponse(response_code=result.response_code, body=json.dumps(body)) class AsyncMetadataTablePostgres(AsyncPostgresTable): metadata_dict = {} run_to_metadata_dict = {} _current_count = 0 _row_type = MetadataRow table_name = METADATA_TABLE_NAME keys = ["flow_id", "run_number", "run_id", "step_name", "task_id", "task_name", "id", "field_name", "value", "type", "user_name", "ts_epoch", "tags", "system_tags"] primary_keys = ["flow_id", "run_number", "step_name", "task_id", "field_name"] trigger_keys = ["flow_id", "run_number", "step_name", "task_id", "field_name", "value", "tags"] trigger_operations = ["INSERT"] select_columns = keys async def add_metadata( self, flow_id, run_number, run_id, step_name, task_id, task_name, field_name, value, type, user_name, tags, system_tags, ): dict = { "flow_id": flow_id, "run_number": str(run_number), "run_id": run_id, "step_name": step_name, "task_id": str(task_id), "task_name": task_name, "field_name": field_name, "value": value, "type": type, "user_name": user_name, "tags": json.dumps(tags), "system_tags": json.dumps(system_tags), } return await self.create_record(dict) async def get_metadata_in_runs(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_metadata( self, flow_id: str, run_id: int, step_name: str, task_id: str ): run_id_key, run_id_value = translate_run_key(run_id) task_id_key, task_id_value = translate_task_key(task_id) filter_dict = { "flow_id": flow_id, run_id_key: run_id_value, "step_name": step_name, task_id_key: task_id_value, } return await self.get_records(filter_dict=filter_dict) async def get_filtered_task_pathspecs(self, flow_id: str, run_id: str, step_name: str, field_name: str, pattern: str): """ Returns a list of task pathspecs that match the given field_name and regexp pattern for the value """ 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, } conditions = [f"{k} = %s" for k, v in filter_dict.items() if v is not None] values = [v for k, v in filter_dict.items() if v is not None] if field_name: conditions.append("field_name = %s") values.append(field_name) if pattern: conditions.append("regexp_match(value, %s) IS NOT NULL") values.append(pattern) # We must return distinct task pathspecs, so we construct the select statement by hand sql_template = """ SELECT DISTINCT {select_columns} FROM ( SELECT {keys} FROM {table_name} ) T {where} {order_by} """ select_sql = sql_template.format( keys=",".join(self.select_columns), table_name=self.table_name, where="WHERE {}".format(" AND ".join(conditions)), order_by="ORDER BY task_id", select_columns=",".join(["flow_id, run_number, run_id, step_name, task_name, task_id"]) ).strip() db_response, pagination = await self.execute_sql(select_sql=select_sql, values=values, serialize=False) # flatten the ids in the response def _format_id(row): flow_id, run_number, run_id, step_name, task_name, task_id = row # pathspec return f"{flow_id}/{run_id or run_number}/{step_name}/{task_name or task_id}" flattened_response = DBResponse(body=[_format_id(row) for row in db_response.body], response_code=db_response.response_code) return flattened_response, pagination class AsyncArtifactTablePostgres(AsyncPostgresTable): artifact_dict = {} run_to_artifact_dict = {} step_to_artifact_dict = {} task_to_artifact_dict = {} current_count = 0 _row_type = ArtifactRow table_name = ARTIFACT_TABLE_NAME ordering = ["attempt_id DESC"] keys = ["flow_id", "run_number", "run_id", "step_name", "task_id", "task_name", "name", "location", "ds_type", "sha", "type", "content_type", "user_name", "attempt_id", "ts_epoch", "tags", "system_tags"] primary_keys = ["flow_id", "run_number", "step_name", "task_id", "attempt_id", "name"] trigger_keys = primary_keys trigger_operations = ["INSERT"] select_columns = keys async def add_artifact( self, flow_id, run_number, run_id, step_name, task_id, task_name, name, location, ds_type, sha, type, content_type, user_name, attempt_id, tags, system_tags, ): dict = { "flow_id": flow_id, "run_number": str(run_number), "run_id": run_id, "step_name": step_name, "task_id": str(task_id), "task_name": task_name, "name": name, "location": location, "ds_type": ds_type, "sha": sha, "type": type, "content_type": content_type, "user_name": user_name, "attempt_id": str(attempt_id), "tags": json.dumps(tags), "system_tags": json.dumps(system_tags), } return await self.create_record(dict) async def get_artifacts_in_runs(self, flow_id: str, run_id: int): 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, ordering=self.ordering) async def get_artifact_in_steps(self, flow_id: str, run_id: int, step_name: str): 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, ordering=self.ordering) async def get_artifact_in_task( self, flow_id: str, run_id: int, step_name: str, task_id: int ): run_id_key, run_id_value = translate_run_key(run_id) task_id_key, task_id_value = translate_task_key(task_id) filter_dict = { "flow_id": flow_id, run_id_key: run_id_value, "step_name": step_name, task_id_key: task_id_value, } return await self.get_records(filter_dict=filter_dict, ordering=self.ordering) async def get_artifact( self, flow_id: str, run_id: int, step_name: str, task_id: int, name: str ): # Return the artifact metadata for the latest attempt of the task. # # The quirk here is that different attempts may have different sets of # artifacts. That is, if artifact "foo" was set in attempt N, that # doesn't mean it was set in attempt N+1, and vice versa. # # To get the artifact value for the "latest" attempt, we first find # the latest attempt_id by querying the artifacts table for the # artifact that always exists for every attempt (the artifact called # 'name', containing the flow name), then use that attempt_id to get # the artifact we're interested in. run_id_key, run_id_value = translate_run_key(run_id) task_id_key, task_id_value = translate_task_key(task_id) filter_dict = { "flow_id": flow_id, run_id_key: run_id_value, "step_name": step_name, task_id_key: task_id_value, '"name"': "name" } name_record = await self.get_records(filter_dict=filter_dict, fetch_single=True, ordering=self.ordering) return await self.get_artifact_by_attempt( flow_id, run_id, step_name, task_id, name, name_record.body.get('attempt_id', 0)) async def get_artifact_by_attempt( self, flow_id: str, run_id: int, step_name: str, task_id: int, name: str, attempt: int): run_id_key, run_id_value = translate_run_key(run_id) task_id_key, task_id_value = translate_task_key(task_id) filter_dict = { "flow_id": flow_id, run_id_key: run_id_value, "step_name": step_name, task_id_key: task_id_value, '"name"': name, '"attempt_id"': attempt } return await self.get_records(filter_dict=filter_dict, fetch_single=True, ordering=self.ordering)