diff --git a/src/onegov/agency/api.py b/src/onegov/agency/api.py index 4c35a00861..1c937c528c 100644 --- a/src/onegov/agency/api.py +++ b/src/onegov/agency/api.py @@ -7,6 +7,7 @@ from onegov.agency.collections import PaginatedAgencyCollection from onegov.agency.collections import PaginatedMembershipCollection from onegov.agency.forms.person import AuthenticatedPersonMutationForm +from onegov.api import AdjacencyListApiEndpoint from onegov.api import ApiEndpoint, ApiInvalidParamException from onegov.api.utils import is_authorized from onegov.gis import Coordinates @@ -213,7 +214,9 @@ def apply_changes( do_report_person_change(item, form.meta.request, form) -class AgencyApiEndpoint(ApiEndpoint['ExtendedAgency', int], ApisMixin): +class AgencyApiEndpoint( + AdjacencyListApiEndpoint['ExtendedAgency', int], ApisMixin +): request: CoreRequest app: AgencyApp endpoint = 'agencies' diff --git a/src/onegov/api/__init__.py b/src/onegov/api/__init__.py index e80fec1e72..c3af6608b8 100644 --- a/src/onegov/api/__init__.py +++ b/src/onegov/api/__init__.py @@ -1,10 +1,12 @@ from __future__ import annotations from onegov.api.integration import ApiApp +from onegov.api.models import AdjacencyListApiEndpoint from onegov.api.models import ApiEndpoint, ApiInvalidParamException from onegov.api.models import log __all__ = [ + 'AdjacencyListApiEndpoint', 'ApiApp', 'ApiInvalidParamException', 'ApiEndpoint', diff --git a/src/onegov/api/models.py b/src/onegov/api/models.py index af9b153c9c..1c47dad9eb 100644 --- a/src/onegov/api/models.py +++ b/src/onegov/api/models.py @@ -25,6 +25,7 @@ from collections.abc import Callable, Collection, Iterator, Mapping from onegov.core import Framework from onegov.core.collection import PKType + from onegov.core.orm.abstract import AdjacencyList from onegov.core.request import CoreRequest from onegov.form import Form from sqlalchemy.orm import DeclarativeBase, Query, Session @@ -492,6 +493,26 @@ def __link_alias__(self) -> str: ) +class AdjacencyListApiEndpoint[L: AdjacencyList, IdT: PKType]( + ApiEndpoint[L, IdT] +): + """ An API endpoint for models deriving from :class:`AdjacencyList`. + + Preloads the ancestors of the whole batch, so building each item's link + (which renders its path by walking the parent chain) doesn't emit a query + per ancestor (N+1). + + """ + + @property + def batch(self) -> dict[ApiEndpointItem[L, IdT], L]: + result = super().batch + items = tuple(result.values()) + if items: + items[0].preload_ancestors(self.session, items) + return result + + class ApiEndpointCollection: """ A collection of all available API endpoints. """ diff --git a/src/onegov/core/orm/abstract/adjacency_list.py b/src/onegov/core/orm/abstract/adjacency_list.py index 85f64d8191..0d1d8e2e78 100644 --- a/src/onegov/core/orm/abstract/adjacency_list.py +++ b/src/onegov/core/orm/abstract/adjacency_list.py @@ -16,14 +16,14 @@ validates, Mapped ) -from sqlalchemy.orm.attributes import get_history +from sqlalchemy.orm.attributes import get_history, set_committed_value from sqlalchemy.schema import Index from sqlalchemy.sql.expression import column, nullsfirst from typing import overload, Any, Literal, TYPE_CHECKING if TYPE_CHECKING: - from collections.abc import Callable, Iterator, Sequence + from collections.abc import Callable, Iterable, Iterator, Sequence from sqlalchemy.orm.query import Query from sqlalchemy.orm.session import Session from typing import Self @@ -254,6 +254,48 @@ def ancestors(self) -> Iterator[AdjacencyList]: yield from self.parent.ancestors yield self.parent + @classmethod + def preload_ancestors( + cls, + session: Session, + items: Iterable[Self] + ) -> None: + """ Bulk loads all ancestors of the given items and wires up their + ``parent`` relationship, so walking the parent chain afterwards (e.g. + via :attr:`ancestors` or :attr:`path`) doesn't emit a query per + ancestor (N+1). + + Ancestors that are already loaded (e.g. eager loaded direct parents) + are reused instead of being queried again. + + """ + by_id: dict[int, Self] = {} + for item in items: + by_id[item.id] = item + # reuse an already-loaded parent without triggering a lazy load + parent = item.__dict__.get('parent') + if parent is not None: + by_id[parent.id] = parent + + # load the missing ancestors one generation at a time (a handful of + # queries bound by the tree depth, instead of one query per ancestor) + while True: + missing = { + item.parent_id for item in by_id.values() + if item.parent_id is not None and item.parent_id not in by_id + } + if not missing: + break + for ancestor in session.query(cls).filter(cls.id.in_(missing)): + by_id[ancestor.id] = ancestor + + # populate the parent relationship from the loaded set, so accessing + # it later resolves in-memory rather than triggering a lazy load + for item in by_id.values(): + parent_id = item.parent_id + if parent_id is not None and parent_id in by_id: + set_committed_value(item, 'parent', by_id[parent_id]) + @property def siblings(self) -> Query[Self]: """ A query including all siblings and the item itself. """ diff --git a/src/onegov/org/api.py b/src/onegov/org/api.py index 7a6d05d877..398d939ccb 100644 --- a/src/onegov/org/api.py +++ b/src/onegov/org/api.py @@ -4,6 +4,7 @@ from datetime import date from functools import cached_property +from onegov.api.models import AdjacencyListApiEndpoint from onegov.api.models import ApiEndpoint, ApiEndpointItem from onegov.api.models import ApiInvalidParamException from onegov.core.collection import Pagination @@ -449,6 +450,8 @@ def item_links(self, item: Occurrence) -> dict[str, Any]: } +# NOTE: News is a flat two-level tree (one shared root), so it's not an N+1 +# like topics/agencies -- hence plain ApiEndpoint here. class NewsApiEndpoint(ApiEndpoint[News, int]): app: OrgApp request: OrgRequest @@ -510,7 +513,7 @@ def item_links(self, item: News) -> dict[str, Any]: } -class TopicApiEndpoint(ApiEndpoint[Topic, int]): +class TopicApiEndpoint(AdjacencyListApiEndpoint[Topic, int]): request: OrgRequest app: OrgApp endpoint = 'topics' diff --git a/tests/onegov/agency/test_api.py b/tests/onegov/agency/test_api.py index 41cc1d1c5c..c1bd574478 100644 --- a/tests/onegov/agency/test_api.py +++ b/tests/onegov/agency/test_api.py @@ -4,6 +4,11 @@ from base64 import b64encode from collection_json import Collection, Template # type: ignore[import-untyped] from freezegun import freeze_time +from onegov.agency.api import AgencyApiEndpoint +from onegov.agency.collections import ExtendedAgencyCollection +from onegov.core.utils import Bunch +from sqlalchemy import event +from sqlalchemy.engine import Engine from tests.onegov.api.test_views import patch_collection_json # noqa: F401 from unittest.mock import patch @@ -11,6 +16,7 @@ if TYPE_CHECKING: from onegov.agency import AgencyApp + from sqlalchemy.orm import Session from tests.shared.client import Client from unittest.mock import MagicMock @@ -743,3 +749,60 @@ def template(item: Any) -> Any: '/api/memberships?updated_gt=2023-05-08T01:02').items } assert set(memberships) == {'Teacher'} + + +def test_agency_api_preloads_ancestors(session: Session) -> None: + """ + Rendering the ``html`` link of an agency builds its full path, which + walks up the parent chain. Only the direct parent is eager loaded by the + API collection, so without preloading each ancestor above it is loaded + with its own ``SELECT ... WHERE agencies.id = ...`` query (N+1). + """ + + agencies = ExtendedAgencyCollection(session) + root = agencies.add_root(title='Root') + a = agencies.add(root, title='A') + b = agencies.add(a, title='B') + agencies.add(b, title='Leaf 1') + agencies.add(b, title='Leaf 2') + session.flush() + b_id = b.id + + # Drop everything from the identity map so the ancestors above the page's + # direct parents actually have to be (pre)loaded. + session.expunge_all() + + pk_lookups: list[str] = [] + + def after_cursor_execute( + conn: Any, cursor: Any, statement: str, *args: Any + ) -> None: + if ( + 'FROM agencies' in statement + and 'agencies.id = ' in statement + and 'JOIN' not in statement + ): + pk_lookups.append(statement) + + # request only the children of ``b``, so their grandparents (``a``, + # ``root``) are not part of the batch and would otherwise be lazy loaded + # per item while building each agency's ``html`` link. + request: Any = Bunch(app=Bunch(session=lambda: session)) + endpoint = AgencyApiEndpoint( + request, extra_parameters={'parent': [str(b_id)]} + ) + + event.listen(Engine, 'after_cursor_execute', after_cursor_execute) + try: + leaves = list(endpoint.batch.values()) + assert len(leaves) == 2 + # walking the ancestors (as ``request.link`` does when building the + # ``html`` link) must not emit a query per ancestor + for leaf in leaves: + assert leaf.path.startswith('root/a/b/') + finally: + event.remove(Engine, 'after_cursor_execute', after_cursor_execute) + + assert pk_lookups == [], ( + f'Expected no per-ancestor queries, got {len(pk_lookups)}' + ) diff --git a/tests/onegov/core/test_adjacency_list.py b/tests/onegov/core/test_adjacency_list.py index 3df7849a14..c77b67d17d 100644 --- a/tests/onegov/core/test_adjacency_list.py +++ b/tests/onegov/core/test_adjacency_list.py @@ -10,9 +10,11 @@ sort_siblings, ) from onegov.core.orm.abstract.adjacency_list import numeric_priority +from sqlalchemy import event +from sqlalchemy.engine import Engine -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING if TYPE_CHECKING: from sqlalchemy.orm import Session @@ -399,3 +401,56 @@ def test_add_uses_binary_gap(session: Session) -> None: # Verify final list order based on calculated numeric orders assert [item.title for item in root.children] == ['a', 'b', 'c'] + + +def test_preload_ancestors(session: Session) -> None: + family = FamilyMemberCollection(session) + adam = family.add_root('Adam') + cain = family.add(parent=adam, title='Cain') + enoch = family.add(parent=cain, title='Enoch') + family.add(parent=enoch, title='Irad') + family.add(parent=enoch, title='Mehujael') + session.flush() + enoch_id = enoch.id + + # Drop everything from the identity map so the ancestors above the direct + # parents actually have to be (pre)loaded. + session.expunge_all() + + # a "batch" that only contains the leaves; their ancestors (enoch, cain, + # adam) are not part of it and would otherwise be lazy loaded per item. + leaves = session.query(FamilyMember).filter( + FamilyMember.parent_id == enoch_id + ).all() + assert len(leaves) == 2 + + FamilyMember.preload_ancestors(session, leaves) + + pk_lookups: list[str] = [] + + def after_cursor_execute( + conn: Any, cursor: Any, statement: str, *args: Any + ) -> None: + if ( + 'FROM familymembers' in statement + and 'familymembers.id = ' in statement + and 'JOIN' not in statement + ): + pk_lookups.append(statement) + + event.listen(Engine, 'after_cursor_execute', after_cursor_execute) + try: + # walking the ancestors (as building a path/link does) must not emit + # a query per ancestor now that they've been preloaded + for leaf in leaves: + names = [a.name for a in leaf.ancestors] + assert names == ['adam', 'cain', 'enoch'] + assert leaf.path.startswith('adam/cain/enoch/') + # the shared ancestors are also wired up in-memory + assert leaves[0].parent is leaves[1].parent + finally: + event.remove(Engine, 'after_cursor_execute', after_cursor_execute) + + assert pk_lookups == [], ( + f'Expected no per-ancestor queries, got {len(pk_lookups)}' + ) diff --git a/tests/onegov/org/test_api.py b/tests/onegov/org/test_api.py index ae7e5487bb..37514ca316 100644 --- a/tests/onegov/org/test_api.py +++ b/tests/onegov/org/test_api.py @@ -5,12 +5,18 @@ from base64 import b64encode from datetime import timedelta from collection_json import Collection # type: ignore[import-untyped] +from freezegun import freeze_time +from onegov.core.utils import Bunch from onegov.directory import DirectoryCollection, DirectoryConfiguration from onegov.form import FormCollection +from onegov.org.api import TopicApiEndpoint from onegov.org.models.external_link import ( ExternalFormLink, ExternalResourceLink) +from onegov.page import PageCollection from onegov.people import Person from sedate import utcnow +from sqlalchemy import event +from sqlalchemy.engine import Engine from tests.onegov.api.test_views import patch_collection_json # noqa: F401 from unittest.mock import patch from uuid import uuid4 @@ -21,6 +27,7 @@ from .conftest import Client from onegov.agency import AgencyApp from onegov.org.models import ExtendedDirectory + from sqlalchemy.orm import Session from unittest.mock import MagicMock @@ -711,3 +718,66 @@ def test_api_directory_content_hash(client: Client) -> None: items = api_items(client, '/api/clubs') item_data = api_item_data(items[0]) assert item_data['content_hash'] != first_hash + + +def test_topic_api_preloads_ancestors(session: Session) -> None: + """ + Rendering the ``html`` link of a topic builds its full path, which walks + up the parent chain. The topics API collection doesn't eager load the + parent, so without preloading each ancestor that isn't part of the batch + is loaded with its own ``SELECT ... WHERE pages.id = ...`` query (N+1). + """ + + pages = PageCollection(session) + # the ancestors are created first (oldest), the leaves last (newest); the + # API orders topics by ``published_or_created`` desc, so page 0 of the + # batch is filled with leaves while their ancestors land on a later page + # and would otherwise be lazy loaded per item while building each link. + with freeze_time('2020-01-01'): + root = pages.add_root(title='Root', type='topic') + a = pages.add(root, title='A', type='topic') + b = pages.add(a, title='B', type='topic') + session.flush() + with freeze_time('2020-06-01'): + for i in range(30): + pages.add(b, title=f'Leaf {i}', type='topic') + session.flush() + + # Drop everything from the identity map so the ancestors that aren't part + # of the batch actually have to be (pre)loaded. + session.expunge_all() + + pk_lookups: list[str] = [] + + def after_cursor_execute( + conn: Any, cursor: Any, statement: str, *args: Any + ) -> None: + if ( + 'FROM pages' in statement + and 'pages.id = ' in statement + and 'JOIN' not in statement + ): + pk_lookups.append(statement) + + request: Any = Bunch( + app=Bunch(session=lambda: session), + session=session, + identity=Bunch(role='admin'), + ) + endpoint = TopicApiEndpoint(request) + + event.listen(Engine, 'after_cursor_execute', after_cursor_execute) + try: + topics = list(endpoint.batch.values()) + # page 0 holds the 25 newest topics, i.e. only leaves + assert len(topics) == 25 + # walking the ancestors (as ``request.link`` does when building the + # ``html`` link) must not emit a query per ancestor + for topic in topics: + assert topic.path.startswith('root/a/b/') + finally: + event.remove(Engine, 'after_cursor_execute', after_cursor_execute) + + assert pk_lookups == [], ( + f'Expected no per-ancestor queries, got {len(pk_lookups)}' + )