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.
This commit is contained in:
GotthardG 2025-01-08 14:11:34 +01:00
parent 6d67d02259
commit 4f73f41717
3 changed files with 185 additions and 88 deletions

View File

@ -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:
- `puck_name` (str): The cleaned name of the puck.
- `segment` (str): The segment in the dewar (e.g., "A-F").
- `puck_in_segment` (int): The position within the segment (1-5).
pucks (List[SetTellPosition]): A list of puck definitions
with potential new positions:
- `puck_name` (Optional[str]): The name of the puck to update.
- `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
puck_name = puck_data.get("puckname")
segment = puck_data.get("segment")
puck_in_segment = puck_data.get("puck_in_segment")
# Step 1: Extract data from the Pydantic model
puckname = puck_data.puckname
tell_position = (
puck_data.tell_position
) # Combines `segment` + `puck_in_segment`
# 1. Validate `segment` and `puck_in_segment`
if not segment or segment not in "ABCDEF":
raise ValueError(
f"Invalid segment '{segment}'. Must be A, B, C, D, E, or F."
)
if not (1 <= puck_in_segment <= 5):
raise ValueError(
f"Invalid puck_in_segment "
f"'{puck_in_segment}'. Must be in range 1-5."
if not puckname:
# Step 3: If `puckname` is missing, clear tell positions for ALL pucks
db.query(PuckEventModel).filter(
PuckEventModel.tell_position.isnot(None)
).update({"tell_position": None})
db.commit()
# Add the result for the clean-up case
results.append(
{
"action": "clear",
"message": "Tell positions cleared for all pucks.",
}
)
break # No need to process further since all tell positions are cleared
# Generate tell_position
tell_position = f"{segment}{puck_in_segment}"
# Step 2: Normalize puck name for database lookups
normalized_name = normalize_puck_name(puckname)
# 2. Find the puck by its cleaned name
normalized_name = normalize_puck_name(puck_name)
# Use SQLAlchemy to match normalized names
# Query the puck from the database
puck = (
db.query(PuckModel)
.filter(
@ -94,76 +110,86 @@ async def set_tell_positions(
)
if not puck:
raise ValueError(
f"Puck with cleaned name '{puck_name}' not found in the database."
)
raise ValueError(f"Puck with name '{puckname}' not found.")
# 3. Find the dewar associated with the puck
dewar = db.query(DewarModel).filter(DewarModel.id == puck.dewar_id).first()
if not dewar:
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)
# Query the most recent tell_position for this puck
last_tell_position_event = (
db.query(PuckEventModel)
.filter(
LogisticsEventModel.dewar_id == dewar.id,
LogisticsEventModel.event_type == "beamline",
PuckEventModel.puck_id == puck.id,
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
slot = (
db.query(SlotModel)
.filter(SlotModel.id == logistics_event.slot_id)
.first()
)
if not slot:
raise ValueError(
f"No slot associated with the most recent 'beamline' "
f"logistics event "
f"for dewar '{dewar.dewar_name}'."
if tell_position:
# Step 4: Compare `tell_position` in the payload with the current one
if (
last_tell_position_event
and last_tell_position_event.tell_position == tell_position
):
# Step 4.1: If positions match, do nothing
results.append(
{
"puck_name": puck.puck_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'
new_puck_event = PuckEventModel(
# Step 4.2: If the position is different, update it
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.commit()
db.refresh(new_puck_event)
db.add(new_event)
# Add success result to the results list
results.append(
{
"puck_id": puck.id,
"puck_name": puck.puck_name,
"dewar_id": dewar.id,
"dewar_name": dewar.dewar_name,
"slot_id": slot.id,
"slot_label": slot.label,
"tell_position": tell_position,
"event_timestamp": new_puck_event.timestamp,
"previous_position": last_tell_position_event.tell_position
if last_tell_position_event
else None,
"new_position": tell_position,
"status": "updated",
"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(
{
"puck_name": puck_name,
"error": str(e),
}
)
results.append({"puck_name": puckname, "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

View File

@ -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

View File

@ -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"}