Agency: Preload agency ancestors in the API to fix N+1 query

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'

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',

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. """
 

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. """

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'

tests/onegov/agency/test_api.py

@@ -4,13 +4,19 @@
 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
 
 from typing import Any, TYPE_CHECKING
 
 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)}'
+    )

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)}'
+    )