fixup: do not crash when encountering an bad loc

This will allow us to add new kinds of location without doing schema
changes and migrations. If you add a location and downgrade to a version
that doesn't understand it, you'll get an
"UnknownLocationSequenceComponent", just like a bad run.

Author: sfoster1

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,

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,
     )
 

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):