add docstring etc

Author: jsstevenson

src/anyvar/storage/base.py

@@ -1,5 +1,7 @@
 """Provide PostgreSQL-based storage implementation."""
 
+import base64
+import json
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from dataclasses import dataclass
@@ -181,6 +183,30 @@ def delete_annotation(self, annotation: metadata.Annotation) -> None:
             depended upon by another object
         """
 
+    @staticmethod
+    def _encode_search_cursor(start: int, allele_id: str) -> str:
+        """Create cursor for search
+
+        :param start: start value for next row
+        :param allele_id: ID for next row
+        :return: cursor to use to fetch next page
+        """
+        raw = json.dumps(
+            {"start": start, "id": allele_id}, separators=(",", ":")
+        ).encode()
+        return base64.urlsafe_b64encode(raw).decode()
+
+    @staticmethod
+    def _decode_search_cursor(cursor: str) -> tuple[int, str]:
+        """Decode cursor for getting next page during search
+
+        :param cursor: opaque key included with previous result
+        :return: start and ID values indicating the first row of the next page
+        """
+        raw = base64.urlsafe_b64decode(cursor.encode())
+        obj = json.loads(raw)
+        return int(obj["start"]), str(obj["id"])
+
     @abstractmethod
     def search_alleles(
         self,
@@ -196,6 +222,9 @@ def search_alleles(
         the RefGet SequenceReference accession (`SQ.*`). Both `start` and `stop` are
         inclusive and represent inter-residue positions.
 
+        Uses keyset pagination, meaning that altering the page size while looping through
+        successive cursors will effectively nullify the search loop.
+
         Currently, any variation which overlaps the queried region is returned.
 
         Todo (see Issue #338):
@@ -209,8 +238,8 @@ def search_alleles(
         :param refget_accession: refget accession (e.g. `"SQ.IW78mgV5Cqf6M24hy52hPjyyo5tCCd86"`)
         :param start: Inclusive, inter-residue start position of the interval
         :param stop: Inclusive, inter-residue end position of the interval
-        :param page_size:
-        :param cursor:
-        :return:
+        :param page_size: Max # of results to return
+        :param cursor: Opaque key indicating start location for query in pagination
+        :return: Results page including variants and a cursor for next result page, if available
         :raise InvalidSearchParamsError: if above search param requirements are violated
         """

src/anyvar/storage/postgres.py

@@ -1,6 +1,5 @@
 """Provide PostgreSQL-based storage implementation."""
 
-import base64
 import json
 import logging
 from collections import defaultdict
@@ -149,7 +148,6 @@ def get_objects(
                         )
                     )
                     .where(orm.Allele.id.in_(object_ids_list))
-                    .limit(self.MAX_ROWS)
                 )
                 db_objects = session.scalars(stmt).all()
             elif object_type is vrs_models.SequenceLocation:
@@ -158,15 +156,12 @@ def get_objects(
                     select(orm.Location)
                     .options(joinedload(orm.Location.sequence_reference))
                     .where(orm.Location.id.in_(object_ids_list))
-                    .limit(self.MAX_ROWS)
                 )
                 db_objects = session.scalars(stmt).all()
             elif object_type is vrs_models.SequenceReference:
                 # Get sequence references
-                stmt = (
-                    select(orm.SequenceReference)
-                    .where(orm.SequenceReference.id.in_(object_ids_list))
-                    .limit(self.MAX_ROWS)
+                stmt = select(orm.SequenceReference).where(
+                    orm.SequenceReference.id.in_(object_ids_list)
                 )
 
                 db_objects = session.scalars(stmt).all()
@@ -382,29 +377,17 @@ def search_alleles(
         :param refget_accession: refget accession (e.g. `"SQ.IW78mgV5Cqf6M24hy52hPjyyo5tCCd86"`)
         :param start: Inclusive, inter-residue start position of the interval
         :param stop: Inclusive, inter-residue end position of the interval
-        :param page_size:
-        :param cursor:
-        :return:
+        :param page_size: Max # of results to return
+        :param cursor: Opaque key indicating start location for query in pagination
+        :return: Results page including variants and a cursor for next result page, if available
         :raise InvalidSearchParamsError: if above search param requirements are violated
         """
-
-        def _encode_cursor(start: int, allele_id: str) -> str:
-            raw = json.dumps(
-                {"start": start, "id": allele_id}, separators=(",", ":")
-            ).encode()
-            return base64.urlsafe_b64encode(raw).decode()
-
-        def _decode_cursor(cursor: str) -> tuple[int, str]:
-            raw = base64.urlsafe_b64decode(cursor.encode())
-            obj = json.loads(raw)
-            return int(obj["start"]), str(obj["id"])
-
         if start < 0 or stop < 0 or start > stop:
             raise InvalidSearchParamsError
         seek_start: int | None = None
         seek_id: str | None = None
         if cursor:
-            seek_start, seek_id = _decode_cursor(cursor)
+            seek_start, seek_id = self._decode_search_cursor(cursor)
 
         with self.session_factory() as session:
             stmt = (
@@ -440,5 +423,5 @@ def _decode_cursor(cursor: str) -> tuple[int, str]:
             if not page_db:
                 return AlleleSearchPage(items=[], next_cursor=None)
             last = page_db[-1]
-            next_cursor = _encode_cursor(last.location.start, last.id)
+            next_cursor = self._encode_search_cursor(last.location.start, last.id)
             return AlleleSearchPage(items=items, next_cursor=next_cursor)

src/anyvar/storage/snowflake.py

@@ -14,7 +14,17 @@
 from ga4gh.vrs import models as vrs_models
 from snowflake.sqlalchemy import MergeInto
 from snowflake.sqlalchemy.snowdialect import SnowflakeDialect
-from sqlalchemy import String, column, create_engine, delete, insert, select, text
+from sqlalchemy import (
+    String,
+    and_,
+    column,
+    create_engine,
+    delete,
+    insert,
+    or_,
+    select,
+    text,
+)
 from sqlalchemy.engine.url import URL
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.compiler import compiles
@@ -27,6 +37,7 @@
 from anyvar.core import objects as anyvar_objects
 from anyvar.storage import orm
 from anyvar.storage.base import (
+    AlleleSearchPage,
     DataIntegrityError,
     IncompleteVrsObjectError,
     InvalidSearchParamsError,
@@ -616,13 +627,18 @@ def search_alleles(
         refget_accession: str,
         start: int,
         stop: int,
-    ) -> Iterable[vrs_models.Allele]:
+        page_size: int = 1000,
+        cursor: str | None = None,
+    ) -> AlleleSearchPage:
         """Find all Alleles that are located within the specified interval.
 
         The interval is the closed range [start, stop] on the sequence identified by
         the RefGet SequenceReference accession (`SQ.*`). Both `start` and `stop` are
         inclusive and represent inter-residue positions.
 
+        Uses keyset pagination, meaning that altering the page size while looping through
+        successive cursors will effectively nullify the search loop.
+
         Currently, any variation which overlaps the queried region is returned.
 
         Todo (see Issue #338):
@@ -636,12 +652,19 @@ def search_alleles(
         :param refget_accession: refget accession (e.g. `"SQ.IW78mgV5Cqf6M24hy52hPjyyo5tCCd86"`)
         :param start: Inclusive, inter-residue start position of the interval
         :param stop: Inclusive, inter-residue end position of the interval
-        :return: an iterable of matching VRS alleles
+        :param page_size: Max # of results to return
+        :param cursor: Opaque key indicating start location for query in pagination
+        :return: Results page including variants and a cursor for next result page, if available
         :raise InvalidSearchParamsError: if above search param requirements are violated
         """
         if start < 0 or stop < 0 or start > stop:
             raise InvalidSearchParamsError
 
+        seek_start: int | None = None
+        seek_id: str | None = None
+        if cursor:
+            seek_start, seek_id = self._decode_search_cursor(cursor)
+
         with self.session_factory() as session:
             # Query alleles with overlapping locations
             # NOTE: this is any overlap, not containment.
@@ -659,13 +682,26 @@ def search_alleles(
                     orm.Location.start <= stop,
                     orm.Location.end >= start,
                 )
-                .limit(self.MAX_ROWS)
+                .order_by(orm.Location.start, orm.Allele.id)
+                .limit(page_size)
             )
-            db_alleles = session.scalars(stmt).all()
 
-            return [
-                mapper_registry.from_db_entity(db_allele) for db_allele in db_alleles
-            ]
+            # seek predicate -- assumes ORDER BY location.start ASC, allele.id ASC
+            if seek_start is not None and seek_id is not None:
+                stmt = stmt.where(
+                    or_(
+                        orm.Location.start > seek_start,
+                        and_(orm.Location.start == seek_start, orm.Allele.id > seek_id),
+                    )
+                )
+
+            page_db = session.scalars(stmt).all()
+            items = [mapper_registry.from_db_entity(a) for a in page_db]
+            if not page_db:
+                return AlleleSearchPage(items=[], next_cursor=None)
+            last = page_db[-1]
+            next_cursor = self._encode_search_cursor(last.location.start, last.id)
+            return AlleleSearchPage(items=items, next_cursor=next_cursor)
 
 
 @compiles(Insert, "snowflake")

tests/unit/storage/storage_test_funcs.py

@@ -425,30 +425,30 @@ def run_search_alleles(
     result = storage.search_alleles(
         rle.location.sequenceReference.refgetAccession, 36561660, 36561665
     )
-    assert result == [rle]
+    assert result.items == [rle]
     result = storage.search_alleles(
         egfr_variant.location.sequenceReference.refgetAccession, 55174010, 140753340
     )
-    sorted(result, key=lambda a: a.id)
-    assert result == [egfr_variant, braf_variant]
+    sorted(result.items, key=lambda a: a.id)
+    assert result.items == [egfr_variant, braf_variant]
 
     # result partially overlaps with interval
     result = storage.search_alleles(
         rle.location.sequenceReference.refgetAccession, 36561662, 36561665
     )
-    assert result == [rle]
+    assert result.items == [rle]
     assert storage.search_alleles(
         rle.location.sequenceReference.refgetAccession, 36561662, 36561663
-    ) == [rle]
+    ).items == [rle]
 
     # position ranges are inclusive
     result = storage.search_alleles(
         braf_variant.location.sequenceReference.refgetAccession, 140753335, 140753336
     )
-    assert result == [braf_variant]
+    assert result.items == [braf_variant]
 
     # handle unrecognized accession
-    assert storage.search_alleles("SQ.unknown-sequence", 1, 10) == []
+    assert storage.search_alleles("SQ.unknown-sequence", 1, 10).items == []
 
     # handle invalid params
     with pytest.raises(InvalidSearchParamsError):
@@ -469,15 +469,32 @@ def run_search_alleles(
     assert (
         storage.search_alleles(
             long_ins.location.sequenceReference.refgetAccession, 10599292, 10599295
-        )
+        ).items
         == []
     )
     assert (
         storage.search_alleles(
             rle_del.location.sequenceReference.refgetAccession, 905, 910
-        )
+        ).items
         == []
     )
     assert storage.search_alleles(
         rle_del.location.sequenceReference.refgetAccession, 904, 910
-    ) == [rle_del]
+    ).items == [rle_del]
+
+    # test pagination
+    result = storage.search_alleles(
+        egfr_variant.location.sequenceReference.refgetAccession,
+        55174010,
+        140753340,
+        page_size=1,
+    )
+    assert result.items == [egfr_variant]
+    result = storage.search_alleles(
+        egfr_variant.location.sequenceReference.refgetAccession,
+        55174010,
+        140753340,
+        page_size=1,
+        cursor=result.next_cursor,
+    )
+    assert result.items == [braf_variant]

tests/unit/storage/test_postgres.py

@@ -14,7 +14,6 @@
     run_incomplete_objects_error,
     run_mappings_crud,
     run_objects_raises_integrityerror,
-    run_query_max_rows,
     run_search_alleles,
     run_sequencelocations_crud,
     run_sequencereferences_crud,
@@ -49,15 +48,6 @@ def test_db_lifecycle(
     run_db_lifecycle(storage, validated_vrs_alleles)
 
 
-@pytest.mark.ci_ok
-def test_query_max_rows(
-    monkeypatch,
-    postgres_storage: PostgresObjectStore,
-    focus_alleles: tuple[models.Allele, models.Allele, models.Allele],
-):
-    run_query_max_rows(monkeypatch, postgres_storage, focus_alleles)
-
-
 @pytest.mark.ci_ok
 def test_alleles_crud(
     postgres_storage: PostgresObjectStore,