Migrated 'processing_job' feedback-callback block to its own workflows module

tieneupin · tieneupin · commit 4591e76d0a2c · 2025-10-28T14:46:24.000Z
diff --git a/pyproject.toml b/pyproject.toml
@@ -112,6 +112,7 @@ GitHub = "https://github.com/DiamondLightSource/python-murfey"
 "pato" = "murfey.workflows.notifications:notification_setup"
 "picked_particles" = "murfey.workflows.spa.picking:particles_picked"
 "picked_tomogram" = "murfey.workflows.tomo.picking:picked_tomogram"
+"processing_job" = "murfey.workflows.register_processing_job:run"
 "spa.flush_spa_preprocess" = "murfey.workflows.spa.flush_spa_preprocess:flush_spa_preprocess"
 
 [tool.setuptools]
diff --git a/src/murfey/server/feedback.py b/src/murfey/server/feedback.py
@@ -26,7 +26,6 @@
     DataCollection,
     DataCollectionGroup,
     ProcessingJob,
-    ProcessingJobParameter,
 )
 from sqlalchemy import func
 from sqlalchemy.exc import (
@@ -41,7 +40,6 @@
 import murfey.server
 import murfey.server.prometheus as prom
 import murfey.util.db as db
-from murfey.server.ispyb import ISPyBSession
 from murfey.server.murfey_db import url  # murfey_db
 from murfey.util import sanitise
 from murfey.util.config import (
@@ -1901,7 +1899,6 @@ def _save_bfactor(message: dict, _db, demo: bool = False):
 
 def feedback_callback(header: dict, message: dict, _db=murfey_db) -> None:
     try:
-        record = None
         if "environment" in message:
             params = message["recipe"][str(message["recipe-pointer"])].get(
                 "parameters", {}
@@ -2013,81 +2010,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"] == "processing_job":
-            murfey_session_id = message["session_id"]
-            logger.info("registering processing job")
-            dc = _db.exec(
-                select(db.DataCollection, db.DataCollectionGroup)
-                .where(db.DataCollection.dcg_id == db.DataCollectionGroup.id)
-                .where(db.DataCollectionGroup.session_id == murfey_session_id)
-                .where(db.DataCollectionGroup.tag == message["source"])
-                .where(db.DataCollection.tag == message["tag"])
-            ).all()
-            if dc:
-                _dcid = dc[0][0].id
-            else:
-                logger.warning(
-                    f"No data collection ID found for {sanitise(message['tag'])}"
-                )
-                if murfey.server._transport_object:
-                    murfey.server._transport_object.transport.nack(header, requeue=True)
-                return None
-            if pj_murfey := _db.exec(
-                select(db.ProcessingJob)
-                .where(db.ProcessingJob.recipe == message["recipe"])
-                .where(db.ProcessingJob.dc_id == _dcid)
-            ).all():
-                pid = pj_murfey[0].id
-            else:
-                if ISPyBSession() is None:
-                    murfey_pj = db.ProcessingJob(recipe=message["recipe"], dc_id=_dcid)
-                else:
-                    record = ProcessingJob(
-                        dataCollectionId=_dcid, recipe=message["recipe"]
-                    )
-                    run_parameters = message.get("parameters", {})
-                    assert isinstance(run_parameters, dict)
-                    if message.get("job_parameters"):
-                        job_parameters = [
-                            ProcessingJobParameter(parameterKey=k, parameterValue=v)
-                            for k, v in message["job_parameters"].items()
-                        ]
-                        pid = _register(ExtendedRecord(record, job_parameters), header)
-                    else:
-                        pid = _register(record, header)
-                    murfey_pj = db.ProcessingJob(
-                        id=pid, recipe=message["recipe"], dc_id=_dcid
-                    )
-                _db.add(murfey_pj)
-                _db.commit()
-                pid = murfey_pj.id
-                _db.close()
-            if pid is None and murfey.server._transport_object:
-                murfey.server._transport_object.transport.nack(header, requeue=True)
-                return None
-            prom.preprocessed_movies.labels(processing_job=pid)
-            if not _db.exec(
-                select(db.AutoProcProgram).where(db.AutoProcProgram.pj_id == pid)
-            ).all():
-                if ISPyBSession() is None:
-                    murfey_app = db.AutoProcProgram(pj_id=pid)
-                else:
-                    record = AutoProcProgram(
-                        processingJobId=pid, processingStartTime=datetime.now()
-                    )
-                    appid = _register(record, header)
-                    if appid is None and murfey.server._transport_object:
-                        murfey.server._transport_object.transport.nack(
-                            header, requeue=True
-                        )
-                        return None
-                    murfey_app = db.AutoProcProgram(id=appid, pj_id=pid)
-                _db.add(murfey_app)
-                _db.commit()
-                _db.close()
-            if murfey.server._transport_object:
-                murfey.server._transport_object.transport.ack(header)
-            return None
         elif message["register"] == "flush_tomography_preprocess":
             _flush_tomography_preprocessing(message, _db)
             if murfey.server._transport_object:
diff --git a/src/murfey/workflows/register_processing_job.py b/src/murfey/workflows/register_processing_job.py
@@ -0,0 +1,99 @@
+import logging
+from datetime import datetime
+
+import ispyb.sqlalchemy._auto_db_schema as ISPyBDB
+from sqlmodel import select
+from sqlmodel.orm.session import Session as SQLModelSession
+
+import murfey.server.prometheus as prom
+import murfey.util.db as MurfeyDB
+from murfey.server import _transport_object
+from murfey.server.ispyb import ISPyBSession
+from murfey.util import sanitise
+
+logger = logging.getLogger("murfey.workflows.register_processing_job")
+
+
+def run(message: dict, murfey_db: SQLModelSession, demo: bool = False):
+    # Faill immediately if not transport manager is set
+    if _transport_object is None:
+        logger.error("Unable to find transport manager")
+        return {"success": False, "requeue": False}
+
+    logger.info(f"Registering the following processing job: \n{message}")
+
+    murfey_session_id = message["session_id"]
+    dc = murfey_db.exec(
+        select(MurfeyDB.DataCollection, MurfeyDB.DataCollectionGroup)
+        .where(MurfeyDB.DataCollection.dcg_id == MurfeyDB.DataCollectionGroup.id)
+        .where(MurfeyDB.DataCollectionGroup.session_id == murfey_session_id)
+        .where(MurfeyDB.DataCollectionGroup.tag == message["source"])
+        .where(MurfeyDB.DataCollection.tag == message["tag"])
+    ).all()
+
+    if dc:
+        _dcid = dc[0][0].id
+    else:
+        logger.warning(f"No data collection ID found for {sanitise(message['tag'])}")
+        return {"success": False, "requeue": True}
+    if pj_murfey := murfey_db.exec(
+        select(MurfeyDB.ProcessingJob)
+        .where(MurfeyDB.ProcessingJob.recipe == message["recipe"])
+        .where(MurfeyDB.ProcessingJob.dc_id == _dcid)
+    ).all():
+        pid = pj_murfey[0].id
+    else:
+        if ISPyBSession() is None:
+            murfey_pj = MurfeyDB.ProcessingJob(recipe=message["recipe"], dc_id=_dcid)
+        else:
+            record = ISPyBDB.ProcessingJob(
+                dataCollectionId=_dcid, recipe=message["recipe"]
+            )
+            run_parameters = message.get("parameters", {})
+            assert isinstance(run_parameters, dict)
+            if message.get("job_parameters"):
+                job_parameters = [
+                    ISPyBDB.ProcessingJobParameter(parameterKey=k, parameterValue=v)
+                    for k, v in message["job_parameters"].items()
+                ]
+                pid = _transport_object.do_create_ispyb_job(
+                    record, params=job_parameters
+                ).get("return_value", None)
+            else:
+                pid = _transport_object.do_create_ispyb_job(record).get(
+                    "return_value", None
+                )
+            murfey_pj = MurfeyDB.ProcessingJob(
+                id=pid, recipe=message["recipe"], dc_id=_dcid
+            )
+        murfey_db.add(murfey_pj)
+        murfey_db.commit()
+        pid = murfey_pj.id
+        murfey_db.close()
+
+    if pid is None:
+        return {"success": False, "requeue": True}
+
+    # Update Prometheus counter for preprocessed movies
+    prom.preprocessed_movies.labels(processing_job=pid)
+
+    # Register AutoProcProgram database entry if it doesn't already exist
+    if not murfey_db.exec(
+        select(MurfeyDB.AutoProcProgram).where(MurfeyDB.AutoProcProgram.pj_id == pid)
+    ).all():
+        if ISPyBSession() is None:
+            murfey_app = MurfeyDB.AutoProcProgram(pj_id=pid)
+        else:
+            record = ISPyBDB.AutoProcProgram(
+                processingJobId=pid, processingStartTime=datetime.now()
+            )
+            appid = _transport_object.do_update_processing_status(record).get(
+                "return_value", None
+            )
+            if appid is None:
+                return {"success": False, "requeue": True}
+            murfey_app = MurfeyDB.AutoProcProgram(id=appid, pj_id=pid)
+        murfey_db.add(murfey_app)
+        murfey_db.commit()
+        murfey_db.close()
+    return {"success": True}
