diff --git a/src/murfey/workflows/clem/register_preprocessing_results.py b/src/murfey/workflows/clem/register_preprocessing_results.py index 68dc9c16..e8e20bf5 100644 --- a/src/murfey/workflows/clem/register_preprocessing_results.py +++ b/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" diff --git a/tests/workflows/clem/test_register_preprocessing_results.py b/tests/workflows/clem/test_register_preprocessing_results.py index 33153231..752aea1b 100644 --- a/tests/workflows/clem/test_register_preprocessing_results.py +++ b/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