feat(gateway): WriteEngine (status_code/timestamp/flags, safe converter, rate-limit)
This commit is contained in:
178
ingestor/engines/write.py
Normal file
178
ingestor/engines/write.py
Normal file
@@ -0,0 +1,178 @@
|
||||
"""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 datetime import datetime
|
||||
|
||||
from asyncua import ua
|
||||
from asyncua.ua.uaerrors import 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.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',
|
||||
})
|
||||
|
||||
|
||||
@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,
|
||||
metadata: dict | None = None,
|
||||
write_timestamp: bool = True,
|
||||
write_quality: bool = True,
|
||||
min_pause_s: float = 0.0,
|
||||
):
|
||||
self.pool = pool
|
||||
self.sink = sink
|
||||
self.metadata = metadata or {}
|
||||
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
|
||||
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
async def write(self, cmd: dict) -> WriteResult:
|
||||
await self._rate_limit()
|
||||
|
||||
try:
|
||||
typed = to_typed_value(cmd['value'], cmd['data_type'])
|
||||
except ConversionError as e:
|
||||
return await self._emit_error(cmd, 'conversion', str(e), 'error')
|
||||
|
||||
dv = self._build_datavalue(cmd, typed)
|
||||
|
||||
start = self.clock.elapsed()
|
||||
try:
|
||||
client = await self.pool.get(cmd['server_id'], self._server_cfg(cmd))
|
||||
node = client.get_node(cmd['addr'])
|
||||
await node.write_value(dv)
|
||||
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
|
||||
|
||||
def _server_cfg(self, cmd: dict) -> dict:
|
||||
# ponytail: minimal cfg built straight from the C2 message. Resolving
|
||||
# the full connector config (url/security) from the connectors DB by
|
||||
# server_id is F4 wiring, not this task — pool is mocked in tests.
|
||||
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:
|
||||
return _VARIANT_TYPES.get(data_type.strip().lower(), ua.VariantType.String)
|
||||
|
||||
def _build_datavalue(self, cmd: dict, typed_value) -> ua.DataValue:
|
||||
variant = ua.Variant(typed_value, self._variant_type(cmd['data_type']))
|
||||
write_quality = cmd.get('write_quality', self.write_quality)
|
||||
write_timestamp = cmd.get('write_timestamp', self.write_timestamp)
|
||||
# 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 = ua.StatusCode(ua.StatusCodes.Good) if write_quality else None
|
||||
source_ts = datetime.fromisoformat(cmd['timestamp']) if write_timestamp else None
|
||||
return ua.DataValue(variant, StatusCode_=status_code, SourceTimestamp=source_ts)
|
||||
|
||||
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):
|
||||
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'},
|
||||
)
|
||||
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['server_id']}
|
||||
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
|
||||
220
tests/unit/engines/test_write.py
Normal file
220
tests/unit/engines/test_write.py
Normal file
@@ -0,0 +1,220 @@
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from asyncua import ua
|
||||
|
||||
import ingestor.metrics as metrics
|
||||
from ingestor.engines.write import WriteEngine
|
||||
|
||||
|
||||
def make_engine(**kwargs):
|
||||
eng = WriteEngine(
|
||||
pool=AsyncMock(),
|
||||
sink=AsyncMock(),
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
metadata={},
|
||||
**kwargs,
|
||||
)
|
||||
eng.emit_metric = AsyncMock()
|
||||
return eng
|
||||
|
||||
|
||||
def make_client(node=None):
|
||||
node = node or make_node()
|
||||
client = MagicMock()
|
||||
client.get_node = MagicMock(return_value=node)
|
||||
return client
|
||||
|
||||
|
||||
def make_node():
|
||||
node = MagicMock()
|
||||
node.write_value = AsyncMock()
|
||||
return node
|
||||
|
||||
|
||||
def base_cmd(**overrides):
|
||||
cmd = {
|
||||
'server_id': 'srv1',
|
||||
'addr': 'ns=2;i=5',
|
||||
'value': 42.0,
|
||||
'data_type': 'float',
|
||||
'tag_type': 'prediction',
|
||||
'model_id': 'm1',
|
||||
'model_name': 'model-1',
|
||||
'timestamp': '2026-07-13T12:00:00Z',
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'predictions_m1',
|
||||
}
|
||||
cmd.update(overrides)
|
||||
return cmd
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_success_writes_datavalue_and_returns_ok():
|
||||
eng = make_engine()
|
||||
node = make_node()
|
||||
client = make_client(node)
|
||||
eng.pool.get.return_value = client
|
||||
|
||||
result = await eng.write(base_cmd())
|
||||
|
||||
client.get_node.assert_called_once_with('ns=2;i=5')
|
||||
node.write_value.assert_awaited_once()
|
||||
dv = node.write_value.await_args.args[0]
|
||||
assert isinstance(dv, ua.DataValue)
|
||||
assert dv.Value.Value == 42.0
|
||||
assert dv.Value.VariantType == ua.VariantType.Float
|
||||
assert result.ok is True
|
||||
assert result.status == 'ok'
|
||||
assert result.opc_status_code == 'Good'
|
||||
eng.sink.update_write_status.assert_awaited_once_with(base_cmd(), result)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_success_emits_metrics():
|
||||
eng = make_engine()
|
||||
eng.pool.get.return_value = make_client()
|
||||
|
||||
await eng.write(base_cmd())
|
||||
|
||||
eng.emit_metric.assert_any_await(
|
||||
metric_object=metrics.OPC_WRITE_TOTAL,
|
||||
tags={'pod_id': eng.pod_id, 'server_name': 'srv1', 'model_id': 'm1', 'result': 'success'},
|
||||
)
|
||||
status_call = next(
|
||||
c
|
||||
for c in eng.emit_metric.await_args_list
|
||||
if c.kwargs['metric_object'] is metrics.OPC_WRITE_STATUS_BY_MODEL
|
||||
)
|
||||
assert status_call.kwargs['value'] == 1
|
||||
assert status_call.kwargs['tags'] == {'pod_id': eng.pod_id, 'model_id': 'm1', 'status': 'ok'}
|
||||
latency_call = next(
|
||||
c
|
||||
for c in eng.emit_metric.await_args_list
|
||||
if c.kwargs['metric_object'] is metrics.OPC_WRITE_LATENCY
|
||||
)
|
||||
assert latency_call.kwargs['method'] == 'observe'
|
||||
assert latency_call.kwargs['value'] >= 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_defaults_carry_statuscode_and_timestamp():
|
||||
eng = make_engine(write_timestamp=True, write_quality=True)
|
||||
node = make_node()
|
||||
eng.pool.get.return_value = make_client(node)
|
||||
|
||||
await eng.write(base_cmd())
|
||||
|
||||
dv = node.write_value.await_args.args[0]
|
||||
assert dv.StatusCode_ is not None
|
||||
assert dv.StatusCode_.is_good()
|
||||
assert dv.SourceTimestamp == datetime.fromisoformat('2026-07-13T12:00:00Z')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_engine_flags_off_omit_statuscode_and_timestamp():
|
||||
eng = make_engine(write_timestamp=False, write_quality=False)
|
||||
node = make_node()
|
||||
eng.pool.get.return_value = make_client(node)
|
||||
|
||||
await eng.write(base_cmd())
|
||||
|
||||
dv = node.write_value.await_args.args[0]
|
||||
assert dv.StatusCode_ is None
|
||||
assert dv.SourceTimestamp is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_message_flag_override_disables_timestamp():
|
||||
# Engine default is True, but the command explicitly turns it off.
|
||||
eng = make_engine(write_timestamp=True, write_quality=True)
|
||||
node = make_node()
|
||||
eng.pool.get.return_value = make_client(node)
|
||||
|
||||
await eng.write(base_cmd(write_timestamp=False))
|
||||
|
||||
dv = node.write_value.await_args.args[0]
|
||||
assert dv.SourceTimestamp is None
|
||||
assert dv.StatusCode_ is not None # write_quality untouched, stays on
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_bad_type_returns_error_without_raising():
|
||||
eng = make_engine()
|
||||
|
||||
result = await eng.write(base_cmd(value='abc', data_type='int'))
|
||||
|
||||
assert result.ok is False
|
||||
assert result.status == 'error'
|
||||
eng.sink.update_write_status.assert_awaited_once()
|
||||
eng.pool.get.assert_not_awaited()
|
||||
eng.emit_metric.assert_any_await(
|
||||
metric_object=metrics.OPC_ERROR_BY_TYPE,
|
||||
tags={'pod_id': eng.pod_id, 'server_name': 'srv1', 'error_type': 'conversion'},
|
||||
)
|
||||
eng.emit_metric.assert_any_await(
|
||||
metric_object=metrics.OPC_WRITE_TOTAL,
|
||||
tags={'pod_id': eng.pod_id, 'server_name': 'srv1', 'model_id': 'm1', 'result': 'error'},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_value_raises_maps_to_error_result_not_raised():
|
||||
eng = make_engine()
|
||||
node = make_node()
|
||||
node.write_value.side_effect = RuntimeError('boom')
|
||||
eng.pool.get.return_value = make_client(node)
|
||||
|
||||
result = await eng.write(base_cmd())
|
||||
|
||||
assert result.ok is False
|
||||
assert result.status == 'error'
|
||||
assert result.opc_status_code == 'RuntimeError'
|
||||
eng.sink.update_write_status.assert_awaited_once()
|
||||
eng.emit_metric.assert_any_await(
|
||||
metric_object=metrics.OPC_ERROR_BY_TYPE,
|
||||
tags={'pod_id': eng.pod_id, 'server_name': 'srv1', 'error_type': 'RuntimeError'},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_session_bad_status_code_error_maps_to_session_bad():
|
||||
from asyncua.ua.uaerrors import BadSessionClosed
|
||||
|
||||
eng = make_engine()
|
||||
node = make_node()
|
||||
node.write_value.side_effect = BadSessionClosed('closed')
|
||||
eng.pool.get.return_value = make_client(node)
|
||||
|
||||
result = await eng.write(base_cmd())
|
||||
|
||||
assert result.ok is False
|
||||
assert result.status == 'session_bad'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limit_sleeps_between_back_to_back_writes():
|
||||
eng = make_engine(min_pause_s=0.05)
|
||||
eng.pool.get.return_value = make_client()
|
||||
|
||||
with patch('ingestor.engines.write.asyncio.sleep', new_callable=AsyncMock) as sleep_mock:
|
||||
await eng.write(base_cmd())
|
||||
sleep_mock.assert_not_awaited()
|
||||
await eng.write(base_cmd())
|
||||
sleep_mock.assert_awaited_once()
|
||||
remaining = sleep_mock.await_args.args[0]
|
||||
assert 0 < remaining <= 0.05
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_rate_limit_when_min_pause_is_zero():
|
||||
eng = make_engine(min_pause_s=0.0)
|
||||
eng.pool.get.return_value = make_client()
|
||||
|
||||
with patch('ingestor.engines.write.asyncio.sleep', new_callable=AsyncMock) as sleep_mock:
|
||||
await eng.write(base_cmd())
|
||||
await eng.write(base_cmd())
|
||||
sleep_mock.assert_not_awaited()
|
||||
Reference in New Issue
Block a user