Skip to content

Commit

Permalink
fixup: do not crash when encountering an bad loc
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
sfoster1 committed Feb 4, 2025
1 parent 7df44a2 commit 443a8df
Show file tree
Hide file tree
Showing 4 changed files with 167 additions and 34 deletions.
21 changes: 18 additions & 3 deletions robot-server/robot_server/labware_offsets/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
13 changes: 11 additions & 2 deletions robot-server/robot_server/labware_offsets/router.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
from .store import (
LabwareOffsetNotFoundError,
LabwareOffsetStore,
IncomingStoredLabwareOffset,
)
from .fastapi_dependencies import get_labware_offset_store
from .models import (
Expand Down Expand Up @@ -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,
Expand All @@ -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,
)

Expand Down
54 changes: 43 additions & 11 deletions robot-server/robot_server/labware_offsets/store.py
Original file line number Diff line number Diff line change
@@ -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 (
Expand All @@ -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."""
Expand All @@ -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(
Expand Down Expand Up @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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,
Expand All @@ -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):
Expand Down
Loading

0 comments on commit 443a8df

Please sign in to comment.