Migrated 'data_collection_group' feedback-callback block into its own workflows module

Author: tieneupin

pyproject.toml

@@ -62,6 +62,7 @@ server = [
     "aiohttp",
     "cryptography",
     "fastapi[standard]<0.116.0",
+    "graypy",
     "ispyb>=10.2.4", # Responsible for setting requirements for SQLAlchemy and mysql-connector-python;
     "jinja2",
     "mrcfile",
@@ -73,7 +74,7 @@ server = [
     "python-jose[cryptography]",
     "sqlalchemy[postgresql]", # Add as explicit dependency
     "sqlmodel",
-    "stomp-py<=8.1.0", # 8.1.1 (released 2024-04-06) doesn't work with our project
+    "stomp-py>8.1.0", # 8.1.1 (released 2024-04-06) doesn't work with our project
     "zocalo>=1",
 ]
 [project.urls]
@@ -105,6 +106,7 @@ GitHub = "https://github.com/DiamondLightSource/python-murfey"
 "clem.process_raw_tiffs" = "murfey.workflows.clem.process_raw_tiffs:zocalo_cluster_request"
 "clem.register_align_and_merge_result" = "murfey.workflows.clem.register_align_and_merge_results:register_align_and_merge_result"
 "clem.register_preprocessing_result" = "murfey.workflows.clem.register_preprocessing_results:run"
+"data_collection_group" = "murfey.workflows.register_data_collection_group:run"
 "pato" = "murfey.workflows.notifications:notification_setup"
 "picked_particles" = "murfey.workflows.spa.picking:particles_picked"
 "spa.flush_spa_preprocess" = "murfey.workflows.spa.flush_spa_preprocess:flush_spa_preprocess"

src/murfey/server/feedback.py

@@ -21,7 +21,6 @@
 import numpy as np
 from backports.entry_points_selectable import entry_points
 from ispyb.sqlalchemy._auto_db_schema import (
-    Atlas,
     AutoProcProgram,
     Base,
     DataCollection,
@@ -2012,75 +2011,6 @@ def feedback_callback(header: dict, message: dict, _db=murfey_db) -> None:
             if murfey.server._transport_object:
                 murfey.server._transport_object.transport.ack(header)
             return None
-        elif message["register"] == "data_collection_group":
-            ispyb_session_id = get_session_id(
-                microscope=message["microscope"],
-                proposal_code=message["proposal_code"],
-                proposal_number=message["proposal_number"],
-                visit_number=message["visit_number"],
-                db=ISPyBSession(),
-            )
-            if dcg_murfey := _db.exec(
-                select(db.DataCollectionGroup)
-                .where(db.DataCollectionGroup.session_id == message["session_id"])
-                .where(db.DataCollectionGroup.tag == message.get("tag"))
-            ).all():
-                dcgid = dcg_murfey[0].id
-            else:
-                if ispyb_session_id is None:
-                    murfey_dcg = db.DataCollectionGroup(
-                        session_id=message["session_id"],
-                        tag=message.get("tag"),
-                    )
-                    dcgid = murfey_dcg.id
-                else:
-                    record = DataCollectionGroup(
-                        sessionId=ispyb_session_id,
-                        experimentType=message["experiment_type"],
-                        experimentTypeId=message["experiment_type_id"],
-                    )
-                    dcgid = _register(record, header)
-                    atlas_record = Atlas(
-                        dataCollectionGroupId=dcgid,
-                        atlasImage=message.get("atlas", ""),
-                        pixelSize=message.get("atlas_pixel_size", 0),
-                        cassetteSlot=message.get("sample"),
-                    )
-                    if murfey.server._transport_object:
-                        atlas_id = murfey.server._transport_object.do_insert_atlas(
-                            atlas_record
-                        )["return_value"]
-                    else:
-                        atlas_id = None
-                    murfey_dcg = db.DataCollectionGroup(
-                        id=dcgid,
-                        atlas_id=atlas_id,
-                        atlas=message.get("atlas", ""),
-                        atlas_pixel_size=message.get("atlas_pixel_size"),
-                        sample=message.get("sample"),
-                        session_id=message["session_id"],
-                        tag=message.get("tag"),
-                    )
-                _db.add(murfey_dcg)
-                _db.commit()
-                _db.close()
-            if murfey.server._transport_object:
-                if dcgid is None:
-                    time.sleep(2)
-                    murfey.server._transport_object.transport.nack(header, requeue=True)
-                    return None
-                murfey.server._transport_object.transport.ack(header)
-            if dcg_hooks := entry_points().select(
-                group="murfey.hooks", name="data_collection_group"
-            ):
-                try:
-                    for hook in dcg_hooks:
-                        hook.load()(dcgid, session_id=message["session_id"])
-                except Exception:
-                    logger.error(
-                        "Call to data collection group hook failed", exc_info=True
-                    )
-            return None
         elif message["register"] == "atlas_update":
             if murfey.server._transport_object:
                 murfey.server._transport_object.do_update_atlas(

src/murfey/workflows/register_data_collection_group.py

@@ -0,0 +1,89 @@
+import logging
+import time
+
+import ispyb.sqlalchemy._auto_db_schema as ISPyBDB
+from backports.entry_points_selectable import entry_points
+from sqlmodel import select
+from sqlmodel.orm.session import Session as SQLModelSession
+
+import murfey.util.db as MurfeyDB
+from murfey.server import _transport_object
+from murfey.server.ispyb import ISPyBSession, get_session_id
+
+logger = logging.getLogger("murfey.workflows.register_data_collection_group")
+
+
+def run(
+    message: dict, murfey_db: SQLModelSession, demo: bool = False
+) -> dict[str, bool]:
+    # Fail immediately if no transport wrapper is found
+    if _transport_object is None:
+        logger.error("Unable to find transport manager")
+        return {"success": False, "requeue": False}
+
+    ispyb_session_id = get_session_id(
+        microscope=message["microscope"],
+        proposal_code=message["proposal_code"],
+        proposal_number=message["proposal_number"],
+        visit_number=message["visit_number"],
+        db=ISPyBSession(),
+    )
+    if dcg_murfey := murfey_db.exec(
+        select(MurfeyDB.DataCollectionGroup)
+        .where(MurfeyDB.DataCollectionGroup.session_id == message["session_id"])
+        .where(MurfeyDB.DataCollectionGroup.tag == message.get("tag"))
+    ).all():
+        dcgid = dcg_murfey[0].id
+    else:
+        if ispyb_session_id is None:
+            murfey_dcg = MurfeyDB.DataCollectionGroup(
+                session_id=message["session_id"],
+                tag=message.get("tag"),
+            )
+            dcgid = murfey_dcg.id
+        else:
+            record = ISPyBDB.DataCollectionGroup(
+                sessionId=ispyb_session_id,
+                experimentType=message["experiment_type"],
+                experimentTypeId=message["experiment_type_id"],
+            )
+
+            dcgid = _transport_object.do_insert_data_collection_group(record).get(
+                "return_value", None
+            )
+            atlas_record = ISPyBDB.Atlas(
+                dataCollectionGroupId=dcgid,
+                atlasImage=message.get("atlas", ""),
+                pixelSize=message.get("atlas_pixel_size", 0),
+                cassetteSlot=message.get("sample"),
+            )
+            if _transport_object:
+                atlas_id = _transport_object.do_insert_atlas(atlas_record)[
+                    "return_value"
+                ]
+            else:
+                atlas_id = None
+            murfey_dcg = MurfeyDB.DataCollectionGroup(
+                id=dcgid,
+                atlas_id=atlas_id,
+                atlas=message.get("atlas", ""),
+                atlas_pixel_size=message.get("atlas_pixel_size"),
+                sample=message.get("sample"),
+                session_id=message["session_id"],
+                tag=message.get("tag"),
+            )
+        murfey_db.add(murfey_dcg)
+        murfey_db.commit()
+        murfey_db.close()
+    if dcgid is None:
+        time.sleep(2)
+        return {"success": False, "requeue": True}
+    if dcg_hooks := entry_points().select(
+        group="murfey.hooks", name="data_collection_group"
+    ):
+        try:
+            for hook in dcg_hooks:
+                hook.load()(dcgid, session_id=message["session_id"])
+        except Exception:
+            logger.error("Call to data collection group hook failed", exc_info=True)
+    return {"success": True}