Safer registration of data collections (#726)

Ensure that dcg, dc and pjid exist before inserting into murfey db. Add a sleep for the case where they cannot be registered to allow the database to settle.

Author: stephen-riggs

.github/workflows/test.yml

@@ -3,10 +3,11 @@ name: Build and test
 on: [push, pull_request]
 
 env:
-  ISPYB_DATABASE_SCHEMA: 4.8.0
+  ISPYB_DATABASE_SCHEMA: 4.11.0
   # Installs from GitHub
   # Versions: https://github.com/DiamondLightSource/ispyb-database/tags
   # Previous version(s):
+  # 4.8.0
   # 4.2.1 # released 2024-08-19
   # 4.1.0 # released 2024-03-26
 

src/murfey/server/feedback.py

@@ -2160,7 +2160,7 @@ def feedback_callback(header: dict, message: dict, _db=murfey_db) -> None:
                 murfey_db=_db,
             )
             if murfey.server._transport_object:
-                if result.get("success", False):
+                if result.get("success"):
                     murfey.server._transport_object.transport.ack(header)
                 else:
                     # Send it directly to DLQ without trying to rerun it

src/murfey/workflows/register_data_collection.py

@@ -1,4 +1,5 @@
 import logging
+import time
 
 import ispyb.sqlalchemy._auto_db_schema as ISPyBDB
 from sqlmodel import select
@@ -37,6 +38,7 @@ def run(message: dict, murfey_db: SQLModelSession) -> dict[str, bool]:
         dcgid = dcg[0].id
         # flush_data_collections(message["source"], murfey_db)
     else:
+        time.sleep(2)
         logger.warning(
             "No data collection group ID was found for image directory "
             f"{sanitise(message['image_directory'])} and source "
@@ -82,21 +84,20 @@ def run(message: dict, murfey_db: SQLModelSession) -> dict[str, bool]:
                     else ""
                 ),
             ).get("return_value", None)
+            if dcid is None:
+                time.sleep(2)
+                logger.error(
+                    "Failed to register the following data collection: \n"
+                    f"{message} \n"
+                    "Requeueing message"
+                )
+                return {"success": False, "requeue": True}
             murfey_dc = MurfeyDB.DataCollection(
                 id=dcid,
                 tag=message.get("tag"),
                 dcg_id=dcgid,
             )
         murfey_db.add(murfey_dc)
         murfey_db.commit()
-        dcid = murfey_dc.id
         murfey_db.close()
-
-    if dcid is None:
-        logger.error(
-            "Failed to register the following data collection: \n"
-            f"{message} \n"
-            "Requeueing message"
-        )
-        return {"success": False, "requeue": True}
     return {"success": True}

src/murfey/workflows/register_data_collection_group.py

@@ -52,6 +52,15 @@ def run(message: dict, murfey_db: SQLModelSession) -> dict[str, bool]:
                 "return_value", None
             )
 
+            if dcgid is None:
+                time.sleep(2)
+                logger.error(
+                    "Failed to register the following data collection group: \n"
+                    f"{message} \n"
+                    "Requeuing message"
+                )
+                return {"success": False, "requeue": True}
+
             atlas_record = ISPyBDB.Atlas(
                 dataCollectionGroupId=dcgid,
                 atlasImage=message.get("atlas", ""),
@@ -75,15 +84,6 @@ def run(message: dict, murfey_db: SQLModelSession) -> dict[str, bool]:
         murfey_db.commit()
         murfey_db.close()
 
-    if dcgid is None:
-        time.sleep(2)
-        logger.error(
-            "Failed to register the following data collection group: \n"
-            f"{message} \n"
-            "Requeuing message"
-        )
-        return {"success": False, "requeue": True}
-
     if dcg_hooks := entry_points(group="murfey.hooks", name="data_collection_group"):
         try:
             for hook in dcg_hooks:

src/murfey/workflows/register_processing_job.py

@@ -63,6 +63,8 @@ def run(message: dict, murfey_db: SQLModelSession):
                 pid = _transport_object.do_create_ispyb_job(record).get(
                     "return_value", None
                 )
+            if pid is None:
+                return {"success": False, "requeue": True}
             murfey_pj = MurfeyDB.ProcessingJob(
                 id=pid, recipe=message["recipe"], dc_id=_dcid
             )
@@ -71,9 +73,6 @@ def run(message: dict, murfey_db: SQLModelSession):
         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)
 

tests/workflows/test_register_data_collection.py

@@ -99,6 +99,6 @@ def test_run(
                     assert result == {"success": False, "requeue": True}
             else:
                 mock_transport_object.do_insert_data_collection.assert_not_called()
-                assert result == {"success": False, "requeue": True}
+                assert result == {"success": True}
         else:
             assert result == {"success": True}

tests/workflows/test_register_data_collection_group.py

@@ -77,10 +77,10 @@ def test_run(
     else:
         if ispyb_session_id is not None:
             mock_transport_object.do_insert_data_collection_group.assert_called_once()
-            mock_transport_object.do_insert_atlas.assert_called_once()
             if insert_dcg is not None:
+                mock_transport_object.do_insert_atlas.assert_called_once()
                 assert result == {"success": True}
             else:
                 assert result == {"success": False, "requeue": True}
         else:
-            assert result == {"success": False, "requeue": True}
+            assert result == {"success": True}

tests/workflows/test_processing_job.py …orkflows/test_register_processing_job.pytests/workflows/test_processing_job.py renamed to tests/workflows/test_register_processing_job.py tests/workflows/test_processing_job.py renamed to tests/workflows/test_register_processing_job.py

@@ -104,6 +104,6 @@ def test_run(
                 else:
                     assert result == {"success": False, "requeue": True}
             else:
-                assert result == {"success": False, "requeue": True}
+                assert result == {"success": True}
     else:
         assert result == {"success": False, "requeue": True}