feature/SIENTIAPDE-2162 - Inject mongo_repository into gates and add best-effort prediction metadata write
This commit is contained in:
@@ -29,6 +29,7 @@ from laborious.workflows.sub_workflows.prediction_process import PredictionProce
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.repository.mongodb_repository_sync import MongoDBRepository
|
||||
|
||||
# Single source of truth for the test database schema. Mirrors the production
|
||||
# DDL for ``sientia_data`` so any production change can be pasted directly into
|
||||
@@ -141,6 +142,24 @@ def notification_inserts(notification_handler):
|
||||
collection.insert_one = original_insert_one
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
def mongo_repository(mock_logger, notification_handler, metrics_controller, mongo_container):
|
||||
"""Real MongoDBRepository for gate metadata (prediction_metadata), same testcontainer as
|
||||
``notification_handler`` above."""
|
||||
mongo_port = mongo_container.get_exposed_port(27017)
|
||||
repository = MongoDBRepository(
|
||||
connection_string=f'mongodb://localhost:{mongo_port}',
|
||||
database_name='test_db',
|
||||
logger=mock_logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
try:
|
||||
yield repository
|
||||
finally:
|
||||
repository.close()
|
||||
|
||||
|
||||
class _FakeModelWrapper:
|
||||
"""External MLflow wrapper double used by repository stub."""
|
||||
|
||||
@@ -233,6 +252,7 @@ async def test_activities(
|
||||
plugin_store_stub,
|
||||
pi_web_api_client_stub,
|
||||
opc_repository_stub,
|
||||
mongo_repository,
|
||||
):
|
||||
"""Activities with real infra and external-system stubs only."""
|
||||
minio_client = minio_container.get_client()
|
||||
@@ -269,6 +289,7 @@ async def test_activities(
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
mlflow_repository=mlflow_repository_stub,
|
||||
mongo_repository=mongo_repository,
|
||||
)
|
||||
activities.pi_web_api_client = pi_web_api_client_stub
|
||||
activities.opc_repository = {'1': opc_repository_stub}
|
||||
@@ -289,6 +310,7 @@ async def test_activities_real_minio(
|
||||
plugin_store_stub,
|
||||
pi_web_api_client_stub,
|
||||
opc_repository_stub,
|
||||
mongo_repository,
|
||||
):
|
||||
"""Compatibility alias for offload tests."""
|
||||
minio_client = minio_container.get_client()
|
||||
@@ -325,6 +347,7 @@ async def test_activities_real_minio(
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
mlflow_repository=mlflow_repository_stub,
|
||||
mongo_repository=mongo_repository,
|
||||
)
|
||||
activities.pi_web_api_client = pi_web_api_client_stub
|
||||
activities.opc_repository = {'1': opc_repository_stub}
|
||||
@@ -535,6 +558,7 @@ async def test_activities_real_opc(
|
||||
plugin_store_stub,
|
||||
pi_web_api_client_stub,
|
||||
opc_e2e_server: OpcE2ETestServer,
|
||||
mongo_repository,
|
||||
):
|
||||
"""Activities with real OpcRepository connected to the in-process OPC UA server."""
|
||||
minio_client = minio_container.get_client()
|
||||
@@ -571,6 +595,7 @@ async def test_activities_real_opc(
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
mlflow_repository=mlflow_repository_stub,
|
||||
mongo_repository=mongo_repository,
|
||||
)
|
||||
activities.pi_web_api_client = pi_web_api_client_stub
|
||||
await asyncio.to_thread(_connect_activities_to_opc_server, activities, opc_e2e_server.url)
|
||||
|
||||
@@ -7,6 +7,8 @@ with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.repository.minio_repository_sync import MinioRepository
|
||||
from sientia_do.repository.mongodb_repository_sync import MongoDBRepository
|
||||
from sientia_do.utils.connectors_config import build_mongodb_config
|
||||
from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
|
||||
from sientia_model.model_repository.plugin_store import PluginStore
|
||||
|
||||
@@ -48,6 +50,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController | None = None,
|
||||
mlflow_repository: SientiaMLflowRepository | None = None,
|
||||
mongo_repository: MongoDBRepository | None = None,
|
||||
):
|
||||
"""
|
||||
Wire Postgres, MinIO, MLflow, OPC, gates, metrics, and PI Web API into a single object.
|
||||
@@ -67,6 +70,9 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
- metrics_controller: Optional shared controller; if ``None``, a new one is created.
|
||||
- mlflow_repository: Optional ``SientiaMLflowRepository`` for unit/e2e tests; in production
|
||||
leave unset so the repository is built from environment via ``build_mlflow_config()``.
|
||||
- mongo_repository: Optional ``MongoDBRepository`` (gate metadata, best-effort) for
|
||||
unit/e2e tests; in production leave unset so the repository is built from environment
|
||||
via ``build_mongodb_config()``, same pattern as ``mlflow_repository`` above.
|
||||
|
||||
Raises:
|
||||
Exception: If any parent ``__init__`` fails (e.g. invalid config keys).
|
||||
@@ -89,6 +95,16 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
metrics_controller=mc,
|
||||
)
|
||||
|
||||
if mongo_repository is None:
|
||||
mongo_cfg = build_mongodb_config()
|
||||
mongo_repository = MongoDBRepository(
|
||||
connection_string=mongo_cfg['connection_string'],
|
||||
database_name=mongo_cfg['database_name'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mc,
|
||||
)
|
||||
|
||||
minio_repository = MinioRepository(
|
||||
endpoint=minio_config['endpoint_url'],
|
||||
access_key=minio_config['access_key'],
|
||||
@@ -129,6 +145,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
Gates.__init__(
|
||||
self,
|
||||
minio_repository=minio_repository,
|
||||
mongo_repository=mongo_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mc,
|
||||
|
||||
@@ -5,6 +5,7 @@ with workflow.unsafe.imports_passed_through():
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Any
|
||||
|
||||
from bson import ObjectId
|
||||
from pandas import DataFrame
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
@@ -12,6 +13,7 @@ with workflow.unsafe.imports_passed_through():
|
||||
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.repository.mongodb_repository_sync import MongoDBRepository
|
||||
from sientia_do.utils.formatters import create_sample_dict
|
||||
|
||||
from laborious import metrics
|
||||
@@ -86,11 +88,14 @@ class Gates(SientiaMonitoring):
|
||||
"""
|
||||
|
||||
minio_repository: MinioRepository | None = None
|
||||
mongo_repository: MongoDBRepository | None = None
|
||||
_MAX_DEBUG_DATAFRAME_ROWS = 100
|
||||
_PREDICTION_METADATA_COLLECTION = 'prediction_metadata'
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
minio_repository: MinioRepository | None = None,
|
||||
mongo_repository: MongoDBRepository | None = None,
|
||||
logger: Logger | None = None,
|
||||
notification_handler: NotificationHandler | None = None,
|
||||
metrics_controller: MetricsController | None = None,
|
||||
@@ -101,11 +106,14 @@ class Gates(SientiaMonitoring):
|
||||
Args:
|
||||
logger: Logger instance for observability and debugging
|
||||
notification_handler: Notification handler for alerts and monitoring
|
||||
mongo_repository: MongoDB, where gate metadata (triggered filter, policy, counter) is
|
||||
written best-effort. `None` is a valid state (see `_write_prediction_metadata`).
|
||||
|
||||
Raises:
|
||||
Exception: If BaseActivity initialization fails
|
||||
"""
|
||||
self.minio_repository = minio_repository
|
||||
self.mongo_repository = mongo_repository
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
@@ -123,6 +131,11 @@ class Gates(SientiaMonitoring):
|
||||
self.minio_repository.close()
|
||||
finally:
|
||||
self.minio_repository = None
|
||||
if self.mongo_repository is not None:
|
||||
try:
|
||||
self.mongo_repository.close()
|
||||
finally:
|
||||
self.mongo_repository = None
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
def __del__(self):
|
||||
@@ -146,6 +159,56 @@ class Gates(SientiaMonitoring):
|
||||
metadata,
|
||||
)
|
||||
|
||||
def _write_prediction_metadata(
|
||||
self, document: dict[str, Any] | None, metadata: dict[str, Any]
|
||||
) -> str | None:
|
||||
"""
|
||||
Best-effort write of gate metadata (triggered filter, policy, counter) to Mongo.
|
||||
|
||||
Implements both business rules for this feature in one place:
|
||||
- The prediction write must never be blocked or failed by a metadata write. A failed
|
||||
`insert(...)` is swallowed, never re-raised, logged here with gate-level context (which
|
||||
`mongo_repository.insert(...)` itself can't know) — matching how `storage.py`'s
|
||||
`cleanup_minio_objects_expired` also logs again at the call site even though the
|
||||
repository already logs internally and emits its own Prometheus metric
|
||||
(`MONGODB_WRITE_ERROR_COUNT`); the generic library log is for the metric/trace, this
|
||||
one is for whoever is reading logs for this specific gate.
|
||||
- Nothing is written when there's nothing relevant to report. A falsy `document` (e.g.
|
||||
no filter triggered) skips Mongo entirely, so callers can call this unconditionally
|
||||
and get `None` back for "no metadata this time" instead of deciding upfront whether
|
||||
there's something worth saving.
|
||||
|
||||
`insert(...)` doesn't hand back the inserted document's `_id`, so it is generated here,
|
||||
client-side, before the insert, as a plain string (not a BSON ObjectId) — matching this
|
||||
codebase's convention of string `_id`s (e.g. the `models` collection) and the 24-hex-char
|
||||
size of the `predictions.metadata_id` Postgres column that will reference it.
|
||||
|
||||
Args:
|
||||
document: Metadata payload (filter, policy, counter, etc.); `None`/empty skips the
|
||||
write entirely. No `_id` key expected.
|
||||
metadata: Workflow metadata for contextual logging.
|
||||
|
||||
Return:
|
||||
The generated id, or `None` when there's nothing to write, Mongo isn't configured, or
|
||||
the insert failed.
|
||||
"""
|
||||
if not document or self.mongo_repository is None:
|
||||
return None
|
||||
|
||||
generated_id = str(ObjectId())
|
||||
|
||||
try:
|
||||
self.mongo_repository.insert(
|
||||
self._PREDICTION_METADATA_COLLECTION,
|
||||
{**document, '_id': generated_id},
|
||||
metadata,
|
||||
)
|
||||
except Exception as e:
|
||||
self.error(f'Failed to write prediction metadata: {e}', metadata)
|
||||
return None
|
||||
|
||||
return generated_id
|
||||
|
||||
@staticmethod
|
||||
def _read_filter_entry(config: dict[str, Any]) -> tuple[str, dict[str, Any]]:
|
||||
"""
|
||||
|
||||
@@ -47,6 +47,7 @@ def test___init__(
|
||||
}
|
||||
|
||||
mlflow_repository = MagicMock()
|
||||
mongo_repository = MagicMock()
|
||||
plugin_store = MagicMock()
|
||||
|
||||
opc_config = {
|
||||
@@ -73,6 +74,7 @@ def test___init__(
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
mlflow_repository=mlflow_repository,
|
||||
mongo_repository=mongo_repository,
|
||||
)
|
||||
|
||||
assert isinstance(activities, Activities)
|
||||
@@ -120,6 +122,7 @@ def test___init__(
|
||||
mock_gates_init.assert_called_once_with(
|
||||
ANY,
|
||||
minio_repository=mock_minio_repository.return_value,
|
||||
mongo_repository=mongo_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mock_metrics_controller.return_value,
|
||||
@@ -217,6 +220,7 @@ def test_shutdown(
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
mlflow_repository=mlflow_repository,
|
||||
mongo_repository=MagicMock(),
|
||||
)
|
||||
|
||||
activities.shutdown()
|
||||
@@ -286,6 +290,7 @@ def test___init___builds_mlflow_repository_when_not_provided(
|
||||
pi_web_api_config=pi_web_api_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
mongo_repository=MagicMock(),
|
||||
)
|
||||
|
||||
mock_build_mlflow_config.assert_called_once()
|
||||
@@ -297,3 +302,83 @@ def test___init___builds_mlflow_repository_when_not_provided(
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mock_metrics_controller.return_value,
|
||||
)
|
||||
|
||||
|
||||
@patch('laborious.activities.activities.MongoDBRepository')
|
||||
@patch('laborious.activities.activities.build_mongodb_config')
|
||||
@patch('laborious.activities.activities.Storage.__init__')
|
||||
@patch('laborious.activities.activities.MLFlow.__init__')
|
||||
@patch('laborious.activities.activities.OPC.__init__')
|
||||
@patch('laborious.activities.activities.Gates.__init__')
|
||||
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
||||
@patch('laborious.activities.activities.API.__init__')
|
||||
@patch('laborious.activities.activities.MinioRepository')
|
||||
@patch('laborious.activities.activities.MetricsController')
|
||||
def test___init___builds_mongo_repository_when_not_provided(
|
||||
mock_metrics_controller,
|
||||
mock_minio_repository,
|
||||
_mock_api_init,
|
||||
_mock_model_metrics_init,
|
||||
mock_gates_init,
|
||||
_mock_opc_init,
|
||||
_mock_mlflow_init,
|
||||
_mock_storage_init,
|
||||
mock_build_mongodb_config,
|
||||
mock_mongo_repository_cls,
|
||||
):
|
||||
postgres_config = {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'user': 'postgres',
|
||||
'password': 'postgres',
|
||||
'dbname': 'postgres',
|
||||
'min_connections': 1,
|
||||
'max_connections': 10,
|
||||
}
|
||||
minio_config = {
|
||||
'endpoint_url': 'localhost:9000',
|
||||
'access_key': 'minio',
|
||||
'secret_key': 'minio123',
|
||||
'default_bucket': 'test',
|
||||
'retention_hours': 24,
|
||||
'secure': False,
|
||||
}
|
||||
opc_config = {'bootstrap_servers': 'localhost:9092', 'polling_time': 1000, 'group_id': 'test'}
|
||||
pi_web_api_config = {'base_url': 'https://pi', 'auth_type': 'bearer', 'auth_token': 'token'}
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
plugin_store = MagicMock()
|
||||
mlflow_repository = MagicMock()
|
||||
mock_build_mongodb_config.return_value = {
|
||||
'connection_string': 'mongodb://root:pass@localhost:27018',
|
||||
'database_name': 'sientia',
|
||||
}
|
||||
|
||||
Activities(
|
||||
postgres_config=postgres_config,
|
||||
plugin_store=plugin_store,
|
||||
minio_config=minio_config,
|
||||
opc_config=opc_config,
|
||||
pi_web_api_config=pi_web_api_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
mlflow_repository=mlflow_repository,
|
||||
)
|
||||
|
||||
mock_build_mongodb_config.assert_called_once()
|
||||
mock_mongo_repository_cls.assert_called_once_with(
|
||||
connection_string='mongodb://root:pass@localhost:27018',
|
||||
database_name='sientia',
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mock_metrics_controller.return_value,
|
||||
)
|
||||
# The same repository built here is the one handed to Gates, not a second instance.
|
||||
mock_gates_init.assert_called_once_with(
|
||||
ANY,
|
||||
minio_repository=mock_minio_repository.return_value,
|
||||
mongo_repository=mock_mongo_repository_cls.return_value,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mock_metrics_controller.return_value,
|
||||
)
|
||||
|
||||
@@ -35,8 +35,14 @@ def _minio_payload(retrieve_return, status=None):
|
||||
|
||||
|
||||
@fixture
|
||||
def gates_activity():
|
||||
def mongo() -> MagicMock:
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@fixture
|
||||
def gates_activity(mongo: MagicMock):
|
||||
gates = Gates(
|
||||
mongo_repository=mongo,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(spec=MetricsController),
|
||||
@@ -61,6 +67,75 @@ metadata = {
|
||||
}
|
||||
|
||||
|
||||
def test_write_prediction_metadata_no_document_skips_mongo(gates_activity, mongo: MagicMock):
|
||||
# Act
|
||||
result = gates_activity._write_prediction_metadata(None, metadata)
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
mongo.insert.assert_not_called()
|
||||
|
||||
|
||||
def test_write_prediction_metadata_empty_document_skips_mongo(gates_activity, mongo: MagicMock):
|
||||
# Act
|
||||
result = gates_activity._write_prediction_metadata({}, metadata)
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
mongo.insert.assert_not_called()
|
||||
|
||||
|
||||
def test_write_prediction_metadata_not_configured_returns_none():
|
||||
# Arrange
|
||||
gates = Gates(
|
||||
mongo_repository=None,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(spec=MetricsController),
|
||||
)
|
||||
gates.error = MagicMock()
|
||||
|
||||
# Act
|
||||
result = gates._write_prediction_metadata({'filter_triggered': 'EMPTY_DATA'}, metadata)
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
gates.error.assert_not_called()
|
||||
|
||||
|
||||
def test_write_prediction_metadata_success_generates_id_and_inserts(
|
||||
gates_activity, mongo: MagicMock
|
||||
):
|
||||
# Arrange
|
||||
document = {'filter_triggered': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 3}
|
||||
|
||||
# Act
|
||||
result = gates_activity._write_prediction_metadata(document, metadata)
|
||||
|
||||
# Assert: a plain 24-hex-char string id (matches predictions.metadata_id's varchar(24)),
|
||||
# not a BSON ObjectId — `insert(...)` never hands one back, so it's generated client-side.
|
||||
assert result is not None
|
||||
assert len(result) == 24
|
||||
int(result, 16) # raises ValueError if not hex
|
||||
|
||||
collection, inserted_document, inserted_metadata = mongo.insert.call_args.args
|
||||
assert collection == Gates._PREDICTION_METADATA_COLLECTION
|
||||
assert inserted_document == {**document, '_id': result}
|
||||
assert inserted_metadata == metadata
|
||||
|
||||
|
||||
def test_write_prediction_metadata_insert_failure_swallowed(gates_activity, mongo: MagicMock):
|
||||
# Arrange
|
||||
mongo.insert.side_effect = RuntimeError('connection reset')
|
||||
|
||||
# Act
|
||||
result = gates_activity._write_prediction_metadata({'filter_triggered': 'EMPTY_DATA'}, metadata)
|
||||
|
||||
# Assert: the prediction write must never be blocked or failed by this
|
||||
assert result is None
|
||||
gates_activity.error.assert_called_once()
|
||||
|
||||
|
||||
def test_input_gate_invalid_filter(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
|
||||
Reference in New Issue
Block a user