fixups

Author: sfoster1

api/src/opentrons/protocol_engine/types/__init__.py

@@ -94,6 +94,7 @@
     OnLabwareOffsetLocationSequenceComponent,
     OnModuleOffsetLocationSequenceComponent,
     OnAddressableAreaOffsetLocationSequenceComponent,
+    LabwareOffsetLocationSequenceComponents,
 )
 from .labware_offset_vector import LabwareOffsetVector
 from .well_position import (
@@ -204,6 +205,7 @@
     # Labware offset location
     "LegacyLabwareOffsetLocation",
     "LabwareOffsetLocationSequence",
+    "LabwareOffsetLocationSequenceComponents",
     "OnLabwareOffsetLocationSequenceComponent",
     "OnModuleOffsetLocationSequenceComponent",
     "OnAddressableAreaOffsetLocationSequenceComponent",

robot-server/robot_server/labware_offsets/models.py

@@ -10,6 +10,21 @@
 from robot_server.errors.error_responses import ErrorDetails
 
 
+class StoredLabwareOffsetCreate(BaseModel):
+    """Create an offset for storage."""
+
+    definitionUri: str = Field(..., description="The URI for the labware's definition.")
+
+    locationSequence: LabwareOffsetLocationSequence = Field(
+        ...,
+        description="Where the labware is located on the robot. Can represent all locations, but may not be present for older runs.",
+    )
+    vector: LabwareOffsetVector = Field(
+        ...,
+        description="The offset applied to matching labware.",
+    )
+
+
 class StoredLabwareOffset(BaseModel):
     """An offset that the robot adds to a pipette's position when it moves to labware."""
 
@@ -20,8 +35,8 @@ 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: LabwareOffsetLocationSequence | None = Field(
-        default=None,
+    locationSequence: LabwareOffsetLocationSequence = Field(
+        ...,
         description="Where the labware is located on the robot. Can represent all locations, but may not be present for older runs.",
     )
     vector: LabwareOffsetVector = Field(

robot-server/robot_server/labware_offsets/router.py

@@ -9,7 +9,7 @@
 from pydantic.json_schema import SkipJsonSchema
 from server_utils.fastapi_utils.light_router import LightRouter
 
-from opentrons.protocol_engine import LabwareOffsetCreate, ModuleModel
+from opentrons.protocol_engine import ModuleModel
 from opentrons.types import DeckSlotName
 
 from robot_server.labware_offsets.models import LabwareOffsetNotFound
@@ -30,7 +30,7 @@
     LabwareOffsetStore,
 )
 from .fastapi_dependencies import get_labware_offset_store
-from .models import StoredLabwareOffset
+from .models import StoredLabwareOffset, StoredLabwareOffsetCreate
 
 
 router = LightRouter()
@@ -54,13 +54,13 @@ async def post_labware_offset(  # noqa: D103
     store: Annotated[LabwareOffsetStore, fastapi.Depends(get_labware_offset_store)],
     new_offset_id: Annotated[str, fastapi.Depends(get_unique_id)],
     new_offset_created_at: Annotated[datetime, fastapi.Depends(get_current_time)],
-    request_body: Annotated[RequestModel[LabwareOffsetCreate], fastapi.Body()],
+    request_body: Annotated[RequestModel[StoredLabwareOffsetCreate], fastapi.Body()],
 ) -> PydanticResponse[SimpleBody[StoredLabwareOffset]]:
     new_offset = StoredLabwareOffset.model_construct(
         id=new_offset_id,
         createdAt=new_offset_created_at,
         definitionUri=request_body.data.definitionUri,
-        location=request_body.data.location,
+        locationSequence=request_body.data.locationSequence,
         vector=request_body.data.vector,
     )
     store.add(new_offset)

robot-server/robot_server/labware_offsets/store.py

@@ -6,9 +6,9 @@
 from opentrons.protocol_engine.types import (
     LabwareOffsetVector,
     ModuleModel,
-    OnAddressableAreaOffsetSequenceComponent,
-    OnModuleOffsetSequenceComponent,
-    OnLabwareOffsetSequenceComponent,
+    OnAddressableAreaOffsetLocationSequenceComponent,
+    OnModuleOffsetLocationSequenceComponent,
+    OnLabwareOffsetLocationSequenceComponent,
     LabwareOffsetLocationSequenceComponents,
 )
 from opentrons.types import DeckSlotName
@@ -42,8 +42,6 @@ class _DoNotFilter(enum.Enum):
 """
 
 
-# todo(mm, 2024-12-06): Convert to be SQL-based and persistent instead of in-memory.
-# https://opentrons.atlassian.net/browse/EXEC-1015
 class LabwareOffsetStore:
     """A persistent store for labware offsets, to support the `/labwareOffsets` endpoints."""
 
@@ -69,7 +67,9 @@ def add(self, offset: StoredLabwareOffset) -> None:
                 .row_id
             )
             transaction.execute(
-                sqlalchemy.insert(labware_offset_table).values(
+                sqlalchemy.insert(
+                    labware_offset_location_sequence_components_table
+                ).values(
                     list(
                         _pydantic_to_sql_location_sequence_iterator(
                             offset, offset_row_id
@@ -98,7 +98,7 @@ def search(
             .join(
                 labware_offset_location_sequence_components_table,
                 labware_offset_table.c.row_id
-                == labware_offset_sequence_table.c.offset_id,
+                == labware_offset_location_sequence_components_table.c.offset_id,
             )
             .where(labware_offset_table.c.active == True)  # noqa: E712
         )
@@ -112,7 +112,7 @@ def search(
                 labware_offset_table.c.definition_uri == definition_uri_filter
             )
         if location_slot_name_filter is not DO_NOT_FILTER:
-            filter_statement = (
+            filter_statement = filter_statement.where(
                 filter_statement.where(
                     labware_offset_location_sequence_components_table.c.component_kind
                     == "onAddressableArea"
@@ -121,15 +121,15 @@ def search(
                     labware_offset_location_sequence_components_table.c.primary_component_value
                     == location_slot_name_filter.value
                 )
-                .exists(0)
+                .exists()
             )
         if location_module_model_filter is not DO_NOT_FILTER:
             location_module_model_filter_value = (
                 location_module_model_filter.value
                 if location_module_model_filter is not None
                 else None
             )
-            filter_statement = (
+            filter_statement = filter_statement.where(
                 filter_statement.where(
                     labware_offset_location_sequence_components_table.c.component_kind
                     == "onModule"
@@ -141,7 +141,7 @@ def search(
                 .exists()
             )
         if location_definition_uri_filter is not DO_NOT_FILTER:
-            filter_statement = (
+            filter_statement = filter_statement.where(
                 filter_statement.where(
                     labware_offset_location_sequence_components_table.c.component_kind
                     == "onLabware"
@@ -217,15 +217,15 @@ def _sql_sequence_component_to_pydantic_sequence_component(
     component_row: sqlalchemy.engine.Row,
 ) -> LabwareOffsetLocationSequenceComponents:
     if component_row.component_kind == "onLabware":
-        yield OnLabwareOffsetSequenceComponent(
+        yield OnLabwareOffsetLocationSequenceComponent(
             labwareUri=component_row.primary_component_value
         )
     elif component_row.component_kind == "onModule":
-        yield OnModuleOffsetSequenceComponent(
+        yield OnModuleOffsetLocationSequenceComponent(
             moduleModel=ModuleModel(component_row.primary_component_value)
         )
     elif component_row.component_kind == "onAddressableArea":
-        yield OnAddressableAreaOffsetSequenceComponent(
+        yield OnAddressableAreaOffsetLocationSequenceComponent(
             addressableAreaName=component_row.primary_component_value
         )
     else:
@@ -297,23 +297,23 @@ def _pydantic_to_sql_location_sequence_iterator(
     labware_offset: StoredLabwareOffset, offset_row_id: int
 ) -> Iterator[dict[str, object]]:
     for index, component in labware_offset.locationSequence:
-        if isinstance(component, OnLabwareOffsetSequenceComponent):
+        if isinstance(component, OnLabwareOffsetLocationSequenceComponent):
             yield dict(
                 offset_id=offset_row_id,
                 sequence_ordinal=index,
                 component_kind=component.kind,
                 primary_component_value=component.labwareUri,
                 component_value_json=component.model_dump(),
             )
-        elif isinstance(component, OnModuleOffsetSequenceComponent):
+        elif isinstance(component, OnModuleOffsetLocationSequenceComponent):
             yield dict(
                 offset_id=offset_row_id,
                 sequence_ordinal=index,
                 component_kind=component.kind,
                 primary_component_value=component.moduleModel.value,
                 component_value_json=component.model_dump(),
             )
-        elif isinstance(component, OnAddressableAreaOffsetSequenceComponent):
+        elif isinstance(component, OnAddressableAreaOffsetLocationSequenceComponent):
             yield dict(
                 offset_id=offset_row_id,
                 sequence_ordinal=index,

robot-server/robot_server/persistence/_migrations/v9_to_v10.py

@@ -62,7 +62,8 @@ def _upmigrate_stored_offsets(connection: sqlalchemy.engine.Connection) -> None:
         DeckType(guess_deck_type_from_global_config()), version=5
     )
 
-    offsets = sqlalchemy.select(schema_9.labware_offset_table)
+    offsets = connection.execute(sqlalchemy.select(schema_9.labware_offset_table))
+
     for offset in offsets:
         new_row = (
             connection.execute(
@@ -100,25 +101,28 @@ def _v9_offset_to_v10_offset_locations(
     location_sequence = legacy_offset_location_to_offset_location_sequence(
         LegacyLabwareOffsetLocation(
             slotName=DeckSlotName(v9_offset.location_slot_name),
-            moduleModel=ModuleModel(v9_offset.location_module_model)
-            if v9_offset.location_module_model is not None
-            else None,
+            moduleModel=(
+                ModuleModel(v9_offset.location_module_model)
+                if v9_offset.location_module_model is not None
+                else None
+            ),
             definitionUri=v9_offset.location_definition_uri,
         ),
+        deck_definition,
     )
     values: list[dict[str, object]] = []
-    for index, sequence_component in location_sequence.enumerate():
+    for index, sequence_component in enumerate(location_sequence):
         primary_component_value = ""
         component_value_json = ""
         if sequence_component.kind == "onLabware":
             primary_component_value = sequence_component.labwareUri
-            component_value_json = sequence_component.model_dump()
+            component_value_json = sequence_component.model_dump_json()
         elif sequence_component.kind == "onModule":
             primary_component_value = sequence_component.moduleModel.value
-            component_value_json = sequence_component.model_dump()
+            component_value_json = sequence_component.model_dump_json()
         elif sequence_component.kind == "onAddressableArea":
             primary_component_value = sequence_component.addressableAreaName
-            component_value_json = sequence_component.model_dump()
+            component_value_json = sequence_component.model_dump_json()
         else:
             # This should never happen since we're exhaustively checking kinds here
             continue