diff --git a/robot-server/robot_server/labware_offsets/models.py b/robot-server/robot_server/labware_offsets/models.py index 33bae92e6b3..7b3f523fccd 100644 --- a/robot-server/robot_server/labware_offsets/models.py +++ b/robot-server/robot_server/labware_offsets/models.py @@ -2,7 +2,7 @@ from datetime import datetime import enum -from typing import Literal, Annotated, Final, TypeAlias +from typing import Literal, Annotated, Final, TypeAlias, Sequence from pydantic import BaseModel, Field @@ -35,18 +35,33 @@ class _DoNotFilter(enum.Enum): """ +class UnknownLabwareOffsetLocationSequenceComponent(BaseModel): + """A labware offset location sequence component from the future.""" + + kind: Literal["unknown"] = "unknown" + storedKind: str + primaryValue: str + + # This is redefined here so we can add stuff to it easily StoredLabwareOffsetLocationSequenceComponents = Annotated[ LabwareOffsetLocationSequenceComponentsUnion, Field(discriminator="kind") ] +ReturnedLabwareOffsetLocationSequenceComponents = Annotated[ + LabwareOffsetLocationSequenceComponentsUnion + | UnknownLabwareOffsetLocationSequenceComponent, + Field(discriminator="kind"), +] + + class StoredLabwareOffsetCreate(BaseModel): """Create an offset for storage.""" definitionUri: str = Field(..., description="The URI for the labware's definition.") - locationSequence: list[StoredLabwareOffsetLocationSequenceComponents] = Field( + locationSequence: Sequence[StoredLabwareOffsetLocationSequenceComponents] = Field( ..., description="Where the labware is located on the robot. Can represent all locations, but may not be present for older runs.", min_length=1, @@ -67,7 +82,7 @@ class StoredLabwareOffset(BaseModel): createdAt: datetime = Field(..., description="When this labware offset was added.") definitionUri: str = Field(..., description="The URI for the labware's definition.") - locationSequence: list[StoredLabwareOffsetLocationSequenceComponents] = Field( + locationSequence: Sequence[ReturnedLabwareOffsetLocationSequenceComponents] = Field( ..., description="Where the labware is located on the robot. Can represent all locations, but may not be present for older runs.", min_length=1, diff --git a/robot-server/robot_server/labware_offsets/router.py b/robot-server/robot_server/labware_offsets/router.py index 6c44cd554db..2a545093165 100644 --- a/robot-server/robot_server/labware_offsets/router.py +++ b/robot-server/robot_server/labware_offsets/router.py @@ -25,6 +25,7 @@ from .store import ( LabwareOffsetNotFoundError, LabwareOffsetStore, + IncomingStoredLabwareOffset, ) from .fastapi_dependencies import get_labware_offset_store from .models import ( @@ -58,7 +59,7 @@ async def post_labware_offset( # noqa: D103 new_offset_created_at: Annotated[datetime, fastapi.Depends(get_current_time)], request_body: Annotated[RequestModel[StoredLabwareOffsetCreate], fastapi.Body()], ) -> PydanticResponse[SimpleBody[StoredLabwareOffset]]: - new_offset = StoredLabwareOffset.model_construct( + new_offset = IncomingStoredLabwareOffset( id=new_offset_id, createdAt=new_offset_created_at, definitionUri=request_body.data.definitionUri, @@ -67,7 +68,15 @@ async def post_labware_offset( # noqa: D103 ) store.add(new_offset) return await PydanticResponse.create( - content=SimpleBody.model_construct(data=new_offset), + content=SimpleBody.model_construct( + data=StoredLabwareOffset( + id=new_offset_id, + createdAt=new_offset_created_at, + definitionUri=request_body.data.definitionUri, + locationSequence=request_body.data.locationSequence, + vector=request_body.data.vector, + ) + ), status_code=201, ) diff --git a/robot-server/robot_server/labware_offsets/store.py b/robot-server/robot_server/labware_offsets/store.py index 6863928e86a..e17c1a30686 100644 --- a/robot-server/robot_server/labware_offsets/store.py +++ b/robot-server/robot_server/labware_offsets/store.py @@ -1,6 +1,8 @@ # noqa: D100 -from typing import Iterator +from datetime import datetime +from dataclasses import dataclass +from typing import Iterator, Sequence from typing_extensions import assert_never from opentrons.protocol_engine.types import ( @@ -9,21 +11,41 @@ OnAddressableAreaOffsetLocationSequenceComponent, OnModuleOffsetLocationSequenceComponent, OnLabwareOffsetLocationSequenceComponent, - LabwareOffsetLocationSequenceComponents, - LabwareOffsetLocationSequence, ) from robot_server.persistence.tables import ( labware_offset_table, labware_offset_location_sequence_components_table, ) -from .models import StoredLabwareOffset, DoNotFilterType, DO_NOT_FILTER +from .models import ( + StoredLabwareOffset, + DoNotFilterType, + DO_NOT_FILTER, + StoredLabwareOffsetLocationSequenceComponents, + ReturnedLabwareOffsetLocationSequenceComponents, + UnknownLabwareOffsetLocationSequenceComponent, +) import sqlalchemy import sqlalchemy.exc from ._search_query_builder import SearchQueryBuilder +ReturnedLabwareOffsetLocationSequence = Sequence[ + ReturnedLabwareOffsetLocationSequenceComponents +] + + +@dataclass +class IncomingStoredLabwareOffset: + """Internal class for representing valid incoming offsets.""" + + id: str + createdAt: datetime + definitionUri: str + locationSequence: Sequence[StoredLabwareOffsetLocationSequenceComponents] + vector: LabwareOffsetVector + class LabwareOffsetStore: """A persistent store for labware offsets, to support the `/labwareOffsets` endpoints.""" @@ -37,7 +59,10 @@ def __init__(self, sql_engine: sqlalchemy.engine.Engine) -> None: """ self._sql_engine = sql_engine - def add(self, offset: StoredLabwareOffset) -> None: + def add( + self, + offset: IncomingStoredLabwareOffset, + ) -> None: """Store a new labware offset.""" with self._sql_engine.begin() as transaction: offset_row_id = transaction.execute( @@ -131,7 +156,7 @@ def __init__(self, bad_offset_id: str) -> None: def _sql_sequence_component_to_pydantic_sequence_component( component_row: sqlalchemy.engine.Row, -) -> LabwareOffsetLocationSequenceComponents: +) -> ReturnedLabwareOffsetLocationSequenceComponents: if component_row.component_kind == "onLabware": return OnLabwareOffsetLocationSequenceComponent( labwareUri=component_row.primary_component_value @@ -145,14 +170,19 @@ def _sql_sequence_component_to_pydantic_sequence_component( addressableAreaName=component_row.primary_component_value ) else: - raise KeyError(component_row.component_kind) + return UnknownLabwareOffsetLocationSequenceComponent( + storedKind=component_row.component_kind, + primaryValue=component_row.primary_component_value, + ) def _collate_sql_locations( first_row: sqlalchemy.engine.Row, rest_rows: Iterator[sqlalchemy.engine.Row] -) -> tuple[LabwareOffsetLocationSequence, sqlalchemy.engine.Row | None]: +) -> tuple[ + list[ReturnedLabwareOffsetLocationSequenceComponents], sqlalchemy.engine.Row | None +]: offset_id = first_row.offset_id - location_sequence: list[LabwareOffsetLocationSequenceComponents] = [ + location_sequence: list[ReturnedLabwareOffsetLocationSequenceComponents] = [ _sql_sequence_component_to_pydantic_sequence_component(first_row) ] while True: @@ -197,7 +227,9 @@ def _collate_sql_to_pydantic( yield result -def _pydantic_to_sql_offset(labware_offset: StoredLabwareOffset) -> dict[str, object]: +def _pydantic_to_sql_offset( + labware_offset: IncomingStoredLabwareOffset, +) -> dict[str, object]: return dict( offset_id=labware_offset.id, definition_uri=labware_offset.definitionUri, @@ -210,7 +242,7 @@ def _pydantic_to_sql_offset(labware_offset: StoredLabwareOffset) -> dict[str, ob def _pydantic_to_sql_location_sequence_iterator( - labware_offset: StoredLabwareOffset, offset_row_id: int + labware_offset: IncomingStoredLabwareOffset, offset_row_id: int ) -> Iterator[dict[str, object]]: for index, component in enumerate(labware_offset.locationSequence): if isinstance(component, OnLabwareOffsetLocationSequenceComponent): diff --git a/robot-server/tests/labware_offsets/test_store.py b/robot-server/tests/labware_offsets/test_store.py index 81bbd641ea1..cc2f132665b 100644 --- a/robot-server/tests/labware_offsets/test_store.py +++ b/robot-server/tests/labware_offsets/test_store.py @@ -16,6 +16,7 @@ from robot_server.labware_offsets.store import ( LabwareOffsetStore, LabwareOffsetNotFoundError, + IncomingStoredLabwareOffset, ) from robot_server.labware_offsets.models import ( StoredLabwareOffset, @@ -165,7 +166,7 @@ def test_filter_fields( ) -> None: """Test each filterable field to make sure it returns only matching entries.""" offsets = { - "a": StoredLabwareOffset( + "a": IncomingStoredLabwareOffset( id="a", createdAt=datetime.now(timezone.utc), definitionUri="definitionUri a", @@ -182,7 +183,7 @@ def test_filter_fields( ], vector=LabwareOffsetVector(x=1, y=2, z=3), ), - "b": StoredLabwareOffset( + "b": IncomingStoredLabwareOffset( id="b", createdAt=datetime.now(timezone.utc), definitionUri="definitionUri b", @@ -199,7 +200,7 @@ def test_filter_fields( ], vector=LabwareOffsetVector(x=2, y=4, z=6), ), - "c": StoredLabwareOffset( + "c": IncomingStoredLabwareOffset( id="c", createdAt=datetime.now(timezone.utc), definitionUri="definitionUri a", @@ -210,7 +211,7 @@ def test_filter_fields( ], vector=LabwareOffsetVector(x=3, y=6, z=9), ), - "d": StoredLabwareOffset( + "d": IncomingStoredLabwareOffset( id="d", createdAt=datetime.now(timezone.utc), definitionUri="definitionUri a", @@ -224,7 +225,7 @@ def test_filter_fields( ], vector=LabwareOffsetVector(x=4, y=8, z=12), ), - "e": StoredLabwareOffset( + "e": IncomingStoredLabwareOffset( id="e", createdAt=datetime.now(timezone.utc), definitionUri="definitionUri a", @@ -248,10 +249,19 @@ def test_filter_fields( location_module_model_filter=location_module_model_filter, location_definition_uri_filter=location_labware_uri_filter, ) - assert sorted( - results, + assert sorted(results, key=lambda o: o.id,) == sorted( + [ + StoredLabwareOffset( + id=offsets[id_].id, + createdAt=offsets[id_].createdAt, + definitionUri=offsets[id_].definitionUri, + locationSequence=offsets[id_].locationSequence, + vector=offsets[id_].vector, + ) + for id_ in returned_ids + ], key=lambda o: o.id, - ) == sorted([offsets[id_] for id_ in returned_ids], key=lambda o: o.id) + ) def test_filter_combinations(subject: LabwareOffsetStore) -> None: @@ -265,7 +275,7 @@ def test_filter_combinations(subject: LabwareOffsetStore) -> None: ("id-6", "definition-uri-b"), ] labware_offsets = [ - StoredLabwareOffset( + IncomingStoredLabwareOffset( id=id, createdAt=datetime.now(timezone.utc), definitionUri=definition_uri, @@ -278,18 +288,28 @@ def test_filter_combinations(subject: LabwareOffsetStore) -> None: ) for (id, definition_uri) in ids_and_definition_uris ] + outgoing_offsets = [ + StoredLabwareOffset( + id=offset.id, + createdAt=offset.createdAt, + definitionUri=offset.definitionUri, + locationSequence=offset.locationSequence, + vector=offset.vector, + ) + for offset in labware_offsets + ] for labware_offset in labware_offsets: subject.add(labware_offset) # No filters: - assert subject.search() == labware_offsets + assert subject.search() == outgoing_offsets # Filter on one thing: result = subject.search(definition_uri_filter="definition-uri-b") assert len(result) == 3 assert result == [ - entry for entry in labware_offsets if entry.definitionUri == "definition-uri-b" + entry for entry in outgoing_offsets if entry.definitionUri == "definition-uri-b" ] # Filter on two things: @@ -297,7 +317,7 @@ def test_filter_combinations(subject: LabwareOffsetStore) -> None: id_filter="id-2", definition_uri_filter="definition-uri-b", ) - assert result == [labware_offsets[1]] + assert result == [outgoing_offsets[1]] # Filters should be ANDed, not ORed, together: result = subject.search( @@ -309,8 +329,8 @@ def test_filter_combinations(subject: LabwareOffsetStore) -> None: def test_delete(subject: LabwareOffsetStore) -> None: """Test the `delete()` and `delete_all()` methods.""" - a, b, c = [ - StoredLabwareOffset( + incoming_offsets = [ + IncomingStoredLabwareOffset( id=id, createdAt=datetime.now(timezone.utc), definitionUri="", @@ -323,6 +343,18 @@ def test_delete(subject: LabwareOffsetStore) -> None: ) for id in ["id-a", "id-b", "id-c"] ] + outgoing_offsets = [ + StoredLabwareOffset( + id=offset.id, + createdAt=offset.createdAt, + definitionUri=offset.definitionUri, + locationSequence=offset.locationSequence, + vector=offset.vector, + ) + for offset in incoming_offsets + ] + a, b, c = incoming_offsets + out_a, out_b, out_c = outgoing_offsets with pytest.raises(LabwareOffsetNotFoundError): subject.delete("b") @@ -330,10 +362,11 @@ def test_delete(subject: LabwareOffsetStore) -> None: subject.add(a) subject.add(b) subject.add(c) - assert subject.delete(b.id) == b - assert _get_all(subject) == [a, c] + + assert subject.delete(b.id) == out_b + assert _get_all(subject) == [out_a, out_c] with pytest.raises(LabwareOffsetNotFoundError): - subject.delete(b.id) + subject.delete(out_b.id) subject.delete_all() assert _get_all(subject) == []