Modify GraphQL upstream/downstream endpoints to use dynamic field loading, which reduces database load for GraphQL requests that only need basic fields. (#1605)

Author: shangyian

datajunction-server/datajunction_server/api/graphql/queries/dag.py

@@ -8,8 +8,12 @@
 from strawberry.types import Info
 
 from datajunction_server.database.node import Node
-from datajunction_server.api.graphql.resolvers.nodes import find_nodes_by
+from datajunction_server.api.graphql.resolvers.nodes import (
+    find_nodes_by,
+    load_node_options,
+)
 from datajunction_server.api.graphql.scalars.node import DimensionAttribute
+from datajunction_server.api.graphql.utils import extract_fields
 from datajunction_server.sql.dag import (
     get_common_dimensions,
     get_downstream_nodes,
@@ -75,13 +79,19 @@ async def downstream_nodes(
     fanout threshold check and BFS fallback work better with single nodes.
     """
     session = info.context["session"]
+
+    # Build load options based on requested GraphQL fields
+    fields = extract_fields(info)
+    options = load_node_options(fields)
+
     all_downstreams: dict[int, Node] = {}
     for node_name in node_names:
         downstreams = await get_downstream_nodes(
             session,
             node_name=node_name,
             node_type=node_type,
             include_deactivated=include_deactivated,
+            options=options,
         )
         for node in downstreams:
             if node.id not in all_downstreams:  # pragma: no cover
@@ -116,9 +126,15 @@ async def upstream_nodes(
     Results are deduplicated by node ID.
     """
     session = info.context["session"]
+
+    # Build load options based on requested GraphQL fields
+    fields = extract_fields(info)
+    options = load_node_options(fields)
+
     return await get_upstream_nodes(  # type: ignore
         session,
         node_name=node_names,
         node_type=node_type,
         include_deactivated=include_deactivated,
+        options=options,
     )

datajunction-server/datajunction_server/api/helpers.py

@@ -104,15 +104,17 @@ async def get_node_by_name(
     """
     Get a node by name
     """
+    from datajunction_server.models.node import NodeOutput
+
     statement = select(Node).where(Node.name == name)
     if not include_inactive:
         statement = statement.where(is_(Node.deactivated_at, None))
     if node_type:
         statement = statement.where(Node.type == node_type)
     if with_current:
-        statement = statement.options(joinedload(Node.current)).options(
-            joinedload(Node.tags),
-        )
+        # Use full NodeOutput load options to ensure all required fields
+        # (like dimension_links) are eagerly loaded for serialization
+        statement = statement.options(*NodeOutput.load_options())
         result = await session.execute(statement)
         node = result.unique().scalar_one_or_none()
     else:

datajunction-server/datajunction_server/sql/dag.py

@@ -8,6 +8,7 @@
 from typing import Dict, List, Tuple, Union, cast
 
 from sqlalchemy import and_, func, join, literal, or_, select, distinct
+from sqlalchemy.sql.base import ExecutableOption
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import aliased, joinedload, selectinload
 from sqlalchemy.sql.operators import is_
@@ -65,15 +66,20 @@ async def get_downstream_nodes(
     include_deactivated: bool = True,
     include_cubes: bool = True,
     depth: int = -1,
-) -> List[Node]:
+    options: list[ExecutableOption] = None,
+) -> list[Node]:
     """
     Gets all downstream children of the given node, filterable by node type.
     Uses a recursive CTE query to build out all descendants from the node.
     """
+    # Use full options if none provided (for REST API DAGNodeOutput compatibility)
+    result_options = options if options is not None else _node_output_options()
+
+    # Initial lookup always uses light options (only need node.id)
     node = await Node.get_by_name(
         session,
         node_name,
-        options=_node_output_options(),
+        options=[joinedload(Node.current)],
     )
     if not node:
         return []
@@ -160,7 +166,7 @@ async def get_downstream_nodes(
         final_select = final_select.where(max_depths.c.max_depth < depth)
 
     statement = final_select.order_by(max_depths.c.max_depth, Node.id).options(
-        *_node_output_options(),
+        *result_options,
     )
     results = (await session.execute(statement)).unique().scalars().all()
     return [
@@ -308,6 +314,7 @@ async def get_upstream_nodes(
     node_name: Union[str, List[str]],
     node_type: NodeType = None,
     include_deactivated: bool = True,
+    options: List = None,
 ) -> List[Node]:
     """
     Gets all upstreams of the given node(s), filterable by node type.
@@ -320,10 +327,14 @@ async def get_upstream_nodes(
     # Normalize to list
     node_names = [node_name] if isinstance(node_name, str) else node_name
 
+    # Use full options if none provided (for REST API DAGNodeOutput compatibility)
+    result_options = options if options is not None else _node_output_options()
+
+    # Initial lookup always uses light options (only need type and current.id)
     nodes = await Node.get_by_names(
         session,
         node_names,
-        options=_node_output_options(),
+        options=[joinedload(Node.current)],
     )
 
     if not nodes:
@@ -417,19 +428,19 @@ async def get_upstream_nodes(
             (Node.current_version == NodeRevision.version)
             & (Node.id == NodeRevision.node_id),
         )
-        .options(*_node_output_options())
+        .options(*result_options)
     )
 
     results = list((await session.execute(statement)).unique().scalars().all())
 
     # For metrics, include the immediate parents in the results
     # (they are the starting point for the CTE, so not included by default)
     if immediate_parent_ids:
-        # Load parents with full options for consistent output
+        # Load parents with same options for consistent output
         parent_query = (
             select(Node)
             .where(Node.id.in_(immediate_parent_ids))
-            .options(*_node_output_options())
+            .options(*result_options)
         )
         if not include_deactivated:
             parent_query = parent_query.where(is_(Node.deactivated_at, None))

datajunction-server/tests/api/graphql/downstream_nodes_test.py

@@ -174,3 +174,75 @@ async def test_downstream_nodes_deactivated(
         {"name": "default.avg_repair_order_discounts", "type": "METRIC"},
         {"name": "default.avg_time_to_dispatch", "type": "METRIC"},
     ]
+
+
+@pytest.mark.asyncio
+async def test_downstream_nodes_with_nested_fields(
+    module__client_with_roads: AsyncClient,
+) -> None:
+    """
+    Test downstream nodes query with nested fields that require database joins.
+    This tests that load_node_options correctly builds options based on the
+    requested GraphQL fields (tags, owners, current.columns, etc.).
+    """
+    # Query with many nested fields that require joins
+    query = """
+    {
+        downstreamNodes(nodeNames: ["default.repair_order_details"], nodeType: TRANSFORM) {
+            name
+            type
+            tags {
+                name
+                tagType
+            }
+            owners {
+                username
+            }
+            current {
+                displayName
+                status
+                description
+                columns {
+                    name
+                    type
+                }
+                parents {
+                    name
+                }
+            }
+        }
+    }
+    """
+
+    response = await module__client_with_roads.post("/graphql", json={"query": query})
+    assert response.status_code == 200
+    data = response.json()
+
+    # Verify we got results
+    downstreams = data["data"]["downstreamNodes"]
+    assert len(downstreams) > 0
+
+    # Find the repair_orders_fact transform
+    repair_orders_fact = next(
+        (n for n in downstreams if n["name"] == "default.repair_orders_fact"),
+        None,
+    )
+    assert repair_orders_fact is not None
+
+    # Verify nested fields are populated
+    assert repair_orders_fact["type"] == "TRANSFORM"
+    assert repair_orders_fact["current"] is not None
+    assert repair_orders_fact["current"]["status"] == "VALID"
+    assert repair_orders_fact["current"]["displayName"] == "Repair Orders Fact"
+
+    # Verify columns are loaded (requires selectinload)
+    columns = repair_orders_fact["current"]["columns"]
+    assert columns is not None
+    assert len(columns) > 0
+    column_names = {c["name"] for c in columns}
+    assert "repair_order_id" in column_names
+
+    # Verify parents are loaded (requires selectinload)
+    parents = repair_orders_fact["current"]["parents"]
+    assert parents is not None
+    assert len(parents) > 0

datajunction-server/tests/api/graphql/upstream_nodes_test.py

@@ -237,3 +237,85 @@ async def test_upstream_nodes_deactivated(
     data = response.json()
     upstream_names = {node["name"] for node in data["data"]["upstreamNodes"]}
     assert "default.repair_orders_fact" in upstream_names
+
+
+@pytest.mark.asyncio
+async def test_upstream_nodes_with_nested_fields(
+    module__client_with_roads: AsyncClient,
+) -> None:
+    """
+    Test upstream nodes query with nested fields that require database joins.
+    This tests that load_node_options correctly builds options based on the
+    requested GraphQL fields (tags, owners, current.columns, etc.).
+    """
+    # Query with many nested fields that require joins
+    # Use includeDeactivated: true since earlier tests may have deactivated nodes
+    query = """
+    {
+        upstreamNodes(nodeNames: ["default.num_repair_orders"], includeDeactivated: true) {
+            name
+            type
+            tags {
+                name
+                tagType
+            }
+            owners {
+                username
+            }
+            current {
+                displayName
+                status
+                description
+                columns {
+                    name
+                    type
+                }
+                parents {
+                    name
+                }
+            }
+        }
+    }
+    """
+
+    response = await module__client_with_roads.post("/graphql", json={"query": query})
+    assert response.status_code == 200
+    data = response.json()
+
+    # Verify we got results
+    upstreams = data["data"]["upstreamNodes"]
+    assert len(upstreams) > 0
+
+    # Find the repair_orders_fact transform (immediate parent of the metric)
+    repair_orders_fact = next(
+        (n for n in upstreams if n["name"] == "default.repair_orders_fact"),
+        None,
+    )
+    assert repair_orders_fact is not None
+
+    # Verify nested fields are populated
+    assert repair_orders_fact["type"] == "TRANSFORM"
+    assert repair_orders_fact["current"] is not None
+    assert repair_orders_fact["current"]["status"] == "VALID"
+    assert repair_orders_fact["current"]["displayName"] == "Repair Orders Fact"
+
+    # Verify columns are loaded (requires selectinload)
+    columns = repair_orders_fact["current"]["columns"]
+    assert columns is not None
+    assert len(columns) > 0
+    column_names = {c["name"] for c in columns}
+    assert "repair_order_id" in column_names
+
+    # Verify parents are loaded (requires selectinload)
+    parents = repair_orders_fact["current"]["parents"]
+    assert parents is not None
+    assert len(parents) > 0
+
+    # Find a source node and verify its fields
+    source_node = next(
+        (n for n in upstreams if n["type"] == "SOURCE"),
+        None,
+    )
+    assert source_node is not None
+    assert source_node["current"] is not None
+    assert source_node["current"]["columns"] is not None