"""Keyset and offset pagination helpers.
Keyset pagination here is *composite*: the cursor encodes the
boundary row's value for every ORDER BY key (the sort fields, then
the pk tiebreak), and the continuation predicate is the
lexicographic expansion over those keys::
(k1 > v1) OR (k1 = v1 AND k2 > v2) OR ...
A cursor on the pk alone is only correct when the pk *is* the sort
order; with ``ORDER BY name`` a bare ``WHERE id > cursor`` selects a
window unrelated to the name ordering and pages skip / duplicate
rows. :func:`run_keyset_query` therefore takes the
:class:`~fsh_lib.ordering.OrderKey` list that
:func:`~fsh_lib.ordering.apply_ordering` resolved, appends the pk
tiebreak (making the order total), and builds the matching
predicate.
NULL boundary values follow Postgres's default placement (ASC =
NULLS LAST, DESC = NULLS FIRST), which is what the plain
``.asc()`` / ``.desc()`` ordering emits.
"""
from __future__ import annotations
import base64
import binascii
import datetime
import decimal
import enum
import json
import uuid
from typing import TYPE_CHECKING, Any
from sqlalchemy import and_, false, func, or_, select
from fsh_lib.errors import FieldError
from fsh_lib.ordering import OrderKey
if TYPE_CHECKING:
from collections.abc import Sequence
from sqlalchemy import ColumnElement, Select
from sqlalchemy.ext.asyncio import AsyncSession
__all__ = [
"InvalidCursorError",
"apply_offset_pagination",
"decode_cursor",
"encode_cursor",
"keyset_predicate",
"run_keyset_query",
]
[docs]
class InvalidCursorError(FieldError):
"""A cursor that doesn't decode to this route's key shape.
Only arises from tampering or a cursor minted by an older
deploy with a different sort surface -- either way the right
response is a 422 on the ``cursor`` field, not a 500.
"""
def __init__(self) -> None:
"""Build the field-error envelope for the cursor field."""
super().__init__(("body", "cursor"), "Invalid pagination cursor.")
def _to_wire(value: Any) -> Any:
"""JSON-encodable form of one boundary value."""
if value is None or isinstance(value, (bool, int, float, str)):
return value
if isinstance(value, enum.Enum):
return _to_wire(value.value)
if isinstance(value, (datetime.datetime, datetime.date, datetime.time)):
return value.isoformat()
# uuid.UUID, decimal.Decimal, and anything else with a faithful
# string form -- _coerce parses it back via the column type.
return str(value)
[docs]
def encode_cursor(values: Sequence[Any]) -> str:
"""Encode boundary *values* (sort keys + pk) as an opaque token."""
payload = json.dumps(
[_to_wire(v) for v in values],
separators=(",", ":"),
)
return base64.urlsafe_b64encode(payload.encode()).decode().rstrip("=")
#: Wire-string parsers for types whose constructor isn't the parse
#: (e.g. ``datetime("2026-...")`` is not a thing); anything else
#: falls back to calling the type on the wire value.
_PARSERS: dict[type, Any] = {
datetime.datetime: datetime.datetime.fromisoformat,
datetime.date: datetime.date.fromisoformat,
datetime.time: datetime.time.fromisoformat,
uuid.UUID: uuid.UUID,
decimal.Decimal: decimal.Decimal,
}
def _coerce(value: Any, expr: ColumnElement) -> Any:
"""Parse one wire value back to *expr*'s python type."""
if value is None:
return None
try:
py_type = expr.type.python_type
except NotImplementedError:
return value
if isinstance(value, py_type):
return value
parser = _PARSERS.get(py_type, py_type)
try:
return parser(value)
except (TypeError, ValueError, decimal.InvalidOperation) as err:
raise InvalidCursorError from err
[docs]
def decode_cursor(cursor: str, keys: Sequence[OrderKey]) -> list[Any]:
"""Decode *cursor* into one boundary value per key.
Raises:
InvalidCursorError: The token isn't base64 JSON, or its
arity doesn't match this route's key shape.
"""
try:
padded = cursor + "=" * (-len(cursor) % 4)
raw = json.loads(base64.urlsafe_b64decode(padded.encode()))
except (ValueError, binascii.Error) as err:
raise InvalidCursorError from err
if not isinstance(raw, list) or len(raw) != len(keys):
raise InvalidCursorError
return [
_coerce(value, key.expr) for value, key in zip(raw, keys, strict=True)
]
def _is_nullable(col: ColumnElement) -> bool:
"""Whether NULLs can appear under *col* -- assume yes if unknown."""
try:
return bool(col.nullable)
except AttributeError:
return True
def _level_after(key: OrderKey, value: Any) -> ColumnElement:
"""Rows strictly after *value* at this key, Postgres NULL placement."""
col = key.expr
if key.direction == "desc":
# DESC = NULLS FIRST: after a null boundary comes every
# non-null; after a value, only smaller values.
if value is None:
return col.is_not(None)
return col < value
# ASC = NULLS LAST: nothing follows a null boundary at this
# level; after a value come greater values and then the nulls.
if value is None:
return false()
if not _is_nullable(col):
return col > value
return or_(col > value, col.is_(None))
def _level_equal(key: OrderKey, value: Any) -> ColumnElement:
"""Rows tied with *value* at this key (NULL-safe)."""
col = key.expr
return col.is_(None) if value is None else col == value
[docs]
def keyset_predicate(
keys: Sequence[OrderKey],
values: Sequence[Any],
) -> ColumnElement:
"""Lexicographic continuation predicate over *keys*.
``OR`` over key positions: tied on every earlier key AND
strictly after at this one. The final key is the non-null
unique pk tiebreak, so the predicate always makes progress.
"""
branches = []
for index, (key, value) in enumerate(zip(keys, values, strict=True)):
level = [
_level_equal(k, v)
for k, v in zip(keys[:index], values[:index], strict=True)
]
level.append(_level_after(key, value))
branches.append(and_(*level))
return or_(*branches)
[docs]
async def run_keyset_query(
db: AsyncSession,
stmt: Select,
model: type,
cursor: str | None,
cursor_field: str,
page_size: int,
max_page_size: int,
order_keys: Sequence[OrderKey] = (),
) -> tuple[list[Any], str | None, int]:
"""Execute *stmt* as one keyset page and mint the next cursor.
Appends the pk tiebreak to the ORDER BY (totalizing whatever
*order_keys* applied), applies the continuation predicate when
a cursor is given, selects the boundary expressions alongside
the entity (so a relationship sort's value is available without
a lazy load), and fetches ``page_size + 1`` rows to detect
whether more exist.
Args:
db: Async session to execute on.
stmt: SELECT with filters / search / ordering applied.
model: The listed model class (provides the pk column).
cursor: Opaque continuation token, or ``None`` for page 1.
cursor_field: Name of the pk tiebreak column.
page_size: Requested page size.
max_page_size: Hard ceiling on *page_size*.
order_keys: The :class:`~fsh_lib.ordering.OrderKey` list
:func:`~fsh_lib.ordering.apply_ordering` returned --
must match the ORDER BY already on *stmt*.
Returns:
``(items, next_cursor, effective_page_size)``.
Raises:
InvalidCursorError: *cursor* doesn't decode against this
route's key shape.
"""
pk_col = getattr(model, cursor_field)
keys = [
*order_keys,
OrderKey(field=cursor_field, direction="asc", expr=pk_col),
]
effective_page_size = min(page_size, max_page_size)
stmt = stmt.order_by(pk_col.asc())
if cursor is not None:
stmt = stmt.where(keyset_predicate(keys, decode_cursor(cursor, keys)))
stmt = stmt.add_columns(*[k.expr for k in keys])
stmt = stmt.limit(effective_page_size + 1)
result = await db.execute(stmt)
rows = result.all()
has_more = len(rows) > effective_page_size
rows = rows[:effective_page_size]
items = [row[0] for row in rows]
next_cursor = (
encode_cursor(list(rows[-1])[1:]) if has_more and rows else None
)
return items, next_cursor, effective_page_size