feat(gateway): DataSink.update_write_status Postgres by timestamp+model_id (C3); drop unused metadata param
This commit is contained in:
@@ -1,16 +1,18 @@
|
|||||||
|
import re
|
||||||
|
|
||||||
|
from sqlalchemy import text
|
||||||
|
|
||||||
import ingestor.metrics as metrics
|
import ingestor.metrics as metrics
|
||||||
|
|
||||||
|
_IDENTIFIER_RE = re.compile(r'^[A-Za-z0-9_]+$')
|
||||||
|
|
||||||
|
|
||||||
class DataSink:
|
class DataSink:
|
||||||
"""Gateway egress wrapper: fans a reading out to every configured name/collection.
|
"""Gateway egress wrapper: fans a reading out to every configured name/collection."""
|
||||||
|
|
||||||
`postgres` is unused for now — task 12 adds `update_write_status` on top of it.
|
def __init__(self, data_manager, postgres) -> None:
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, data_manager, postgres, metadata: dict | None = None) -> None:
|
|
||||||
self.data_manager = data_manager
|
self.data_manager = data_manager
|
||||||
self.postgres = postgres
|
self.postgres = postgres
|
||||||
self.metadata = metadata
|
|
||||||
|
|
||||||
async def emit_metric(self, **kwargs) -> None:
|
async def emit_metric(self, **kwargs) -> None:
|
||||||
# ponytail: delegate to data_manager, it's already a SientiaMonitoring
|
# ponytail: delegate to data_manager, it's already a SientiaMonitoring
|
||||||
@@ -37,3 +39,30 @@ class DataSink:
|
|||||||
'model_id': name['model_id'],
|
'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 _IDENTIFIER_RE.match(name):
|
||||||
|
raise ValueError(f'invalid SQL identifier: {name!r}')
|
||||||
|
return name
|
||||||
|
|||||||
72
tests/unit/io/test_data_sink_status.py
Normal file
72
tests/unit/io/test_data_sink_status.py
Normal file
@@ -0,0 +1,72 @@
|
|||||||
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from ingestor.engines.write import WriteResult
|
||||||
|
from ingestor.io.data_sink import DataSink
|
||||||
|
|
||||||
|
|
||||||
|
def make_sink():
|
||||||
|
postgres = MagicMock()
|
||||||
|
connection = MagicMock()
|
||||||
|
postgres.engine.begin.return_value.__enter__.return_value = connection
|
||||||
|
sink = DataSink(data_manager=MagicMock(), postgres=postgres)
|
||||||
|
return sink, connection
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_write_status_targets_model_row():
|
||||||
|
sink, connection = make_sink()
|
||||||
|
cmd = {
|
||||||
|
'schema': 'sientia_data',
|
||||||
|
'table_name': 'pred_m1',
|
||||||
|
'model_id': '1',
|
||||||
|
'timestamp': '2026-07-13T00:00:00Z',
|
||||||
|
}
|
||||||
|
result = WriteResult(ok=True, status='ok', opc_status_code='Good', message='')
|
||||||
|
|
||||||
|
await sink.update_write_status(cmd, result)
|
||||||
|
|
||||||
|
assert connection.execute.call_count == 1
|
||||||
|
stmt, params = connection.execute.call_args[0]
|
||||||
|
sql = str(stmt)
|
||||||
|
assert 'update' in sql.lower() and 'write_status' in sql.lower()
|
||||||
|
assert 'sientia_data.pred_m1' in sql
|
||||||
|
# values must be bound params, never interpolated into the SQL text itself
|
||||||
|
assert params == {'status': 'ok', 'timestamp': '2026-07-13T00:00:00Z', 'model_id': '1'}
|
||||||
|
assert 'ok' not in sql
|
||||||
|
assert '2026-07-13T00:00:00Z' not in sql
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_write_status_rejects_bad_schema():
|
||||||
|
sink, connection = make_sink()
|
||||||
|
cmd = {
|
||||||
|
'schema': 'a; DROP TABLE x',
|
||||||
|
'table_name': 'pred_m1',
|
||||||
|
'model_id': '1',
|
||||||
|
'timestamp': '2026-07-13T00:00:00Z',
|
||||||
|
}
|
||||||
|
result = WriteResult(ok=False, status='error', opc_status_code='Bad', message='x')
|
||||||
|
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
await sink.update_write_status(cmd, result)
|
||||||
|
|
||||||
|
connection.execute.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_write_status_rejects_bad_table_name():
|
||||||
|
sink, connection = make_sink()
|
||||||
|
cmd = {
|
||||||
|
'schema': 'sientia_data',
|
||||||
|
'table_name': 'pred_m1; DROP TABLE x',
|
||||||
|
'model_id': '1',
|
||||||
|
'timestamp': '2026-07-13T00:00:00Z',
|
||||||
|
}
|
||||||
|
result = WriteResult(ok=False, status='error', opc_status_code='Bad', message='x')
|
||||||
|
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
await sink.update_write_status(cmd, result)
|
||||||
|
|
||||||
|
connection.execute.assert_not_called()
|
||||||
Reference in New Issue
Block a user