Refactor puck handling for responses

Revised and optimized the `get_pucks_with_tell_position` endpoint to streamline logic and enhance clarity.

backend/app/routers/puck.py

@@ -12,6 +12,7 @@ from app.schemas import (
     PuckWithTellPosition,
     Sample,
     SetTellPosition,
+    DataCollectionParameters,
 )
 from app.models import (
     Puck as PuckModel,
@@ -189,6 +190,86 @@ async def set_tell_positions(
     return results
 
 
+@router.get("/with-tell-position", response_model=List[PuckWithTellPosition])
+async def get_pucks_with_tell_position(db: Session = Depends(get_db)):
+    """
+    Retrieve all pucks with a `tell_position` set (not null),
+    their associated samples, and the latest `tell_position` value (if any).
+    """
+    # Step 1: Prepare a subquery to fetch the latest event timestamp for each
+    # puck with a non-null tell_position
+    latest_event_subquery = (
+        db.query(
+            PuckEventModel.puck_id,
+            func.max(PuckEventModel.timestamp).label("latest_timestamp"),
+        )
+        .filter(
+            PuckEventModel.tell_position.isnot(None)
+        )  # Filter non-null tell_positions
+        .group_by(PuckEventModel.puck_id)  # Group by puck
+        .subquery()
+    )
+
+    # Step 2: Query the pucks and their latest `tell_position` by joining the subquery
+    pucks_with_events = (
+        db.query(PuckModel, PuckEventModel, DewarModel)
+        .join(PuckEventModel, PuckModel.id == PuckEventModel.puck_id)
+        .join(
+            latest_event_subquery,
+            (PuckEventModel.puck_id == latest_event_subquery.c.puck_id)
+            & (PuckEventModel.timestamp == latest_event_subquery.c.latest_timestamp),
+        )
+        .outerjoin(
+            DewarModel, PuckModel.dewar_id == DewarModel.id
+        )  # Outer join with DewarModel
+        .all()
+    )
+
+    if not pucks_with_events:
+        return []
+
+    # Step 3: Construct the response with pucks and their latest tell_position
+    results = []
+    for puck, event, dewar in pucks_with_events:
+        # Fetch associated samples for this puck
+        samples = db.query(SampleModel).filter(SampleModel.puck_id == puck.id).all()
+
+        # Construct the response model
+        results.append(
+            PuckWithTellPosition(
+                id=int(puck.id),
+                puck_name=str(puck.puck_name),
+                puck_type=str(puck.puck_type),
+                puck_location_in_dewar=str(puck.puck_location_in_dewar)
+                if puck.puck_location_in_dewar
+                else None,
+                dewar_id=int(puck.dewar_id) if puck.dewar_id else None,
+                dewar_name=str(dewar.dewar_name)
+                if dewar and dewar.dewar_name
+                else None,
+                samples=[
+                    Sample(
+                        id=sample.id,
+                        sample_name=sample.sample_name,
+                        position=sample.position,
+                        puck_id=sample.puck_id,
+                        data_collection_parameters=(
+                            DataCollectionParameters(
+                                **sample.data_collection_parameters
+                            )
+                            if isinstance(sample.data_collection_parameters, dict)
+                            else sample.data_collection_parameters
+                        ),
+                    )
+                    for sample in samples
+                ],
+                tell_position=str(event.tell_position) if event else None,
+            )
+        )
+
+    return results
+
+
 @router.get("/{puck_id}", response_model=PuckSchema)
 async def get_puck(puck_id: str, db: Session = Depends(get_db)):
     puck = db.query(PuckModel).filter(PuckModel.id == puck_id).first()
@@ -403,60 +484,3 @@ async def get_pucks_by_slot(slot_identifier: str, db: Session = Depends(get_db))
 
     logger.info(f"Final response for slot {slot_identifier}: {results}")
     return results
-
-
-@router.get("/with-tell-position", response_model=List[PuckWithTellPosition])
-async def get_pucks_with_tell_position(db: Session = Depends(get_db)):
-    """
-    Retrieve all pucks with a `tell_position` set (not null),
-    their associated samples, and the latest `tell_position` value (if any).
-    """
-
-    pucks_with_events = (
-        db.query(PuckModel, PuckEventModel)
-        .join(PuckEventModel, PuckModel.id == PuckEventModel.puck_id)
-        .filter(
-            PuckEventModel.tell_position.isnot(None)
-        )  # Ensure only non-null tell_positions
-        .order_by(PuckEventModel.timestamp.desc())  # Get the most recent event
-        .distinct(PuckModel.id)  # Ensure one row per puck (latest event is prioritized)
-        .all()
-    )
-
-    if not pucks_with_events:
-        raise HTTPException(
-            status_code=404,
-            detail="No pucks with a `tell_position` found.",
-        )
-
-    # Construct the response with pucks and their latest tell_position
-    results = []
-    for puck, event in pucks_with_events:
-        # Retrieve associated samples for this puck
-        samples = db.query(SampleModel).filter(SampleModel.puck_id == puck.id).all()
-
-        # Construct the response model
-        results.append(
-            PuckWithTellPosition(
-                id=int(puck.id),  # Explicit casting
-                puck_name=str(puck.puck_name),
-                puck_type=str(puck.puck_type),
-                puck_location_in_dewar=str(puck.puck_location_in_dewar)
-                if puck.puck_location_in_dewar
-                else None,
-                dewar_id=int(puck.dewar_id) if puck.dewar_id else None,
-                samples=[
-                    Sample(
-                        id=sample.id,
-                        sample_name=sample.sample_name,
-                        position=sample.position,
-                    )
-                    for sample in samples
-                ],
-                tell_position=str(event.tell_position)
-                if event
-                else None,  # Include tell_position
-            )
-        )
-
-    return results