diff --git a/robot-server/robot_server/labware_offsets/_search_query_builder.py b/robot-server/robot_server/labware_offsets/_search_query_builder.py
new file mode 100644
index 00000000000..e6c8d2bc058
--- /dev/null
+++ b/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
diff --git a/robot-server/robot_server/labware_offsets/models.py b/robot-server/robot_server/labware_offsets/models.py
index bc49a2a37c9..33bae92e6b3 100644
--- a/robot-server/robot_server/labware_offsets/models.py
+++ b/robot-server/robot_server/labware_offsets/models.py
@@ -1,7 +1,8 @@
"""Request/response models for the `/labwareOffsets` endpoints."""
from datetime import datetime
-from typing import Literal, Annotated
+import enum
+from typing import Literal, Annotated, Final, TypeAlias
from pydantic import BaseModel, Field
@@ -14,6 +15,26 @@
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")
diff --git a/robot-server/robot_server/labware_offsets/router.py b/robot-server/robot_server/labware_offsets/router.py
index 8ef930d4ddf..6c44cd554db 100644
--- a/robot-server/robot_server/labware_offsets/router.py
+++ b/robot-server/robot_server/labware_offsets/router.py
@@ -23,13 +23,16 @@
)
from .store import (
- DO_NOT_FILTER,
- DoNotFilterType,
LabwareOffsetNotFoundError,
LabwareOffsetStore,
)
from .fastapi_dependencies import get_labware_offset_store
-from .models import StoredLabwareOffset, StoredLabwareOffsetCreate
+from .models import (
+ StoredLabwareOffset,
+ StoredLabwareOffsetCreate,
+ DO_NOT_FILTER,
+ DoNotFilterType,
+)
router = LightRouter()
diff --git a/robot-server/robot_server/labware_offsets/store.py b/robot-server/robot_server/labware_offsets/store.py
index 5b97f799f36..7e3da8d3522 100644
--- a/robot-server/robot_server/labware_offsets/store.py
+++ b/robot-server/robot_server/labware_offsets/store.py
@@ -1,7 +1,6 @@
# noqa: D100
-import enum
-from typing import Final, Literal, TypeAlias, Iterator
+from typing import Iterator
from opentrons.protocol_engine.types import (
LabwareOffsetVector,
@@ -17,29 +16,12 @@
labware_offset_table,
labware_offset_location_sequence_components_table,
)
-from .models import StoredLabwareOffset
+from .models import StoredLabwareOffset, DoNotFilterType, DO_NOT_FILTER
import sqlalchemy
import sqlalchemy.exc
-
-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.
-"""
+from ._search_query_builder import SearchQueryBuilder
class LabwareOffsetStore:
@@ -89,114 +71,19 @@ def search(
# making it worse.
) -> list[StoredLabwareOffset]:
"""Return all matching labware offsets in order from oldest-added to newest."""
- filter_statement = (
- 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_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,
- )
- )
- .where(labware_offset_table.c.active == True) # noqa: E712
+ builder = (
+ SearchQueryBuilder()
+ .do_active_filter(True)
+ .do_id_filter(id_filter)
+ .do_definition_uri_filter(definition_uri_filter)
+ .do_on_addressable_area_filter(location_addressable_area_filter)
+ .do_on_module_filter(location_module_model_filter)
+ .do_on_labware_filter(location_definition_uri_filter)
)
- location_positive_filter_subquery: sqlalchemy.sql.selectable.Exists | None = (
- None
- )
- location_negative_filter_subqueries: list[sqlalchemy.sql.selectable.Exists] = []
-
- locations_2 = labware_offset_location_sequence_components_table.alias()
-
- def _query_or_build(
- query: sqlalchemy.sql.selectable.Exists | None,
- ) -> sqlalchemy.sql.selectable.Exists:
- if query is not None:
- return query
- return sqlalchemy.exists().where(
- locations_2.c.offset_id
- == labware_offset_location_sequence_components_table.c.offset_id
- )
-
- if id_filter is not DO_NOT_FILTER:
- filter_statement = filter_statement.where(
- labware_offset_table.c.offset_id == id_filter
- )
- if definition_uri_filter is not DO_NOT_FILTER:
- filter_statement = filter_statement.where(
- labware_offset_table.c.definition_uri == definition_uri_filter
- )
- if location_addressable_area_filter is not DO_NOT_FILTER:
- location_positive_filter_subquery = (
- _query_or_build(location_positive_filter_subquery)
- .where(locations_2.c.component_kind == "onAddressableArea")
- .where(
- locations_2.c.primary_component_value
- == location_addressable_area_filter
- )
- )
-
- if location_module_model_filter is not DO_NOT_FILTER:
- if location_module_model_filter is not None:
- location_positive_filter_subquery = (
- _query_or_build(location_positive_filter_subquery)
- .where(locations_2.c.component_kind == "onModule")
- .where(
- locations_2.c.primary_component_value
- == location_module_model_filter.value
- )
- )
- else:
- location_negative_filter_subqueries.append(
- sqlalchemy.exists()
- .where(
- locations_2.c.offset_id
- == labware_offset_location_sequence_components_table.c.offset_id
- )
- .where(locations_2.c.component_kind == "onModule")
- )
- if location_definition_uri_filter is not DO_NOT_FILTER:
- if location_definition_uri_filter is not None:
- location_positive_filter_subquery = (
- _query_or_build(location_positive_filter_subquery)
- .where(locations_2.c.component_kind == "onLabware")
- .where(
- locations_2.c.primary_component_value
- == location_definition_uri_filter
- )
- )
- else:
- location_negative_filter_subqueries.append(
- sqlalchemy.exists()
- .where(
- locations_2.c.offset_id
- == labware_offset_location_sequence_components_table.c.offset_id
- )
- .where(locations_2.c.component_kind == "onLabware")
- )
-
- if location_positive_filter_subquery is not None:
- filter_statement = filter_statement.where(location_positive_filter_subquery)
- for subq in location_negative_filter_subqueries:
- filter_statement = filter_statement.where(sqlalchemy.not_(subq))
-
- filter_statement = filter_statement.order_by(
- labware_offset_table.c.row_id
- ).order_by(labware_offset_location_sequence_components_table.c.sequence_ordinal)
+ query = builder.build_query()
with self._sql_engine.begin() as transaction:
- result = transaction.execute(filter_statement).all()
+ result = transaction.execute(query).all()
if len(result) == 0:
return []
@@ -204,32 +91,11 @@ def _query_or_build(
def delete(self, offset_id: str) -> StoredLabwareOffset:
"""Delete a labware offset by its ID. Return what was just deleted."""
+ builder = SearchQueryBuilder().do_id_filter(offset_id)
+ query = builder.build_query()
with self._sql_engine.begin() as transaction:
try:
- offset_rows = transaction.execute(
- 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,
- )
- )
- .where(labware_offset_table.c.offset_id == offset_id)
- ).all()
+ offset_rows = transaction.execute(query).all()
except sqlalchemy.exc.NoResultFound:
raise LabwareOffsetNotFoundError(bad_offset_id=offset_id) from None
if len(offset_rows) == 0:
diff --git a/robot-server/tests/labware_offsets/test_store.py b/robot-server/tests/labware_offsets/test_store.py
index 7f8c0a2406e..81bbd641ea1 100644
--- a/robot-server/tests/labware_offsets/test_store.py
+++ b/robot-server/tests/labware_offsets/test_store.py
@@ -16,10 +16,12 @@
from robot_server.labware_offsets.store import (
LabwareOffsetStore,
LabwareOffsetNotFoundError,
+)
+from robot_server.labware_offsets.models import (
+ StoredLabwareOffset,
DoNotFilterType,
DO_NOT_FILTER,
)
-from robot_server.labware_offsets.models import StoredLabwareOffset
@pytest.fixture