diff --git a/datajunction-server/datajunction_server/api/data.py b/datajunction-server/datajunction_server/api/data.py index a6511ce7e..1b30867b6 100644 --- a/datajunction-server/datajunction_server/api/data.py +++ b/datajunction-server/datajunction_server/api/data.py @@ -25,6 +25,7 @@ ) from datajunction_server.construction.build_v3.types import GeneratedSQL from datajunction_server.database.availabilitystate import AvailabilityState +from datajunction_server.database.catalog import Catalog from datajunction_server.database.history import History from datajunction_server.database.node import Node, NodeRevision from datajunction_server.database.user import User @@ -364,7 +365,16 @@ async def get_data( node = cast( Node, - await Node.get_by_name(session, node_name, raise_if_not_exists=True), + await Node.get_by_name( + session, + node_name, + options=[ + joinedload(Node.current).options( + joinedload(NodeRevision.catalog).joinedload(Catalog.engines), + ), + ], + raise_if_not_exists=True, + ), ) engine = await resolve_engine( session=session, @@ -418,7 +428,16 @@ async def get_data_stream_for_node( request_headers = dict(request.headers) node = cast( Node, - await Node.get_by_name(session, node_name, raise_if_not_exists=True), + await Node.get_by_name( + session, + node_name, + options=[ + joinedload(Node.current).options( + joinedload(NodeRevision.catalog).joinedload(Catalog.engines), + ), + ], + raise_if_not_exists=True, + ), ) engine = await resolve_engine( session=session, diff --git a/datajunction-server/datajunction_server/api/metrics.py b/datajunction-server/datajunction_server/api/metrics.py index ac4f047fb..46b9e143d 100644 --- a/datajunction-server/datajunction_server/api/metrics.py +++ b/datajunction-server/datajunction_server/api/metrics.py @@ -7,7 +7,7 @@ from fastapi import BackgroundTasks, Depends, HTTPException, Query from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import selectinload +from sqlalchemy.orm import load_only, noload, selectinload from sqlalchemy.sql.operators import is_ from datajunction_server.api.nodes import list_nodes @@ -130,6 +130,11 @@ async def get_common_dimensions( input_errors = [] statement = ( select(Node) + .options( + load_only(Node.id, Node.name, Node.type, Node.current_version), + noload(Node.created_by), + noload(Node.tags), + ) .where(Node.name.in_(metric)) # type: ignore .where(is_(Node.deactivated_at, None)) ) diff --git a/datajunction-server/datajunction_server/api/sql.py b/datajunction-server/datajunction_server/api/sql.py index 3e1fca938..3b6583152 100644 --- a/datajunction-server/datajunction_server/api/sql.py +++ b/datajunction-server/datajunction_server/api/sql.py @@ -9,6 +9,7 @@ from typing import cast from fastapi import BackgroundTasks, Depends, Query, Request +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from datajunction_server.construction.build_v3 import ( @@ -702,11 +703,17 @@ async def get_sql_for_metrics( # 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=[]) + # Fetch only name/type for validation here — no need to hydrate full + # Node entities (which would trigger the mapper-level created_by/tags + # selectin loads on each row). + node_rows = ( + await session.execute( + select(Node.name, Node.type).where(Node.name.in_(metrics)), + ) + ).all() # Check if all requested nodes exist - found_names = {node.name for node in nodes} + found_names = {row.name for row in node_rows} missing_nodes = set(metrics) - found_names if missing_nodes: raise DJInvalidInputException( @@ -715,11 +722,11 @@ async def get_sql_for_metrics( ) # Validate node types - non_metric_nodes = [node for node in nodes if node and node.type != NodeType.METRIC] + non_metric_nodes = [row for row in node_rows if row.type != NodeType.METRIC] if non_metric_nodes: raise DJInvalidInputException( message="All nodes must be of metric type, but some are not: " - f"{', '.join([f'{n.name} ({n.type})' for n in non_metric_nodes])} .", + f"{', '.join([f'{row.name} ({row.type})' for row in non_metric_nodes])} .", http_status_code=HTTPStatus.UNPROCESSABLE_ENTITY, ) diff --git a/datajunction-server/datajunction_server/internal/sql.py b/datajunction-server/datajunction_server/internal/sql.py index 26689d09e..b02e92717 100644 --- a/datajunction-server/datajunction_server/internal/sql.py +++ b/datajunction-server/datajunction_server/internal/sql.py @@ -136,7 +136,16 @@ async def build_node_sql( node = cast( Node, - await Node.get_by_name(session, node_name, raise_if_not_exists=True), + await Node.get_by_name( + session, + node_name, + options=[ + joinedload(Node.current).options( + joinedload(NodeRevision.catalog).joinedload(Catalog.engines), + ), + ], + raise_if_not_exists=True, + ), ) if not engine: # pragma: no cover engine = node.current.catalog.engines[0]