stash

Author: jsstevenson

src/anyvar/restapi/schema.py

@@ -296,6 +296,7 @@ class SearchResponse(BaseModel):
     """Describe response for the GET /search endpoint"""
 
     variations: list[models.Variation]
+    next_cursor: str | None
 
 
 class RunStatusResponse(BaseModel):

src/anyvar/restapi/search_router.py

@@ -36,6 +36,8 @@ def search_variations(
         int,
         Query(..., description="End position for genomic region", examples=[2781758]),
     ],
+    page_size: int = Query(1000, ge=1, le=10000),
+    cursor: str | None = Query(None, description="Opaque pagination cursor"),
 ) -> SearchResponse:
     """Perform genomic coordinate-based search over all registered variations."""
     av: AnyVar = request.app.state.anyvar
@@ -50,15 +52,18 @@ def search_variations(
             detail="Unable to dereference provided accession ID",
         ) from e
 
-    alleles = []
-    if ga4gh_id:
-        try:
-            refget_accession = ga4gh_id.split("ga4gh:")[-1]
-            alleles = av.object_store.search_alleles(refget_accession, start, end)
-        except NotImplementedError as e:
-            raise HTTPException(
-                status_code=HTTPStatus.NOT_IMPLEMENTED,
-                detail="Search not implemented for current storage backend",
-            ) from e
+    if not ga4gh_id:
+        return SearchResponse(variations=[], next_cursor=None)
 
-    return SearchResponse(variations=alleles)
+    try:
+        refget_accession = ga4gh_id.split("ga4gh:")[-1]
+        page = av.object_store.search_alleles(
+            refget_accession, start, end, page_size=page_size, cursor=cursor
+        )
+    except NotImplementedError as e:
+        raise HTTPException(
+            status_code=HTTPStatus.NOT_IMPLEMENTED,
+            detail="Search not implemented for current storage backend",
+        ) from e
+
+    return SearchResponse(variations=page.items, next_cursor=page.next_cursor)

src/anyvar/storage/base.py

@@ -2,6 +2,7 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
+from dataclasses import dataclass
 
 from ga4gh.vrs import models as vrs_models
 
@@ -29,6 +30,14 @@ class InvalidSearchParamsError(StorageError):
     """Raise if search params violate specified logical constraints"""
 
 
+@dataclass(frozen=True)
+class AlleleSearchPage:
+    """Return object for implementing keyset pagination in allele search"""
+
+    items: list[vrs_models.Allele]
+    next_cursor: str | None
+
+
 class Storage(ABC):
     """Abstract base class for interacting with storage backends."""
 
@@ -174,7 +183,9 @@ def search_alleles(
         refget_accession: str,
         start: int,
         stop: int,
-    ) -> list[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
@@ -194,7 +205,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
-        :return: a list of matching VRS alleles
+        :param page_size:
+        :param cursor:
+        :return:
         :raise InvalidSearchParamsError: if above search param requirements are violated
-
         """

src/anyvar/storage/postgres.py

@@ -1,12 +1,13 @@
 """Provide PostgreSQL-based storage implementation."""
 
+import base64
 import json
 import logging
 from collections import defaultdict
 from collections.abc import Iterable
 
 from ga4gh.vrs import models as vrs_models
-from sqlalchemy import create_engine, delete, func, select
+from sqlalchemy import and_, create_engine, delete, func, or_, select
 from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import joinedload, sessionmaker
@@ -15,6 +16,7 @@
 from anyvar.core import objects as anyvar_objects
 from anyvar.storage import orm
 from anyvar.storage.base import (
+    AlleleSearchPage,
     DataIntegrityError,
     IncompleteVrsObjectError,
     InvalidSearchParamsError,
@@ -33,10 +35,6 @@ class PostgresObjectStore(Storage):
     with object mapping to convert between VRS models and database entities.
     """
 
-    # temporary cap on max # of rows that can be returned by a single SQL query
-    # issue 295 should convert this to a batch size parameter
-    MAX_ROWS = 100
-
     _VRS_OBJECT_INSERT_ORDER: list[str] = [  # noqa: RUF012
         orm.SequenceReference.__name__,
         orm.Location.__name__,
@@ -51,6 +49,7 @@ def __init__(self, db_url: str, *args, **kwargs) -> None:
         self.db_url = db_url
         self.engine = create_engine(db_url)
         self.session_factory = sessionmaker(bind=self.engine)
+        self.batch_size = kwargs.get("batch_size", 1000)
         create_tables(self.db_url)
 
     def close(self) -> None:
@@ -133,8 +132,10 @@ def get_objects(
         :return: iterable collection of VRS objects matching given IDs
         """
         object_ids_list = list(object_ids)
-        results = []
+        if not object_ids_list:
+            return []
 
+        results = []
         with self.session_factory() as session:
             if object_type is vrs_models.Allele:
                 # Get alleles with eager loading
@@ -146,7 +147,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:
@@ -155,15 +155,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()
             else:
@@ -353,7 +350,9 @@ def search_alleles(
         refget_accession: str,
         start: int,
         stop: int,
-    ) -> list[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
@@ -373,16 +372,29 @@ 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: a list of matching VRS alleles
+        :return:
         :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)
 
         with self.session_factory() as session:
-            # Query alleles with overlapping locations
-            # NOTE: this is any overlap, not containment.
             stmt = (
                 select(orm.Allele)
                 .options(
@@ -398,10 +410,25 @@ def search_alleles(
                         func.int8range(start, stop, "[]")
                     ),
                 )
-                .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
-            ]
+            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 = _encode_cursor(last.location.start, last.id)
+            return AlleleSearchPage(items=items, next_cursor=next_cursor)

src/anyvar/storage/snowflake.py

@@ -616,7 +616,7 @@ def search_alleles(
         refget_accession: str,
         start: int,
         stop: int,
-    ) -> list[vrs_models.Allele]:
+    ) -> Iterable[vrs_models.Allele]:
         """Find all Alleles that are located within the specified interval.
 
         The interval is the closed range [start, stop] on the sequence identified by
@@ -636,9 +636,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
-        :return: a list of matching VRS alleles
+        :return: an iterable of matching VRS alleles
         :raise InvalidSearchParamsError: if above search param requirements are violated
-
         """
         if start < 0 or stop < 0 or start > stop:
             raise InvalidSearchParamsError