75 lines
3.0 KiB
Python
75 lines
3.0 KiB
Python
import re
|
|
|
|
from sqlalchemy import text
|
|
|
|
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'[A-Za-z0-9_]+')
|
|
|
|
|
|
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."""
|
|
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) -> None:
|
|
"""UPDATE the per-model prediction row's write_status (C3), by timestamp+model_id.
|
|
|
|
schema/table_name are per-model SQL identifiers, not values — they can't be
|
|
bound params, so they're checked against an allowlist before interpolation.
|
|
status/timestamp/model_id ARE bound params, never f-stringed.
|
|
"""
|
|
schema = self._validate_identifier(cmd['schema'])
|
|
table = self._validate_identifier(cmd['table_name'])
|
|
# schema/table are allowlist-validated identifiers above; status/timestamp/model_id
|
|
# are bound params below, never interpolated — not a real injection vector.
|
|
stmt = text(
|
|
f'UPDATE {schema}.{table} SET write_status = :status ' # noqa: S608
|
|
'WHERE timestamp = :timestamp AND model_id = :model_id'
|
|
)
|
|
params = {
|
|
'status': result.status,
|
|
'timestamp': cmd['timestamp'],
|
|
'model_id': cmd['model_id'],
|
|
}
|
|
# ponytail: sync engine call inside async def — same pattern sientia_do's own
|
|
# Postgres activities use (e.g. export_data_to_postgres); no async pg driver here.
|
|
with self.postgres.engine.begin() as connection:
|
|
connection.execute(stmt, params)
|
|
|
|
@staticmethod
|
|
def _validate_identifier(name: str) -> str:
|
|
if not isinstance(name, str) or not _IDENTIFIER_RE.fullmatch(name):
|
|
raise ValueError(f'invalid SQL identifier: {name!r}')
|
|
return name
|