Simplified imports of classes and functions in 'murfey.server.ispyb'

Author: tieneupin

src/murfey/server/ispyb.py

@@ -5,9 +5,6 @@
 from typing import Callable, Generator, List, Literal, Optional
 
 import ispyb
-
-# import ispyb.sqlalchemy
-import sqlalchemy.orm
 import workflows.transport
 from fastapi import Depends
 from ispyb.sqlalchemy import (
@@ -29,6 +26,8 @@
     ZcZocaloBuffer,
     url,
 )
+from sqlalchemy import create_engine
+from sqlalchemy.orm import Session, sessionmaker
 
 from murfey.util import sanitise
 from murfey.util.config import get_security_config
@@ -38,16 +37,16 @@
 security_config = get_security_config()
 
 try:
-    Session = sqlalchemy.orm.sessionmaker(
-        bind=sqlalchemy.create_engine(
+    ISPyBSession = sessionmaker(
+        bind=create_engine(
             url(credentials=security_config.ispyb_credentials),
             connect_args={"use_pure": True},
         )
     )
     log.info("Loaded ISPyB database session")
 except AttributeError:
     log.error("Error loading ISPyB session", exc_info=True)
-    Session = lambda: None
+    ISPyBSession = lambda: None
 
 
 def _send_using_new_connection(transport_type: str, queue: str, message: dict) -> None:
@@ -94,7 +93,7 @@ def do_insert_data_collection_group(
         **kwargs,
     ):
         try:
-            with Session() as db:
+            with ISPyBSession() as db:
                 db.add(record)
                 db.commit()
                 log.info(f"Created DataCollectionGroup {record.dataCollectionGroupId}")
@@ -109,7 +108,7 @@ def do_insert_data_collection_group(
 
     def do_insert_atlas(self, record: Atlas):
         try:
-            with Session() as db:
+            with ISPyBSession() as db:
                 db.add(record)
                 db.commit()
                 log.info(f"Created Atlas {record.atlasId}")
@@ -126,7 +125,7 @@ def do_update_atlas(
         self, atlas_id: int, atlas_image: str, pixel_size: float, slot: int
     ):
         try:
-            with Session() as db:
+            with ISPyBSession() as db:
                 atlas = db.query(Atlas).filter(Atlas.atlasId == atlas_id).one()
                 atlas.atlasImage = atlas_image or atlas.atlasImage
                 atlas.pixelSize = pixel_size or atlas.pixelSize
@@ -194,7 +193,7 @@ def do_insert_grid_square(
             pixelSize=grid_square_parameters.pixel_size,
         )
         try:
-            with Session() as db:
+            with ISPyBSession() as db:
                 db.add(record)
                 db.commit()
                 log.info(f"Created GridSquare {record.gridSquareId}")
@@ -211,7 +210,7 @@ def do_update_grid_square(
         self, grid_square_id: int, grid_square_parameters: GridSquareParameters
     ):
         try:
-            with Session() as db:
+            with ISPyBSession() as db:
                 grid_square = (
                     db.query(GridSquare)
                     .filter(GridSquare.gridSquareId == grid_square_id)
@@ -302,7 +301,7 @@ def do_insert_foil_hole(
             pixelSize=foil_hole_parameters.pixel_size,
         )
         try:
-            with Session() as db:
+            with ISPyBSession() as db:
                 db.add(record)
                 db.commit()
                 log.info(f"Created FoilHole {record.foilHoleId}")
@@ -322,7 +321,7 @@ def do_update_foil_hole(
         foil_hole_parameters: FoilHoleParameters,
     ):
         try:
-            with Session() as db:
+            with ISPyBSession() as db:
                 foil_hole = (
                     db.query(FoilHole).filter(FoilHole.foilHoleId == foil_hole_id).one()
                 )
@@ -380,7 +379,7 @@ def do_insert_data_collection(self, record: DataCollection, message=None, **kwar
             else "Created for Murfey"
         )
         try:
-            with Session() as db:
+            with ISPyBSession() as db:
                 record.comments = comment
                 db.add(record)
                 db.commit()
@@ -396,7 +395,7 @@ def do_insert_data_collection(self, record: DataCollection, message=None, **kwar
 
     def do_insert_sample_group(self, record: BLSampleGroup) -> dict:
         try:
-            with Session() as db:
+            with ISPyBSession() as db:
                 db.add(record)
                 db.commit()
                 log.info(f"Created BLSampleGroup {record.blSampleGroupId}")
@@ -411,7 +410,7 @@ def do_insert_sample_group(self, record: BLSampleGroup) -> dict:
 
     def do_insert_sample(self, record: BLSample, sample_group_id: int) -> dict:
         try:
-            with Session() as db:
+            with ISPyBSession() as db:
                 db.add(record)
                 db.commit()
                 log.info(f"Created BLSample {record.blSampleId}")
@@ -434,7 +433,7 @@ def do_insert_sample(self, record: BLSample, sample_group_id: int) -> dict:
 
     def do_insert_subsample(self, record: BLSubSample) -> dict:
         try:
-            with Session() as db:
+            with ISPyBSession() as db:
                 db.add(record)
                 db.commit()
                 log.info(f"Created BLSubSample {record.blSubSampleId}")
@@ -449,7 +448,7 @@ def do_insert_subsample(self, record: BLSubSample) -> dict:
 
     def do_insert_sample_image(self, record: BLSampleImage) -> dict:
         try:
-            with Session() as db:
+            with ISPyBSession() as db:
                 db.add(record)
                 db.commit()
                 log.info(f"Created BLSampleImage {record.blSampleImageId}")
@@ -532,7 +531,7 @@ def do_update_processing_status(self, record: AutoProcProgram, **kwargs):
             return {"success": False, "return_value": None}
 
     def do_buffer_lookup(self, app_id: int, uuid: int) -> Optional[int]:
-        with Session() as db:
+        with ISPyBSession() as db:
             buffer_objects = (
                 db.query(ZcZocaloBuffer)
                 .filter_by(AutoProcProgramID=app_id, UUID=uuid)
@@ -543,8 +542,8 @@ def do_buffer_lookup(self, app_id: int, uuid: int) -> Optional[int]:
         return reference
 
 
-def _get_session() -> Generator[Optional[sqlalchemy.orm.Session], None, None]:
-    db = Session()
+def _get_session() -> Generator[Optional[Session], None, None]:
+    db = ISPyBSession()
     if db is None:
         yield None
         return
@@ -563,7 +562,7 @@ def get_session_id(
     proposal_code: str,
     proposal_number: str,
     visit_number: str,
-    db: sqlalchemy.orm.Session | None,
+    db: Session | None,
 ) -> int | None:
 
     # Log received lookup parameters
@@ -596,9 +595,7 @@ def get_session_id(
     return res
 
 
-def get_proposal_id(
-    proposal_code: str, proposal_number: str, db: sqlalchemy.orm.Session
-) -> int:
+def get_proposal_id(proposal_code: str, proposal_number: str, db: Session) -> int:
     query = (
         db.query(Proposal)
         .filter(
@@ -610,7 +607,7 @@ def get_proposal_id(
     return query[0].proposalId
 
 
-def get_sub_samples_from_visit(visit: str, db: sqlalchemy.orm.Session) -> List[Sample]:
+def get_sub_samples_from_visit(visit: str, db: Session) -> List[Sample]:
     proposal_id = get_proposal_id(visit[:2], visit.split("-")[0][2:], db)
     samples = (
         db.query(BLSampleGroup, BLSampleGroupHasBLSample, BLSample, BLSubSample)
@@ -635,9 +632,7 @@ def get_sub_samples_from_visit(visit: str, db: sqlalchemy.orm.Session) -> List[S
     return res
 
 
-def get_all_ongoing_visits(
-    microscope: str, db: sqlalchemy.orm.Session | None
-) -> list[Visit]:
+def get_all_ongoing_visits(microscope: str, db: Session | None) -> list[Visit]:
     if db is None:
         print("No database found")
         return []
@@ -676,7 +671,7 @@ def get_all_ongoing_visits(
 
 def get_data_collection_group_ids(session_id):
     query = (
-        Session()
+        ISPyBSession()
         .query(DataCollectionGroup)
         .filter(
             DataCollectionGroup.sessionId == session_id,