lint and persistence test fixes

sfoster1 · sfoster1 · commit abd1564b7f49 · 2025-01-29T11:09:41.000-05:00
diff --git a/api/src/opentrons/protocol_engine/__init__.py b/api/src/opentrons/protocol_engine/__init__.py
@@ -32,6 +32,10 @@
     LabwareOffsetVector,
     LegacyLabwareOffsetLocation,
     LabwareOffsetLocationSequence,
+    OnLabwareOffsetLocationSequenceComponent,
+    OnModuleOffsetLocationSequenceComponent,
+    OnAddressableAreaOffsetLocationSequenceComponent,
+    LabwareOffsetLocationSequenceComponents,
     LabwareMovementStrategy,
     AddressableOffsetVector,
     DeckPoint,
@@ -101,6 +105,10 @@
     "LegacyLabwareOffsetCreate",
     "LabwareOffsetLocationSequence",
     "LabwareOffsetVector",
+    "OnLabwareOffsetLocationSequenceComponent",
+    "OnModuleOffsetLocationSequenceComponent",
+    "OnAddressableAreaOffsetLocationSequenceComponent",
+    "LabwareOffsetLocationSequenceComponents",
     "LegacyLabwareOffsetCreate",
     "LegacyLabwareOffsetLocation",
     "LabwareMovementStrategy",
diff --git a/robot-server/robot_server/labware_offsets/router.py b/robot-server/robot_server/labware_offsets/router.py
@@ -10,7 +10,6 @@
 from server_utils.fastapi_utils.light_router import LightRouter
 
 from opentrons.protocol_engine import ModuleModel
-from opentrons.types import DeckSlotName
 
 from robot_server.labware_offsets.models import LabwareOffsetNotFound
 from robot_server.service.dependencies import get_current_time, get_unique_id
@@ -96,8 +95,8 @@ async def get_labware_offsets(  # noqa: D103
             ),
         ),
     ] = DO_NOT_FILTER,
-    location_slot_name: Annotated[
-        Json[DeckSlotName] | SkipJsonSchema[DoNotFilterType],
+    location_addressable_area_name: Annotated[
+        Json[str] | SkipJsonSchema[DoNotFilterType],
         fastapi.Query(
             alias="locationSlotName",
             description="Filter for exact matches on the `location.slotName` field.",
@@ -146,7 +145,7 @@ async def get_labware_offsets(  # noqa: D103
     result_data = store.search(
         id_filter=id,
         definition_uri_filter=definition_uri,
-        location_slot_name_filter=location_slot_name,
+        location_addressable_area_filter=location_addressable_area_name,
         location_definition_uri_filter=location_definition_uri,
         location_module_model_filter=location_module_model,
     )
diff --git a/robot-server/robot_server/labware_offsets/store.py b/robot-server/robot_server/labware_offsets/store.py
@@ -10,8 +10,8 @@
     OnModuleOffsetLocationSequenceComponent,
     OnLabwareOffsetLocationSequenceComponent,
     LabwareOffsetLocationSequenceComponents,
+    LabwareOffsetLocationSequence,
 )
-from opentrons.types import DeckSlotName
 
 from robot_server.persistence.tables import (
     labware_offset_table,
@@ -59,9 +59,9 @@ def add(self, offset: StoredLabwareOffset) -> None:
         with self._sql_engine.begin() as transaction:
             offset_row_id = (
                 transaction.execute(
-                    sqlalchemy.insert(labware_offset_table)
-                    .values(_pydantic_to_sql_offset(offset))
-                    .returning(labware_offset_table.c.row_id)
+                    sqlalchemy.insert(labware_offset_table).values(
+                        _pydantic_to_sql_offset(offset)
+                    )
                 )
                 .one()
                 .row_id
@@ -82,7 +82,7 @@ def search(
         self,
         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_addressable_area_filter: str | DoNotFilterType = DO_NOT_FILTER,
         location_module_model_filter: (
             ModuleModel | None | DoNotFilterType
         ) = DO_NOT_FILTER,
@@ -111,48 +111,72 @@ def search(
             filter_statement = filter_statement.where(
                 labware_offset_table.c.definition_uri == definition_uri_filter
             )
-        if location_slot_name_filter is not DO_NOT_FILTER:
-            filter_statement = filter_statement.where(
-                filter_statement.where(
-                    labware_offset_location_sequence_components_table.c.component_kind
-                    == "onAddressableArea"
+        if location_addressable_area_filter is not DO_NOT_FILTER:
+            if location_addressable_area_filter is not None:
+                filter_statement = filter_statement.where(
+                    filter_statement.where(
+                        labware_offset_location_sequence_components_table.c.component_kind
+                        == "onAddressableArea"
+                    )
+                    .where(
+                        labware_offset_location_sequence_components_table.c.primary_component_value
+                        == location_addressable_area_filter
+                    )
+                    .exists()
                 )
-                .where(
-                    labware_offset_location_sequence_components_table.c.primary_component_value
-                    == location_slot_name_filter.value
+            else:
+                filter_statement = filter_statement.where(
+                    sqlalchemy.not_(
+                        filter_statement.where(
+                            labware_offset_location_sequence_components_table.c.component_kind
+                            == "onAddressableArea"
+                        ).exists()
+                    )
                 )
-                .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.where(
-                filter_statement.where(
-                    labware_offset_location_sequence_components_table.c.component_kind
-                    == "onModule"
+            if location_module_model_filter is not None:
+                filter_statement = filter_statement.where(
+                    filter_statement.where(
+                        labware_offset_location_sequence_components_table.c.component_kind
+                        == "onModule"
+                    )
+                    .where(
+                        labware_offset_location_sequence_components_table.c.primary_component_value
+                        == location_module_model_filter.value
+                    )
+                    .exists()
                 )
-                .where(
-                    labware_offset_location_sequence_components_table.c.primary_component_value
-                    == location_module_model_filter.value
+            else:
+                filter_statement = filter_statement.where(
+                    sqlalchemy.not_(
+                        filter_statement.where(
+                            labware_offset_location_sequence_components_table.c.component_kind
+                            == "onModule"
+                        ).exists()
+                    )
                 )
-                .exists()
-            )
         if location_definition_uri_filter is not DO_NOT_FILTER:
-            filter_statement = filter_statement.where(
-                filter_statement.where(
-                    labware_offset_location_sequence_components_table.c.component_kind
-                    == "onLabware"
+            if location_definition_uri_filter is not None:
+                filter_statement = filter_statement.where(
+                    filter_statement.where(
+                        labware_offset_location_sequence_components_table.c.component_kind
+                        == "onLabware"
+                    )
+                    .where(
+                        labware_offset_location_sequence_components_table.c.primary_component_value
+                        == location_definition_uri_filter
+                    )
+                    .exists()
                 )
-                .where(
-                    labware_offset_location_sequence_components_table.c.primary
-                    - component_value
-                    == location_definition_uri_filter
+            else:
+                filter_statement = filter_statement.where(
+                    sqlalchemy.not_(
+                        filter_statement.where(
+                            labware_offset_location_sequence_components_table.c.component_kind
+                            == "onLabware"
+                        ).exists()
+                    )
                 )
-                .exists()
-            )
 
         filter_statement = filter_statement.order_by(
             labware_offset_table.c.row_id
@@ -167,14 +191,20 @@ def delete(self, offset_id: str) -> StoredLabwareOffset:
         """Delete a labware offset by its ID. Return what was just deleted."""
         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
+                offset_rows = transaction.execute(
+                    sqlalchemy.select(labware_offset_table)
+                    .join(
+                        labware_offset_location_sequence_components_table,
+                        labware_offset_table.c.row_id
+                        == labware_offset_location_sequence_components_table.c.offset_id,
                     )
-                ).one()
+                    .where(labware_offset_table.c.offset_id == offset_id)
+                ).all()
             except sqlalchemy.exc.NoResultFound:
                 raise LabwareOffsetNotFoundError(bad_offset_id=offset_id) from None
-            if not row_to_delete.active:
+            if len(offset_rows) == 0:
+                raise LabwareOffsetNotFoundError(bad_offset_id=offset_id)
+            if not offset_rows[0].active:
                 # Already soft-deleted.
                 raise LabwareOffsetNotFoundError(bad_offset_id=offset_id)
 
@@ -184,18 +214,7 @@ def delete(self, offset_id: str) -> StoredLabwareOffset:
                 .values(active=False)
             )
 
-            location_sequence = transaction.execute(
-                sqlalchemy.get(labware_offset_location_sequence_components_table)
-                .where(
-                    labware_offset_location_sequence_components_table.c.offset_id
-                    == row_to_delete.row_id
-                )
-                .order_by(
-                    labware_offset_location_sequence_components_table.c.sequence_ordinal
-                )
-            ).all()
-
-        return _sql_to_pydantic(row_to_delete, location_sequence)
+        return next(_collate_sql_to_pydantic(offset_rows))
 
     def delete_all(self) -> None:
         """Delete all labware offsets."""
@@ -217,15 +236,15 @@ def _sql_sequence_component_to_pydantic_sequence_component(
     component_row: sqlalchemy.engine.Row,
 ) -> LabwareOffsetLocationSequenceComponents:
     if component_row.component_kind == "onLabware":
-        yield OnLabwareOffsetLocationSequenceComponent(
+        return OnLabwareOffsetLocationSequenceComponent(
             labwareUri=component_row.primary_component_value
         )
     elif component_row.component_kind == "onModule":
-        yield OnModuleOffsetLocationSequenceComponent(
+        return OnModuleOffsetLocationSequenceComponent(
             moduleModel=ModuleModel(component_row.primary_component_value)
         )
     elif component_row.component_kind == "onAddressableArea":
-        yield OnAddressableAreaOffsetLocationSequenceComponent(
+        return OnAddressableAreaOffsetLocationSequenceComponent(
             addressableAreaName=component_row.primary_component_value
         )
     else:
@@ -234,7 +253,7 @@ def _sql_sequence_component_to_pydantic_sequence_component(
 
 def _collate_sql_locations(
     first_row: sqlalchemy.engine.Row, rest_rows: Iterator[sqlalchemy.engine.Row]
-) -> tuple[LabwareOffsetLocationSequence, Optional[sqlalchemy.engine.Row]]:
+) -> tuple[LabwareOffsetLocationSequence, sqlalchemy.engine.Row | None]:
     offset_id = first_row.offset_id
     location_sequence: list[LabwareOffsetLocationSequenceComponents] = [
         _sql_sequence_component_to_pydantic_sequence_component(first_row)
@@ -244,7 +263,7 @@ def _collate_sql_locations(
             row = next(rest_rows)
         except StopIteration:
             return location_sequence, None
-        if row.offset_id != first_row.offset_id:
+        if row.offset_id != offset_id:
             return location_sequence, row
         location_sequence.append(
             _sql_sequence_component_to_pydantic_sequence_component(row)
@@ -253,7 +272,7 @@ def _collate_sql_locations(
 
 def _sql_to_pydantic(
     first_row: sqlalchemy.engine.Row, rest_rows: Iterator[sqlalchemy.engine.Row]
-) -> tuple[StoredLabwareOffset, Optional[sqlalchemy.engine.Row]]:
+) -> tuple[StoredLabwareOffset, sqlalchemy.engine.Row | None]:
     location_sequence, next_row = _collate_sql_locations(first_row, rest_rows)
     return (
         StoredLabwareOffset(
@@ -275,7 +294,7 @@ def _collate_sql_to_pydantic(
     query_results: list[sqlalchemy.engine.Row],
 ) -> Iterator[StoredLabwareOffset]:
     row_iter = iter(query_results)
-    row = next(row_iter)
+    row: sqlalchemy.engine.Row | None = next(row_iter)
     while row:
         result, row = _sql_to_pydantic(row, row_iter)
         yield result
diff --git a/robot-server/robot_server/persistence/_migrations/v9_to_v10.py b/robot-server/robot_server/persistence/_migrations/v9_to_v10.py
@@ -54,7 +54,7 @@ def migrate(self, source_dir: Path, dest_dir: Path) -> None:
         with sql_engine_ctx(
             dest_dir / DB_FILE
         ) as engine, engine.begin() as transaction:
-            transaction.execute("DROP TABLE ?", schema_9.labware_offset_table.name)
+            schema_9.labware_offset_table.drop(transaction)
 
 
 def _upmigrate_stored_offsets(connection: sqlalchemy.engine.Connection) -> None:
@@ -67,9 +67,9 @@ def _upmigrate_stored_offsets(connection: sqlalchemy.engine.Connection) -> None:
     for offset in offsets:
         new_row = (
             connection.execute(
-                sqlalchemy.insert(schema_10.labware_offset_table)
-                .values(_v9_offset_to_v10_offset(offset))
-                .returning(schema_10.labware_offset_table.c.row_id)
+                sqlalchemy.insert(schema_10.labware_offset_table).values(
+                    _v9_offset_to_v10_offset(offset)
+                )
             )
             .one()
             .row_id
diff --git a/robot-server/robot_server/persistence/persistence_directory.py b/robot-server/robot_server/persistence/persistence_directory.py
@@ -1,6 +1,5 @@
 """Create or reset the server's persistence directory."""
 
-
 from pathlib import Path
 from logging import getLogger
 from shutil import rmtree
@@ -71,7 +70,7 @@ def make_migration_orchestrator(prepared_root: Path) -> MigrationOrchestrator:
             v6_to_v7.Migration6to7(subdirectory="7.1"),
             v7_to_v8.Migration7to8(subdirectory="8"),
             v8_to_v9.Migration8to9(subdirectory="9"),
-            v9_to_v10.Migration9to10(subdirection=LATEST_VERSION_DIRECTORY),
+            v9_to_v10.Migration9to10(subdirectory=LATEST_VERSION_DIRECTORY),
         ],
         temp_file_prefix="temp-",
     )
diff --git a/robot-server/robot_server/persistence/tables/schema_10.py b/robot-server/robot_server/persistence/tables/schema_10.py
@@ -397,7 +397,11 @@ class BooleanSettingKey(enum.Enum):
     sqlalchemy.Column("row_id", sqlalchemy.Integer, primary_key=True),
     # Which offset this belongs to
     sqlalchemy.Column(
-        "offset_id", sqlalchemy.ForeignKey("labware_offset_with_sequence.row_id")
+        "offset_id",
+        sqlalchemy.ForeignKey(
+            "labware_offset_with_sequence.row_id",
+        ),
+        nullable=False,
     ),
     # Its position within the sequence
     sqlalchemy.Column("sequence_ordinal", sqlalchemy.Integer, nullable=False),
diff --git a/robot-server/tests/labware_offsets/test_store.py b/robot-server/tests/labware_offsets/test_store.py
diff --git a/robot-server/tests/persistence/test_tables.py b/robot-server/tests/persistence/test_tables.py
