refactor: batch facet queries into a single round trip

This commit is contained in:
2026-07-26 15:49:58 +02:00
committed by d3vyce
parent c15aee22b0
commit 19a823dc8b
5 changed files with 298 additions and 95 deletions
+67 -27
View File
@@ -44,6 +44,7 @@ from ..types import (
)
from .search import (
SearchConfig,
apply_search_joins,
build_facets,
build_filter_by,
build_search_filters,
@@ -128,17 +129,6 @@ def _apply_joins(q: Any, joins: JoinType | None, outer_join: bool) -> Any:
return q
def _apply_search_joins(q: Any, search_joins: list[Any]) -> Any:
"""Apply relationship-based outer joins (from search/filter_by) to a query."""
seen: set[str] = set()
for join_rel in search_joins:
key = str(join_rel)
if key not in seen:
seen.add(key)
q = q.outerjoin(join_rel)
return q
class AsyncCrud(Generic[ModelType]):
"""Generic async CRUD operations for SQLAlchemy models.
@@ -265,12 +255,12 @@ class AsyncCrud(Generic[ModelType]):
cls: type[Self],
filter_by: dict[str, Any] | BaseModel | None,
facet_fields: Sequence[FacetFieldType] | None,
) -> tuple[list[Any], list[Any]]:
"""Normalize filter_by and return (filters, joins) to apply to the query."""
) -> tuple[dict[str, Any], list[Any]]:
"""Normalize filter_by and return ({facet_key: filter}, joins) to apply to the query."""
if isinstance(filter_by, BaseModel):
filter_by = filter_by.model_dump(exclude_none=True)
if not filter_by:
return [], []
return {}, []
resolved = cls._resolve_facet_fields(facet_fields)
return build_filter_by(filter_by, resolved or [])
@@ -281,8 +271,13 @@ class AsyncCrud(Generic[ModelType]):
facet_fields: Sequence[FacetFieldType] | None,
filters: list[Any],
search_joins: list[Any],
*,
include_facets: bool = True,
own_filters: dict[str, Any] | None = None,
) -> dict[str, list[Any]] | None:
"""Build facet filter_attributes, or return None if no facet fields configured."""
"""Build facet filter_attributes, or None if disabled/no facet fields configured."""
if not include_facets:
return None
resolved = cls._resolve_facet_fields(facet_fields)
if not resolved:
return None
@@ -292,6 +287,7 @@ class AsyncCrud(Generic[ModelType]):
resolved,
base_filters=filters,
base_joins=search_joins,
own_filters=own_filters,
)
@classmethod
@@ -475,6 +471,7 @@ class AsyncCrud(Generic[ModelType]):
default_page_size: int = 20,
max_page_size: int = 100,
include_total: bool = True,
include_facets: bool = True,
search: bool = True,
filter: bool = True,
order: bool = True,
@@ -490,6 +487,7 @@ class AsyncCrud(Generic[ModelType]):
default_page_size: Default ``items_per_page`` value.
max_page_size: Maximum ``items_per_page`` value.
include_total: Whether to include total count (not a query param).
include_facets: Whether to run facet queries (not a query param).
search: Enable search query parameters.
filter: Enable facet filter query parameters.
order: Enable order query parameters.
@@ -519,7 +517,10 @@ class AsyncCrud(Generic[ModelType]):
]
return cls._build_paginate_params(
pagination_params=pagination_params,
pagination_fixed={"include_total": include_total},
pagination_fixed={
"include_total": include_total,
"include_facets": include_facets,
},
dep_name=f"{cls.model.__name__}OffsetPaginateParams",
search=search,
filter=filter,
@@ -537,6 +538,7 @@ class AsyncCrud(Generic[ModelType]):
*,
default_page_size: int = 20,
max_page_size: int = 100,
include_facets: bool = True,
search: bool = True,
filter: bool = True,
order: bool = True,
@@ -551,6 +553,7 @@ class AsyncCrud(Generic[ModelType]):
Args:
default_page_size: Default ``items_per_page`` value.
max_page_size: Maximum ``items_per_page`` value.
include_facets: Whether to run facet queries (not a query param).
search: Enable search query parameters.
filter: Enable facet filter query parameters.
order: Enable order query parameters.
@@ -582,7 +585,7 @@ class AsyncCrud(Generic[ModelType]):
]
return cls._build_paginate_params(
pagination_params=pagination_params,
pagination_fixed={},
pagination_fixed={"include_facets": include_facets},
dep_name=f"{cls.model.__name__}CursorPaginateParams",
search=search,
filter=filter,
@@ -602,6 +605,7 @@ class AsyncCrud(Generic[ModelType]):
max_page_size: int = 100,
default_pagination_type: PaginationType = PaginationType.OFFSET,
include_total: bool = True,
include_facets: bool = True,
search: bool = True,
filter: bool = True,
order: bool = True,
@@ -618,6 +622,7 @@ class AsyncCrud(Generic[ModelType]):
max_page_size: Maximum ``items_per_page`` value.
default_pagination_type: Default pagination strategy.
include_total: Whether to include total count (not a query param).
include_facets: Whether to run facet queries (not a query param).
search: Enable search query parameters.
filter: Enable facet filter query parameters.
order: Enable order query parameters.
@@ -666,7 +671,10 @@ class AsyncCrud(Generic[ModelType]):
]
return cls._build_paginate_params(
pagination_params=pagination_params,
pagination_fixed={"include_total": include_total},
pagination_fixed={
"include_total": include_total,
"include_facets": include_facets,
},
dep_name=f"{cls.model.__name__}PaginateParams",
search=search,
filter=filter,
@@ -1271,6 +1279,7 @@ class AsyncCrud(Generic[ModelType]):
search_column: str | None = None,
order_fields: Sequence[OrderFieldType] | None = None,
facet_fields: Sequence[FacetFieldType] | None = None,
include_facets: bool = True,
filter_by: dict[str, Any] | BaseModel | None = None,
schema: type[BaseModel],
) -> OffsetPaginatedResponse[Any]:
@@ -1292,6 +1301,9 @@ class AsyncCrud(Generic[ModelType]):
search_column: Restrict search to a single column key.
order_fields: Fields allowed for sorting (overrides class default).
facet_fields: Columns to compute distinct values for (overrides class default)
include_facets: When ``False``, skip facet queries entirely;
``filter_attributes`` will be ``None``. Useful on pages 2..N
where the facet counts were already fetched on page 1.
filter_by: Dict of {column_key: value} to filter by declared facet fields.
Keys must match the column.key of a facet field. Scalar → equality,
list → IN clause. Raises InvalidFacetFilterError for unknown keys.
@@ -1304,7 +1316,6 @@ class AsyncCrud(Generic[ModelType]):
offset = (page - 1) * items_per_page
fb_filters, search_joins = cls._prepare_filter_by(filter_by, facet_fields)
filters.extend(fb_filters)
# Build search filters
if search:
@@ -1318,6 +1329,11 @@ class AsyncCrud(Generic[ModelType]):
filters.extend(search_filters)
search_joins.extend(new_search_joins)
# Facets combine these with each facet's own filter individually, so
# fb_filters is applied to the query below but excluded here.
facet_base_filters = list(filters)
filters.extend(fb_filters.values())
# Build query with joins
q = select(cls.model)
@@ -1325,11 +1341,11 @@ class AsyncCrud(Generic[ModelType]):
q = _apply_joins(q, joins, outer_join)
# Apply search joins (always outer joins for search)
q = _apply_search_joins(q, search_joins)
q = apply_search_joins(q, search_joins)
# Apply order joins (relation joins required for order_by field)
if order_joins:
q = _apply_search_joins(q, order_joins)
q = apply_search_joins(q, order_joins)
if filters:
q = q.where(and_(*filters))
@@ -1352,7 +1368,7 @@ class AsyncCrud(Generic[ModelType]):
count_q = _apply_joins(count_q, joins, outer_join)
# Apply search joins to count query
count_q = _apply_search_joins(count_q, search_joins)
count_q = apply_search_joins(count_q, search_joins)
if filters:
count_q = count_q.where(and_(*filters))
@@ -1372,7 +1388,12 @@ class AsyncCrud(Generic[ModelType]):
items: list[Any] = [schema.model_validate(item) for item in raw_items]
filter_attributes = await cls._build_filter_attributes(
session, facet_fields, filters, search_joins
session,
facet_fields,
facet_base_filters,
search_joins,
include_facets=include_facets,
own_filters=fb_filters,
)
search_columns = cls._resolve_search_columns(search_fields)
order_columns = cls._resolve_order_columns(order_fields)
@@ -1408,6 +1429,7 @@ class AsyncCrud(Generic[ModelType]):
search_column: str | None = None,
order_fields: Sequence[OrderFieldType] | None = None,
facet_fields: Sequence[FacetFieldType] | None = None,
include_facets: bool = True,
filter_by: dict[str, Any] | BaseModel | None = None,
schema: type[BaseModel],
) -> CursorPaginatedResponse[Any]:
@@ -1430,6 +1452,8 @@ class AsyncCrud(Generic[ModelType]):
search_column: Restrict search to a single column key.
order_fields: Fields allowed for sorting (overrides class default).
facet_fields: Columns to compute distinct values for (overrides class default).
include_facets: When ``False``, skip facet queries entirely;
``filter_attributes`` will be ``None``.
filter_by: Dict of {column_key: value} to filter by declared facet fields.
Keys must match the column.key of a facet field. Scalar → equality,
list → IN clause. Raises InvalidFacetFilterError for unknown keys.
@@ -1441,7 +1465,6 @@ class AsyncCrud(Generic[ModelType]):
filters = list(filters) if filters else []
fb_filters, search_joins = cls._prepare_filter_by(filter_by, facet_fields)
filters.extend(fb_filters)
if cls.cursor_column is None:
raise ValueError(
@@ -1473,6 +1496,11 @@ class AsyncCrud(Generic[ModelType]):
filters.extend(search_filters)
search_joins.extend(new_search_joins)
# Facets combine these with each facet's own filter individually, so
# fb_filters is applied to the query below but excluded here.
facet_base_filters = list(filters)
filters.extend(fb_filters.values())
# Build query
q = select(cls.model)
@@ -1480,11 +1508,11 @@ class AsyncCrud(Generic[ModelType]):
q = _apply_joins(q, joins, outer_join)
# Apply search joins (always outer joins)
q = _apply_search_joins(q, search_joins)
q = apply_search_joins(q, search_joins)
# Apply order joins (relation joins required for order_by field)
if order_joins:
q = _apply_search_joins(q, order_joins)
q = apply_search_joins(q, order_joins)
if filters:
q = q.where(and_(*filters))
@@ -1541,7 +1569,12 @@ class AsyncCrud(Generic[ModelType]):
items: list[Any] = [schema.model_validate(item) for item in items_page]
filter_attributes = await cls._build_filter_attributes(
session, facet_fields, filters, search_joins
session,
facet_fields,
facet_base_filters,
search_joins,
include_facets=include_facets,
own_filters=fb_filters,
)
search_columns = cls._resolve_search_columns(search_fields)
order_columns = cls._resolve_order_columns(order_fields)
@@ -1581,6 +1614,7 @@ class AsyncCrud(Generic[ModelType]):
search_column: str | None = ...,
order_fields: Sequence[OrderFieldType] | None = ...,
facet_fields: Sequence[FacetFieldType] | None = ...,
include_facets: bool = ...,
filter_by: dict[str, Any] | BaseModel | None = ...,
schema: type[BaseModel],
) -> OffsetPaginatedResponse[Any]: ...
@@ -1607,6 +1641,7 @@ class AsyncCrud(Generic[ModelType]):
search_column: str | None = ...,
order_fields: Sequence[OrderFieldType] | None = ...,
facet_fields: Sequence[FacetFieldType] | None = ...,
include_facets: bool = ...,
filter_by: dict[str, Any] | BaseModel | None = ...,
schema: type[BaseModel],
) -> CursorPaginatedResponse[Any]: ...
@@ -1632,6 +1667,7 @@ class AsyncCrud(Generic[ModelType]):
search_column: str | None = None,
order_fields: Sequence[OrderFieldType] | None = None,
facet_fields: Sequence[FacetFieldType] | None = None,
include_facets: bool = True,
filter_by: dict[str, Any] | BaseModel | None = None,
schema: type[BaseModel],
) -> OffsetPaginatedResponse[Any] | CursorPaginatedResponse[Any]:
@@ -1662,6 +1698,8 @@ class AsyncCrud(Generic[ModelType]):
order_fields: Fields allowed for sorting (overrides class default).
facet_fields: Columns to compute distinct values for (overrides
class default).
include_facets: When ``False``, skip facet queries entirely;
``filter_attributes`` will be ``None``.
filter_by: Dict of ``{column_key: value}`` to filter by declared
facet fields. Keys must match the ``column.key`` of a facet
field. Scalar → equality, list → IN clause. Raises
@@ -1692,6 +1730,7 @@ class AsyncCrud(Generic[ModelType]):
search_column=search_column,
order_fields=order_fields,
facet_fields=facet_fields,
include_facets=include_facets,
filter_by=filter_by,
schema=schema,
)
@@ -1714,6 +1753,7 @@ class AsyncCrud(Generic[ModelType]):
search_column=search_column,
order_fields=order_fields,
facet_fields=facet_fields,
include_facets=include_facets,
filter_by=filter_by,
schema=schema,
)
+83 -59
View File
@@ -1,12 +1,12 @@
"""Search utilities for AsyncCrud."""
import asyncio
import functools
from collections.abc import Sequence
from dataclasses import dataclass, replace
from typing import TYPE_CHECKING, Any, Literal
from sqlalchemy import String, and_, func, or_, select
from sqlalchemy import String, and_, distinct, func, or_, select
from sqlalchemy.dialects.postgresql import aggregate_order_by
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.orm.attributes import InstrumentedAttribute
@@ -181,6 +181,21 @@ def search_field_keys(fields: Sequence[SearchFieldType]) -> list[str]:
return facet_keys(fields)
def apply_search_joins(q: Any, joins: Sequence[Any]) -> Any:
"""Apply relationship-based outer joins (from search/filter_by/facets) to a query.
Deduplicates by relationship identity so a join used by several fields
(e.g. search + a facet on the same relation) is only applied once.
"""
seen: set[str] = set()
for rel in joins:
rel_key = str(rel)
if rel_key not in seen:
seen.add(rel_key)
q = q.outerjoin(rel)
return q
def facet_keys(facet_fields: Sequence[FacetFieldType]) -> list[str]:
"""Return a key for each facet field.
@@ -207,6 +222,7 @@ async def build_facets(
*,
base_filters: "list[ColumnElement[bool]] | None" = None,
base_joins: list[InstrumentedAttribute[Any]] | None = None,
own_filters: "dict[str, ColumnElement[bool]] | None" = None,
) -> dict[str, list[Any]]:
"""Return distinct values for each facet field, respecting current filters.
@@ -216,15 +232,24 @@ async def build_facets(
facet_fields: Columns or relationship tuples to facet on
base_filters: Filter conditions already applied to the main query (search + caller filters)
base_joins: Relationship joins already applied to the main query
own_filters: Map of facet key -> the ``filter_by`` condition for that
same key (if any). Excluded from that facet's own subquery so
filtering on a facet doesn't collapse its own value list down to
just the filtered value.
Returns:
Dict mapping column key to sorted list of distinct non-None values
"""
existing_join_keys: set[str] = {str(j) for j in (base_joins or [])}
if not facet_fields:
return {}
keys = facet_keys(facet_fields)
own_filters = own_filters or {}
async def _query_facet(field: FacetFieldType, key: str) -> tuple[str, list[Any]]:
scalars: list[Any] = []
enum_classes: dict[str, Any] = {}
for field, key in zip(facet_fields, keys):
if isinstance(field, tuple):
# Relationship chain: (User.role, Role.name) — last element is the column
rels = field[:-1]
@@ -235,51 +260,48 @@ async def build_facets(
col_type = column.property.columns[0].type
is_array = isinstance(col_type, ARRAY)
enum_classes[key] = getattr(col_type, "enum_class", None)
if is_array:
unnested = func.unnest(column).label(column.key)
q = select(unnested).select_from(model).distinct()
else:
q = select(column).select_from(model).distinct()
# Apply base joins (deduplicated) — needed here independently
seen_joins: set[str] = set()
for rel in base_joins or []:
rel_key = str(rel)
if rel_key not in seen_joins:
seen_joins.add(rel_key)
q = q.outerjoin(rel)
# Add any extra joins required by this facet field that aren't already applied
for rel in rels:
rel_key = str(rel)
if rel_key not in existing_join_keys and rel_key not in seen_joins:
seen_joins.add(rel_key)
q = q.outerjoin(rel)
if base_filters:
q = q.where(and_(*base_filters))
if is_array:
q = q.order_by(unnested)
else:
q = q.order_by(column)
result = await session.execute(q)
col_type = column.property.columns[0].type
enum_class = getattr(col_type, "enum_class", None)
values = [
row[0].name
if (enum_class is not None and isinstance(row[0], enum_class))
else row[0]
for row in result.all()
if row[0] is not None
filters = [
*(base_filters or []),
*(f for k, f in own_filters.items() if k != key),
]
return key, values
joins = [*(base_joins or []), *rels]
pairs = await asyncio.gather(
*[_query_facet(f, k) for f, k in zip(facet_fields, keys)]
)
return dict(pairs)
if is_array:
unnested = apply_search_joins(
select(func.unnest(column).label("v")).select_from(model), joins
)
if filters:
unnested = unnested.where(and_(*filters))
unnested_sq = unnested.subquery()
v = unnested_sq.c.v
agg = (
select(func.array_agg(aggregate_order_by(distinct(v), v)))
.select_from(unnested_sq)
.where(v.isnot(None))
)
else:
agg = apply_search_joins(
select(
func.array_agg(aggregate_order_by(distinct(column), column))
).select_from(model),
joins,
)
agg = agg.where(and_(*filters, column.isnot(None)))
scalars.append(agg.scalar_subquery().label(key))
row = (await session.execute(select(*scalars))).one()
facets: dict[str, list[Any]] = {}
for key, values in zip(keys, row):
enum_class = enum_classes[key]
facets[key] = [
v.name if (enum_class is not None and isinstance(v, enum_class)) else v
for v in (values or [])
]
return facets
_EQUALITY_TYPES = (String, Integer, Numeric, Date, DateTime, Time, Enum, Uuid)
@@ -301,7 +323,7 @@ def _coerce_bool(value: Any) -> bool:
def build_filter_by(
filter_by: dict[str, Any],
facet_fields: Sequence[FacetFieldType],
) -> tuple["list[ColumnElement[bool]]", list[InstrumentedAttribute[Any]]]:
) -> tuple["dict[str, ColumnElement[bool]]", list[InstrumentedAttribute[Any]]]:
"""Translate a {column_key: value} dict into SQLAlchemy filter conditions.
Args:
@@ -309,7 +331,9 @@ def build_filter_by(
facet_fields: Declared facet fields to validate keys against
Returns:
Tuple of (filter_conditions, joins_needed)
Tuple of ({facet_key: filter_condition}, joins_needed). One filter
condition per key, so callers can identify (and exclude) a facet's
own filter when computing that facet's distinct values.
Raises:
InvalidFacetFilterError: If a key in filter_by is not a declared facet field
@@ -327,7 +351,7 @@ def build_filter_by(
index[key] = (column, rels)
valid_keys = set(index)
filters: list[ColumnElement[bool]] = []
filters: dict[str, ColumnElement[bool]] = {}
joins: list[InstrumentedAttribute[Any]] = []
added_join_keys: set[str] = set()
@@ -347,14 +371,14 @@ def build_filter_by(
if isinstance(col_type, Boolean):
coerce = _coerce_bool
if isinstance(value, list):
filters.append(column.in_([coerce(v) for v in value]))
filters[key] = column.in_([coerce(v) for v in value])
else:
filters.append(column == coerce(value))
filters[key] = column == coerce(value)
elif isinstance(col_type, ARRAY):
if isinstance(value, list):
filters.append(column.overlap(value))
filters[key] = column.overlap(value)
else:
filters.append(column.any(value))
filters[key] = column.any(value)
elif isinstance(col_type, Enum):
enum_class = col_type.enum_class
if enum_class is not None:
@@ -365,19 +389,19 @@ def build_filter_by(
return enum_class[v] # lookup by name: "PENDING", "RED"
if isinstance(value, list):
filters.append(column.in_([_coerce_enum(v) for v in value]))
filters[key] = column.in_([_coerce_enum(v) for v in value])
else:
filters.append(column == _coerce_enum(value))
filters[key] = column == _coerce_enum(value)
else: # pragma: no cover
if isinstance(value, list):
filters.append(column.in_(value))
filters[key] = column.in_(value)
else:
filters.append(column == value)
filters[key] = column == value
elif isinstance(col_type, _EQUALITY_TYPES):
if isinstance(value, list):
filters.append(column.in_(value))
filters[key] = column.in_(value)
else:
filters.append(column == value)
filters[key] = column == value
else:
raise UnsupportedFacetTypeError(key, type(col_type).__name__)