Refactor beamtime relationships in models and related APIs

Updated relationships for beamtime in models to support many-to-many associations with pucks, samples, and dewars. Refactored API endpoints to accommodate these changes, ensuring accurate assignment and retrieval of data. Improved sample data generation logic and incremented the application version for the new updates.

backend/app/routers/dewar.py

@@ -36,6 +36,7 @@ from app.models import (
     LogisticsEvent,
     PuckEvent,
     SampleEvent,
+    Beamtime as BeamtimeModel,
 )
 from app.dependencies import get_db
 import qrcode
@@ -598,18 +599,37 @@ async def assign_beamtime_to_dewar(
     if not dewar:
         raise HTTPException(status_code=404, detail="Dewar not found")
 
-    dewar.beamtime_id = None if beamtime_id == 0 else beamtime_id
+    # Find the Beamtime instance, if not unassigning
+    beamtime = (
+        db.query(BeamtimeModel).filter(BeamtimeModel.id == beamtime_id).first()
+        if beamtime_id
+        else None
+    )
+
+    if beamtime_id == 0:
+        dewar.beamtimes = []
+    else:
+        dewar.beamtimes = [
+            beamtime
+        ]  # assign one; append if you want to support multiple
+
     db.commit()
     db.refresh(dewar)
     for puck in dewar.pucks:
-        puck.beamtime_id = None if beamtime_id == 0 else beamtime_id
+        if beamtime_id == 0:
+            puck.beamtimes = []
+        else:
+            puck.beamtimes = [beamtime]
         for sample in puck.samples:
             has_sample_event = (
                 db.query(SampleEvent).filter(SampleEvent.sample_id == sample.id).count()
                 > 0
             )
             if not has_sample_event:
-                sample.beamtime_id = None if beamtime_id == 0 else beamtime_id
+                if beamtime_id == 0:
+                    sample.beamtimes = []
+                else:
+                    sample.beamtimes = [beamtime]
 
     db.commit()
     return {"status": "success", "dewar_id": dewar.id, "beamtime_id": beamtime_id}
@@ -726,5 +746,12 @@ async def get_single_shipment(id: int, db: Session = Depends(get_db)):
     operation_id="get_dewars_by_beamtime",
 )
 async def get_dewars_by_beamtime(beamtime_id: int, db: Session = Depends(get_db)):
-    """List all dewars assigned to a given beamtime."""
-    return db.query(DewarModel).filter(DewarModel.beamtime_id == beamtime_id).all()
+    beamtime = (
+        db.query(BeamtimeModel)
+        .options(joinedload(BeamtimeModel.dewars))
+        .filter(BeamtimeModel.id == beamtime_id)
+        .first()
+    )
+    if not beamtime:
+        raise HTTPException(status_code=404, detail="Beamtime not found")
+    return beamtime.dewars

backend/app/routers/puck.py

@@ -1,6 +1,6 @@
 from datetime import datetime
 from fastapi import APIRouter, HTTPException, status, Depends
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, joinedload
 from sqlalchemy.sql import func
 from typing import List
 import uuid
@@ -21,6 +21,7 @@ from app.models import (
     LogisticsEvent as LogisticsEventModel,
     Dewar as DewarModel,
     SampleEvent,
+    Beamtime as BeamtimeModel,
 )
 from app.dependencies import get_db
 import logging
@@ -664,23 +665,38 @@ async def get_pucks_by_slot(slot_identifier: str, db: Session = Depends(get_db))
 @router.patch("/puck/{puck_id}/assign-beamtime", operation_id="assignPuckToBeamtime")
 async def assign_beamtime_to_puck(
     puck_id: int,
-    beamtime_id: int,  # If you want ?beamtime_id=123 in the query
+    beamtime_id: int,  # expects ?beamtime_id=123 in the query
     db: Session = Depends(get_db),
 ):
     puck = db.query(PuckModel).filter(PuckModel.id == puck_id).first()
     if not puck:
         raise HTTPException(status_code=404, detail="Puck not found")
 
-    puck.beamtime_id = None if beamtime_id == 0 else beamtime_id
+    beamtime = (
+        db.query(BeamtimeModel).filter(BeamtimeModel.id == beamtime_id).first()
+        if beamtime_id
+        else None
+    )
+
+    if beamtime_id == 0:
+        puck.beamtimes = []
+    else:
+        puck.beamtimes = [
+            beamtime
+        ]  # or use .append(beamtime) if you want to support multiple
+
     db.commit()
     db.refresh(puck)
-    # Update samples
+    # Update samples as well
     for sample in puck.samples:
         has_sample_event = (
             db.query(SampleEvent).filter(SampleEvent.sample_id == sample.id).count() > 0
         )
         if not has_sample_event:
-            sample.beamtime_id = None if beamtime_id == 0 else beamtime_id
+            if beamtime_id == 0:
+                sample.beamtimes = []
+            else:
+                sample.beamtimes = [beamtime]
     db.commit()
     return {"status": "success", "puck_id": puck.id, "beamtime_id": beamtime_id}
 
@@ -691,5 +707,12 @@ async def assign_beamtime_to_puck(
     operation_id="get_pucks_by_beamtime",
 )
 async def get_pucks_by_beamtime(beamtime_id: int, db: Session = Depends(get_db)):
-    """List all pucks assigned to a given beamtime."""
-    return db.query(PuckModel).filter(PuckModel.beamtime_id == beamtime_id).all()
+    beamtime = (
+        db.query(BeamtimeModel)
+        .options(joinedload(BeamtimeModel.pucks))  # eager load pucks
+        .filter(BeamtimeModel.id == beamtime_id)
+        .first()
+    )
+    if not beamtime:
+        raise HTTPException(status_code=404, detail="Beamtime not found")
+    return beamtime.pucks

backend/app/routers/sample.py

@@ -32,6 +32,7 @@ from app.models import (
     Results as ResultsModel,
     Jobs as JobModel,
     JobStatus,
+    Beamtime as BeamtimeModel,
 )
 from app.dependencies import get_db
 import logging
@@ -463,6 +464,7 @@ async def get_results_for_run_and_sample(
     formatted_results = [
         ResultResponse(
             id=result.id,
+            status=result.status,
             sample_id=result.sample_id,
             run_id=result.run_id,
             result=ProcessingResults(**result.result),
@@ -479,5 +481,12 @@ async def get_results_for_run_and_sample(
     operation_id="get_samples_by_beamtime",
 )
 async def get_samples_by_beamtime(beamtime_id: int, db: Session = Depends(get_db)):
-    """List all samples assigned to a given beamtime."""
-    return db.query(SampleModel).filter(SampleModel.beamtime_id == beamtime_id).all()
+    beamtime = (
+        db.query(BeamtimeModel)
+        .options(joinedload(BeamtimeModel.samples))
+        .filter(BeamtimeModel.id == beamtime_id)
+        .first()
+    )
+    if not beamtime:
+        raise HTTPException(status_code=404, detail="Beamtime not found")
+    return beamtime.samples