feat(robot-server): handle new offset locations in /labwareOffsets

api/src/opentrons/protocol_engine/__init__.py

@@ -30,6 +30,11 @@
     LabwareOffsetCreate,
     LabwareOffsetVector,
     LegacyLabwareOffsetLocation,
+    LabwareOffsetLocationSequence,
+    OnLabwareOffsetLocationSequenceComponent,
+    OnModuleOffsetLocationSequenceComponent,
+    OnAddressableAreaOffsetLocationSequenceComponent,
+    LabwareOffsetLocationSequenceComponents,
     LabwareMovementStrategy,
     AddressableOffsetVector,
     DeckPoint,
@@ -96,7 +101,13 @@
     # public value interfaces and models
     "LabwareOffset",
     "LabwareOffsetCreate",
+    "LegacyLabwareOffsetCreate",
+    "LabwareOffsetLocationSequence",
     "LabwareOffsetVector",
+    "OnLabwareOffsetLocationSequenceComponent",
+    "OnModuleOffsetLocationSequenceComponent",
+    "OnAddressableAreaOffsetLocationSequenceComponent",
+    "LabwareOffsetLocationSequenceComponents",
     "LegacyLabwareOffsetCreate",
     "LegacyLabwareOffsetLocation",
     "LabwareMovementStrategy",

api/src/opentrons/protocol_engine/labware_offset_standardization.py

@@ -37,9 +37,10 @@ def standardize_labware_offset_create(
     )
 
 
-def _legacy_offset_location_to_offset_location_sequence(
+def legacy_offset_location_to_offset_location_sequence(
     location: LegacyLabwareOffsetLocation, deck_definition: DeckDefinitionV5
 ) -> LabwareOffsetLocationSequence:
+    """Convert a legacy location to a new-style sequence."""
     sequence: LabwareOffsetLocationSequence = []
     if location.definitionUri:
         sequence.append(
@@ -165,7 +166,7 @@ def _locations_for_create(
             }
         )
         return (
-            _legacy_offset_location_to_offset_location_sequence(
+            legacy_offset_location_to_offset_location_sequence(
                 normalized, deck_definition
             ),
             normalized,

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

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

@@ -3,7 +3,7 @@
 This is its own module to fix circular imports.
 """
 
-from typing import Optional, Literal
+from typing import Optional, Literal, Annotated
 
 from pydantic import BaseModel, Field
 
@@ -48,12 +48,16 @@ class OnAddressableAreaOffsetLocationSequenceComponent(BaseModel):
     )
 
 
-LabwareOffsetLocationSequenceComponents = (
+LabwareOffsetLocationSequenceComponentsUnion = (
     OnLabwareOffsetLocationSequenceComponent
     | OnModuleOffsetLocationSequenceComponent
     | OnAddressableAreaOffsetLocationSequenceComponent
 )
 
+LabwareOffsetLocationSequenceComponents = Annotated[
+    LabwareOffsetLocationSequenceComponentsUnion, Field(discriminator="kind")
+]
+
 LabwareOffsetLocationSequence = list[LabwareOffsetLocationSequenceComponents]
 
 

robot-server/robot_server/labware_offsets/_search_query_builder.py

@@ -0,0 +1,177 @@
+"""Helper to build a search query."""
+
+from __future__ import annotations
+from typing import Final, TYPE_CHECKING
+
+import sqlalchemy
+
+from opentrons.protocol_engine import ModuleModel
+
+from robot_server.persistence.tables import (
+    labware_offset_table,
+    labware_offset_location_sequence_components_table,
+)
+from .models import DoNotFilterType, DO_NOT_FILTER
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+
+class SearchQueryBuilder:
+    """Helper class to build a search query.
+
+    This object is stateful, and should be kept around just long enough to have the parameters
+    of a single search injected.
+    """
+
+    def __init__(self) -> None:
+        """Build the object."""
+        super().__init__()
+        self._filter_original: Final = sqlalchemy.select(
+            labware_offset_table.c.row_id,
+            labware_offset_table.c.offset_id,
+            labware_offset_table.c.definition_uri,
+            labware_offset_table.c.vector_x,
+            labware_offset_table.c.vector_y,
+            labware_offset_table.c.vector_z,
+            labware_offset_table.c.created_at,
+            labware_offset_table.c.active,
+            labware_offset_location_sequence_components_table.c.sequence_ordinal,
+            labware_offset_location_sequence_components_table.c.component_kind,
+            labware_offset_location_sequence_components_table.c.primary_component_value,
+        ).select_from(
+            sqlalchemy.join(
+                labware_offset_table,
+                labware_offset_location_sequence_components_table,
+                labware_offset_table.c.row_id
+                == labware_offset_location_sequence_components_table.c.offset_id,
+            )
+        )
+        self._offset_location_alias: Final = (
+            labware_offset_location_sequence_components_table.alias()
+        )
+        self._current_base_filter_statement = self._filter_original
+        self._current_positive_location_filter: (
+            sqlalchemy.sql.selectable.Exists | None
+        ) = None
+        self._current_negative_filter_subqueries: list[
+            sqlalchemy.sql.selectable.Exists
+        ] = []
+
+    def _positive_query(self) -> sqlalchemy.sql.selectable.Exists:
+        if self._current_positive_location_filter is not None:
+            return self._current_positive_location_filter
+        return sqlalchemy.exists().where(
+            self._offset_location_alias.c.offset_id
+            == labware_offset_location_sequence_components_table.c.offset_id
+        )
+
+    def build_query(self) -> sqlalchemy.sql.selectable.Selectable:
+        """Render the query into a sqlalchemy object suitable for passing to the database."""
+        statement = self._current_base_filter_statement
+        if self._current_positive_location_filter is not None:
+            statement = statement.where(self._current_positive_location_filter)
+        for subq in self._current_negative_filter_subqueries:
+            statement = statement.where(sqlalchemy.not_(subq))
+        statement = statement.order_by(labware_offset_table.c.row_id).order_by(
+            labware_offset_location_sequence_components_table.c.sequence_ordinal
+        )
+        return statement
+
+    def do_active_filter(self, active: bool) -> Self:
+        """Filter to only active=active rows."""
+        self._current_base_filter_statement = self._current_base_filter_statement.where(
+            labware_offset_table.c.active == True  # noqa: E712
+        )
+        return self
+
+    def do_id_filter(self, id_filter: str | DoNotFilterType) -> Self:
+        """Filter to rows with only the given offset ID."""
+        if id_filter is DO_NOT_FILTER:
+            return self
+
+        self._current_base_filter_statement = self._current_base_filter_statement.where(
+            labware_offset_table.c.offset_id == id_filter
+        )
+        return self
+
+    def do_definition_uri_filter(
+        self, definition_uri_filter: str | DoNotFilterType
+    ) -> Self:
+        """Filter to rows of an offset that apply to a definition URI."""
+        if definition_uri_filter is DO_NOT_FILTER:
+            return self
+        self._current_base_filter_statement = self._current_base_filter_statement.where(
+            labware_offset_table.c.definition_uri == definition_uri_filter
+        )
+        return self
+
+    def do_on_addressable_area_filter(
+        self,
+        addressable_area_filter: str | DoNotFilterType,
+    ) -> Self:
+        """Filter to rows of an offset that applies to the given addressable area."""
+        if addressable_area_filter is DO_NOT_FILTER:
+            return self
+        self._current_positive_location_filter = (
+            self._positive_query()
+            .where(self._offset_location_alias.c.component_kind == "onAddressableArea")
+            .where(
+                self._offset_location_alias.c.primary_component_value
+                == addressable_area_filter
+            )
+        )
+        return self
+
+    def do_on_labware_filter(
+        self, labware_uri_filter: str | DoNotFilterType | None
+    ) -> Self:
+        """Filter to the rows of an offset located on the given labware (or no labware)."""
+        if labware_uri_filter is DO_NOT_FILTER:
+            return self
+        if labware_uri_filter is None:
+            self._current_negative_filter_subqueries.append(
+                sqlalchemy.exists()
+                .where(
+                    self._offset_location_alias.c.offset_id
+                    == labware_offset_location_sequence_components_table.c.offset_id
+                )
+                .where(self._offset_location_alias.c.component_kind == "onLabware")
+            )
+            return self
+        self._current_positive_location_filter = (
+            self._positive_query()
+            .where(self._offset_location_alias.c.component_kind == "onLabware")
+            .where(
+                self._offset_location_alias.c.primary_component_value
+                == labware_uri_filter
+            )
+        )
+        return self
+
+    def do_on_module_filter(
+        self,
+        module_model_filter: ModuleModel | DoNotFilterType | None,
+    ) -> Self:
+        """Filter to the rows of an offset located on the given module (or no module)."""
+        if module_model_filter is DO_NOT_FILTER:
+            return self
+        if module_model_filter is None:
+            self._current_negative_filter_subqueries.append(
+                sqlalchemy.exists()
+                .where(
+                    self._offset_location_alias.c.offset_id
+                    == labware_offset_location_sequence_components_table.c.offset_id
+                )
+                .where(self._offset_location_alias.c.component_kind == "onModule")
+            )
+            return self
+        self._current_positive_location_filter = (
+            self._positive_query()
+            .where(self._offset_location_alias.c.component_kind == "onModule")
+            .where(
+                self._offset_location_alias.c.primary_component_value
+                == module_model_filter.value
+            )
+        )
+        return self

robot-server/robot_server/labware_offsets/models.py

@@ -1,11 +1,83 @@
 """Request/response models for the `/labwareOffsets` endpoints."""
 
+from datetime import datetime
+import enum
+from typing import Literal, Annotated, Final, TypeAlias
 
-from typing import Literal
+from pydantic import BaseModel, Field
+
+from opentrons.protocol_engine import (
+    LabwareOffsetVector,
+)
+from opentrons.protocol_engine.types.labware_offset_location import (
+    LabwareOffsetLocationSequenceComponentsUnion,
+)
 
 from robot_server.errors.error_responses import ErrorDetails
 
 
+class _DoNotFilter(enum.Enum):
+    DO_NOT_FILTER = enum.auto()
+
+
+DO_NOT_FILTER: Final = _DoNotFilter.DO_NOT_FILTER
+"""A sentinel value for when a filter should not be applied.
+
+This is different from filtering on `None`, which returns only entries where the
+value is equal to `None`.
+"""
+
+
+DoNotFilterType: TypeAlias = Literal[_DoNotFilter.DO_NOT_FILTER]
+"""The type of `DO_NOT_FILTER`, as `NoneType` is to `None`.
+
+Unfortunately, mypy doesn't let us write `Literal[DO_NOT_FILTER]`. Use this instead.
+"""
+
+
+# This is redefined here so we can add stuff to it easily
+StoredLabwareOffsetLocationSequenceComponents = Annotated[
+    LabwareOffsetLocationSequenceComponentsUnion, 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(
+        ...,
+        description="Where the labware is located on the robot. Can represent all locations, but may not be present for older runs.",
+        min_length=1,
+    )
+    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."""
+
+    # This is a separate thing from the model defined in protocol engine because as a new API it does
+    # not have to handle legacy locations. There is probably a better way to do this than to copy the model
+    # contents, but I'm not sure what it is.
+    id: str = Field(..., description="Unique labware offset record identifier.")
+    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(
+        ...,
+        description="Where the labware is located on the robot. Can represent all locations, but may not be present for older runs.",
+        min_length=1,
+    )
+    vector: LabwareOffsetVector = Field(
+        ...,
+        description="The offset applied to matching labware.",
+    )
+
+
 class LabwareOffsetNotFound(ErrorDetails):
     """An error returned when a requested labware offset does not exist."""