Add beamtime relationships and enhance sample handling

This commit adds relationships to link Pucks and Samples to Beamtime in the models, enabling better data association. Includes changes to assign beamtime IDs during data generation and updates in API response models for improved data loading. Removed redundant code in testfunctions.ipynb to clean up the notebook.

backend/app/data/data.py

@@ -407,9 +407,9 @@ beamtimes = [
     ),
     Beamtime(
         id=2,
-        pgroups="p20002",
+        pgroups="p20003",
         shift="afternoon",
-        beamtime_name="p20001-test",
+        beamtime_name="p20003-test",
         beamline="X06DA",
         start_date=datetime.strptime("07.05.2025", "%d.%m.%Y").date(),
         end_date=datetime.strptime("08.05.2025", "%d.%m.%Y").date(),
@@ -677,8 +677,15 @@ pucks = [
 # Define samples
 samples = []
 sample_id_counter = 1
+# Assign a beamtime to each dewar
+dewar_to_beamtime = {
+    dewar.id: random.choice([1, 2]) for dewar in dewars  # Or use actual beamtime ids
+}
 
 for puck in pucks:
+    dewar_id = puck.dewar_id  # Assuming puck has dewar_id
+    assigned_beamtime = dewar_to_beamtime[dewar_id]
+
     positions_with_samples = random.randint(1, 16)
     occupied_positions = random.sample(range(1, 17), positions_with_samples)
 
@@ -689,6 +696,7 @@ for puck in pucks:
                 sample_name=f"Sample{sample_id_counter:03}",
                 position=pos,
                 puck_id=puck.id,
+                beamtime_id=assigned_beamtime,  # IMPORTANT: Use the dewar's beamtime
             )
             samples.append(sample)
             sample_id_counter += 1

backend/app/models.py

@@ -154,6 +154,10 @@ class Puck(Base):
     dewar = relationship("Dewar", back_populates="pucks")
     samples = relationship("Sample", back_populates="puck")
     events = relationship("PuckEvent", back_populates="puck")
+    beamtime_id = Column(Integer, ForeignKey("beamtimes.id"), nullable=True)
+    beamtime = relationship(
+        "Beamtime", back_populates="pucks", foreign_keys=[beamtime_id]
+    )
 
 
 class Sample(Base):
@@ -173,6 +177,8 @@ class Sample(Base):
     puck = relationship("Puck", back_populates="samples")
     events = relationship("SampleEvent", back_populates="sample", lazy="joined")
     images = relationship("Image", back_populates="sample", lazy="joined")
+    beamtime_id = Column(Integer, ForeignKey("beamtimes.id"), nullable=True)
+    beamtime = relationship("Beamtime", back_populates="samples")
 
     @property
     def mount_count(self) -> int:
@@ -256,6 +262,8 @@ class Beamtime(Base):
 
     local_contact = relationship("LocalContact")
     dewars = relationship("Dewar", back_populates="beamtime")
+    pucks = relationship("Puck", back_populates="beamtime")
+    samples = relationship("Sample", back_populates="beamtime")
 
 
 class Image(Base):

backend/app/routers/beamtime.py

@@ -1,9 +1,14 @@
 from fastapi import APIRouter, HTTPException, status, Depends
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, joinedload
 from sqlalchemy import or_
 
 from app.models import Beamtime as BeamtimeModel
-from app.schemas import Beamtime as BeamtimeSchema, BeamtimeCreate, loginData
+from app.schemas import (
+    Beamtime as BeamtimeSchema,
+    BeamtimeCreate,
+    loginData,
+    BeamtimeResponse,
+)
 from app.dependencies import get_db
 from app.routers.auth import get_current_user
 
@@ -60,7 +65,7 @@ async def create_beamtime(
 
 @beamtime_router.get(
     "/my-beamtimes",
-    response_model=list[BeamtimeSchema],
+    response_model=list[BeamtimeResponse],
 )
 async def get_my_beamtimes(
     db: Session = Depends(get_db),
@@ -68,5 +73,10 @@ async def get_my_beamtimes(
 ):
     user_pgroups = current_user.pgroups
     filters = [BeamtimeModel.pgroups.like(f"%{pgroup}%") for pgroup in user_pgroups]
-    beamtimes = db.query(BeamtimeModel).filter(or_(*filters)).all()
+    beamtimes = (
+        db.query(BeamtimeModel)
+        .options(joinedload(BeamtimeModel.local_contact))
+        .filter(or_(*filters))
+        .all()
+    )
     return beamtimes

backend/app/routers/sample.py

@@ -425,6 +425,7 @@ def create_result(payload: ResultCreate, db: Session = Depends(get_db)):
 
     result_entry = ResultsModel(
         sample_id=payload.sample_id,
+        status=payload.status,
         run_id=payload.run_id,
         result=payload.result.model_dump(),  # Serialize entire result to JSON
     )
@@ -435,6 +436,7 @@ def create_result(payload: ResultCreate, db: Session = Depends(get_db)):
 
     return ResultResponse(
         id=result_entry.id,
+        status=result_entry.status,
         sample_id=result_entry.sample_id,
         run_id=result_entry.run_id,
         result=payload.result,  # return original payload directly

backend/app/schemas.py

@@ -534,6 +534,7 @@ class PuckCreate(BaseModel):
     puck_type: str
     puck_location_in_dewar: int
     samples: List[SampleCreate] = []
+    beamtime_id: Optional[int] = None
 
 
 class PuckUpdate(BaseModel):
@@ -541,6 +542,7 @@ class PuckUpdate(BaseModel):
     puck_type: Optional[str] = None
     puck_location_in_dewar: Optional[int] = None
     dewar_id: Optional[int] = None
+    beamtime_id: Optional[int] = None
 
 
 class Puck(BaseModel):
@@ -549,6 +551,7 @@ class Puck(BaseModel):
     puck_type: str
     puck_location_in_dewar: int
     dewar_id: int
+    beamtime_id: Optional[int] = None
     events: List[PuckEvent] = []
     samples: List[Sample] = []
 
@@ -800,6 +803,24 @@ class BeamtimeCreate(BaseModel):
     local_contact_id: Optional[int]
 
 
+class BeamtimeResponse(BaseModel):
+    id: int
+    pgroups: str
+    shift: str
+    beamtime_name: str
+    beamline: str
+    start_date: date
+    end_date: date
+    status: str
+    comments: Optional[str] = None
+    proposal_id: Optional[int]
+    local_contact_id: Optional[int]
+    local_contact: Optional[LocalContact]
+
+    class Config:
+        from_attributes = True
+
+
 class ImageCreate(BaseModel):
     pgroup: str
     sample_id: int