Refactor set_tell_positions logic with updated rules.

Revised the set_tell_positions endpoint to handle updated business rules for puck positioning. Improved event handling to ensure proper nullification, updates, and removal of tell_positions based on the provided payload. Enhanced query performance and normalized puck name processing for consistency.
This commit is contained in:
GotthardG 2025-01-08 17:01:42 +01:00
parent fb5e930096
commit 0fdd5e54b7
2 changed files with 162 additions and 175 deletions

View File

@ -1,6 +1,6 @@
from datetime import datetime
from fastapi import APIRouter, HTTPException, status, Depends
from sqlalchemy.orm import Session, load_only
from sqlalchemy.orm import Session
from sqlalchemy.sql import func
from typing import List
import uuid
@ -16,6 +16,7 @@ from app.schemas import (
from app.models import (
Puck as PuckModel,
PuckEvent as PuckEventModel,
Sample as SampleModel,
Slot as SlotModel,
LogisticsEvent as LogisticsEventModel,
Dewar as DewarModel,
@ -45,31 +46,8 @@ 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[SetTellPosition], # Accept a validated Pydantic model
db: Session = Depends(get_db),
pucks: List[SetTellPosition], db: Session = Depends(get_db)
):
"""
Set tell positions for multiple pucks with updated rules.
Args:
pucks (List[SetTellPosition]): A list including puck_name,
segment, and puck_in_segment.
Rules:
1. If a puck already has a tell_position matching the payload,
it is ignored (no new timestamps).
2. If a puck is assigned a different position,
set the new position while nullifying the previous.
3. Pucks that have a last `tell_position_set` event with
a non-null `tell_position` but
are not in the payload will get
a `"puck_removed"` event nullifying their position.
4. If the last event for a puck is already null
or it has no events, nothing is added or updated.
Returns:
List[dict]: Status information for processed and ignored pucks.
"""
results = []
# Helper function to normalize puck names for database querying
@ -82,10 +60,8 @@ async def set_tell_positions(
detail="Payload cannot be empty. Provide at least one puck.",
)
# Normalize payload puck names
input_puck_names = {normalize_puck_name(p.puck_name): p for p in pucks}
# Retrieve all pucks in the database
# Retrieve all pucks in the database with their most recent
# `tell_position_set` event
all_pucks_with_last_event = (
db.query(PuckModel, PuckEventModel)
.outerjoin(PuckEventModel, PuckEventModel.puck_id == PuckModel.id)
@ -97,9 +73,12 @@ async def set_tell_positions(
# Dictionary mapping each puck's ID to its latest event
last_events = {}
for puck, last_event in all_pucks_with_last_event:
if puck.id not in last_events:
if puck.id not in last_events: # Only store the latest event for each puck
last_events[puck.id] = last_event
# Track processed puck IDs to avoid double-processing
processed_pucks = set()
# Process pucks provided in the payload
for puck_data in pucks:
try:
@ -108,8 +87,6 @@ async def set_tell_positions(
new_position = (
puck_data.tell_position
) # Combined from segment + puck_in_segment
# Normalize puck name
normalized_name = normalize_puck_name(puck_name)
# Find puck in the database
@ -125,10 +102,13 @@ async def set_tell_positions(
if not puck:
raise ValueError(f"Puck with name '{puck_name}' not found.")
# Mark this puck as processed
processed_pucks.add(puck.id)
# Query the last event for this puck
last_event = last_events.get(puck.id)
# Rule 1: Skip if the last event's tell_position matches the new position
# Rule 1: Skip if the last event's `tell_position` matches the new position
if last_event and last_event.tell_position == new_position:
results.append(
{
@ -140,18 +120,17 @@ async def set_tell_positions(
)
continue
# Rule 2: If the last tell_position is not None, nullify it with a
# puck_removed event
# Rule 2: Add a "puck_removed" event if the last tell_position is not None
if last_event and last_event.tell_position is not None:
remove_event = PuckEventModel(
puck_id=puck.id,
tell_position=None,
event_type="puck_removed", # Set event_type to "puck_removed"
event_type="puck_removed", # Event type set to "puck_removed"
timestamp=datetime.utcnow(),
)
db.add(remove_event)
# Add a new event with the updated tell_position
# Add a new "tell_position_set" event
if new_position:
new_event = PuckEventModel(
puck_id=puck.id,
@ -160,6 +139,7 @@ async def set_tell_positions(
timestamp=datetime.utcnow(),
)
db.add(new_event)
results.append(
{
"puck_name": puck.puck_name,
@ -172,50 +152,51 @@ async def set_tell_positions(
}
)
# Commit changes after processing each puck
db.commit()
except Exception as e:
# Handle individual puck errors
results.append({"puck_name": puck_data.puck_name, "error": str(e)})
# Process pucks not included in the payload but present in the DB
# Process pucks not included in the payload but present in the database
for puck_id, last_event in last_events.items():
# Skip pucks already processed in the previous loop
if puck_id in processed_pucks:
continue
puck = db.query(PuckModel).filter(PuckModel.id == puck_id).first()
normalized_name = normalize_puck_name(puck.puck_name)
if not puck:
continue
# Check if the puck is missing from the payload
if (
normalized_name not in input_puck_names
and last_event
and last_event.tell_position is not None
):
try:
# Add a puck_removed event
remove_event = PuckEventModel(
puck_id=puck.id,
tell_position=None,
event_type="puck_removed", # Set event_type to "puck_removed"
timestamp=datetime.utcnow(),
)
db.add(remove_event)
# Skip if the last event's tell_position is already null
if not last_event or last_event.tell_position is None:
continue
# Append to results
results.append(
{
"puck_name": puck.puck_name,
"removed_position": last_event.tell_position,
"status": "removed",
"message": "Puck is not in payload"
" and has been marked as removed from tell_position.",
}
)
try:
# Add a "puck_removed" event
remove_event = PuckEventModel(
puck_id=puck.id,
tell_position=None,
event_type="puck_removed", # Event type set to "puck_removed"
timestamp=datetime.utcnow(),
)
db.add(remove_event)
db.commit()
results.append(
{
"puck_name": puck.puck_name,
"removed_position": last_event.tell_position,
"status": "removed",
"message": "Puck is not in payload and"
" has been marked as removed from tell_position.",
}
)
except Exception as e:
# Handle errors for individual puck removal
results.append({"puck_name": puck.puck_name, "error": str(e)})
db.commit()
except Exception as e:
# Handle errors for individual puck removal
results.append({"puck_name": puck.puck_name, "error": str(e)})
return results
@ -299,14 +280,10 @@ async def get_last_tell_position(puck_id: str, db: Session = Depends(get_db)):
}
@router.get("/slot/{slot_identifier}", response_model=List[dict])
@router.get("/slot/{slot_identifier}", response_model=List[PuckWithTellPosition])
async def get_pucks_by_slot(slot_identifier: str, db: Session = Depends(get_db)):
"""
Retrieve all pucks associated with all dewars linked to the given slot
(by ID or keyword) via 'beamline' events.
- Accepts slot keywords like PXI, PXII, PXIII.
- Retrieves all dewars (and their names) associated with the slot.
Retrieve all pucks in a slot with their latest `tell_position`.
"""
# Map keywords to slot IDs
slot_aliases = {
@ -318,152 +295,162 @@ async def get_pucks_by_slot(slot_identifier: str, db: Session = Depends(get_db))
"X06DA": 49,
}
# Check if the slot identifier is an alias or ID
# Resolve slot ID or alias
try:
slot_id = int(slot_identifier) # If the user provided a numeric ID
alias = next(
(k for k, v in slot_aliases.items() if v == slot_id), slot_identifier
)
slot_id = int(slot_identifier)
except ValueError:
slot_id = slot_aliases.get(slot_identifier.upper()) # Try mapping alias
alias = slot_identifier.upper() # Keep alias as-is for error messages
slot_id = slot_aliases.get(slot_identifier.upper())
if not slot_id:
logger.error(f"Invalid slot alias: {slot_identifier}")
raise HTTPException(
status_code=400,
detail="Invalid slot identifier."
"Must be an ID or one of the following:"
"PXI, PXII, PXIII, X06SA, X10SA, X06DA.",
status_code=400, detail=f"Invalid slot identifier: {slot_identifier}"
)
# Verify that the slot exists
logger.info(f"Resolved slot identifier: {slot_identifier} to Slot ID: {slot_id}")
# Verify slot existence
slot = db.query(SlotModel).filter(SlotModel.id == slot_id).first()
if not slot:
logger.error(f"Slot not found: {slot_identifier}")
raise HTTPException(
status_code=404, detail=f"Slot not found for identifier '{alias}'."
status_code=404, detail=f"Slot not found for identifier {slot_identifier}"
)
logger.info(f"Slot found: ID={slot.id}, Label={slot.label}")
# Retrieve all beamline events associated with the slot
beamline_events = (
db.query(LogisticsEventModel)
# Fetch dewars in the slot
dewars = (
db.query(DewarModel)
.join(LogisticsEventModel, DewarModel.id == LogisticsEventModel.dewar_id)
.filter(
LogisticsEventModel.slot_id == slot_id,
LogisticsEventModel.event_type == "beamline",
)
.order_by(LogisticsEventModel.timestamp.desc())
.all()
)
if not beamline_events:
logger.warning(f"No dewars associated to this beamline '{alias}'.")
raise HTTPException(
status_code=404, detail=f"No dewars found for the given beamline '{alias}'."
)
logger.info(f"Found {len(beamline_events)} beamline events for slot_id={slot_id}.")
# Use the beamline events to find all associated dewars
dewar_ids = {event.dewar_id for event in beamline_events if event.dewar_id}
dewars = db.query(DewarModel).filter(DewarModel.id.in_(dewar_ids)).all()
if not dewars:
logger.warning(f"No dewars found for beamline '{alias}'.")
logger.warning(f"No dewars found for slot: {slot_identifier}")
raise HTTPException(
status_code=404, detail=f"No dewars found for beamline '{alias}'."
status_code=404, detail=f"No dewars found for slot {slot_identifier}"
)
logger.info(f"Found {len(dewars)} dewars for beamline '{alias}'.")
logger.info(
f"Found dewars for slot {slot_identifier}: {[dewar.id for dewar in dewars]}"
)
# Create a mapping of dewar_id to dewar_name
dewar_mapping = {dewar.id: dewar.dewar_name for dewar in dewars}
dewar_ids = [dewar.id for dewar in dewars]
# Retrieve all pucks associated with the dewars
puck_list = (
db.query(PuckModel)
.filter(PuckModel.dewar_id.in_([dewar.id for dewar in dewars]))
# Subquery to fetch the latest tell_position for each puck
subquery = (
db.query(
PuckEventModel.puck_id,
func.max(PuckEventModel.timestamp).label("latest_event_time"),
)
.filter(PuckEventModel.event_type == "tell_position_set")
.group_by(PuckEventModel.puck_id)
.subquery()
)
# Fetch pucks with their latest tell_position
pucks_with_positions = (
db.query(PuckModel, PuckEventModel.tell_position)
.outerjoin(subquery, subquery.c.puck_id == PuckModel.id)
.outerjoin(
PuckEventModel,
(PuckEventModel.puck_id == PuckModel.id)
& (PuckEventModel.timestamp == subquery.c.latest_event_time),
)
.filter(PuckModel.dewar_id.in_(dewar_ids))
.all()
)
if not puck_list:
logger.warning(f"No pucks found for dewars associated with beamline '{alias}'.")
# Log the results of the subquery and pucks fetched:
logger.debug(f"Results from subquery (tell_position): {pucks_with_positions}")
if not pucks_with_positions:
logger.warning(f"No pucks found for slot: {slot_identifier}")
raise HTTPException(
status_code=404,
detail=f"No pucks found for dewars associated with beamline '{alias}'.",
status_code=404, detail=f"No pucks found for slot '{slot_identifier}'"
)
logger.info(f"Found {len(puck_list)} pucks for beamline '{alias}'.")
# Prepare results:
results = []
for puck, tell_position in pucks_with_positions:
logger.debug(
f"Puck ID: {puck.id}, Name: {puck.puck_name},"
f" Tell Position: {tell_position}"
)
# Add the dewar_name to the output for each puck
puck_output = [
{
"id": puck.id,
"puck_name": puck.puck_name,
"puck_type": puck.puck_type,
"dewar_id": puck.dewar_id,
"dewar_name": dewar_mapping.get(puck.dewar_id), # Link dewar_name
}
for puck in puck_list
]
# Prepare the PuckWithTellPosition instance
results.append(
PuckWithTellPosition(
id=puck.id,
puck_name=puck.puck_name,
puck_type=puck.puck_type,
puck_location_in_dewar=str(puck.puck_location_in_dewar)
if puck.puck_location_in_dewar
else None,
dewar_id=puck.dewar_id,
tell_position=tell_position, # Latest tell_position from subquery
)
)
# Return the list of pucks with their associated dewar names
return puck_output
class SampleModel:
pass
logger.info(f"Final response prepared 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)
and their associated samples.
Retrieve all pucks with a `tell_position` set (not null),
their associated samples, and the latest `tell_position` value (if any).
"""
# Query pucks with events where `tell_position` is not null
pucks = (
db.query(PuckModel)
# Query pucks with their latest `tell_position_set` event where
# `tell_position` is not null
pucks_with_events = (
db.query(PuckModel, PuckEventModel)
.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,
)
)
.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:
if not pucks_with_events:
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()
],
# 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
)
)
for puck in pucks
]
return results

View File

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