Moved common correlative workflow functions into 'murfey.server.api.shared' and updated function name

Author: tieneupin

src/murfey/server/api/session_control.py

@@ -9,7 +9,6 @@
 from pydantic import BaseModel
 from sqlalchemy import func
 from sqlmodel import select
-from werkzeug.utils import secure_filename
 
 import murfey.server.prometheus as prom
 from murfey.server import _transport_object
@@ -18,18 +17,20 @@
     validate_instrument_token,
 )
 from murfey.server.api.shared import (
+    find_upstream_visits as _find_upstream_visits,
+    gather_upstream_tiffs as _gather_upstream_tiffs,
     get_foil_hole as _get_foil_hole,
     get_foil_holes_from_grid_square as _get_foil_holes_from_grid_square,
     get_grid_squares as _get_grid_squares,
     get_grid_squares_from_dcg as _get_grid_squares_from_dcg,
     get_machine_config_for_instrument,
-    get_upstream_tiff_dirs,
+    get_tiff_file as _get_tiff_file,
     remove_session_by_id,
 )
 from murfey.server.ispyb import DB as ispyb_db, 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
+from murfey.util.config import MachineConfig
 from murfey.util.db import (
     AutoProcProgram,
     ClientEnvironment,
@@ -418,62 +419,21 @@ def register_batch_position(
 
 @correlative_router.get("/sessions/{session_id}/upstream_visits")
 async def find_upstream_visits(session_id: MurfeySessionID, db=murfey_db):
-    murfey_session = db.exec(select(Session).where(Session.id == session_id)).one()
-    visit_name = murfey_session.visit
-    instrument_name = murfey_session.instrument_name
-    machine_config = get_machine_config(instrument_name=instrument_name)[
-        instrument_name
-    ]
-    upstream_visits = {}
-    # Iterates through provided upstream directories
-    for p in machine_config.upstream_data_directories:
-        # Looks for visit name in file path
-        for v in Path(p).glob(f"{visit_name.split('-')[0]}-*"):
-            upstream_visits[v.name] = v / machine_config.processed_directory_name
-    return upstream_visits
+    return _find_upstream_visits(session_id=session_id, db=db)
 
 
 @correlative_router.get(
     "/visits/{visit_name}/sessions/{session_id}/upstream_tiff_paths"
 )
 async def gather_upstream_tiffs(visit_name: str, session_id: int, db=murfey_db):
-    """
-    Looks for TIFF files associated with the current session in the permitted storage
-    servers, and returns their relative file paths as a list.
-    """
-    instrument_name = (
-        db.exec(select(Session).where(Session.id == session_id)).one().instrument_name
-    )
-    upstream_tiff_paths = []
-    tiff_dirs = get_upstream_tiff_dirs(visit_name, instrument_name)
-    if not tiff_dirs:
-        return None
-    for tiff_dir in tiff_dirs:
-        for f in tiff_dir.glob("**/*.tiff"):
-            upstream_tiff_paths.append(str(f.relative_to(tiff_dir)))
-        for f in tiff_dir.glob("**/*.tif"):
-            upstream_tiff_paths.append(str(f.relative_to(tiff_dir)))
-    return upstream_tiff_paths
+    return _gather_upstream_tiffs(visit_name=visit_name, session_id=session_id, db=db)
 
 
 @correlative_router.get(
     "/visits/{visit_name}/sessions/{session_id}/upstream_tiff/{tiff_path:path}"
 )
-async def get_tiff(visit_name: str, session_id: int, tiff_path: str, db=murfey_db):
-    instrument_name = (
-        db.exec(select(Session).where(Session.id == session_id)).one().instrument_name
+async def get_tiff_file(visit_name: str, session_id: int, tiff_path: str, db=murfey_db):
+    tiff_file = _get_tiff_file(
+        visit_name=visit_name, session_id=session_id, tiff_path=tiff_path, db=db
     )
-    tiff_dirs = get_upstream_tiff_dirs(visit_name, instrument_name)
-    if not tiff_dirs:
-        return None
-
-    tiff_path = "/".join(secure_filename(p) for p in tiff_path.split("/"))
-    for tiff_dir in tiff_dirs:
-        test_path = tiff_dir / tiff_path
-        if test_path.is_file():
-            break
-    else:
-        logger.warning(f"TIFF {tiff_path} not found")
-        return None
-
-    return FileResponse(path=test_path)
+    return FileResponse(path=tiff_file) if isinstance(tiff_file, Path) else tiff_file

src/murfey/server/api/session_info.py

@@ -7,7 +7,6 @@
 from fastapi.responses import FileResponse
 from pydantic import BaseModel
 from sqlmodel import select
-from werkzeug.utils import secure_filename
 
 import murfey.server.api.websocket as ws
 from murfey.server import _transport_object
@@ -18,18 +17,20 @@
     validate_token,
 )
 from murfey.server.api.shared import (
+    find_upstream_visits as _find_upstream_visits,
+    gather_upstream_tiffs as _gather_upstream_tiffs,
     get_foil_hole as _get_foil_hole,
     get_foil_holes_from_grid_square as _get_foil_holes_from_grid_square,
     get_grid_squares as _get_grid_squares,
     get_grid_squares_from_dcg as _get_grid_squares_from_dcg,
     get_machine_config_for_instrument,
-    get_upstream_tiff_dirs,
+    get_tiff_file as _get_tiff_file,
     remove_session_by_id,
 )
 from murfey.server.ispyb import DB as ispyb_db, 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
+from murfey.util.config import MachineConfig
 from murfey.util.db import (
     ClientEnvironment,
     DataCollection,
@@ -412,62 +413,21 @@ def get_tilts(
 
 @correlative_router.get("/sessions/{session_id}/upstream_visits")
 async def find_upstream_visits(session_id: MurfeySessionID, db=murfey_db):
-    murfey_session = db.exec(select(Session).where(Session.id == session_id)).one()
-    visit_name = murfey_session.visit
-    instrument_name = murfey_session.instrument_name
-    machine_config = get_machine_config(instrument_name=instrument_name)[
-        instrument_name
-    ]
-    upstream_visits = {}
-    # Iterates through provided upstream directories
-    for p in machine_config.upstream_data_directories:
-        # Looks for visit name in file path
-        for v in Path(p).glob(f"{visit_name.split('-')[0]}-*"):
-            upstream_visits[v.name] = v / machine_config.processed_directory_name
-    return upstream_visits
+    return _find_upstream_visits(session_id=session_id, db=db)
 
 
 @correlative_router.get(
     "/visits/{visit_name}/sessions/{session_id}/upstream_tiff_paths"
 )
 async def gather_upstream_tiffs(visit_name: str, session_id: int, db=murfey_db):
-    """
-    Looks for TIFF files associated with the current session in the permitted storage
-    servers, and returns their relative file paths as a list.
-    """
-    instrument_name = (
-        db.exec(select(Session).where(Session.id == session_id)).one().instrument_name
-    )
-    upstream_tiff_paths = []
-    tiff_dirs = get_upstream_tiff_dirs(visit_name, instrument_name)
-    if not tiff_dirs:
-        return None
-    for tiff_dir in tiff_dirs:
-        for f in tiff_dir.glob("**/*.tiff"):
-            upstream_tiff_paths.append(str(f.relative_to(tiff_dir)))
-        for f in tiff_dir.glob("**/*.tif"):
-            upstream_tiff_paths.append(str(f.relative_to(tiff_dir)))
-    return upstream_tiff_paths
+    return _gather_upstream_tiffs(visit_name=visit_name, session_id=session_id, db=db)
 
 
 @correlative_router.get(
     "/visits/{visit_name}/sessions/{session_id}/upstream_tiff/{tiff_path:path}"
 )
-async def get_tiff(visit_name: str, session_id: int, tiff_path: str, db=murfey_db):
-    instrument_name = (
-        db.exec(select(Session).where(Session.id == session_id)).one().instrument_name
+async def get_tiff_file(visit_name: str, session_id: int, tiff_path: str, db=murfey_db):
+    tiff_file = _get_tiff_file(
+        visit_name=visit_name, session_id=session_id, tiff_path=tiff_path, db=db
     )
-    tiff_dirs = get_upstream_tiff_dirs(visit_name, instrument_name)
-    if not tiff_dirs:
-        return None
-
-    tiff_path = "/".join(secure_filename(p) for p in tiff_path.split("/"))
-    for tiff_dir in tiff_dirs:
-        test_path = tiff_dir / tiff_path
-        if test_path.is_file():
-            break
-    else:
-        logger.warning(f"TIFF {tiff_path} not found")
-        return None
-
-    return FileResponse(path=test_path)
+    return FileResponse(path=tiff_file) if isinstance(tiff_file, Path) else tiff_file

src/murfey/server/api/shared.py

@@ -4,6 +4,7 @@
 from typing import Dict, List, Optional
 
 from sqlmodel import select
+from sqlmodel.orm.session import Session as SQLModelSession
 from werkzeug.utils import secure_filename
 
 import murfey.server.prometheus as prom
@@ -16,7 +17,7 @@
     GridSquare,
     ProcessingJob,
     RsyncInstance,
-    Session,
+    Session as MurfeySession,
 )
 
 logger = logging.getLogger("murfey.server.api.shared")
@@ -32,9 +33,9 @@ def get_machine_config_for_instrument(instrument_name: str) -> Optional[MachineC
 
 
 def remove_session_by_id(session_id: int, db):
-    session = db.exec(select(Session).where(Session.id == session_id)).one()
+    session = db.exec(select(MurfeySession).where(MurfeySession.id == session_id)).one()
     sessions_for_visit = db.exec(
-        select(Session).where(Session.visit == session.visit)
+        select(MurfeySession).where(MurfeySession.visit == session.visit)
     ).all()
     # Don't remove prometheus metrics if there are other sessions using them
     if len(sessions_for_visit) == 1:
@@ -145,6 +146,24 @@ def get_foil_hole(session_id: int, fh_name: int, db) -> Dict[str, int]:
     return {f[1].tag: f[0].id for f in foil_holes}
 
 
+def find_upstream_visits(session_id: int, db: SQLModelSession):
+    murfey_session = db.exec(
+        select(MurfeySession).where(MurfeySession.id == session_id)
+    ).one()
+    visit_name = murfey_session.visit
+    instrument_name = murfey_session.instrument_name
+    machine_config = get_machine_config(instrument_name=instrument_name)[
+        instrument_name
+    ]
+    upstream_visits = {}
+    # Iterates through provided upstream directories
+    for p in machine_config.upstream_data_directories:
+        # Looks for visit name in file path
+        for v in Path(p).glob(f"{visit_name.split('-')[0]}-*"):
+            upstream_visits[v.name] = v / machine_config.processed_directory_name
+    return upstream_visits
+
+
 def get_upstream_tiff_dirs(visit_name: str, instrument_name: str) -> List[Path]:
     tiff_dirs = []
     machine_config = get_machine_config(instrument_name=instrument_name)[
@@ -161,3 +180,49 @@ def get_upstream_tiff_dirs(visit_name: str, instrument_name: str) -> List[Path]:
             f"No candidate directory found for upstream download from visit {sanitise(visit_name)}"
         )
     return tiff_dirs
+
+
+def gather_upstream_tiffs(visit_name: str, session_id: int, db: SQLModelSession):
+    """
+    Looks for TIFF files associated with the current session in the permitted storage
+    servers, and returns their relative file paths as a list.
+    """
+    instrument_name = (
+        db.exec(select(MurfeySession).where(MurfeySession.id == session_id))
+        .one()
+        .instrument_name
+    )
+    upstream_tiff_paths = []
+    tiff_dirs = get_upstream_tiff_dirs(visit_name, instrument_name)
+    if not tiff_dirs:
+        return None
+    for tiff_dir in tiff_dirs:
+        for f in tiff_dir.glob("**/*.tiff"):
+            upstream_tiff_paths.append(str(f.relative_to(tiff_dir)))
+        for f in tiff_dir.glob("**/*.tif"):
+            upstream_tiff_paths.append(str(f.relative_to(tiff_dir)))
+    return upstream_tiff_paths
+
+
+def get_tiff_file(
+    visit_name: str, session_id: int, tiff_path: str, db: SQLModelSession
+):
+    instrument_name = (
+        db.exec(select(MurfeySession).where(MurfeySession.id == session_id))
+        .one()
+        .instrument_name
+    )
+    tiff_dirs = get_upstream_tiff_dirs(visit_name, instrument_name)
+    if not tiff_dirs:
+        return None
+
+    tiff_path = "/".join(secure_filename(p) for p in tiff_path.split("/"))
+    for tiff_dir in tiff_dirs:
+        tiff_file = tiff_dir / tiff_path
+        if tiff_file.is_file():
+            break
+    else:
+        logger.warning(f"TIFF {tiff_path} not found")
+        return None
+
+    return tiff_file

src/murfey/util/route_manifest.yaml

@@ -763,7 +763,7 @@ murfey.server.api.session_control.correlative_router:
     methods:
       - GET
   - path: /session_control/correlative/visits/{visit_name}/sessions/{session_id}/upstream_tiff/{tiff_path:path}
-    function: get_tiff
+    function: get_tiff_file
     path_params:
       - name: visit_name
         type: str
@@ -986,7 +986,7 @@ murfey.server.api.session_info.correlative_router:
     methods:
       - GET
   - path: /session_info/correlative/visits/{visit_name}/sessions/{session_id}/upstream_tiff/{tiff_path:path}
-    function: get_tiff
+    function: get_tiff_file
     path_params:
       - name: visit_name
         type: str