sientia-dataops-opc-ingestor/ingestor/engines/write.py

"""OPC write execution (F2 capstone) from a NATS C2 command dict.

Fixes 1879 correctness gaps: safe value coercion (never raises on bad
input), StatusCode/SourceTimestamp fidelity on the DataValue actually sent
to the server, a rate-limit floor between writes, and OPC error typing.
"""

import asyncio
from dataclasses import dataclass

from asyncua import ua
from asyncua.ua.uaerrors import BadWriteNotSupported, UaStatusCodeError
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring

import ingestor.metrics as metrics
from ingestor.connection.clock import MonotonicClock
from ingestor.connection.pool import OpcConnectionPool
from ingestor.connection.servers import ServerConfigResolver
from ingestor.converters import ConversionError, to_typed_value
from ingestor.io.data_sink import DataSink

_VARIANT_TYPES = {
    'float': ua.VariantType.Float,
    'double': ua.VariantType.Double,
    'real': ua.VariantType.Double,
    'int': ua.VariantType.Int32,
    'integer': ua.VariantType.Int32,
    'bool': ua.VariantType.Boolean,
    'boolean': ua.VariantType.Boolean,
    'str': ua.VariantType.String,
    'string': ua.VariantType.String,
}

# ponytail: same reconnect-worthy Bad*-status subset as laborious's
# opc_repository.py (different repo, not importable here); keep in sync by
# hand or promote to a shared lib if it drifts.
_SESSION_BAD_NAMES = frozenset(
    {
        'BadSessionIdInvalid',
        'BadSessionClosed',
        'BadSessionNotActivated',
        'BadSecureChannelIdInvalid',
        'BadSecureChannelClosed',
        'BadSecureChannelTokenUnknown',
        'BadTcpSecureChannelUnknown',
        'BadServerNotConnected',
        'BadConnectionClosed',
        'BadConnectionRejected',
        'BadNotConnected',
        'BadTimeout',
        'BadRequestTimeout',
        'BadTooManySessions',
    }
)


# `quality` of the C2 command (the Laborious sends its prediction_status) -> OPC quality.
_QUALITY_CODES = {
    'good': ua.StatusCodes.Good,
    'uncertain': ua.StatusCodes.Uncertain,
    'bad': ua.StatusCodes.Bad,
}


class _InvalidFieldError(Exception):
    """A command field needed by an option that is on (timestamp, quality) is bad."""

    def __init__(self, error_type: str, message: str):
        super().__init__(message)
        self.error_type = error_type
        self.message = message


@dataclass
class WriteResult:
    ok: bool
    status: str  # 'ok' | 'error' | 'session_bad'
    opc_status_code: str
    message: str


class WriteEngine(SientiaMonitoring):
    """Executes one OPC write from a NATS C2 command."""

    def __init__(
        self,
        pool: OpcConnectionPool,
        sink: DataSink,
        logger: Logger,
        notification_handler: NotificationHandler,
        metrics_controller: MetricsController,
        write_timestamp: bool = True,
        write_quality: bool = True,
        min_pause_s: float = 0.0,
        servers: ServerConfigResolver | None = None,
    ):
        self.pool = pool
        self.sink = sink
        # Resolves the server connection by server_id (SIENTIAPDE-1975 GW1). Without it,
        # only the connection carried by the command itself is used.
        self.servers = servers
        # Defaults (env); the server document and the command override them (SIENTIAPDE-2202).
        self.write_timestamp = write_timestamp
        self.write_quality = write_quality
        self.min_pause_s = min_pause_s
        self.clock = MonotonicClock()
        self._last_write_at: float | None = None
        # Servers that refused timestamp/quality: written with the value alone.
        self._value_only_servers: set[str] = set()

        SientiaMonitoring.__init__(
            self,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

    async def write(self, cmd: dict) -> WriteResult:
        await self._rate_limit()

        missing = [k for k in ('value', 'data_type', 'server_id', 'addr') if k not in cmd]
        if missing:
            return await self._emit_error(cmd, 'bad_command', f'missing fields: {missing}', 'error')

        try:
            typed = to_typed_value(cmd['value'], cmd['data_type'])
        except ConversionError as e:
            return await self._emit_error(cmd, 'conversion', str(e), 'error')

        try:
            server_cfg = await self._server_cfg(cmd)
        except Exception as e:
            return await self._emit_error(cmd, 'server_lookup', str(e), 'error')
        if server_cfg is None:
            return await self._emit_error(
                cmd, 'server_not_found', f'OPC server {cmd["server_id"]} is not configured', 'error'
            )

        server_id = str(cmd['server_id'])
        write_timestamp, write_quality = await self._write_flags(server_id, cmd)
        try:
            dv = self._build_datavalue(cmd, typed, write_timestamp, write_quality)
        except _InvalidFieldError as e:
            return await self._emit_error(cmd, e.error_type, e.message, 'error')

        start = self.clock.elapsed()
        try:
            client = await self.pool.get(cmd['server_id'], server_cfg)
            node = client.get_node(cmd['addr'])
            await self._write_node(node, dv, server_id)
        except Exception as e:
            status = 'session_bad' if self._is_session_bad(e) else 'error'
            return await self._emit_error(cmd, type(e).__name__, str(e), status)

        return await self._emit_success(cmd, start)

    async def _rate_limit(self) -> None:
        now = self.clock.elapsed()
        if self.min_pause_s > 0 and self._last_write_at is not None:
            remaining = self.min_pause_s - (now - self._last_write_at)
            if remaining > 0:
                await asyncio.sleep(remaining)
                now = self.clock.elapsed()
        self._last_write_at = now

    async def _server_cfg(self, cmd: dict) -> dict | None:
        """Connection settings of the target server: resolved by server_id (connection
        carried by the command, live pooled connection, connectors, opc-servers), or only
        the ones carried by the command when no resolver is wired."""
        if self.servers is not None:
            return await self.servers.resolve(str(cmd['server_id']), cmd)
        if not cmd.get('url'):
            return None
        return {
            'url': cmd.get('url'),
            'use_security': cmd.get('use_security', False),
            'server_uri': cmd.get('server_uri'),
        }

    def _variant_type(self, data_type: str) -> ua.VariantType:
        # direct index: unknown dtypes already became ConversionError in to_typed_value
        return _VARIANT_TYPES[data_type.strip().lower()]

    async def _write_flags(self, server_id: str, cmd: dict) -> tuple[bool, bool]:
        """`write_timestamp` and `write_quality` of this write (SIENTIAPDE-2202): the
        command, then the server document, then the gateway env. A server that already
        refused them (BadWriteNotSupported) gets the value alone."""
        if server_id in self._value_only_servers:
            return False, False
        options = await self.servers.write_options(server_id) if self.servers is not None else {}

        def flag(key: str, default: bool) -> bool:
            value = cmd.get(key)
            return bool(options.get(key, default) if value is None else value)

        return flag('write_timestamp', self.write_timestamp), flag(
            'write_quality', self.write_quality
        )

    def _build_datavalue(
        self, cmd: dict, typed_value, write_timestamp: bool, write_quality: bool
    ) -> ua.DataValue:
        """The value, with the prediction timestamp and quality when they are on.

        Raises:
            _InvalidFieldError: When an option is on and its command field is missing or bad.
        """
        variant = ua.Variant(typed_value, self._variant_type(cmd['data_type']))
        # NOTE: DataValue.StatusCode default_factory populates a Good status
        # even when omitted, so "off" must pass StatusCode=None explicitly
        # or the flag silently no-ops.
        status_code = self._status_code(cmd) if write_quality else None
        try:
            source_ts = ua.DateTime.fromisoformat(cmd['timestamp']) if write_timestamp else None
        except (KeyError, TypeError, ValueError) as e:
            raise _InvalidFieldError('bad_timestamp', repr(e)) from e
        return ua.DataValue(variant, StatusCode=status_code, SourceTimestamp=source_ts)

    @staticmethod
    def _status_code(cmd: dict) -> ua.StatusCode:
        """The OPC quality of the command's `quality` (the prediction status). Commands
        without it (before SIENTIAPDE-2202) are written as Good."""
        quality = cmd.get('quality') or 'good'
        code = _QUALITY_CODES.get(str(quality).strip().lower())
        if code is None:
            raise _InvalidFieldError('bad_quality', f'unknown quality {quality!r}')
        return ua.StatusCode(ua.UInt32(code))

    async def _write_node(self, node, dv: ua.DataValue, server_id: str) -> None:
        """Write `dv`. When the server refuses the timestamp/quality it carries
        (BadWriteNotSupported, e.g. Kepware), write the value alone and remember the
        server, so its next writes go without them."""
        try:
            await node.write_value(dv)
        except BadWriteNotSupported:
            if dv.StatusCode is None and dv.SourceTimestamp is None:
                raise
            await node.write_value(ua.DataValue(dv.Value, StatusCode=None))
            await self._remember_value_only(server_id)

    async def _remember_value_only(self, server_id: str) -> None:
        self._value_only_servers.add(server_id)
        self.logger.warning(
            f'OPC server {server_id} refused the timestamp/quality of a write '
            f'(BadWriteNotSupported): its writes go with the value alone until the gateway '
            f'restarts. Set write_timestamp/write_quality to false on its document to skip '
            f'the refused attempt.'
        )
        await self.emit_metric(
            metric_object=metrics.OPC_WRITE_VALUE_ONLY_FALLBACK,
            tags={'pod_id': self.pod_id, 'server_name': server_id},
        )

    def _is_session_bad(self, exc: Exception) -> bool:
        return isinstance(exc, UaStatusCodeError) and type(exc).__name__ in _SESSION_BAD_NAMES

    async def _emit_error(self, cmd: dict, error_type: str, message: str, status: str):
        # The reason only lives here: the row gets write_status alone (SIENTIAPDE-1975).
        self.logger.error(
            f'OPC write failed ({status}, {error_type}): server={cmd.get("server_id")} '
            f'addr={cmd.get("addr")} model_id={cmd.get("model_id")} '
            f'timestamp={cmd.get("timestamp")}: {message}'
        )
        tags = {'pod_id': self.pod_id, 'server_name': cmd.get('server_id', 'unknown')}
        await self.emit_metric(
            metric_object=metrics.OPC_ERROR_BY_TYPE, tags={**tags, 'error_type': error_type}
        )
        await self.emit_metric(
            metric_object=metrics.OPC_WRITE_TOTAL,
            tags={**tags, 'model_id': cmd.get('model_id', 'unknown'), 'result': 'error'},
        )
        # Same series as a success (status='ok'), so it always shows the last outcome.
        await self.emit_metric(
            metric_object=metrics.OPC_WRITE_STATUS_BY_MODEL,
            method='set',
            value=0,
            tags={
                'pod_id': self.pod_id,
                'model_id': cmd.get('model_id', 'unknown'),
                'status': 'ok',
            },
        )
        result = WriteResult(ok=False, status=status, opc_status_code=error_type, message=message)
        await self.sink.update_write_status(cmd, result)
        return result

    async def _emit_success(self, cmd: dict, start: float):
        tags = {'pod_id': self.pod_id, 'server_name': cmd.get('server_id', 'unknown')}
        await self.emit_metric(
            metric_object=metrics.OPC_WRITE_TOTAL,
            tags={**tags, 'model_id': cmd.get('model_id', 'unknown'), 'result': 'success'},
        )
        await self.emit_metric(
            metric_object=metrics.OPC_WRITE_LATENCY,
            method='observe',
            value=self.clock.elapsed() - start,
            tags=tags,
        )
        await self.emit_metric(
            metric_object=metrics.OPC_WRITE_STATUS_BY_MODEL,
            method='set',
            value=1,
            tags={
                'pod_id': self.pod_id,
                'model_id': cmd.get('model_id', 'unknown'),
                'status': 'ok',
            },
        )
        result = WriteResult(ok=True, status='ok', opc_status_code='Good', message='')
        await self.sink.update_write_status(cmd, result)
        return result