From d715f4c3c9680f0865bbce6d91aeaacfb2aec198 Mon Sep 17 00:00:00 2001 From: Seth Foster Date: Thu, 30 Jan 2025 15:49:53 -0500 Subject: [PATCH] refactor for complexity lint --- .../labware_offsets/_search_query_builder.py | 177 ++++++++++++++++++ .../robot_server/labware_offsets/models.py | 23 ++- .../robot_server/labware_offsets/router.py | 9 +- .../robot_server/labware_offsets/store.py | 166 ++-------------- .../tests/labware_offsets/test_store.py | 4 +- 5 files changed, 224 insertions(+), 155 deletions(-) create mode 100644 robot-server/robot_server/labware_offsets/_search_query_builder.py 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