datajunction-server 0.0.1a37.dev4__tar.gz → 0.0.1a37.dev6__tar.gz

{datajunction_server-0.0.1a37.dev4 → datajunction_server-0.0.1a37.dev6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datajunction-server
-Version: 0.0.1a37.dev4
+Version: 0.0.1a37.dev6
 Summary: DataJunction server library for running to a DataJunction server
 Project-URL: repository, https://github.com/DataJunction/dj
 Author-email: DataJunction Authors <yian.shang@gmail.com>

datajunction_server-0.0.1a37.dev6/datajunction_server/__about__.py

@@ -0,0 +1,4 @@
+"""
+Version for Hatch
+"""
+__version__ = "0.0.1a37.dev6"

{datajunction_server-0.0.1a37.dev4 → datajunction_server-0.0.1a37.dev6}/datajunction_server/api/helpers.py

@@ -3,7 +3,6 @@
 Helpers for API endpoints
 """
 import asyncio
-import collections
 import http.client
 import json
 import logging
@@ -57,7 +56,6 @@ from datajunction_server.models.node import (
     Node,
     NodeMissingParents,
     NodeNamespace,
-    NodeRelationship,
     NodeRevision,
     NodeRevisionBase,
     NodeStatus,
@@ -65,7 +63,7 @@ from datajunction_server.models.node import (
 from datajunction_server.models.node_type import NodeType
 from datajunction_server.models.query import ColumnMetadata, QueryWithResults
 from datajunction_server.service_clients import QueryServiceClient
-from datajunction_server.sql.dag import get_nodes_with_dimension
+from datajunction_server.sql.dag import get_downstream_nodes, get_nodes_with_dimension
 from datajunction_server.sql.parsing import ast
 from datajunction_server.sql.parsing.backends.antlr4 import SqlSyntaxError, parse
 from datajunction_server.sql.parsing.backends.exceptions import DJParseException
@@ -235,81 +233,6 @@ def get_query(  # pylint: disable=too-many-arguments
     return query_ast
 
 
-def get_downstream_nodes(
-    session: Session,
-    node_name: str,
-    node_type: NodeType = 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.
-    """
-    node = get_node_by_name(session=session, name=node_name, include_inactive=True)
-
-    dag = (
-        select(
-            NodeRelationship.parent_id,
-            NodeRevision.node_id,
-        )
-        .where(NodeRelationship.parent_id == node.id)
-        .join(NodeRevision, NodeRelationship.child_id == NodeRevision.id)
-        .join(
-            Node,
-            (Node.id == NodeRevision.node_id)
-            & (Node.current_version == NodeRevision.version),
-        )
-    ).cte("dag", recursive=True)
-
-    paths = dag.union_all(
-        select(
-            dag.c.parent_id,
-            NodeRevision.node_id,
-        )
-        .join(NodeRelationship, dag.c.node_id == NodeRelationship.parent_id)
-        .join(NodeRevision, NodeRelationship.child_id == NodeRevision.id)
-        .join(Node, Node.id == NodeRevision.node_id),
-    )
-
-    statement = (
-        select(Node)
-        .join(paths, paths.c.node_id == Node.id)
-        .options(joinedload(Node.current))
-    )
-
-    results = session.exec(statement).unique().all()
-
-    return [
-        downstream
-        for downstream in results
-        if downstream.type == node_type or node_type is None
-    ]
-
-
-def get_upstream_nodes(
-    session: Session,
-    node_name: str,
-    node_type: NodeType = None,
-) -> List[Node]:
-    """
-    Gets all upstream parents of the given node, filterable by node type.
-    """
-    node = get_node_by_name(session=session, name=node_name)
-    queue = collections.deque([node])
-    all_parents = set()
-    while queue:
-        node = queue.pop()
-        for parent in node.current.parents:
-            all_parents.add(parent)
-        for parent in node.current.parents:
-            queue.append(parent)
-
-    return [
-        upstream
-        for upstream in all_parents
-        if upstream.type == node_type or node_type is None
-    ]
-
-
 def find_bound_dimensions(
     validated_node: NodeRevision,
     dependencies_map: Dict[NodeRevision, List[ast.Table]],

{datajunction_server-0.0.1a37.dev4 → datajunction_server-0.0.1a37.dev6}/datajunction_server/api/nodes.py

@@ -5,7 +5,7 @@ Node related APIs.
 import logging
 import os
 from http import HTTPStatus
-from typing import List, Optional, Union, cast
+from typing import List, Optional, Union
 
 from fastapi import BackgroundTasks, Depends, Query, Response
 from fastapi.responses import JSONResponse
@@ -20,10 +20,8 @@ from datajunction_server.api.helpers import (
     deactivate_node,
     get_catalog_by_name,
     get_column,
-    get_downstream_nodes,
     get_node_by_name,
     get_node_namespace,
-    get_upstream_nodes,
     hard_delete_node,
     raise_if_node_exists,
     revalidate_node,
@@ -64,6 +62,7 @@ from datajunction_server.models.node import (
     CreateCubeNode,
     CreateNode,
     CreateSourceNode,
+    DAGNodeOutput,
     DimensionAttributeOutput,
     LineageColumn,
     Node,
@@ -86,7 +85,12 @@ from datajunction_server.models.partition import (
 )
 from datajunction_server.models.user import User
 from datajunction_server.service_clients import QueryServiceClient
-from datajunction_server.sql.dag import get_dimensions, get_nodes_with_dimension
+from datajunction_server.sql.dag import (
+    get_dimensions,
+    get_downstream_nodes,
+    get_nodes_with_dimension,
+    get_upstream_nodes,
+)
 from datajunction_server.sql.parsing.backends.antlr4 import parse
 from datajunction_server.utils import (
     Version,
@@ -966,12 +970,12 @@ def calculate_node_similarity(
 
 @router.get(
     "/nodes/{name}/downstream/",
-    response_model=List[NodeOutput],
+    response_model=List[DAGNodeOutput],
     name="List Downstream Nodes For A Node",
 )
 def list_downstream_nodes(
     name: str, *, node_type: NodeType = None, session: Session = Depends(get_session)
-) -> List[NodeOutput]:
+) -> List[DAGNodeOutput]:
     """
     List all nodes that are downstream from the given node, filterable by type.
     """
@@ -980,41 +984,46 @@ def list_downstream_nodes(
 
 @router.get(
     "/nodes/{name}/upstream/",
-    response_model=List[NodeOutput],
+    response_model=List[DAGNodeOutput],
     name="List Upstream Nodes For A Node",
 )
 def list_upstream_nodes(
     name: str, *, node_type: NodeType = None, session: Session = Depends(get_session)
-) -> List[NodeOutput]:
+) -> List[DAGNodeOutput]:
     """
     List all nodes that are upstream from the given node, filterable by type.
     """
-    return get_upstream_nodes(session, name, node_type)  # type: ignore
+    return get_upstream_nodes(session, name, node_type)
 
 
 @router.get(
     "/nodes/{name}/dag/",
-    response_model=List[NodeOutput],
+    response_model=List[DAGNodeOutput],
     name="List All Connected Nodes (Upstreams + Downstreams)",
 )
 def list_node_dag(
     name: str, *, session: Session = Depends(get_session)
-) -> List[NodeOutput]:
+) -> List[DAGNodeOutput]:
     """
     List all nodes that are part of the DAG of the given node. This means getting all upstreams,
     downstreams, and linked dimension nodes.
     """
-    node = get_node_by_name(session, name)
+    node = get_node_by_name(session, name, with_current=True)
     dimension_nodes = (
-        get_dimensions(session, node.current.parents[0], attributes=False)
+        get_dimensions(session, node.current.parents[0], with_attributes=False)
         + node.current.parents
         if node.type == NodeType.METRIC
-        else get_dimensions(session, node, attributes=False)
+        else get_dimensions(session, node, with_attributes=False)
     )
     dimension_nodes += [node]
-    downstreams = get_downstream_nodes(session, name)
-    upstreams = get_upstream_nodes(session, name)
-    return list(set(cast(List[Node], dimension_nodes) + downstreams + upstreams))  # type: ignore
+    downstreams = get_downstream_nodes(
+        session,
+        name,
+        include_deactivated=False,
+        include_cubes=False,
+    )
+    upstreams = get_upstream_nodes(session, name, include_deactivated=False)
+    return list(set(dimension_nodes + downstreams + upstreams))  # type: ignore
 
 
 @router.get(
@@ -1029,7 +1038,7 @@ def list_all_dimension_attributes(
     List all available dimension attributes for the given node.
     """
     node = get_node_by_name(session, name)
-    return get_dimensions(session, node, attributes=True)
+    return get_dimensions(session, node, with_attributes=True)
 
 
 @router.get(

{datajunction_server-0.0.1a37.dev4 → datajunction_server-0.0.1a37.dev6}/datajunction_server/internal/nodes.py

@@ -1190,7 +1190,7 @@ def get_cube_revision_metadata(session: Session, name: str):
     )
     result = session.execute(statement).unique().first()
     if not result:
-        raise DJNodeNotFound(
+        raise DJNodeNotFound(  # pragma: no cover
             message=f"A cube node with name `{name}` does not exist.",
             http_status_code=404,
         )

{datajunction_server-0.0.1a37.dev4 → datajunction_server-0.0.1a37.dev6}/datajunction_server/materialization/jobs/materialization_job.py

@@ -120,7 +120,7 @@ class SparkSqlMaterializationJob(  # pylint: disable=too-few-public-methods # pr
                 ),
                 ctes=query_ast.ctes,
             )
-        print("QUERY!!!", str(final_query))
+
         result = query_service_client.materialize(
             GenericMaterializationInput(
                 name=materialization.name,  # type: ignore

{datajunction_server-0.0.1a37.dev4 → datajunction_server-0.0.1a37.dev6}/datajunction_server/models/node.py

@@ -1123,7 +1123,7 @@ class ColumnOutput(BaseSQLModel):
 
     name: str
     display_name: Optional[str]
-    type: ColumnType
+    type: str
     attributes: Optional[List[AttributeOutput]]
     dimension: Optional[NodeNameOutput]
     partition: Optional[PartitionOutput]
@@ -1135,13 +1135,9 @@ class ColumnOutput(BaseSQLModel):
 
         validate_assignment = True
 
-    @root_validator
-    def type_string(cls, values):  # pylint: disable=no-self-argument
-        """
-        Extracts the type as a string
-        """
-        values["type"] = str(values.get("type"))
-        return values
+    _extract_type = validator("type", pre=True, allow_reuse=True)(
+        lambda raw: str(raw),  # pylint: disable=unnecessary-lambda
+    )
 
 
 class SourceColumnOutput(SQLModel):
@@ -1339,6 +1335,43 @@ class NodeOutput(OutputModel):
     current_version: str
 
 
+class DAGNodeRevisionOutput(SQLModel):
+    """
+    Output for a node revision with information about columns and if it is a metric.
+    """
+
+    id: int = Field(alias="node_revision_id")
+    node_id: int
+    type: NodeType
+    name: str
+    display_name: str
+    version: str
+    status: NodeStatus
+    mode: NodeMode
+    catalog: Optional[Catalog]
+    schema_: Optional[str]
+    table: Optional[str]
+    description: str = ""
+    columns: List[ColumnOutput]
+    updated_at: UTCDatetime
+    parents: List[NodeNameOutput]
+
+    class Config:  # pylint: disable=missing-class-docstring,too-few-public-methods
+        allow_population_by_field_name = True
+
+
+class DAGNodeOutput(OutputModel):
+    """
+    Output for a node in another node's DAG
+    """
+
+    namespace: str
+    current: DAGNodeRevisionOutput = PydanticField(flatten=True)
+    created_at: UTCDatetime
+    tags: List["Tag"] = []
+    current_version: str
+
+
 class NodeValidation(SQLModel):
     """
     A validation of a provided node definition

{datajunction_server-0.0.1a37.dev4 → datajunction_server-0.0.1a37.dev6}/datajunction_server/sql/dag.py

@@ -23,15 +23,165 @@ from datajunction_server.utils import SEPARATOR, get_settings
 settings = get_settings()
 
 
+def _node_output_options():
+    """
+    Statement options to retrieve all NodeOutput objects in one query
+    """
+    return [
+        joinedload(Node.current).options(
+            joinedload(NodeRevision.columns).options(
+                joinedload(Column.attributes).joinedload(
+                    ColumnAttribute.attribute_type,
+                ),
+                joinedload(Column.dimension),
+                joinedload(Column.partition),
+            ),
+            joinedload(NodeRevision.catalog),
+            joinedload(NodeRevision.parents),
+        ),
+        joinedload(Node.tags),
+    ]
+
+
+def get_downstream_nodes(
+    session: Session,
+    node_name: str,
+    node_type: NodeType = None,
+    include_deactivated: bool = True,
+    include_cubes: bool = True,
+) -> 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.
+    """
+    node = session.exec(select(Node).where(Node.name == node_name)).one()
+
+    initial_dag = (
+        select(
+            NodeRelationship.parent_id,
+            NodeRevision.node_id,
+        )
+        .where(NodeRelationship.parent_id == node.id)
+        .join(NodeRevision, NodeRelationship.child_id == NodeRevision.id)
+        .join(
+            Node,
+            (Node.id == NodeRevision.node_id)
+            & (Node.current_version == NodeRevision.version),
+        )
+    )
+    if not include_cubes:
+        initial_dag = initial_dag.where((NodeRevision.type != NodeType.CUBE))
+    dag = initial_dag.cte("downstreams", recursive=True)
+
+    next_layer = (
+        select(
+            dag.c.parent_id,
+            NodeRevision.node_id,
+        )
+        .join(NodeRelationship, dag.c.node_id == NodeRelationship.parent_id)
+        .join(NodeRevision, NodeRelationship.child_id == NodeRevision.id)
+        .join(Node, Node.id == NodeRevision.node_id)
+    )
+    if not include_cubes:
+        next_layer = next_layer.where(NodeRevision.type != NodeType.CUBE)
+
+    paths = dag.union_all(next_layer)
+
+    final_select = select(Node)
+    if not include_deactivated:
+        final_select = final_select.where(is_(Node.deactivated_at, None))
+    statement = final_select.join(paths, paths.c.node_id == Node.id).options(
+        *_node_output_options()
+    )
+    results = session.exec(statement).unique().all()
+
+    return [
+        downstream
+        for downstream in results
+        if downstream.type == node_type or node_type is None
+    ]
+
+
+def get_upstream_nodes(
+    session: Session,
+    node_name: str,
+    node_type: NodeType = None,
+    include_deactivated: bool = True,
+) -> List[Node]:
+    """
+    Gets all upstreams of the given node, filterable by node type.
+    Uses a recursive CTE query to build out all parents of the node.
+    """
+    node = (
+        session.exec(
+            select(Node).where(
+                (Node.name == node_name) & (is_(Node.deactivated_at, None)),
+            ),
+        )
+        .unique()
+        .one()
+    )
+
+    dag = (
+        select(
+            NodeRelationship.child_id,
+            NodeRevision.id,
+            NodeRevision.node_id,
+        )
+        .where(NodeRelationship.child_id == node.current.id)
+        .join(Node, NodeRelationship.parent_id == Node.id)
+        .join(
+            NodeRevision,
+            (Node.id == NodeRevision.node_id)
+            & (Node.current_version == NodeRevision.version),
+        )
+    ).cte("upstreams", recursive=True)
+
+    paths = dag.union_all(
+        select(
+            dag.c.child_id,
+            NodeRevision.id,
+            NodeRevision.node_id,
+        )
+        .join(NodeRelationship, dag.c.id == NodeRelationship.child_id)
+        .join(Node, NodeRelationship.parent_id == Node.id)
+        .join(
+            NodeRevision,
+            (Node.id == NodeRevision.node_id)
+            & (Node.current_version == NodeRevision.version),
+        ),
+    )
+
+    node_selector = select(Node)
+    if not include_deactivated:
+        node_selector = node_selector.where(is_(Node.deactivated_at, None))
+    statement = (
+        node_selector.join(paths, paths.c.node_id == Node.id)
+        .join(
+            NodeRevision,
+            (Node.current_version == NodeRevision.version)
+            & (Node.id == NodeRevision.node_id),
+        )
+        .options(*_node_output_options())
+    )
+
+    results = session.exec(statement).unique().all()
+    return [
+        upstream
+        for upstream in results
+        if upstream.type == node_type or node_type is None
+    ]
+
+
 def get_dimensions_dag(
     session: Session,
     node_revision: NodeRevision,
-    attributes: bool = True,
+    with_attributes: bool = True,
 ) -> List[Union[DimensionAttributeOutput, Node]]:
     """
-    Gets the dimensions graph of the given node revision with a single
-    recursive CTE query. This graph is split out into attributes or nodes
-    depending on the `attributes` flag.
+    Gets the dimensions graph of the given node revision with a single recursive CTE query.
+    This graph is split out into dimension attributes or dimension nodes depending on the
+    `with_attributes` flag.
     """
 
     initial_node = aliased(NodeRevision, name="initial_node")
@@ -129,13 +279,13 @@ def get_dimensions_dag(
     # Final SELECT statements
     # ----
     # If attributes was set to False, we only need to return the dimension nodes
-    if not attributes:
+    if not with_attributes:
         return (
             session.exec(
                 select(Node)
-                .options(joinedload(Node.current))
                 .select_from(paths)
-                .join(Node, paths.c.node_name == Node.name),
+                .join(Node, paths.c.node_name == Node.name)
+                .options(*_node_output_options()),
             )
             .unique()
             .all()
@@ -211,26 +361,41 @@ def get_dimensions_dag(
         )
     )
 
+    # Only include a given column it's an attribute on a dimension node or
+    # if the column is tagged with the attribute type 'dimension'
     return sorted(
         [
             DimensionAttributeOutput(
-                name=f"{row[0]}.{row[2]}",
-                node_name=row[0],
-                node_display_name=row[1],
-                is_primary_key=row[4] == "primary_key",
-                type=str(row[3]),
-                path=row[5].split(",")[:-1] if row[5] else [],
+                name=f"{node_name}.{column_name}",
+                node_name=node_name,
+                node_display_name=node_display_name,
+                is_primary_key=(
+                    attribute_types is not None and "primary_key" in attribute_types
+                ),
+                type=str(column_type),
+                path=join_path.split(",")[:-1] if join_path else [],
             )
-            for row in session.exec(final_query).all()
-            if row[5] != ""
-            or (
-                row[5] == ""
-                and row[4] is not None
-                and ("dimension" in row[4] or "primary_key" in row[4])
+            for (
+                node_name,
+                node_display_name,
+                column_name,
+                column_type,
+                attribute_types,
+                join_path,
+            ) in session.exec(
+                final_query,
+            ).all()
+            if (  # column has dimension attribute
+                join_path == ""
+                and attribute_types is not None
+                and ("dimension" in attribute_types or "primary_key" in attribute_types)
             )
-            or (
-                row[0] == node_revision.name
-                and node_revision.type == NodeType.DIMENSION
+            or (  # column is on dimension node
+                join_path != ""
+                or (
+                    node_name == node_revision.name
+                    and node_revision.type == NodeType.DIMENSION
+                )
             )
         ],
         key=lambda x: x.name,
@@ -240,7 +405,7 @@ def get_dimensions_dag(
 def get_dimensions(
     session: Session,
     node: Node,
-    attributes: bool = True,
+    with_attributes: bool = True,
 ) -> List[Union[DimensionAttributeOutput, Node]]:
     """
     Return all available dimensions for a given node.
@@ -248,8 +413,12 @@ def get_dimensions(
     * Setting `attributes` to False will return a list of dimension nodes
     """
     if node.type == NodeType.METRIC:
-        return get_dimensions_dag(session, node.current.parents[0].current, attributes)
-    return get_dimensions_dag(session, node.current, attributes)
+        return get_dimensions_dag(
+            session,
+            node.current.parents[0].current,
+            with_attributes,
+        )
+    return get_dimensions_dag(session, node.current, with_attributes)
 
 
 def check_convergence(path1: List[str], path2: List[str]) -> bool:

{datajunction_server-0.0.1a37.dev4 → datajunction_server-0.0.1a37.dev6}/tests/api/nodes_test.py

@@ -11,7 +11,6 @@ import pytest
 from fastapi.testclient import TestClient
 from sqlmodel import Session, select
 
-from datajunction_server.api.helpers import get_upstream_nodes
 from datajunction_server.internal.materializations import decompose_expression
 from datajunction_server.models import Database, Table
 from datajunction_server.models.column import Column
@@ -24,6 +23,7 @@ from datajunction_server.models.node import (
 from datajunction_server.models.node_type import NodeType
 from datajunction_server.models.partition import PartitionBackfill
 from datajunction_server.service_clients import QueryServiceClient
+from datajunction_server.sql.dag import get_upstream_nodes
 from datajunction_server.sql.parsing import ast, types
 from datajunction_server.sql.parsing.types import IntegerType, StringType, TimestampType
 from tests.sql.utils import compare_query_strings

datajunction_server-0.0.1a37.dev4/datajunction_server/__about__.py

@@ -1,4 +0,0 @@
-"""
-Version for Hatch
-"""
-__version__ = "0.0.1a37.dev4"