Fix create_session API endpoint and perform housekeeping on murfey.server.api.session_info

src/murfey/server/api/session_info.py

@@ -1,7 +1,7 @@
+import re
 from datetime import datetime
 from logging import getLogger
 from pathlib import Path
-from typing import Dict, List, Optional
 
 import requests
 from fastapi import APIRouter, Depends, Request
@@ -12,6 +12,7 @@
 import murfey.server
 import murfey.server.api.websocket as ws
 import murfey.server.prometheus as prom
+import murfey.util.db as MurfeyDB
 from murfey.server.api import templates
 from murfey.server.api.auth import (
     MurfeyInstrumentNameFrontend as MurfeyInstrumentName,
@@ -34,22 +35,6 @@
 from murfey.server.murfey_db import murfey_db
 from murfey.util import sanitise
 from murfey.util.config import get_machine_config
-from murfey.util.db import (
-    ClassificationFeedbackParameters,
-    ClientEnvironment,
-    DataCollection,
-    DataCollectionGroup,
-    FoilHole,
-    GridSquare,
-    Movie,
-    ProcessingJob,
-    RsyncInstance,
-    Session as MurfeySession,
-    SessionProcessingParameters,
-    SPARelionParameters,
-    Tilt,
-    TiltSeries,
-)
 from murfey.util.models import UpstreamFileRequestInfo, Visit
 
 logger = getLogger("murfey.server.api.session_info")
@@ -89,7 +74,7 @@ def check_smartem_availability(instrument_name: str):
     return {"available": bool(machine_config.smartem_api_url)}
 
 
-@router.get("/instruments/{instrument_name}/visits_raw", response_model=List[Visit])
+@router.get("/instruments/{instrument_name}/visits_raw", response_model=list[Visit])
 def get_current_visits(instrument_name: MurfeyInstrumentName, db=ispyb_db):
     logger.debug(
         f"Received request to look up ongoing visits for {sanitise(instrument_name)}"
@@ -130,36 +115,44 @@ def all_visit_info(
         )
 
 
-@router.get("/sessions/{session_id}/rsyncers", response_model=List[RsyncInstance])
+@router.get(
+    "/sessions/{session_id}/rsyncers", response_model=list[MurfeyDB.RsyncInstance]
+)
 def get_rsyncers_for_client(
     session_id: MurfeySessionID, db: SQLModelSession = murfey_db
 ):
     rsync_instances = db.exec(
-        select(RsyncInstance).where(RsyncInstance.session_id == session_id)
+        select(MurfeyDB.RsyncInstance).where(
+            MurfeyDB.RsyncInstance.session_id == session_id
+        )
     )
     return rsync_instances.all()
 
 
 class SessionClients(BaseModel):
-    session: MurfeySession
-    clients: List[ClientEnvironment]
+    session: MurfeyDB.Session
+    clients: list[MurfeyDB.ClientEnvironment]
 
 
 @router.get("/sessions/{session_id}")
 async def get_session(
     session_id: MurfeySessionID, db: SQLModelSession = murfey_db
 ) -> SessionClients:
-    session = db.exec(select(MurfeySession).where(MurfeySession.id == session_id)).one()
+    session = db.exec(
+        select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id)
+    ).one()
     clients = db.exec(
-        select(ClientEnvironment).where(ClientEnvironment.session_id == session_id)
+        select(MurfeyDB.ClientEnvironment).where(
+            MurfeyDB.ClientEnvironment.session_id == session_id
+        )
     ).all()
     return SessionClients(session=session, clients=clients)
 
 
 @router.get("/sessions")
 async def get_sessions(db: SQLModelSession = murfey_db):
-    sessions = db.exec(select(MurfeySession)).all()
-    clients = db.exec(select(ClientEnvironment)).all()
+    sessions = db.exec(select(MurfeyDB.Session)).all()
+    clients = db.exec(select(MurfeyDB.ClientEnvironment)).all()
     res = []
     for sess in sessions:
         r = {"session": sess, "clients": []}
@@ -170,32 +163,57 @@ async def get_sessions(db: SQLModelSession = murfey_db):
     return res
 
 
-class VisitEndTime(BaseModel):
-    end_time: Optional[datetime] = None
+class NewSessionInfo(BaseModel):
+    visit: str  # Alphanumeric hyphenated code
+    name: str  # Human-readable description
+    end_time: datetime | None = None
 
 
-@router.post("/instruments/{instrument_name}/visits/{visit}/sessions/{name}")
+@router.post("/instruments/{instrument_name}/sessions/new")
 def create_session(
     instrument_name: MurfeyInstrumentName,
-    visit: str,
-    name: str,
-    visit_end_time: VisitEndTime,
+    session_info: NewSessionInfo,
     db: SQLModelSession = murfey_db,
 ) -> int:
-    s = MurfeySession(
-        name=name,
-        visit=visit,
+    # For 'visit', keep only alphanumerics, hyphens, and underscores
+    visit_cleaned = re.sub(r"[^A-Za-z0-9_-]", "", sanitise(session_info.visit))
+
+    # For 'name' (human-readable description), keep:
+    # - alphanumerics
+    # - spaces
+    # - hyphens
+    # - underscores
+    # - forward slashes
+    # - round and square brackets
+    # - periods
+    # - commas
+    # - ampersands
+    # - pluses
+    # - single/double quotes
+    name_cleaned = re.sub(
+        r'[^\w\s\-/()\[\].,&+\'"]',
+        "",
+        sanitise(session_info.name),
+    )
+
+    # Add to database
+    session = MurfeyDB.Session(
+        name=name_cleaned,
+        visit=visit_cleaned,
         instrument_name=instrument_name,
-        visit_end_time=visit_end_time.end_time,
+        visit_end_time=session_info.end_time,
     )
-    db.add(s)
+    db.add(session)
     db.commit()
-    sid = s.id
+    session_id = session.id
 
-    if visit_end_time.end_time:
-        prom.alert_end_time.labels(visit=visit).set(visit_end_time.end_time.timestamp())
+    # Add a Prometheus label for the session if an end time is set
+    if session_info.end_time:
+        prom.alert_end_time.labels(visit=session_info.visit).set(
+            session_info.end_time.timestamp()
+        )
 
-    return sid
+    return session_id
 
 
 @router.post("/sessions/{session_id}")
@@ -205,7 +223,9 @@ def update_session(
     smartem_acquisition_uuid: str | None = None,
     db: SQLModelSession = murfey_db,
 ) -> None:
-    session = db.exec(select(MurfeySession).where(MurfeySession.id == session_id)).one()
+    session = db.exec(
+        select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id)
+    ).one()
     session.process = process
     session.smartem_acquisition_uuid = smartem_acquisition_uuid
     db.add(session)
@@ -223,48 +243,52 @@ def get_sessions_with_visit(
     instrument_name: MurfeyInstrumentName,
     visit_name: str,
     db: SQLModelSession = murfey_db,
-) -> List[MurfeySession]:
+) -> list[MurfeyDB.Session]:
     sessions = db.exec(
-        select(MurfeySession)
-        .where(MurfeySession.instrument_name == instrument_name)
-        .where(MurfeySession.visit == visit_name)
+        select(MurfeyDB.Session)
+        .where(MurfeyDB.Session.instrument_name == instrument_name)
+        .where(MurfeyDB.Session.visit == visit_name)
     ).all()
     return sessions
 
 
 @router.get("/instruments/{instrument_name}/sessions")
 async def get_sessions_by_instrument_name(
     instrument_name: MurfeyInstrumentName, db: SQLModelSession = murfey_db
-) -> List[MurfeySession]:
+) -> list[MurfeyDB.Session]:
     sessions = db.exec(
-        select(MurfeySession).where(MurfeySession.instrument_name == instrument_name)
+        select(MurfeyDB.Session).where(
+            MurfeyDB.Session.instrument_name == instrument_name
+        )
     ).all()
     return sessions
 
 
 @router.get("/sessions/{session_id}/data_collection_groups")
 def get_dc_groups(
     session_id: MurfeySessionID, db: SQLModelSession = murfey_db
-) -> Dict[str, DataCollectionGroup]:
+) -> dict[str, MurfeyDB.DataCollectionGroup]:
     data_collection_groups = db.exec(
-        select(DataCollectionGroup).where(DataCollectionGroup.session_id == session_id)
+        select(MurfeyDB.DataCollectionGroup).where(
+            MurfeyDB.DataCollectionGroup.session_id == session_id
+        )
     ).all()
     return {dcg.tag: dcg for dcg in data_collection_groups}
 
 
 @router.get("/sessions/{session_id}/data_collection_groups/{dcgid}/data_collections")
 def get_data_collections(
     session_id: MurfeySessionID, dcgid: int, db: SQLModelSession = murfey_db
-) -> List[DataCollection]:
+) -> list[MurfeyDB.DataCollection]:
     data_collections = db.exec(
-        select(DataCollection).where(DataCollection.dcg_id == dcgid)
+        select(MurfeyDB.DataCollection).where(MurfeyDB.DataCollection.dcg_id == dcgid)
     ).all()
     return data_collections
 
 
 @router.get("/clients")
 async def get_clients(db: SQLModelSession = murfey_db):
-    clients = db.exec(select(ClientEnvironment)).all()
+    clients = db.exec(select(MurfeyDB.ClientEnvironment)).all()
     return clients
 
 
@@ -278,13 +302,15 @@ def update_current_gain_ref(
     new_gain_ref: CurrentGainRef,
     db: SQLModelSession = murfey_db,
 ):
-    session = db.exec(select(MurfeySession).where(MurfeySession.id == session_id)).one()
+    session = db.exec(
+        select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id)
+    ).one()
     session.current_gain_ref = new_gain_ref.path
     db.add(session)
 
     session_processing_parameters = db.exec(
-        select(SessionProcessingParameters).where(
-            SessionProcessingParameters.session_id == session_id
+        select(MurfeyDB.SessionProcessingParameters).where(
+            MurfeyDB.SessionProcessingParameters.session_id == session_id
         )
     ).all()
     if session_processing_parameters:
@@ -390,30 +416,32 @@ def delete_silences(instrument_name: MurfeyInstrumentName):
 
 
 class ProcessingDetails(BaseModel):
-    data_collection_group: DataCollectionGroup
-    data_collections: List[DataCollection]
-    processing_jobs: List[ProcessingJob]
-    relion_params: SPARelionParameters
-    feedback_params: ClassificationFeedbackParameters
+    data_collection_group: MurfeyDB.DataCollectionGroup
+    data_collections: list[MurfeyDB.DataCollection]
+    processing_jobs: list[MurfeyDB.ProcessingJob]
+    relion_params: MurfeyDB.SPARelionParameters
+    feedback_params: MurfeyDB.ClassificationFeedbackParameters
 
 
 @spa_router.get("/sessions/{session_id}/spa_processing_parameters")
 def get_spa_proc_param_details(
     session_id: MurfeySessionID, db: SQLModelSession = murfey_db
-) -> Optional[List[ProcessingDetails]]:
+) -> list[ProcessingDetails] | None:
     params = db.exec(
         select(
-            DataCollectionGroup,
-            DataCollection,
-            ProcessingJob,
-            SPARelionParameters,
-            ClassificationFeedbackParameters,
+            MurfeyDB.DataCollectionGroup,
+            MurfeyDB.DataCollection,
+            MurfeyDB.ProcessingJob,
+            MurfeyDB.SPARelionParameters,
+            MurfeyDB.ClassificationFeedbackParameters,
+        )
+        .where(MurfeyDB.DataCollectionGroup.session_id == session_id)
+        .where(MurfeyDB.DataCollectionGroup.id == MurfeyDB.DataCollection.dcg_id)
+        .where(MurfeyDB.DataCollection.id == MurfeyDB.ProcessingJob.dc_id)
+        .where(MurfeyDB.SPARelionParameters.pj_id == MurfeyDB.ProcessingJob.id)
+        .where(
+            MurfeyDB.ClassificationFeedbackParameters.pj_id == MurfeyDB.ProcessingJob.id
         )
-        .where(DataCollectionGroup.session_id == session_id)
-        .where(DataCollectionGroup.id == DataCollection.dcg_id)
-        .where(DataCollection.id == ProcessingJob.dc_id)
-        .where(SPARelionParameters.pj_id == ProcessingJob.id)
-        .where(ClassificationFeedbackParameters.pj_id == ProcessingJob.id)
     ).all()
     if not params:
         return None
@@ -451,14 +479,19 @@ def get_number_of_movies_from_foil_hole(
     session_id: int, dcgid: int, gsid: int, fhid: int, db: SQLModelSession = murfey_db
 ) -> int:
     movies = db.exec(
-        select(Movie, FoilHole, GridSquare, DataCollectionGroup)
-        .where(Movie.foil_hole_id == FoilHole.id)
-        .where(FoilHole.name == fhid)
-        .where(FoilHole.grid_square_id == GridSquare.id)
-        .where(GridSquare.name == gsid)
-        .where(GridSquare.session_id == session_id)
-        .where(GridSquare.tag == DataCollectionGroup.tag)
-        .where(DataCollectionGroup.id == dcgid)
+        select(
+            MurfeyDB.Movie,
+            MurfeyDB.FoilHole,
+            MurfeyDB.GridSquare,
+            MurfeyDB.DataCollectionGroup,
+        )
+        .where(MurfeyDB.Movie.foil_hole_id == MurfeyDB.FoilHole.id)
+        .where(MurfeyDB.FoilHole.name == fhid)
+        .where(MurfeyDB.FoilHole.grid_square_id == MurfeyDB.GridSquare.id)
+        .where(MurfeyDB.GridSquare.name == gsid)
+        .where(MurfeyDB.GridSquare.session_id == session_id)
+        .where(MurfeyDB.GridSquare.tag == MurfeyDB.DataCollectionGroup.tag)
+        .where(MurfeyDB.DataCollectionGroup.id == dcgid)
     ).all()
     return len(movies)
 
@@ -471,7 +504,7 @@ def get_grid_squares(session_id: MurfeySessionID, db: SQLModelSession = murfey_d
 @spa_router.get("/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares")
 def get_grid_squares_from_dcg(
     session_id: MurfeySessionID, dcgid: int, db: SQLModelSession = murfey_db
-) -> List[GridSquare]:
+) -> list[MurfeyDB.GridSquare]:
     return _get_grid_squares_from_dcg(session_id, dcgid, db)
 
 
@@ -480,14 +513,14 @@ def get_grid_squares_from_dcg(
 )
 def get_foil_holes_from_grid_square(
     session_id: MurfeySessionID, dcgid: int, gsid: int, db: SQLModelSession = murfey_db
-) -> List[FoilHole]:
+) -> list[MurfeyDB.FoilHole]:
     return _get_foil_holes_from_grid_square(session_id, dcgid, gsid, db)
 
 
 @spa_router.get("/sessions/{session_id}/foil_hole/{fh_name}")
 def get_foil_hole(
     session_id: MurfeySessionID, fh_name: int, db: SQLModelSession = murfey_db
-) -> Dict[str, int]:
+) -> dict[str, int]:
     return _get_foil_hole(session_id, fh_name, db)
 
 
@@ -501,14 +534,14 @@ def get_foil_hole(
 @tomo_router.get("/sessions/{session_id}/tilt_series/{tilt_series_tag}/tilts")
 def get_tilts(
     session_id: MurfeySessionID, tilt_series_tag: str, db: SQLModelSession = murfey_db
-) -> Dict[str, List[str]]:
+) -> dict[str, list[str]]:
     res = db.exec(
-        select(TiltSeries, Tilt)
-        .where(TiltSeries.tag == tilt_series_tag)
-        .where(TiltSeries.session_id == session_id)
-        .where(Tilt.tilt_series_id == TiltSeries.id)
+        select(MurfeyDB.TiltSeries, MurfeyDB.Tilt)
+        .where(MurfeyDB.TiltSeries.tag == tilt_series_tag)
+        .where(MurfeyDB.TiltSeries.session_id == session_id)
+        .where(MurfeyDB.Tilt.tilt_series_id == MurfeyDB.TiltSeries.id)
     ).all()
-    tilts: Dict[str, List[str]] = {}
+    tilts: dict[str, list[str]] = {}
     for el in res:
         if tilts.get(el[1].rsync_source):
             tilts[el[1].rsync_source].append(el[2].movie_path)