Optimised database registration logic in CLEM workflow

src/murfey/workflows/clem/register_preprocessing_results.py

@@ -330,16 +330,11 @@ def _register_dcg_and_atlas(
         .where(MurfeyDB.DataCollectionGroup.tag == dcg_name)
     ).one()
 
-    if not (
-        clem_img_site := murfey_db.exec(
-            select(MurfeyDB.ImagingSite)
-            .where(MurfeyDB.ImagingSite.session_id == session_id)
-            .where(MurfeyDB.ImagingSite.site_name == result.site_name)
-        ).one_or_none()
-    ):
-        clem_img_site = MurfeyDB.ImagingSite(
-            session_id=session_id, site_name=result.site_name
-        )
+    clem_img_site = murfey_db.exec(
+        select(MurfeyDB.ImagingSite)
+        .where(MurfeyDB.ImagingSite.session_id == session_id)
+        .where(MurfeyDB.ImagingSite.site_name == result.site_name)
+    ).one()
 
     clem_img_site.dcg_id = dcg_entry.id
     clem_img_site.dcg_name = dcg_entry.tag
@@ -364,17 +359,20 @@ def _register_grid_square(
 
     # Check if an atlas has been registered
     if not (
-        atlas_entry := murfey_db.exec(
+        # Sort by ascending insertion order
+        atlas_results := murfey_db.exec(
             select(MurfeyDB.ImagingSite)
             .where(MurfeyDB.ImagingSite.session_id == session_id)
             .where(MurfeyDB.ImagingSite.dcg_name == dcg_name)
             .where(MurfeyDB.ImagingSite.data_type == "atlas")
-        ).one_or_none()
+            .order_by(MurfeyDB.ImagingSite.id)
+        ).all()
     ):
         logger.info(
             f"No atlas has been registered for data collection group {dcg_name!r} yet"
         )
         return
+    atlas_entry = atlas_results[-1]  # Use the latest registered atlas
 
     # Check if there are CLEM entries to register
     if clem_img_site_to_register := murfey_db.exec(
@@ -531,12 +529,13 @@ def _register_grid_square(
                     image=grid_square_params.image,
                 )
             murfey_db.add(grid_square_entry)
-            murfey_db.commit()
 
             # Add grid square ID to existing CLEM image series entry
             clem_img_site.grid_square_id = grid_square_entry.id
             murfey_db.add(clem_img_site)
-            murfey_db.commit()
+
+        # Do one commit at the end
+        murfey_db.commit()
     else:
         logger.info(
             f"No grid squares to register for data collection group {dcg_name!r} yet"

tests/workflows/clem/test_register_preprocessing_results.py

@@ -69,6 +69,20 @@ def generate_preprocessing_messages(
             for n in range(3)
         ]
     )
+    # Add a second atlas dataset
+    datasets.extend(
+        [
+            (
+                grid_dir / "Overview 2" / "Image 1",
+                False,
+                True,
+                (2400, 2400),
+                1e-6,
+                [0.002, 0.0044, 0.002, 0.0044],
+            )
+        ]
+    )
+    # Add on metadata for denoised datasets
     datasets.extend(
         [
             (
@@ -398,7 +412,7 @@ def test_run_with_db(
             MurfeyDB.GridSquare.session_id == murfey_session.id
         )
     ).all()
-    assert len(murfey_gs_search) == (len(preprocessing_messages) - 1) // 2
+    assert len(murfey_gs_search) == (len(preprocessing_messages) - 2) // 2
 
     # ISPyB's DataCollectionGroup should have an entry
     murfey_dcg = murfey_dcg_search[0]
@@ -449,7 +463,7 @@ def test_run_with_db(
         .scalars()
         .all()
     )
-    assert len(ispyb_gs_search) == (len(preprocessing_messages) - 1) // 2
+    assert len(ispyb_gs_search) == (len(preprocessing_messages) - 2) // 2
     for gs in ispyb_gs_search:
         # Check that all entries point to the denoised images ("_Lng_LVCC")
         assert gs.gridSquareImage is not None and "_Lng_LVCC" in gs.gridSquareImage