"""Snowflake-based ACK processing connector. Replaces the former ows-sound-recordings HTTP connector by connecting directly to Snowflake to manage the full ACK lifecycle. """ import math import typing import config from src.queries import FETCH_AWAITED_DELIVERY_HISTORY from src.queries import FETCH_UNACKED_DELIVERY_HISTORY from src.queries import INSERT_DELIVERY_ACKS from src.queries import UPDATE_DELIVERY_ACK_AS_ERROR from src.queries import UPDATE_DELIVERY_ACKS_AS_MISSING from src.queries import UPDATE_DELIVERY_ACKS_AS_SUCCESS from src.connectors import snowflake as snowflake_connector # Maximum number of ACKs per Snowflake UPDATE batch to avoid query size limits BATCH_SIZE = 500 # Maximum number of rows per Snowflake INSERT VALUES batch. # INSERT statements are cheaper for the query planner than UPDATE tuples; # 5000 rows per statement minimises round-trips while staying within # Snowflake's recommended limits for the classic connector. BATCH_INSERT_SIZE = 5000 def _build_service_filter(service: typing.List[str]) -> str: """Build the SQL service filter fragment. Args: service: List of service name strings. Returns: SQL AND clause string, or empty string when service is empty. """ if not service: return '' values = ', '.join("'{}'".format(s.strip()) for s in service) return 'AND srdh.service IN ({}) '.format(values) def fetch_unacked_delivery_history( service: typing.List[str] = [], limit: typing.Optional[int] = None, ) -> typing.List[typing.Dict[str, typing.Any]]: """Fetch unacknowledged delivery history from the last 3 months. Args: service: Optional list of service names to filter by. limit: Optional cap on the number of (oldest-first) rows returned. Bounds how many rows a single run marks as awaited, keeping the insert in step with the downstream resolution rate. Returns: List of unacked delivery dicts. """ filters_sql = _build_service_filter(service) filters_sql += 'AND srdh.timestamp >= (select dateadd(month, -3, getdate())) ' limit_sql = 'LIMIT {} '.format(limit) if limit else '' sql = FETCH_UNACKED_DELIVERY_HISTORY.format(filters=filters_sql, limit_sql=limit_sql) conn = snowflake_connector.get_connection() try: with conn.cursor() as cur: cur.execute(sql) history = cur.fetchall() finally: conn.close() if not history: return [] return [ { 'sound_recording_id': event[0], 'version_id': event[1], 'sfn_execution_id': event[2] if event[2] else None, 'datetime': event[3], 'ack': None, } for event in history ] def fetch_awaited_delivery_history( service: typing.List[str], filter_awaited_hours: int, limit: typing.Optional[int] = None, offset: typing.Optional[int] = None, ) -> typing.List[typing.Dict[str, typing.Any]]: """Fetch delivery history records currently in awaited state. Args: service: List of service names to filter by. filter_awaited_hours: How many hours back to look for awaited records. limit: Optional maximum number of rows to return. offset: Optional row offset (used together with limit). Returns: List of awaited delivery dicts. """ filters_sql = _build_service_filter(service) filters_sql += 'AND srdh.timestamp >= (select dateadd(hour, -{}, getdate())) '.format(filter_awaited_hours) # noqa: E501 limit_offset_sql = '' if limit: limit_offset_sql += 'LIMIT {} '.format(limit) if offset: limit_offset_sql += 'OFFSET {} '.format(offset) sql = FETCH_AWAITED_DELIVERY_HISTORY.format( filters=filters_sql, limit_offset=limit_offset_sql, ) conn = snowflake_connector.get_connection() try: with conn.cursor() as cur: cur.execute(sql) awaited_history = cur.fetchall() finally: conn.close() if not awaited_history: return [] return [ { 'sound_recording_id': event[0], 'version_id': event[1], 'sfn_execution_id': event[2] if event[2] else None, 'datetime': event[3], 'ack': 'awaited', } for event in awaited_history ] def mark_unacked_as_awaited() -> typing.List[typing.Dict[str, typing.Any]]: """Fetch unacked delivery records and mark them as awaited in Snowflake. Returns: List of dicts with sfn_execution_id, version_id and datetime. """ awaited_deliveries = fetch_awaited_delivery_history( service=[config.ACK_SERVICE], filter_awaited_hours=config.ACKS_FILTER_AWAITED_HOURS, limit=config.LIMIT, ) unacked_deliveries = fetch_unacked_delivery_history( service=[config.ACK_SERVICE], limit=BATCH_INSERT_SIZE, ) deliveries = awaited_deliveries + unacked_deliveries if not deliveries: return [] rows_to_insert = [ d for d in unacked_deliveries if d['version_id'] and d['datetime'] and d['ack'] != 'awaited' ] if rows_to_insert: conn = snowflake_connector.get_connection() try: num_batches = math.ceil(len(rows_to_insert) / BATCH_INSERT_SIZE) for batch_idx in range(num_batches): batch = rows_to_insert[batch_idx * BATCH_INSERT_SIZE: (batch_idx + 1) * BATCH_INSERT_SIZE] first_row = True values_sql = [] for delivery in batch: values_sql.append( ('VALUES' if first_row else ',') + "('{}', '{}', '{}')".format( # noqa: W503 delivery['version_id'], delivery['datetime'], 'awaited', ) ) first_row = False with conn.cursor() as cur: cur.execute(INSERT_DELIVERY_ACKS.format(values=''.join(values_sql))) conn.commit() finally: conn.close() return [ { 'sfn_execution_id': d['sfn_execution_id'], 'version_id': d['version_id'], 'datetime': str(d['datetime']), } for d in deliveries if d['version_id'] and d['datetime'] ] def resolve_awaited_as_success(acks: typing.List[typing.Dict[str, typing.Any]]) -> None: """Mark awaited ACKs as successfully resolved in Snowflake. Args: acks: List of dicts with version_id and datetime. """ _resolve_awaited_in_batches(acks, 'success') def resolve_awaited_as_error(acks: typing.List[typing.Dict[str, typing.Any]]) -> None: """Mark awaited ACKs as resolved with error in Snowflake. Args: acks: List of dicts with version_id, datetime and message. """ _resolve_awaited_in_batches(acks, 'error') def resolve_awaited_as_missing(acks: typing.List[typing.Dict[str, typing.Any]]) -> None: """Mark the given ACK rows as missing in Snowflake. Args: acks: List of dicts with version_id and datetime. """ _resolve_awaited_in_batches(acks, 'missing') def _resolve_awaited_in_batches( acks: typing.List[typing.Dict[str, typing.Any]], status: str, ) -> None: """Execute UPDATE queries in chunks to avoid Snowflake query size limits. Args: acks: List of ACK records to resolve. status: One of 'success', 'missing', or 'error'. """ if not acks: return conn = snowflake_connector.get_connection() try: num_batches = math.ceil(len(acks) / BATCH_SIZE) for batch_idx in range(num_batches): batch = acks[batch_idx * BATCH_SIZE: (batch_idx + 1) * BATCH_SIZE] with conn.cursor() as cur: if status in ('success', 'missing'): acks_tuple = ', '.join( "('{}', '{}')".format(ack['version_id'], ack['datetime']) for ack in batch ) query = ( UPDATE_DELIVERY_ACKS_AS_SUCCESS if status == 'success' else UPDATE_DELIVERY_ACKS_AS_MISSING ) cur.execute(query.format(acks='({})'.format(acks_tuple))) else: for ack in batch: message = ack.get('message', '').replace("'", "''") ack_tuple = "('{}', '{}')".format(ack['version_id'], ack['datetime']) cur.execute( UPDATE_DELIVERY_ACK_AS_ERROR.format( message="'{}'".format(message), ack=ack_tuple, ) ) conn.commit() finally: conn.close()