Adjusted imports of 'murfey.server.ispyb.DB'; imported the object as 'ispyb_db' for greater readability

Author: tieneupin

src/murfey/server/api/session_control.py

@@ -10,13 +10,7 @@
 from sqlmodel import select
 from werkzeug.utils import secure_filename
 
-import murfey.server.ispyb
 import murfey.server.prometheus as prom
-
-try:
-    from murfey.server.ispyb import DB
-except ImportError:
-    DB = None
 from murfey.server import _transport_object
 from murfey.server.api.auth import MurfeySessionID, validate_token
 from murfey.server.api.shared import get_foil_hole as _get_foil_hole
@@ -32,6 +26,8 @@
     get_upstream_tiff_dirs,
     remove_session_by_id,
 )
+from murfey.server.ispyb import DB as ispyb_db
+from murfey.server.ispyb import get_all_ongoing_visits
 from murfey.server.murfey_db import murfey_db
 from murfey.util import sanitise
 from murfey.util.config import MachineConfig, get_machine_config
@@ -88,11 +84,11 @@ async def new_client_id(db=murfey_db):
 
 
 @router.get("/instruments/{instrument_name}/visits_raw", response_model=List[Visit])
-def get_current_visits(instrument_name: str, db=murfey.server.ispyb.DB):
+def get_current_visits(instrument_name: str, db=ispyb_db):
     logger.debug(
         f"Received request to look up ongoing visits for {sanitise(instrument_name)}"
     )
-    return murfey.server.ispyb.get_all_ongoing_visits(instrument_name, db)
+    return get_all_ongoing_visits(instrument_name, db)
 
 
 class SessionInfo(BaseModel):
@@ -142,7 +138,7 @@ def remove_session(session_id: MurfeySessionID, db=murfey_db):
 
 @router.post("/sessions/{session_id}/successful_processing")
 def register_processing_success_in_ispyb(
-    session_id: MurfeySessionID, db=DB, murfey_db=murfey_db
+    session_id: MurfeySessionID, db=ispyb_db, murfey_db=murfey_db
 ):
     collected_ids = murfey_db.exec(
         select(DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram)

src/murfey/server/api/session_info.py

@@ -10,8 +10,8 @@
 from sqlmodel import select
 from werkzeug.utils import secure_filename
 
+import murfey
 import murfey.server.api.websocket as ws
-import murfey.server.ispyb
 from murfey.server import _transport_object
 from murfey.server.api import templates
 from murfey.server.api.auth import MurfeySessionID, validate_token
@@ -28,6 +28,8 @@
     get_upstream_tiff_dirs,
     remove_session_by_id,
 )
+from murfey.server.ispyb import DB as ispyb_db
+from murfey.server.ispyb import get_all_ongoing_visits
 from murfey.server.murfey_db import murfey_db
 from murfey.util import sanitise
 from murfey.util.config import (
@@ -77,7 +79,7 @@ async def root(request: Request):
 
 
 @router.get("/health/")
-def health_check(db=murfey.server.ispyb.DB):
+def health_check(db=ispyb_db):
     conn = db.connection()
     conn.close()
     return {
@@ -97,16 +99,16 @@ def machine_info_by_instrument(instrument_name: str) -> Optional[MachineConfig]:
 
 
 @router.get("/instruments/{instrument_name}/visits_raw", response_model=List[Visit])
-def get_current_visits(instrument_name: str, db=murfey.server.ispyb.DB):
+def get_current_visits(instrument_name: str, db=ispyb_db):
     logger.debug(
         f"Received request to look up ongoing visits for {sanitise(instrument_name)}"
     )
-    return murfey.server.ispyb.get_all_ongoing_visits(instrument_name, db)
+    return get_all_ongoing_visits(instrument_name, db)
 
 
 @router.get("/instruments/{instrument_name}/visits/")
-def all_visit_info(instrument_name: str, request: Request, db=murfey.server.ispyb.DB):
-    visits = murfey.server.ispyb.get_all_ongoing_visits(instrument_name, db)
+def all_visit_info(instrument_name: str, request: Request, db=ispyb_db):
+    visits = get_all_ongoing_visits(instrument_name, db)
 
     if visits:
         return_query = [

src/murfey/server/api/workflow.py

@@ -36,7 +36,8 @@
     get_job_ids,
     get_tomo_proc_params,
 )
-from murfey.server.ispyb import DB, get_proposal_id
+from murfey.server.ispyb import DB as ispyb_db
+from murfey.server.ispyb import get_proposal_id
 from murfey.server.murfey_db import murfey_db
 from murfey.util import sanitise
 from murfey.util.config import get_machine_config
@@ -922,7 +923,7 @@ class Sample(BaseModel):
 
 
 @correlative_router.get("/visit/{visit_name}/samples")
-def get_samples(visit_name: str, db=DB) -> List[Sample]:
+def get_samples(visit_name: str, db=ispyb_db) -> List[Sample]:
     proposal_id = get_proposal_id(visit_name[:2], visit_name.split("-")[0][2:], db)
     samples = (
         db.query(BLSampleGroup, BLSampleGroupHasBLSample, BLSample, BLSubSample)
@@ -948,7 +949,7 @@ def get_samples(visit_name: str, db=DB) -> List[Sample]:
 
 
 @correlative_router.post("/visit/{visit_name}/sample_group")
-def register_sample_group(visit_name: str, db=DB) -> dict:
+def register_sample_group(visit_name: str, db=ispyb_db) -> dict:
     proposal_id = get_proposal_id(visit_name[:2], visit_name.split("-")[0][2:], db=db)
     record = BLSampleGroup(proposalId=proposal_id)
     if _transport_object:

src/murfey/server/ispyb.py

@@ -580,8 +580,8 @@ def _get_session() -> Generator[Optional[Session], None, None]:
         db.close()
 
 
-DB = Depends(_get_session)
 # Shortcut to access the database in a FastAPI endpoint
+DB = Depends(_get_session)
 
 
 def get_session_id(