Refactor and extend puck handling with tell_position updates

Refactored puck updating logic to use Pydantic models for validation and normalization. Added a new endpoint to retrieve all pucks with a tell_position and their associated samples. Updated project version to 0.1.0a11 in pyproject.toml.

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, load_only
 from sqlalchemy.sql import func
 from typing import List
 import uuid
@@ -9,6 +9,9 @@ from app.schemas import (
     Puck as PuckSchema,
     PuckCreate,
     PuckUpdate,
+    PuckWithTellPosition,
+    Sample,
+    SetTellPosition,
 )
 from app.models import (
     Puck as PuckModel,
@@ -42,48 +45,61 @@ async def get_pucks(db: Session = Depends(get_db)):
 
 @router.put("/set-tell-positions", status_code=status.HTTP_200_OK)
 async def set_tell_positions(
-    pucks: List[dict],  # Accept a list of puck definitions from the client
+    pucks: List[SetTellPosition],  # Accept a list of SetTellPosition Pydantic models
     db: Session = Depends(get_db),
 ):
     """
-    Set the tell positions for multiple pucks.
+    Set the tell positions for multiple pucks based on the business rules.
 
     Args:
-        pucks (List[dict]): A list of puck definitions, where each puck contains:
+        pucks (List[SetTellPosition]): A list of puck definitions
-            - `puck_name` (str): The cleaned name of the puck.
+        with potential new positions:
-            - `segment` (str): The segment in the dewar (e.g., "A-F").
+            - `puck_name` (Optional[str]): The name of the puck to update.
-            - `puck_in_segment` (int): The position within the segment (1-5).
+            - `segment` (Optional[str]): The segment (A-F) in the dewar.
+            - `puck_in_segment` (Optional[int]): The position within the segment (1-5).
 
     Returns:
-        List[dict]: A list of responses, one for each successfully processed puck.
+        List[dict]: A list of results indicating the status of each puck update.
     """
     results = []
 
+    # Helper function: Validate and normalize puck names
+    def normalize_puck_name(name: str) -> str:
+        return str(name).strip().replace(" ", "_").upper()
+
+    if not pucks:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Payload cannot be empty. Provide at least one puck.",
+        )
+
     for puck_data in pucks:
         try:
-            # Extract data from input
+            # Step 1: Extract data from the Pydantic model
-            puck_name = puck_data.get("puckname")
+            puckname = puck_data.puckname
-            segment = puck_data.get("segment")
+            tell_position = (
-            puck_in_segment = puck_data.get("puck_in_segment")
+                puck_data.tell_position
+            )  # Combines `segment` + `puck_in_segment`
 
-            # 1. Validate `segment` and `puck_in_segment`
+            if not puckname:
-            if not segment or segment not in "ABCDEF":
+                # Step 3: If `puckname` is missing, clear tell positions for ALL pucks
-                raise ValueError(
+                db.query(PuckEventModel).filter(
-                    f"Invalid segment '{segment}'. Must be A, B, C, D, E, or F."
+                    PuckEventModel.tell_position.isnot(None)
-                )
+                ).update({"tell_position": None})
-            if not (1 <= puck_in_segment <= 5):
+                db.commit()
-                raise ValueError(
+                # Add the result for the clean-up case
-                    f"Invalid puck_in_segment "
+                results.append(
-                    f"'{puck_in_segment}'. Must be in range 1-5."
+                    {
+                        "action": "clear",
+                        "message": "Tell positions cleared for all pucks.",
+                    }
                 )
+                break  # No need to process further since all tell positions are cleared
 
-            # Generate tell_position
+            # Step 2: Normalize puck name for database lookups
-            tell_position = f"{segment}{puck_in_segment}"
+            normalized_name = normalize_puck_name(puckname)
 
-            # 2. Find the puck by its cleaned name
+            # Query the puck from the database
-            normalized_name = normalize_puck_name(puck_name)
-
-            # Use SQLAlchemy to match normalized names
             puck = (
                 db.query(PuckModel)
                 .filter(
@@ -94,76 +110,86 @@ async def set_tell_positions(
             )
 
             if not puck:
-                raise ValueError(
+                raise ValueError(f"Puck with name '{puckname}' not found.")
-                    f"Puck with cleaned name '{puck_name}' not found in the database."
-                )
 
-            # 3. Find the dewar associated with the puck
+            # Query the most recent tell_position for this puck
-            dewar = db.query(DewarModel).filter(DewarModel.id == puck.dewar_id).first()
+            last_tell_position_event = (
-            if not dewar:
+                db.query(PuckEventModel)
-                raise ValueError(f"Dewar associated with puck '{puck_name}' not found.")
-
-            # 4. Find the most recent logistics event for the dewar (type 'beamline')
-            logistics_event = (
-                db.query(LogisticsEventModel)
                 .filter(
-                    LogisticsEventModel.dewar_id == dewar.id,
+                    PuckEventModel.puck_id == puck.id,
-                    LogisticsEventModel.event_type == "beamline",
+                    PuckEventModel.tell_position.isnot(None),
                 )
-                .order_by(LogisticsEventModel.timestamp.desc())
+                .order_by(PuckEventModel.timestamp.desc())
                 .first()
             )
-            if not logistics_event:
-                raise ValueError(
-                    f"No recent 'beamline' logistics event found for dewar "
-                    f"'{dewar.dewar_name}'."
-                )
 
-            # 5. Retrieve the slot from the logistics event
+            if tell_position:
-            slot = (
+                # Step 4: Compare `tell_position` in the payload with the current one
-                db.query(SlotModel)
+                if (
-                .filter(SlotModel.id == logistics_event.slot_id)
+                    last_tell_position_event
-                .first()
+                    and last_tell_position_event.tell_position == tell_position
-            )
+                ):
-            if not slot:
+                    # Step 4.1: If positions match, do nothing
-                raise ValueError(
+                    results.append(
-                    f"No slot associated with the most recent 'beamline' "
+                        {
-                    f"logistics event "
+                            "puck_name": puck.puck_name,
-                    f"for dewar '{dewar.dewar_name}'."
+                            "current_position": tell_position,
+                            "status": "unchanged",
+                            "message": "The tell_position remains the same.",
+                        }
                     )
+                    continue  # Skip to the next puck
 
-            # 6. Create a PuckEvent to set the 'tell_position'
+                # Step 4.2: If the position is different, update it
-            new_puck_event = PuckEventModel(
+                if last_tell_position_event:
+                    # Clear the previous tell_position (set it to null)
+                    last_tell_position_event.tell_position = None
+                    db.add(last_tell_position_event)
+
+                # Add a new event with the updated tell_position
+                new_event = PuckEventModel(
                     puck_id=puck.id,
                     tell_position=tell_position,
                     event_type="tell_position_set",
                     timestamp=datetime.utcnow(),
                 )
-            db.add(new_puck_event)
+                db.add(new_event)
-            db.commit()
-            db.refresh(new_puck_event)
 
-            # Add success result to the results list
                 results.append(
                     {
-                    "puck_id": puck.id,
                         "puck_name": puck.puck_name,
-                    "dewar_id": dewar.id,
+                        "previous_position": last_tell_position_event.tell_position
-                    "dewar_name": dewar.dewar_name,
+                        if last_tell_position_event
-                    "slot_id": slot.id,
+                        else None,
-                    "slot_label": slot.label,
+                        "new_position": tell_position,
-                    "tell_position": tell_position,
+                        "status": "updated",
-                    "event_timestamp": new_puck_event.timestamp,
+                        "message": "The tell_position was updated successfully.",
                     }
                 )
+            else:
+                # Step 5: If the new tell_position is None, clear the current one
+                if last_tell_position_event:
+                    last_tell_position_event.tell_position = None
+                    db.add(last_tell_position_event)
+
+                results.append(
+                    {
+                        "puck_name": puck.puck_name,
+                        "previous_position": last_tell_position_event.tell_position
+                        if last_tell_position_event
+                        else None,
+                        "new_position": None,
+                        "status": "cleared",
+                        "message": "The tell_position was cleared.",
+                    }
+                )
+
+            # Commit transaction for updating tell_position
+            db.commit()
+
         except Exception as e:
             # Handle errors for individual pucks and continue processing others
-            results.append(
+            results.append({"puck_name": puckname, "error": str(e)})
-                {
-                    "puck_name": puck_name,
-                    "error": str(e),
-                }
-            )
 
     return results
 
@@ -356,3 +382,62 @@ async def get_pucks_by_slot(slot_identifier: str, db: Session = Depends(get_db))
 
     # Return the list of pucks with their associated dewar names
     return puck_output
+
+
+class SampleModel:
+    pass
+
+
+@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)
+    and their associated samples.
+    """
+    # Query pucks with events where `tell_position` is not null
+    pucks = (
+        db.query(PuckModel)
+        .join(PuckEventModel, PuckModel.id == PuckEventModel.puck_id)
+        .filter(PuckEventModel.tell_position.isnot(None))
+        .options(
+            load_only(
+                PuckModel.id,
+                PuckModel.puck_name,
+                PuckModel.puck_type,
+                PuckModel.puck_location_in_dewar,
+                PuckModel.dewar_id,
+            )
+        )
+        .all()
+    )
+
+    if not pucks:
+        raise HTTPException(
+            status_code=404,
+            detail="No pucks with a `tell_position` found.",
+        )
+
+    results = [
+        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),
+            samples=[
+                Sample(
+                    id=sample.id,
+                    sample_name=sample.sample_name,
+                    position=sample.position,
+                )
+                for sample in db.query(SampleModel)
+                .filter(SampleModel.puck_id == puck.id)
+                .all()
+            ],
+        )
+        for puck in pucks
+    ]
+
+    return results

backend/app/schemas.py

@@ -613,7 +613,7 @@ class SlotSchema(BaseModel):
 
 
 class SetTellPosition(BaseModel):
-    puckname: str  # The puck name is required.
+    puck_name: str
     segment: Optional[str] = Field(
         None,
         pattern="^[A-F]$",  # Valid segments are A, B, C, D, E, F
@@ -637,3 +637,15 @@ class SetTellPosition(BaseModel):
         if self.segment and self.puck_in_segment:
             return f"{self.segment}{self.puck_in_segment}"
         return None
+
+
+class PuckWithTellPosition(BaseModel):
+    id: int
+    puck_name: str
+    puck_type: str
+    puck_location_in_dewar: Optional[str]
+    dewar_id: int
+    samples: List[Sample]
+
+    class Config:
+        from_attributes = True

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "aareDB"
-version = "0.1.0a10"
+version = "0.1.0a11"
 description = "Backend for next gen sample management system"
 authors = [{name = "Guillaume Gotthard", email = "guillaume.gotthard@psi.ch"}]
 license = {text = "MIT"}