communities: batch parent lookups to fix SearchCommunities/ListAllCommunities N+1

Author: kevinortiz43

app/backend/src/couchers/db.py

@@ -12,7 +12,7 @@
 from alembic.config import Config
 from geoalchemy2 import WKBElement
 from opentelemetry import trace
-from sqlalchemy import Engine, Row, Subquery, create_engine, select, text, true
+from sqlalchemy import Engine, Subquery, create_engine, select, text, true
 from sqlalchemy.dialects import registry
 from sqlalchemy.orm.session import Session
 from sqlalchemy.pool import QueuePool
@@ -197,15 +197,59 @@ def _get_node_parents_recursive_cte_subquery(node_id: int) -> Subquery:
     ).subquery()
 
 
-def get_node_parents_recursively(session: Session, node_id: int) -> Sequence[Row[tuple[int, int, int, Cluster]]]:
-    subquery = _get_node_parents_recursive_cte_subquery(node_id)
-    return session.execute(
-        select(subquery, Cluster)
+def _get_nodes_parents_recursive_cte_subquery(node_ids: Sequence[int]) -> Subquery:
+    parents = (
+        select(
+            Node.id.label("start_node_id"),
+            Node.id,
+            Node.parent_node_id,
+            literal(0).label("level"),
+        )
+        .where(Node.id.in_(node_ids))
+        .cte("parents", recursive=True)
+    )
+
+    return select(
+        parents.union(
+            select(
+                parents.c.start_node_id,
+                Node.id,
+                Node.parent_node_id,
+                (parents.c.level + 1).label("level"),
+            ).join(parents, Node.id == parents.c.parent_node_id)
+        )
+    ).subquery()
+
+
+def get_nodes_parents_recursively(
+    session: Session, node_ids: Sequence[int]
+) -> dict[int, list[tuple[int, int, int, Cluster]]]:
+    if not node_ids:
+        return {}
+
+    subquery = _get_nodes_parents_recursive_cte_subquery(node_ids)
+    rows = session.execute(
+        select(
+            subquery.c.start_node_id,
+            subquery.c.id,
+            subquery.c.parent_node_id,
+            subquery.c.level,
+            Cluster,
+        )
         .join(Cluster, Cluster.parent_node_id == subquery.c.id)
         .where(Cluster.is_official_cluster)
-        .order_by(subquery.c.level.desc())
+        .order_by(subquery.c.start_node_id, subquery.c.level.desc())
     ).all()
 
+    result: dict[int, list[tuple[int, int, int, Cluster]]] = {node_id: [] for node_id in node_ids}
+    for start_node_id, id_, parent_node_id, level, cluster in rows:
+        result[start_node_id].append((id_, parent_node_id, level, cluster))
+    return result
+
+
+def get_node_parents_recursively(session: Session, node_id: int) -> Sequence[tuple[int, int, int, Cluster]]:
+    return get_nodes_parents_recursively(session, [node_id])[node_id]
+
 
 def _can_moderate_any_cluster(session: Session, user_id: int, cluster_ids: list[int]) -> bool:
     query = select(

app/backend/src/couchers/servicers/communities.py

@@ -11,7 +11,12 @@
 from couchers.constants import COMMUNITIES_SEARCH_FUZZY_SIMILARITY_THRESHOLD
 from couchers.context import CouchersContext
 from couchers.crypto import decrypt_page_token, encrypt_page_token
-from couchers.db import can_moderate_node, get_node_parents_recursively, is_user_in_node_geography
+from couchers.db import (
+    can_moderate_node,
+    get_node_parents_recursively,
+    get_nodes_parents_recursively,
+    is_user_in_node_geography,
+)
 from couchers.event_log import log_event
 from couchers.materialized_views import ClusterAdminCount, ClusterSubscriptionCount
 from couchers.models import (
@@ -50,8 +55,7 @@
 }
 
 
-def _parents_to_pb(session: Session, node_id: int) -> list[groups_pb2.Parent]:
-    parents = get_node_parents_recursively(session, node_id)
+def _parents_list_to_pb(parents: Sequence[tuple[int, int, int, Cluster]]) -> list[groups_pb2.Parent]:
     return [
         groups_pb2.Parent(
             community=groups_pb2.CommunityParent(
@@ -65,16 +69,24 @@ def _parents_to_pb(session: Session, node_id: int) -> list[groups_pb2.Parent]:
     ]
 
 
+def _parents_to_pb(session: Session, node_id: int) -> list[groups_pb2.Parent]:
+    return _parents_list_to_pb(get_node_parents_recursively(session, node_id))
+
+
 def _community_summary_to_pb(
-    session: Session, node: Node, cluster: Cluster, member_count: int | None, user_subscription: int | None
+    node: Node,
+    cluster: Cluster,
+    member_count: int | None,
+    user_subscription: int | None,
+    parents: Sequence[tuple[int, int, int, Cluster]],
 ) -> communities_pb2.CommunitySummary:
     return communities_pb2.CommunitySummary(
         community_id=node.id,
         name=cluster.name,
         slug=cluster.slug,
         member=user_subscription is not None,
         member_count=member_count or 1,
-        parents=_parents_to_pb(session, node.id),
+        parents=_parents_list_to_pb(parents),
         created=Timestamp_from_datetime(node.created),
         node_type=nodetype2api[node.node_type],
     )
@@ -207,9 +219,8 @@ def SearchCommunities(
         )
 
         if not raw_query:
-            query = query.order_by(Cluster.name.asc(), Node.id.asc())
+            query = query.order_by(Cluster.name.asc(), Node.id.asc()).limit(page_size)
         elif len(raw_query) < 3:
-            # < 3 chars is too short for a trigram match; prefix-match instead (still uses the trgm index)
             query = (
                 query.where(func.immutable_unaccent(Cluster.name).ilike(func.unaccent(raw_query).concat("%")))
                 .order_by(Cluster.name.asc(), Node.id.asc())
@@ -227,9 +238,11 @@ def SearchCommunities(
 
         rows = session.execute(query).all()
 
+        parents_by_node_id = get_nodes_parents_recursively(session, [node.id for node, _, _, _ in rows])
+
         return communities_pb2.SearchCommunitiesRes(
             results=[
-                _community_summary_to_pb(session, node, cluster, member_count, user_subscription)
+                _community_summary_to_pb(node, cluster, member_count, user_subscription, parents_by_node_id[node.id])
                 for node, cluster, member_count, user_subscription in rows
             ]
         )
@@ -659,9 +672,11 @@ def ListAllCommunities(
             .order_by(Node.id)
         ).all()
 
+        parents_by_node_id = get_nodes_parents_recursively(session, [node.id for node, _, _, _ in results])
+
         return communities_pb2.ListAllCommunitiesRes(
             communities=[
-                _community_summary_to_pb(session, node, cluster, member_count, user_subscription)
+                _community_summary_to_pb(node, cluster, member_count, user_subscription, parents_by_node_id[node.id])
                 for node, cluster, member_count, user_subscription in results
             ],
         )

app/backend/src/tests/test_communities.py

@@ -7,7 +7,12 @@
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
-from couchers.db import is_user_in_node_geography, session_scope
+from couchers.db import (
+    get_node_parents_recursively,
+    get_nodes_parents_recursively,
+    is_user_in_node_geography,
+    session_scope,
+)
 from couchers.helpers.clusters import CHILD_NODE_TYPE
 from couchers.materialized_views import refresh_materialized_views
 from couchers.models import (
@@ -1068,6 +1073,46 @@ def test_ListAllCommunities(testing_communities):
             global_community = next(c for c in res.communities if c.community_id == w_id)
             assert global_community.member  # user6 is a member
 
+    @staticmethod
+    def test_get_nodes_parents_recursively_matches_single_id(testing_communities):
+        """
+        get_nodes_parents_recursively (the batched CTE used by SearchCommunities/ListAllCommunities to
+        avoid an N+1) must return, for every node_id, exactly what get_node_parents_recursively(session,
+        node_id) returns for that node_id individually: same length, same order, same (id, parent_node_id,
+        level, cluster) contents.
+        """
+        with session_scope() as session:
+            w_id = get_community_id(session, "Global")
+            c1_id = get_community_id(session, "Country 1")
+            c1r1_id = get_community_id(session, "Country 1, Region 1")
+            c1r1c1_id = get_community_id(session, "Country 1, Region 1, City 1")
+            c2r1c1_id = get_community_id(session, "Country 2, Region 1, City 1")
+            nonexistent_id = -1
+
+            node_ids = [w_id, c1_id, c1r1_id, c1r1c1_id, c2r1c1_id, nonexistent_id]
+            batched = get_nodes_parents_recursively(session, node_ids)
+
+            assert set(batched.keys()) == set(node_ids)
+
+            for node_id in node_ids:
+                expected = get_node_parents_recursively(session, node_id)
+                actual = batched[node_id]
+                assert len(actual) == len(expected)
+                for (a_id, a_parent_id, a_level, a_cluster), (e_id, e_parent_id, e_level, e_cluster) in zip(
+                    actual, expected
+                ):
+                    assert a_id == e_id
+                    assert a_parent_id == e_parent_id
+                    assert a_level == e_level
+                    assert a_cluster.id == e_cluster.id
+                    assert a_cluster.name == e_cluster.name
+
+            assert [row[0] for row in batched[w_id]] == [w_id]
+
+            assert [row[0] for row in batched[c1r1c1_id]] == [w_id, c1_id, c1r1_id, c1r1c1_id]
+
+            assert batched[nonexistent_id] == []
+
     @staticmethod
     def test_ListRecentCommunities(testing_communities, monkeypatch):
         """

app/web/features/communities/CommunitiesPage/CommunitySearch.test.tsx

@@ -21,7 +21,6 @@ jest.mock("next/router", () => ({
   }),
 }));
 
-// AuthProvider pulls in @sentry/nextjs, which crashes under jsdom
 jest.mock("platform/sentry", () => {
   const mockCaptureException = jest.fn();
   return {

app/web/features/communities/CommunitiesPage/CommunitySearch.tsx

@@ -90,7 +90,6 @@ export default function CommunitySearch() {
       onInputChange={handleInputChange}
       onChange={handleChange}
       getOptionLabel={(option) => option.name}
-      // keep backend order; don't let MUI re-filter
       filterOptions={(x) => x}
       renderOption={(props, option) => {
         const { key, ...optionProps } = props;