feat(robot-server): Persist /labwareOffsets storage across reboots

robot-server/robot_server/labware_offsets/fastapi_dependencies.py

@@ -10,6 +10,9 @@
     AppStateAccessor,
     get_app_state,
 )
+import sqlalchemy
+
+from robot_server.persistence.fastapi_dependencies import get_sql_engine
 from .store import LabwareOffsetStore
 
 
@@ -20,10 +23,11 @@
 
 async def get_labware_offset_store(
     app_state: Annotated[AppState, Depends(get_app_state)],
+    sql_engine: Annotated[sqlalchemy.engine.Engine, Depends(get_sql_engine)],
 ) -> LabwareOffsetStore:
     """Get the server's singleton LabwareOffsetStore."""
     labware_offset_store = _labware_offset_store_accessor.get_from(app_state)
     if labware_offset_store is None:
-        labware_offset_store = LabwareOffsetStore()
+        labware_offset_store = LabwareOffsetStore(sql_engine)
         _labware_offset_store_accessor.set_on(app_state, labware_offset_store)
     return labware_offset_store

robot-server/robot_server/labware_offsets/router.py

@@ -3,7 +3,7 @@
 
 from datetime import datetime
 import textwrap
-from typing import Annotated, Literal, Type
+from typing import Annotated, Literal
 
 import fastapi
 from pydantic import Json
@@ -24,7 +24,12 @@
     SimpleMultiBody,
 )
 
-from .store import DO_NOT_FILTER, LabwareOffsetNotFoundError, LabwareOffsetStore
+from .store import (
+    DO_NOT_FILTER,
+    DoNotFilterType,
+    LabwareOffsetNotFoundError,
+    LabwareOffsetStore,
+)
 from .fastapi_dependencies import get_labware_offset_store
 
 
@@ -78,11 +83,11 @@ async def post_labware_offset(  # noqa: D103
 async def get_labware_offsets(  # noqa: D103
     store: Annotated[LabwareOffsetStore, fastapi.Depends(get_labware_offset_store)],
     id: Annotated[
-        Json[str] | SkipJsonSchema[Type[DO_NOT_FILTER]],
+        Json[str] | SkipJsonSchema[DoNotFilterType],
         fastapi.Query(description="Filter for exact matches on the `id` field."),
     ] = DO_NOT_FILTER,
     definition_uri: Annotated[
-        Json[str] | SkipJsonSchema[Type[DO_NOT_FILTER]],
+        Json[str] | SkipJsonSchema[DoNotFilterType],
         fastapi.Query(
             alias="definitionUri",
             description=(
@@ -92,21 +97,21 @@ async def get_labware_offsets(  # noqa: D103
         ),
     ] = DO_NOT_FILTER,
     location_slot_name: Annotated[
-        Json[DeckSlotName] | SkipJsonSchema[Type[DO_NOT_FILTER]],
+        Json[DeckSlotName] | SkipJsonSchema[DoNotFilterType],
         fastapi.Query(
             alias="locationSlotName",
             description="Filter for exact matches on the `location.slotName` field.",
         ),
     ] = DO_NOT_FILTER,
     location_module_model: Annotated[
-        Json[ModuleModel | None] | SkipJsonSchema[Type[DO_NOT_FILTER]],
+        Json[ModuleModel | None] | SkipJsonSchema[DoNotFilterType],
         fastapi.Query(
             alias="locationModuleModel",
             description="Filter for exact matches on the `location.moduleModel` field.",
         ),
     ] = DO_NOT_FILTER,
     location_definition_uri: Annotated[
-        Json[str | None] | SkipJsonSchema[Type[DO_NOT_FILTER]],
+        Json[str | None] | SkipJsonSchema[DoNotFilterType],
         fastapi.Query(
             alias="locationDefinitionUri",
             description=(

robot-server/robot_server/labware_offsets/store.py

@@ -1,80 +1,137 @@
 # noqa: D100
 
-from typing import Type
-
-from opentrons.protocol_engine import LabwareOffset, ModuleModel
+import enum
+from typing import Final, Literal, TypeAlias
+
+from opentrons.protocol_engine.types import (
+    LabwareOffset,
+    LabwareOffsetLocation,
+    LabwareOffsetVector,
+    ModuleModel,
+)
 from opentrons.types import DeckSlotName
 
+from robot_server.persistence.tables import labware_offset_table
+
+import sqlalchemy
+import sqlalchemy.exc
+
+
+class _DoNotFilter(enum.Enum):
+    DO_NOT_FILTER = enum.auto()
+
 
-class DO_NOT_FILTER:
-    """A sentinel value for when a filter should not be applied.
+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`.
-    """
+This is different from filtering on `None`, which returns only entries where the
+value is equal to `None`.
+"""
 
-    pass
+
+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.
+"""
 
 
 # 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."""
 
-    def __init__(self) -> None:
-        self._offsets_by_id: dict[str, LabwareOffset] = {}
+    def __init__(self, sql_engine: sqlalchemy.engine.Engine) -> None:
+        """Initialize the store.
+
+        Params:
+            sql_engine: The SQL database to use as backing storage. Assumed to already
+                have all the proper tables set up.
+        """
+        self._sql_engine = sql_engine
 
     def add(self, offset: LabwareOffset) -> None:
         """Store a new labware offset."""
-        assert offset.id not in self._offsets_by_id
-        self._offsets_by_id[offset.id] = offset
+        with self._sql_engine.begin() as transaction:
+            transaction.execute(
+                sqlalchemy.insert(labware_offset_table).values(_pydantic_to_sql(offset))
+            )
 
     def search(
         self,
-        id_filter: str | Type[DO_NOT_FILTER] = DO_NOT_FILTER,
-        definition_uri_filter: str | Type[DO_NOT_FILTER] = DO_NOT_FILTER,
-        location_slot_name_filter: DeckSlotName | Type[DO_NOT_FILTER] = DO_NOT_FILTER,
+        id_filter: str | DoNotFilterType = DO_NOT_FILTER,
+        definition_uri_filter: str | DoNotFilterType = DO_NOT_FILTER,
+        location_slot_name_filter: DeckSlotName | DoNotFilterType = DO_NOT_FILTER,
         location_module_model_filter: ModuleModel
         | None
-        | Type[DO_NOT_FILTER] = DO_NOT_FILTER,
-        location_definition_uri_filter: str
-        | None
-        | Type[DO_NOT_FILTER] = DO_NOT_FILTER,
+        | DoNotFilterType = DO_NOT_FILTER,
+        location_definition_uri_filter: str | None | DoNotFilterType = DO_NOT_FILTER,
         # todo(mm, 2024-12-06): Support pagination (cursor & pageLength query params).
         # The logic for that is currently duplicated across several places in
         # robot-server and api. We should try to clean that up, or at least avoid
         # making it worse.
     ) -> list[LabwareOffset]:
         """Return all matching labware offsets in order from oldest-added to newest."""
-
-        def is_match(candidate: LabwareOffset) -> bool:
-            return (
-                id_filter in (DO_NOT_FILTER, candidate.id)
-                and definition_uri_filter in (DO_NOT_FILTER, candidate.definitionUri)
-                and location_slot_name_filter
-                in (DO_NOT_FILTER, candidate.location.slotName)
-                and location_module_model_filter
-                in (DO_NOT_FILTER, candidate.location.moduleModel)
-                and location_definition_uri_filter
-                in (DO_NOT_FILTER, candidate.location.definitionUri)
+        statement = sqlalchemy.select(labware_offset_table).order_by(
+            labware_offset_table.c.row_id
+        )
+
+        if id_filter is not DO_NOT_FILTER:
+            statement = statement.where(labware_offset_table.c.offset_id == id_filter)
+        if definition_uri_filter is not DO_NOT_FILTER:
+            statement = statement.where(
+                labware_offset_table.c.definition_uri == definition_uri_filter
+            )
+        if location_slot_name_filter is not DO_NOT_FILTER:
+            statement = statement.where(
+                labware_offset_table.c.location_slot_name
+                == location_slot_name_filter.value
+            )
+        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
+            )
+            statement = statement.where(
+                labware_offset_table.c.location_module_model
+                == location_module_model_filter_value
+            )
+        if location_definition_uri_filter is not DO_NOT_FILTER:
+            statement = statement.where(
+                labware_offset_table.c.location_definition_uri
+                == location_definition_uri_filter
             )
 
-        return [
-            candidate
-            for candidate in self._offsets_by_id.values()
-            if is_match(candidate)
-        ]
+        with self._sql_engine.begin() as transaction:
+            result = transaction.execute(statement).all()
+
+        return [_sql_to_pydantic(row) for row in result]
 
     def delete(self, offset_id: str) -> LabwareOffset:
         """Delete a labware offset by its ID. Return what was just deleted."""
-        try:
-            return self._offsets_by_id.pop(offset_id)
-        except KeyError:
-            raise LabwareOffsetNotFoundError(bad_offset_id=offset_id) from None
+        with self._sql_engine.begin() as transaction:
+            try:
+                row_to_delete = transaction.execute(
+                    sqlalchemy.select(labware_offset_table).where(
+                        labware_offset_table.c.offset_id == offset_id
+                    )
+                ).one()
+            except sqlalchemy.exc.NoResultFound:
+                raise LabwareOffsetNotFoundError(bad_offset_id=offset_id) from None
+
+            transaction.execute(
+                sqlalchemy.delete(labware_offset_table).where(
+                    labware_offset_table.c.offset_id == offset_id
+                )
+            )
+
+        return _sql_to_pydantic(row_to_delete)
 
     def delete_all(self) -> None:
         """Delete all labware offsets."""
-        self._offsets_by_id.clear()
+        with self._sql_engine.begin() as transaction:
+            transaction.execute(sqlalchemy.delete(labware_offset_table))
 
 
 class LabwareOffsetNotFoundError(KeyError):
@@ -83,3 +140,38 @@ class LabwareOffsetNotFoundError(KeyError):
     def __init__(self, bad_offset_id: str) -> None:
         super().__init__(bad_offset_id)
         self.bad_offset_id = bad_offset_id
+
+
+def _sql_to_pydantic(row: sqlalchemy.engine.Row) -> LabwareOffset:
+    return LabwareOffset(
+        id=row.offset_id,
+        createdAt=row.created_at,
+        definitionUri=row.definition_uri,
+        location=LabwareOffsetLocation(
+            slotName=DeckSlotName(row.location_slot_name),
+            moduleModel=row.location_module_model,
+            definitionUri=row.location_definition_uri,
+        ),
+        vector=LabwareOffsetVector(
+            x=row.vector_x,
+            y=row.vector_y,
+            z=row.vector_z,
+        ),
+    )
+
+
+def _pydantic_to_sql(labware_offset: LabwareOffset) -> dict[str, object]:
+    return dict(
+        offset_id=labware_offset.id,
+        definition_uri=labware_offset.definitionUri,
+        location_slot_name=labware_offset.location.slotName.value,
+        location_module_model=labware_offset.location.moduleModel.value
+        if labware_offset.location.moduleModel is not None
+        else None,
+        location_definition_uri=labware_offset.location.definitionUri,
+        vector_x=labware_offset.vector.x,
+        vector_y=labware_offset.vector.y,
+        vector_z=labware_offset.vector.z,
+        created_at=labware_offset.createdAt,
+        active=True,
+    )

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

@@ -41,6 +41,15 @@ def copy_rows_unmodified(
         dest_connection.execute(insert, row)
 
 
+def copy_contents(source_dir: Path, dest_dir: Path) -> None:
+    """Copy the contents of one directory to another (assumed to be empty)."""
+    for item in source_dir.iterdir():
+        if item.is_dir():
+            shutil.copytree(src=item, dst=dest_dir / item.name)
+        else:
+            shutil.copy(src=item, dst=dest_dir / item.name)
+
+
 def copy_if_exists(src: Path, dst: Path) -> None:
     """Like `shutil.copy()`, but no-op if `src` doesn't exist."""
     try:

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

@@ -8,9 +8,8 @@
 
 from pathlib import Path
 from contextlib import ExitStack
-import shutil
 
-from ._util import add_column
+from ._util import add_column, copy_contents
 from ..database import sql_engine_ctx
 from ..file_and_directory_names import DB_FILE
 from ..tables import schema_4
@@ -21,11 +20,8 @@ class Migration3to4(Migration):  # noqa: D101
     def migrate(self, source_dir: Path, dest_dir: Path) -> None:
         """Migrate the persistence directory from schema 3 to 4."""
         # Copy over all existing directories and files to new version
-        for item in source_dir.iterdir():
-            if item.is_dir():
-                shutil.copytree(src=item, dst=dest_dir / item.name)
-            else:
-                shutil.copy(src=item, dst=dest_dir / item.name)
+        copy_contents(source_dir=source_dir, dest_dir=dest_dir)
+
         dest_db_file = dest_dir / DB_FILE
 
         # Append the new column to existing analyses in v4 database

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

@@ -8,9 +8,8 @@
 
 from pathlib import Path
 from contextlib import ExitStack
-import shutil
 
-from ._util import add_column
+from ._util import add_column, copy_contents
 from ..database import sql_engine_ctx
 from ..file_and_directory_names import DB_FILE
 from ..tables import schema_5
@@ -21,11 +20,8 @@ class Migration4to5(Migration):  # noqa: D101
     def migrate(self, source_dir: Path, dest_dir: Path) -> None:
         """Migrate the persistence directory from schema 4 to 5."""
         # Copy over all existing directories and files to new version
-        for item in source_dir.iterdir():
-            if item.is_dir():
-                shutil.copytree(src=item, dst=dest_dir / item.name)
-            else:
-                shutil.copy(src=item, dst=dest_dir / item.name)
+        copy_contents(source_dir=source_dir, dest_dir=dest_dir)
+
         dest_db_file = dest_dir / DB_FILE
 
         # Append the new column to existing protocols in v4 database

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

@@ -10,11 +10,10 @@
 import json
 from pathlib import Path
 from contextlib import ExitStack
-import shutil
 
 import sqlalchemy
 
-from ._util import add_column
+from ._util import add_column, copy_contents
 from ..database import sql_engine_ctx, sqlite_rowid
 from ..tables import schema_7
 from .._folder_migrator import Migration
@@ -28,11 +27,7 @@ class Migration6to7(Migration):  # noqa: D101
     def migrate(self, source_dir: Path, dest_dir: Path) -> None:
         """Migrate the persistence directory from schema 6 to 7."""
         # Copy over all existing directories and files to new version
-        for item in source_dir.iterdir():
-            if item.is_dir():
-                shutil.copytree(src=item, dst=dest_dir / item.name)
-            else:
-                shutil.copy(src=item, dst=dest_dir / item.name)
+        copy_contents(source_dir=source_dir, dest_dir=dest_dir)
 
         dest_db_file = dest_dir / DB_FILE