Query performance optimizations: reduce eagerloading queries in /sql (#1824)

* Various query optimizations by reducing eager loading

* Add additional test coverage

Author: shangyian

datajunction-server/datajunction_server/api/graphql/dataloaders.py

@@ -2,6 +2,7 @@
 DataLoaders for batching and caching GraphQL queries.
 """
 
+import json
 from typing import Any
 
 from sqlalchemy import select
@@ -109,48 +110,52 @@ def create_node_by_name_loader(request: Request) -> DataLoader[str, DBNode | Non
 
 
 async def batch_load_collection_nodes(
-    collection_ids: list[int],
+    keys: list[tuple[int, str]],
     request: Request,
 ) -> list[list[DBNode]]:
     """
-    Batch load nodes for multiple collections.
+    Batch load nodes for multiple collections with field-aware eager loading.
 
-    This batches multiple collection node lookups into a single query,
-    avoiding N+1 queries when fetching nodes for multiple collections.
+    Keys are (collection_id, fields_json) tuples where fields_json is a
+    JSON-serialized dict of requested GraphQL fields (for load_node_options).
 
     Args:
-        collection_ids: List of collection IDs
+        keys: List of (collection_id, fields_json) tuples
         request: The Starlette request object for creating sessions
 
     Returns:
-        List of node lists, one per collection ID, in the same order
+        List of node lists, one per key, in the same order
     """
+    collection_ids = [cid for cid, _ in keys]
+
+    # Merge all requested fields across all loaders in this batch
+    all_fields: dict[str, Any] = {}
+    for _, fields_json in keys:
+        if fields_json:  # pragma: no branch
+            all_fields.update(json.loads(fields_json))
+
     async with session_context(request) as session:
-        # Load all requested collections with their nodes in one query
+        node_options = load_node_options(all_fields)
         stmt = (
             select(DBCollection)
             .where(DBCollection.id.in_(collection_ids))
-            .options(selectinload(DBCollection.nodes))
+            .options(selectinload(DBCollection.nodes).options(*node_options))
         )
         result = await session.execute(stmt)
         collections = result.unique().scalars().all()
 
-        # Create a lookup map: collection_id -> nodes
         collection_nodes_map = {c.id: c.nodes for c in collections}
-
-        # Return node lists in the same order as requested collection IDs
-        # Return empty list if collection not found
         return [collection_nodes_map.get(cid, []) for cid in collection_ids]
 
 
 def create_collection_nodes_loader(
     request: Request,
-) -> DataLoader[int, list[DBNode]]:
+) -> DataLoader[tuple[int, str], list[DBNode]]:
     """
     Create a DataLoader for loading nodes by collection ID.
 
-    This loader batches multiple collection node lookups within a single request
-    and caches the results to avoid N+1 queries.
+    Keys are (collection_id, fields_json) tuples so the loader can
+    eagerly load only the node relationships the query actually requests.
 
     Args:
         request: The Starlette request object

datajunction-server/datajunction_server/api/graphql/scalars/collection.py

@@ -2,6 +2,7 @@
 Collection GraphQL scalar types.
 """
 
+import json
 from datetime import datetime
 from typing import TYPE_CHECKING
 
@@ -10,6 +11,7 @@
 
 from datajunction_server.api.graphql.scalars.node import Node
 from datajunction_server.api.graphql.scalars.user import User
+from datajunction_server.api.graphql.utils import extract_fields
 
 if TYPE_CHECKING:
     from datajunction_server.database.collection import (
@@ -37,8 +39,8 @@ async def nodes(self, info: Info) -> list[Node]:
         Uses dataloader to batch requests efficiently.
         """
         loader = info.context["collection_nodes_loader"]
-        nodes = await loader.load(self.id)
-        return nodes  # type: ignore
+        node_fields = extract_fields(info)
+        return await loader.load((self.id, json.dumps(node_fields, sort_keys=True)))  # type: ignore
 
     @classmethod
     def from_db_collection(

datajunction-server/datajunction_server/api/helpers.py

@@ -527,11 +527,14 @@ async def validate_cube(
             message=("Metrics and dimensions must be part of a common catalog"),
         )
 
-    await validate_shared_dimensions(
-        session,
-        metric_nodes,
-        dimension_names,
-    )
+    # Only validate shared dimensions if dimensions were actually requested
+    # This avoids expensive dimension graph loading when dimensions=[]
+    if dimension_names:
+        await validate_shared_dimensions(
+            session,
+            metric_nodes,
+            dimension_names,
+        )
     return metrics, metric_nodes, list(dimension_nodes.values()), dimensions, catalog
 
 

datajunction-server/datajunction_server/api/sql.py

@@ -578,13 +578,23 @@ async def get_sql_for_metrics(
     """
     Return SQL for a set of metrics with dimensions and filters
     """
-    # make sure all metrics exist and have correct node type
-    nodes = [
-        await Node.get_by_name(session, node, raise_if_not_exists=True)
-        for node in metrics
-    ]
-    non_metric_nodes = [node for node in nodes if node and node.type != NodeType.METRIC]
+    # Label this session for debugging
+    session.info["session_label"] = "initial node loading"
+
+    # Fetch all metric nodes in a single query (only name/type needed for validation here)
+    nodes = await Node.get_by_names(session, metrics, options=[])
 
+    # Check if all requested nodes exist
+    found_names = {node.name for node in nodes}
+    missing_nodes = set(metrics) - found_names
+    if missing_nodes:
+        raise DJInvalidInputException(
+            message=f"The following nodes do not exist: {', '.join(missing_nodes)}",
+            http_status_code=HTTPStatus.NOT_FOUND,
+        )
+
+    # Validate node types
+    non_metric_nodes = [node for node in nodes if node and node.type != NodeType.METRIC]
     if non_metric_nodes:
         raise DJInvalidInputException(
             message="All nodes must be of metric type, but some are not: "
@@ -596,6 +606,7 @@ async def get_sql_for_metrics(
         cache=cache,
         query_type=QueryBuildType.METRICS,
     )
+
     return await query_cache_manager.get_or_load(
         background_tasks,
         request,
@@ -611,4 +622,5 @@ async def get_sql_for_metrics(
             use_materialized=use_materialized,
             ignore_errors=ignore_errors,
         ),
+        session=session,  # Pass the session to reuse it
     )

datajunction-server/datajunction_server/construction/build_v2.py

@@ -16,7 +16,7 @@
 
 from sqlalchemy import text, bindparam, select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload, selectinload
+from sqlalchemy.orm import joinedload, selectinload, noload
 
 from datajunction_server.internal.access.authorization import (
     AccessChecker,
@@ -723,6 +723,8 @@ async def find_join_paths_batch(
 
         This is O(1) database calls instead of O(nodes * depth) individual queries.
         """
+        # Filter out empty strings and check if we have any valid dimension names
+        target_dimension_names = {name for name in target_dimension_names if name}
         if not target_dimension_names:
             return {}  # pragma: no cover
 
@@ -800,18 +802,24 @@ async def load_dimension_links_and_nodes(
             .where(DimensionLink.id.in_(link_ids))
             .options(
                 joinedload(DimensionLink.dimension).options(
+                    noload(Node.created_by),
                     joinedload(Node.current).options(
+                        noload(NodeRevision.created_by),
                         selectinload(NodeRevision.columns).options(
                             joinedload(Column.attributes).joinedload(
                                 ColumnAttribute.attribute_type,
                             ),
-                            joinedload(Column.dimension),
+                            joinedload(Column.dimension).options(
+                                noload(Node.created_by),
+                            ),
                             joinedload(Column.partition),
                         ),
                         joinedload(NodeRevision.catalog),
                         selectinload(NodeRevision.availability),
                         selectinload(NodeRevision.dimension_links).options(
-                            joinedload(DimensionLink.dimension),
+                            joinedload(DimensionLink.dimension).options(
+                                noload(Node.created_by),
+                            ),
                         ),
                     ),
                 ),
@@ -1323,6 +1331,8 @@ async def build(self) -> ast.Query:
         Builds SQL for multiple metrics with the requested set of dimensions,
         filter expressions, order by, and limit clauses.
         """
+        # Always add dimensions referenced in the metric queries themselves
+        # (e.g., if a metric references a joinable dimension in its SQL definition)
         self.add_dimensions(get_dimensions_referenced_in_metrics(self.metric_nodes))
 
         measures_queries = await self.build_measures_queries()

datajunction-server/datajunction_server/construction/build_v3/cube_matcher.py

@@ -13,7 +13,7 @@
 
 from sqlalchemy import and_, select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload, selectinload
+from sqlalchemy.orm import joinedload, selectinload, noload
 
 from datajunction_server.construction.build_v3.decomposition import is_derived_metric
 from datajunction_server.models.dialect import Dialect
@@ -87,9 +87,13 @@ async def find_matching_cube(
             ),
         )
         .options(
+            noload(Node.created_by),  # Prevent User N+1 queries
             joinedload(Node.current).options(
-                selectinload(NodeRevision.cube_elements).selectinload(
-                    Column.node_revision,
+                noload(NodeRevision.created_by),  # Prevent User N+1 queries
+                selectinload(NodeRevision.cube_elements).options(
+                    selectinload(Column.node_revision).options(
+                        noload(NodeRevision.created_by),  # Prevent User N+1 queries
+                    ),
                 ),
                 joinedload(NodeRevision.availability),
                 selectinload(NodeRevision.materializations),
@@ -225,7 +229,18 @@ async def resolve_dialect_and_engine_for_metrics(
                 )
 
     # Fallback: use first metric's catalog's default engine
-    node = await Node.get_by_name(session, metrics[0], raise_if_not_exists=True)
+    node = await Node.get_by_name(
+        session,
+        metrics[0],
+        raise_if_not_exists=True,
+        options=[
+            joinedload(Node.current).options(
+                noload(NodeRevision.created_by),  # Prevent User N+1 queries
+                joinedload(NodeRevision.catalog),
+            ),
+            noload(Node.created_by),  # Prevent User N+1 queries
+        ],
+    )
     if not node:  # pragma: no cover
         raise ValueError(f"Metric not found: {metrics[0]}")
 

datajunction-server/datajunction_server/construction/build_v3/loaders.py

@@ -9,7 +9,7 @@
 
 from sqlalchemy import select, text, bindparam
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import selectinload, joinedload, load_only
+from sqlalchemy.orm import selectinload, joinedload, load_only, noload
 
 from datajunction_server.database.dimensionlink import DimensionLink
 from datajunction_server.database.node import Node, NodeRevision, Column
@@ -275,7 +275,9 @@ async def load_dimension_links_batch(
         .where(DimensionLink.id.in_(link_ids))
         .options(
             joinedload(DimensionLink.dimension).options(
+                noload(Node.created_by),  # Prevent User N+1 queries
                 joinedload(Node.current).options(
+                    noload(NodeRevision.created_by),  # Prevent User N+1 queries
                     # Load what's needed for table references, parsing, and type lookups
                     joinedload(NodeRevision.catalog),
                     joinedload(NodeRevision.availability),
@@ -375,6 +377,7 @@ async def load_nodes(ctx: BuildContext) -> None:
                 Node.current_version,
             ),
             joinedload(Node.current).options(
+                noload(NodeRevision.created_by),  # Prevent User N+1 queries
                 load_only(
                     NodeRevision.name,
                     NodeRevision.query,
@@ -391,13 +394,18 @@ async def load_nodes(ctx: BuildContext) -> None:
                 selectinload(NodeRevision.required_dimensions).options(
                     # Load the node_revision and node to reconstruct full dimension path
                     joinedload(Column.node_revision).options(
-                        joinedload(NodeRevision.node),
+                        noload(NodeRevision.created_by),  # Prevent User N+1 queries
+                        joinedload(NodeRevision.node).options(
+                            noload(Node.created_by),  # Prevent User N+1 queries
+                        ),
                     ),
                 ),
                 joinedload(NodeRevision.availability),  # For materialization support
                 selectinload(NodeRevision.dimension_links).options(
                     # Load dimension node for link matching in temporal filters
-                    joinedload(DimensionLink.dimension),
+                    joinedload(DimensionLink.dimension).options(
+                        noload(Node.created_by),  # Prevent User N+1 queries
+                    ),
                 ),
             ),
         )

datajunction-server/datajunction_server/database/node.py

@@ -31,6 +31,7 @@
     Mapped,
     joinedload,
     mapped_column,
+    noload,
     relationship,
     selectinload,
     MappedColumn,
@@ -309,6 +310,7 @@ class Node(Base):
         secondary="tagnoderelationship",
         primaryjoin="TagNodeRelationship.node_id==Node.id",
         secondaryjoin="TagNodeRelationship.tag_id==Tag.id",
+        lazy="selectin",
     )
 
     namespace_obj: Mapped[Optional["NodeNamespace"]] = relationship(
@@ -550,9 +552,6 @@ async def get_by_name(
             joinedload(Node.current).options(
                 *NodeRevision.default_load_options(),
             ),
-            selectinload(Node.tags),
-            selectinload(Node.created_by),
-            selectinload(Node.owners),
         ]
         statement = statement.options(*options)
         if not include_inactive:
@@ -581,7 +580,12 @@ async def get_by_names(
         """
         Get nodes by names
         """
+        # Early return if no names provided to avoid useless query
+        if not names:
+            return []
+
         statement = select(Node).where(Node.name.in_(names))
+
         options = options or [
             joinedload(Node.current).options(
                 *NodeRevision.default_load_options(),
@@ -1081,7 +1085,7 @@ class NodeRevision(
         secondary="cube",
         primaryjoin="NodeRevision.id==CubeRelationship.cube_id",
         secondaryjoin="Column.id==CubeRelationship.cube_element_id",
-        lazy="selectin",
+        # No lazy strategy - control via options (selectinload or noload)
         order_by="Column.order",
     )
 
@@ -1188,22 +1192,34 @@ def default_load_options(cls):
                 joinedload(Column.attributes).joinedload(
                     ColumnAttribute.attribute_type,
                 ),
-                joinedload(Column.dimension),
+                joinedload(Column.dimension).options(
+                    noload(Node.created_by),
+                ),
                 joinedload(Column.partition),
             ),
             joinedload(NodeRevision.catalog),
-            selectinload(NodeRevision.parents),
+            selectinload(NodeRevision.parents).options(
+                selectinload(Node.current).options(
+                    noload(NodeRevision.created_by),
+                ),
+                noload(Node.created_by),
+            ),
             selectinload(NodeRevision.materializations),
             selectinload(NodeRevision.metric_metadata),
             selectinload(NodeRevision.availability),
             selectinload(NodeRevision.dimension_links).options(
                 joinedload(DimensionLink.dimension).options(
-                    selectinload(Node.current),
+                    selectinload(Node.current).options(
+                        noload(NodeRevision.created_by),
+                    ),
+                    noload(Node.created_by),
                 ),
                 joinedload(DimensionLink.node_revision),
             ),
             selectinload(NodeRevision.required_dimensions),
             selectinload(NodeRevision.availability),
+            # Load created_by for API responses (but noload in /sql/ endpoint's custom options)
+            selectinload(NodeRevision.created_by),
         )
 
     @classmethod