sientia-dataops-opc-ingestor/ingestor/io/data_sink.py

import asyncio
import re

from sqlalchemy import Select, Update, case, column, select, table, update

import ingestor.metrics as metrics

# fullmatch (no ^/$ anchors) so a trailing newline can't slip through — `$` would
# match before a trailing '\n', which is not what we want for a SQL identifier.
_IDENTIFIER_RE = re.compile(r'\w+', re.ASCII)

# Worst-wins precedence of predictions.write_status (SIENTIAPDE-1975), from best to worst;
# NULL, 'pending' and unknown values rank 0. The gateway is the only writer that updates the
# column: the Laborious only sets the value of the rows it inserts ('pending', PI outcome).
STATUS_RANK = {'ok': 1, 'error': 2, 'session_bad': 3}


def write_status_statements(
    schema: str, table_name: str, status: str, timestamp, model_id
) -> tuple[Update, Select]:
    """The UPDATE that records `status` unless the row has a worse one, and the SELECT that
    tells whether the row exists. Built with SQLAlchemy Core: identifiers and values are
    never formatted into the SQL text."""
    predictions = table(
        table_name,
        column('write_status'),
        column('timestamp'),
        column('model_id'),
        schema=schema,
    )
    row = (predictions.c.timestamp == timestamp) & (predictions.c.model_id == model_id)
    current_rank = case(STATUS_RANK, value=predictions.c.write_status, else_=0)
    return (
        update(predictions)
        .where(row, current_rank <= STATUS_RANK.get(status, 0))
        .values(write_status=status),
        select(1).select_from(predictions).where(row),
    )


class DataSink:
    """Gateway egress wrapper: fans a reading out to every configured name/collection."""

    def __init__(self, data_manager, postgres) -> None:
        self.data_manager = data_manager
        self.postgres = postgres

    async def emit_metric(self, **kwargs) -> None:
        # ponytail: delegate to data_manager, it's already a SientiaMonitoring
        await self.data_manager.emit_metric(**kwargs)

    async def publish_reading(self, tag: str, value, timestamp, names: list[dict]) -> None:
        """Publish one MongoDB document per name entry and record its tag value metric.

        Every name of the slot v2 carries collection, tag_name and model_id (parse_slot
        checks it)."""
        for name in names:
            doc = {
                'tag': tag,
                'name': name['tag_name'],
                'model_id': name['model_id'],
                'timestamp': timestamp,
                'value': value,
            }
            await self.data_manager.publish(name['collection'], doc)
            await self.emit_metric(
                metric_object=metrics.OPC_TAG_VALUE,
                method='set',
                value=value,
                tags={
                    'pod_id': self.data_manager.pod_id,
                    'tag_name': name['tag_name'],
                    'model_id': name['model_id'],
                },
            )

    async def update_write_status(self, cmd: dict, result) -> str:
        """Record a write outcome on the prediction row (C3), keeping the worst one.

        The row is found by timestamp + model_id. A prediction is written to several
        tags, each answered by its own command, so a better outcome never replaces a
        worse one: ``ok`` < ``error`` < ``session_bad``, and ``pending``/NULL are
        replaced by any answer, so the order of the answers does not matter. Only the
        gateway updates the column (SIENTIAPDE-1975): the Laborious publishes the commands
        and sets ``pending`` (merged with its PI Web API outcome) on the row it inserts.

        Never raises: a failure here must not reach the NATS callback nor hide the OPC
        write itself, which already happened. The outcome is logged and counted instead.

        Return:
            str: ``updated``, ``kept`` (the row already has a worse outcome),
            ``row_not_found`` (e.g. the Laborious INSERT failed), ``invalid`` (bad or
            missing schema/table/timestamp/model_id) or ``error``.
        """
        try:
            schema = self._validate_identifier(cmd.get('schema'))
            table = self._validate_identifier(cmd.get('table_name'))
            if cmd.get('timestamp') is None or cmd.get('model_id') is None:
                raise ValueError('missing timestamp or model_id')
        except ValueError as e:
            return await self._status_outcome('invalid', cmd, f'write_status not recorded: {e}')

        statements = write_status_statements(
            schema, table, result.status, cmd['timestamp'], cmd['model_id']
        )
        try:
            # The engine is synchronous (no async pg driver here): run it off the event
            # loop, which also serves the subscriptions of the read path.
            outcome = await asyncio.to_thread(self._update_sync, *statements)
        except Exception as e:
            return await self._status_outcome('error', cmd, f'write_status UPDATE failed: {e}')
        detail = {
            'updated': f'write_status={result.status} recorded',
            'kept': f'write_status={result.status} not recorded: the row has a worse outcome',
            'row_not_found': f'write_status={result.status} not recorded: no prediction row',
        }[outcome]
        return await self._status_outcome(outcome, cmd, detail)

    def _update_sync(self, update_stmt: Update, exists_stmt: Select) -> str:
        with self.postgres.engine.begin() as connection:
            if connection.execute(update_stmt).rowcount:
                return 'updated'
            row = connection.execute(exists_stmt).first()
        return 'kept' if row is not None else 'row_not_found'

    async def _status_outcome(self, outcome: str, cmd: dict, detail: str) -> str:
        message = (
            f'{detail} (server={cmd.get("server_id")} addr={cmd.get("addr")} '
            f'model_id={cmd.get("model_id")} timestamp={cmd.get("timestamp")})'
        )
        if outcome in ('updated', 'kept'):
            self.data_manager.logger.debug(message)
        else:
            self.data_manager.logger.warning(message)
        await self.emit_metric(
            metric_object=metrics.OPC_WRITE_STATUS_UPDATE_TOTAL,
            tags={'pod_id': self.data_manager.pod_id, 'result': outcome},
        )
        return outcome

    @staticmethod
    def _validate_identifier(name: object) -> str:
        if not isinstance(name, str) or not _IDENTIFIER_RE.fullmatch(name):
            raise ValueError(f'invalid SQL identifier: {name!r}')
        return name