diff --git a/e2e/db_schema.sql b/e2e/db_schema.sql index 9057815..2c3d15f 100644 --- a/e2e/db_schema.sql +++ b/e2e/db_schema.sql @@ -40,6 +40,7 @@ CREATE TABLE IF NOT EXISTS sientia_data.predictions ( "timestamp" timestamptz NOT NULL, created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL, "comments" text NULL, + metadata_id varchar(24) NULL, CONSTRAINT unique_model_id_timestamp UNIQUE (model_id, "timestamp") ); diff --git a/laborious/activities/storage.py b/laborious/activities/storage.py index ea284fc..b4ae983 100644 --- a/laborious/activities/storage.py +++ b/laborious/activities/storage.py @@ -2,6 +2,7 @@ from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): # Extend the Temporal Postgres activities for convenient query -> MinIO export + import time import traceback from datetime import timedelta from typing import Any @@ -9,12 +10,14 @@ with workflow.unsafe.imports_passed_through(): import pandas as pd from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.models import NotificationLevel + from sientia_do.observability import metrics from sientia_do.observability.logger import Logger from sientia_do.observability.metrics_controller import MetricsController from sientia_do.observability.sientia_monitoring import SientiaMonitoring from sientia_do.repository.minio_repository_sync import MinioRepository from sientia_do.temporal.activities.postgres_sync import Postgres from sientia_do.temporal.constants import now + from sqlalchemy import text from laborious.utils.messages import MISSING_MINIO_REPOSITORY from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload @@ -26,7 +29,9 @@ _LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_thr class Storage(Postgres, SientiaMonitoring): """ Extensions for Postgres activities with a helper to export query results - directly to MinIO as Parquet and return the object name. + directly to MinIO as Parquet and return the object name, and with + ``repeat_last_prediction``, a query specific to this repo's REPEAT path + (moved here from the shared library — nothing else consumes it). """ minio_repository: MinioRepository | None = None @@ -141,6 +146,110 @@ class Storage(Postgres, SientiaMonitoring): } ) + def _build_repeat_last_prediction_query(self, schema: str, table_name: str): + table_identifier = f'"{schema}".{table_name}' + return text( # noqa: S608 + f""" + INSERT INTO {table_identifier} ( + model_id, + prediction, + timestamp, + response_time, + prediction_status, + prediction_confidence, + metadata_id, + created_at + ) + SELECT + model_id, + prediction, + :last_timestamp, + response_time, + prediction_status, + prediction_confidence, + :metadata_id, + NOW() + FROM {table_identifier} + WHERE model_id = :model + ORDER BY timestamp DESC + LIMIT 1; + """ # nosec B608 + ) + + @activity.defn(name='repeat_last_prediction') + def repeat_last_prediction(self, query_items: dict[str, Any]): + """ + Insert a new `predictions` row that copies the last one for a model forward in time. + + REPEAT persists a new row (unlike STOP, which persists nothing), reusing the previous + row's `prediction`/`response_time`/`prediction_status`/`prediction_confidence` values with + a fresh `timestamp`/`created_at`. `metadata_id` references the `prediction_metadata` + document for this cycle, if any (see `PredictionMetadata.write_prediction_metadata`); + `None` is a valid value — the column is nullable. + + Args (query_items): + metadata (dict): Workflow metadata for logging and metrics. + schema (str): Postgres schema of the `predictions` table. + table_name (str): `predictions` table name. + model (str): `model_id` to repeat the last prediction for. + last_timestamp: Timestamp to stamp the new row with. + metadata_id (str | None): `prediction_metadata` document id, or `None`. + + Raises: + Exception: Re-raised after logging/metrics if the insert fails. + """ + metadata = query_items['metadata'] + schema = query_items['schema'] + table_name = query_items['table_name'] + model = query_items['model'] + last_timestamp = query_items['last_timestamp'] + metadata_id = query_items.get('metadata_id') + + repeat_query = self._build_repeat_last_prediction_query(schema, table_name) + self.info(f'Repeating last prediction for model {model}', metadata) + self.debug( + f'Query parameters: schema={schema}, model={model}, last_timestamp={last_timestamp}', + metadata, + ) + core_labels = { + **self.get_core_labels(metadata, 'execute_sql'), + 'conflict_mode': 'error', + } + + with self.session_factory.begin() as session: + try: + start_time = time.time() + session.execute( + repeat_query, + { + 'last_timestamp': last_timestamp, + 'model': model, + 'metadata_id': metadata_id, + }, + ) + except Exception as e: + trace = traceback.format_exc() + self.send_notification( + metadata=metadata, + notification_id='ERROR_REPEATING_LAST_PREDICTION', + message=f'Error repeating last prediction: {e}', + block='repeat_last_prediction', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + self.error(trace, metadata) + self.emit_metric_sync( + metric_object=metrics.POSTGRES_WRITE_ERROR_COUNT, + tags=core_labels, + ) + raise + else: + self.observe_lag_sync(start_time, metrics.POSTGRES_WRITE_LAG, core_labels) + self.emit_metric_sync( + metric_object=metrics.POSTGRES_WRITE_COUNT, + tags=core_labels, + ) + @activity.defn(name='cleanup_minio_objects_expired') def cleanup_minio_objects_expired(self, input_data: dict[str, Any]) -> dict[str, Any]: """ diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index aba6006..9c59240 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -360,10 +360,9 @@ class PredictionProcess: return True elif path_flag == 'REPEAT': metadata_id = await self._write_prediction_metadata(metadata, document) - # Repeat last prediction if available. `metadata_id` isn't consumed by - # `repeat_last_prediction` yet (its INSERT...SELECT has no `metadata_id` column — - # pending a change in `sientia-dataops-library`); passed through already so no further - # change is needed here once that lands. + # Repeat last prediction, copying the previous row's values forward with a new + # timestamp. `repeat_last_prediction` now lives in `Storage` (laborious), not the + # shared library — its INSERT...SELECT references `metadata_id` directly. await workflow.execute_activity_method( Activities.repeat_last_prediction, { diff --git a/tests/laborious/activities/test_storage.py b/tests/laborious/activities/test_storage.py index f88df3a..51124f6 100644 --- a/tests/laborious/activities/test_storage.py +++ b/tests/laborious/activities/test_storage.py @@ -257,6 +257,86 @@ def test_export_payload_to_postgres(storage): assert result == {'success': True} +def test_build_repeat_last_prediction_query(storage): + query = storage._build_repeat_last_prediction_query('public', 'predictions') + query_str = str(query) + + assert 'INSERT INTO "public".predictions' in query_str + assert 'metadata_id' in query_str + assert ':last_timestamp' in query_str + assert ':metadata_id' in query_str + assert ':model' in query_str + + +@patch('laborious.activities.storage.text') +def test_repeat_last_prediction_success(text_mock, storage): + storage.emit_metric_sync = MagicMock() + storage.observe_lag_sync = MagicMock() + query_items = { + **metadata, + 'schema': 'public', + 'table_name': 'predictions', + 'model': 1, + 'last_timestamp': '2022-01-01', + 'metadata_id': 'abc123', + } + + with patch('sqlalchemy.orm.session.Session.execute') as mock_execute: + storage.repeat_last_prediction(query_items) + + text_mock.assert_called_once() + mock_execute.assert_called_once_with( + text_mock.return_value, + {'last_timestamp': '2022-01-01', 'model': 1, 'metadata_id': 'abc123'}, + ) + storage.observe_lag_sync.assert_called_once() + storage.emit_metric_sync.assert_called_once() + + +@patch('laborious.activities.storage.text') +def test_repeat_last_prediction_without_metadata_id(text_mock, storage): + storage.emit_metric_sync = MagicMock() + storage.observe_lag_sync = MagicMock() + query_items = { + **metadata, + 'schema': 'public', + 'table_name': 'predictions', + 'model': 1, + 'last_timestamp': '2022-01-01', + } + + with patch('sqlalchemy.orm.session.Session.execute') as mock_execute: + storage.repeat_last_prediction(query_items) + + mock_execute.assert_called_once_with( + text_mock.return_value, + {'last_timestamp': '2022-01-01', 'model': 1, 'metadata_id': None}, + ) + + +@patch('laborious.activities.storage.text') +def test_repeat_last_prediction_error(text_mock, storage): + storage.emit_metric_sync = MagicMock() + storage.send_notification = MagicMock() + query_items = { + **metadata, + 'schema': 'public', + 'table_name': 'predictions', + 'model': 1, + 'last_timestamp': '2022-01-01', + 'metadata_id': 'abc123', + } + + with ( + patch('sqlalchemy.orm.session.Session.execute', side_effect=ValueError('db')), + raises(ValueError), + ): + storage.repeat_last_prediction(query_items) + + storage.send_notification.assert_called_once() + storage.emit_metric_sync.assert_called() + + def test_cleanup_minio_objects_expired_minio_not_initialized(storage): storage.minio_repository = None