From d8ed79ddf5e005936d057bbf265f8eeb18334f57 Mon Sep 17 00:00:00 2001 From: Val Redchenko Date: Wed, 8 Jul 2026 10:20:31 +0100 Subject: [PATCH] feat: ingest motion-corrected micrograph image paths (#309) Persist the motion-corrected image artefacts cryoem-services produces so the micrograph leaf can be served (ADR 0021). This is the ingestion half; the serving endpoint is #308. - add nullable micrograph.motion_corrected_image_path and motion_corrected_snapshot_path columns (Alembic migration, no backfill) - extend MotionCorrectionCompletedRequest and MotionCorrectionCompleteBody additively with both optional paths (backward-compatible) - thread the paths through publish_motion_correction_completed and persist whichever are present in handle_motion_correction_complete The two paths arrive via the existing motion_correction/completed contract; callers that omit them are unaffected. --- src/smartem_backend/api_server.py | 2 ++ src/smartem_backend/consumer.py | 4 +++ ...micrograph_motion_corrected_image_paths.py | 25 +++++++++++++++ src/smartem_backend/model/database.py | 2 ++ src/smartem_backend/model/http_request.py | 2 ++ src/smartem_backend/model/mq_event.py | 2 ++ src/smartem_backend/mq_publisher.py | 10 +++++- .../test_consumer_async_handlers.py | 31 +++++++++++++++++++ .../test_processing_feedback_endpoints.py | 16 ++++++++++ 9 files changed, 93 insertions(+), 1 deletion(-) create mode 100644 src/smartem_backend/migrations/versions/2026_07_08_1200-d5e6f7a8b9c0_add_micrograph_motion_corrected_image_paths.py diff --git a/src/smartem_backend/api_server.py b/src/smartem_backend/api_server.py index 860ed4d4..0a00aeb8 100644 --- a/src/smartem_backend/api_server.py +++ b/src/smartem_backend/api_server.py @@ -1389,6 +1389,8 @@ async def publish_micrograph_motion_correction_completed( micrograph_uuid=micrograph_uuid, total_motion=payload.total_motion, average_motion=payload.average_motion, + motion_corrected_image_path=payload.motion_corrected_image_path, + motion_corrected_snapshot_path=payload.motion_corrected_snapshot_path, ) return _publish_or_502(success, "motion_correction_completed", micrograph_uuid) diff --git a/src/smartem_backend/consumer.py b/src/smartem_backend/consumer.py index 74816350..2115778b 100755 --- a/src/smartem_backend/consumer.py +++ b/src/smartem_backend/consumer.py @@ -455,6 +455,10 @@ async def handle_motion_correction_complete(event_data: dict[str, Any]) -> None: .first() ) if micrograph: + if event.motion_corrected_image_path is not None: + micrograph.motion_corrected_image_path = event.motion_corrected_image_path + if event.motion_corrected_snapshot_path is not None: + micrograph.motion_corrected_snapshot_path = event.motion_corrected_snapshot_path micrograph.updated_at = datetime.now() await session.commit() await publish_motion_correction_registered( diff --git a/src/smartem_backend/migrations/versions/2026_07_08_1200-d5e6f7a8b9c0_add_micrograph_motion_corrected_image_paths.py b/src/smartem_backend/migrations/versions/2026_07_08_1200-d5e6f7a8b9c0_add_micrograph_motion_corrected_image_paths.py new file mode 100644 index 00000000..1f4f9ba1 --- /dev/null +++ b/src/smartem_backend/migrations/versions/2026_07_08_1200-d5e6f7a8b9c0_add_micrograph_motion_corrected_image_paths.py @@ -0,0 +1,25 @@ +"""Add micrograph motion-corrected image and snapshot paths + +Revision ID: d5e6f7a8b9c0 +Revises: c4d5e6f7a8b9 +Create Date: 2026-07-08 12:00:00.000000 + +""" + +import sqlalchemy as sa +from alembic import op + +revision = "d5e6f7a8b9c0" +down_revision = "c4d5e6f7a8b9" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.add_column("micrograph", sa.Column("motion_corrected_image_path", sa.String(), nullable=True)) + op.add_column("micrograph", sa.Column("motion_corrected_snapshot_path", sa.String(), nullable=True)) + + +def downgrade() -> None: + op.drop_column("micrograph", "motion_corrected_snapshot_path") + op.drop_column("micrograph", "motion_corrected_image_path") diff --git a/src/smartem_backend/model/database.py b/src/smartem_backend/model/database.py index da54a6d2..cae483ae 100644 --- a/src/smartem_backend/model/database.py +++ b/src/smartem_backend/model/database.py @@ -223,6 +223,8 @@ class Micrograph(SQLModel, table=True, table_name="micrograph"): binning_y: int | None = Field(default=None) total_motion: float | None = Field(default=None) average_motion: float | None = Field(default=None) + motion_corrected_image_path: str | None = Field(default=None) + motion_corrected_snapshot_path: str | None = Field(default=None) ctf_max_resolution_estimate: float | None = Field(default=None) number_of_particles_selected: int | None = Field(default=None) number_of_particles_rejected: int | None = Field(default=None) diff --git a/src/smartem_backend/model/http_request.py b/src/smartem_backend/model/http_request.py index 4cbddc4c..2d3fddd1 100644 --- a/src/smartem_backend/model/http_request.py +++ b/src/smartem_backend/model/http_request.py @@ -300,6 +300,8 @@ class MicrographUpdateRequest(MicrographBaseFields): class MotionCorrectionCompletedRequest(BaseModel): total_motion: float average_motion: float + motion_corrected_image_path: str | None = None + motion_corrected_snapshot_path: str | None = None class MotionCorrectionRegisteredRequest(BaseModel): diff --git a/src/smartem_backend/model/mq_event.py b/src/smartem_backend/model/mq_event.py index 99f231e3..7831e49e 100644 --- a/src/smartem_backend/model/mq_event.py +++ b/src/smartem_backend/model/mq_event.py @@ -331,6 +331,8 @@ class MotionCorrectionCompleteBody(GenericEventMessageBody): micrograph_uuid: str total_motion: float average_motion: float + motion_corrected_image_path: str | None = None + motion_corrected_snapshot_path: str | None = None @model_validator(mode="after") def check_model(self): diff --git a/src/smartem_backend/mq_publisher.py b/src/smartem_backend/mq_publisher.py index 5c05dcb0..f70ddc80 100644 --- a/src/smartem_backend/mq_publisher.py +++ b/src/smartem_backend/mq_publisher.py @@ -367,12 +367,20 @@ async def publish_model_parameter_update( return await _publish(MessageQueueEventType.MODEL_PARAMETER_UPDATE, event) -async def publish_motion_correction_completed(micrograph_uuid: str, total_motion: float, average_motion: float) -> bool: +async def publish_motion_correction_completed( + micrograph_uuid: str, + total_motion: float, + average_motion: float, + motion_corrected_image_path: str | None = None, + motion_corrected_snapshot_path: str | None = None, +) -> bool: event = MotionCorrectionCompleteBody( event_type=MessageQueueEventType.MOTION_CORRECTION_COMPLETE, micrograph_uuid=micrograph_uuid, total_motion=total_motion, average_motion=average_motion, + motion_corrected_image_path=motion_corrected_image_path, + motion_corrected_snapshot_path=motion_corrected_snapshot_path, ) return await _publish(MessageQueueEventType.MOTION_CORRECTION_COMPLETE, event) diff --git a/tests/smartem_backend/test_consumer_async_handlers.py b/tests/smartem_backend/test_consumer_async_handlers.py index 01d25dea..d35b78ec 100644 --- a/tests/smartem_backend/test_consumer_async_handlers.py +++ b/tests/smartem_backend/test_consumer_async_handlers.py @@ -179,6 +179,37 @@ async def _stub_publish(*args, **kwargs): assert db.add.call_count == 1 assert db.commit.await_count == 1 + def test_persists_motion_corrected_image_paths(self, db, monkeypatch, stub_publisher): + grid_row = MagicMock() + grid_row.grid_uuid = "grid-1" + micrograph = MagicMock() + db.execute.return_value.one.return_value = (grid_row,) + db.execute.return_value.scalars.return_value.all.return_value = [] + db.execute.return_value.scalars.return_value.first.return_value = micrograph + + async def _stub_check(*args, **kwargs): + return 0.7 + + async def _stub_prior(*args, **kwargs): + return None + + async def _stub_publish(*args, **kwargs): + return True + + monkeypatch.setattr(consumer, "_check_against_statistics", _stub_check) + monkeypatch.setattr(consumer, "prior_update", _stub_prior) + monkeypatch.setattr(consumer, "publish_motion_correction_registered", _stub_publish) + + import asyncio + + event = dict(self.base_event) + event["motion_corrected_image_path"] = "/dls/mc/mic.mrc" + event["motion_corrected_snapshot_path"] = "/dls/mc/mic.mrc.jpeg" + asyncio.run(consumer.handle_motion_correction_complete(event)) + + assert micrograph.motion_corrected_image_path == "/dls/mc/mic.mrc" + assert micrograph.motion_corrected_snapshot_path == "/dls/mc/mic.mrc.jpeg" + class TestRefreshPredictions: base_event = { diff --git a/tests/smartem_backend/test_processing_feedback_endpoints.py b/tests/smartem_backend/test_processing_feedback_endpoints.py index 8d2ec06e..75704b52 100644 --- a/tests/smartem_backend/test_processing_feedback_endpoints.py +++ b/tests/smartem_backend/test_processing_feedback_endpoints.py @@ -71,8 +71,24 @@ def test_happy_path(self, client, captured): "micrograph_uuid": "abc-123", "total_motion": 1.5, "average_motion": 0.1, + "motion_corrected_image_path": None, + "motion_corrected_snapshot_path": None, } + def test_image_paths_are_forwarded(self, client, captured): + resp = client.post( + self.endpoint, + json={ + "total_motion": 1.5, + "average_motion": 0.1, + "motion_corrected_image_path": "/dls/mc/mic.mrc", + "motion_corrected_snapshot_path": "/dls/mc/mic.mrc.jpeg", + }, + ) + assert resp.status_code == 202 + assert captured[0].kwargs["motion_corrected_image_path"] == "/dls/mc/mic.mrc" + assert captured[0].kwargs["motion_corrected_snapshot_path"] == "/dls/mc/mic.mrc.jpeg" + def test_missing_micrograph_returns_404(self, client, captured): client._db.execute.return_value = make_execute_result(None) resp = client.post(self.endpoint, json={"total_motion": 1.5, "average_motion": 0.1})