fix(lineage): expand each shared node once when rendering lineage trees (#870)

Author: kvlonge

src/sqlbuild/compiler/lineage/_helpers/tree_render.py

@@ -27,30 +27,29 @@ def render_dependency_branch[Node: Hashable](
     branch_style: Callable[[str], str],
     already_shown: Callable[[], str],
 ) -> list[str]:
-    """Render a box-drawing dependency branch for opaque node values."""
+    """Render a box-drawing dependency branch, expanding each node once across the tree."""
 
     lines: list[str] = []
-    children: list[Node] = sorted(deps.get(node, ()), key=sort_key)
-    index: int
-    child: Node
-    for index, child in enumerate(children):
-        is_last: bool = index == len(children) - 1
+    shown: set[Node] = set(seen)
+    pending: list[tuple[Node, str, bool]] = _pending_children(
+        children=sorted(deps.get(node, ()), key=sort_key), prefix=prefix
+    )
+    while pending:
+        child: Node
+        child_prefix: str
+        is_last: bool
+        child, child_prefix, is_last = pending.pop()
         branch: str = _BRANCH_LAST if is_last else _BRANCH_MID
-        continuation: str = _CONTINUATION_LAST if is_last else _CONTINUATION_MID
-        suffix: str = already_shown() if child in seen else ""
-        lines.append(f"{branch_style(prefix + branch)}{format_node(child)}{suffix}")
-        if child in seen:
+        suffix: str = already_shown() if child in shown else ""
+        lines.append(f"{branch_style(child_prefix + branch)}{format_node(child)}{suffix}")
+        if child in shown:
             continue
-        lines.extend(
-            render_dependency_branch(
-                node=child,
-                deps=deps,
-                prefix=prefix + continuation,
-                seen=seen | {child},
-                format_node=format_node,
-                sort_key=sort_key,
-                branch_style=branch_style,
-                already_shown=already_shown,
+        shown.add(child)
+        continuation: str = _CONTINUATION_LAST if is_last else _CONTINUATION_MID
+        pending.extend(
+            _pending_children(
+                children=sorted(deps.get(child, ()), key=sort_key),
+                prefix=child_prefix + continuation,
             )
         )
     return lines
@@ -68,41 +67,59 @@ def render_column_trace_branch[Column, Edge](
     branch_style: Callable[[str], str],
     already_shown: Callable[[], str],
 ) -> list[str]:
-    """Render a box-drawing column-trace branch for opaque column/edge values."""
+    """Render a box-drawing column-trace branch, expanding each column once across the tree."""
 
     lines: list[str] = []
-    edges: list[Edge] = sorted(
-        deps.get(column_id(column), ()),
-        key=lambda edge: column_id(related_column(edge)),
+    shown: set[str] = set(seen)
+    pending: list[tuple[Edge, str, bool]] = _pending_children(
+        children=_sorted_edges(
+            edges=deps.get(column_id(column), ()),
+            column_id=column_id,
+            related_column=related_column,
+        ),
+        prefix=prefix,
     )
-    index: int
-    edge: Edge
-    for index, edge in enumerate(edges):
-        is_last: bool = index == len(edges) - 1
+    while pending:
+        edge: Edge
+        edge_prefix: str
+        is_last: bool
+        edge, edge_prefix, is_last = pending.pop()
         branch: str = _BRANCH_LAST if is_last else _BRANCH_MID
-        continuation: str = _CONTINUATION_LAST if is_last else _CONTINUATION_MID
         related: Column = related_column(edge)
         related_id: str = column_id(related)
-        suffix: str = already_shown() if related_id in seen else ""
-        lines.append(f"{prefix}{branch_style(branch)}{format_related(edge)}{suffix}")
-        if related_id in seen:
+        suffix: str = already_shown() if related_id in shown else ""
+        lines.append(f"{edge_prefix}{branch_style(branch)}{format_related(edge)}{suffix}")
+        if related_id in shown:
             continue
-        lines.extend(
-            render_column_trace_branch(
-                column=related,
-                deps=deps,
-                prefix=prefix + continuation,
-                seen=seen | {related_id},
-                column_id=column_id,
-                related_column=related_column,
-                format_related=format_related,
-                branch_style=branch_style,
-                already_shown=already_shown,
+        shown.add(related_id)
+        continuation: str = _CONTINUATION_LAST if is_last else _CONTINUATION_MID
+        pending.extend(
+            _pending_children(
+                children=_sorted_edges(
+                    edges=deps.get(related_id, ()),
+                    column_id=column_id,
+                    related_column=related_column,
+                ),
+                prefix=edge_prefix + continuation,
             )
         )
     return lines
 
 
+def _pending_children[Item](*, children: list[Item], prefix: str) -> list[tuple[Item, str, bool]]:
+    last: int = len(children) - 1
+    return [(child, prefix, index == last) for index, child in reversed(list(enumerate(children)))]
+
+
+def _sorted_edges[Column, Edge](
+    *,
+    edges: list[Edge] | tuple[()],
+    column_id: Callable[[Column], str],
+    related_column: Callable[[Edge], Column],
+) -> list[Edge]:
+    return sorted(edges, key=lambda edge: column_id(related_column(edge)))
+
+
 def render_column_trace_limit_note(
     *,
     total: int,

tests/e2e/src/sqlbuild/cli/commands/main/lineage/_test_types.py

@@ -31,3 +31,11 @@ class LineageErrorCliTestCase:
     description: str
     command: tuple[str, ...]
     expected_fragments: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class DiamondLineageTreeCliTestCase:
+    description: str
+    command: tuple[str, ...]
+    expected_expanded_names: tuple[str, ...]
+    expected_max_lines: int

tests/e2e/src/sqlbuild/cli/commands/main/lineage/helpers.py

@@ -28,3 +28,40 @@ def prepare_lineage_cache_project(*, tmp_path: Path) -> Path:
 def lineage_node_ids(*, payload: dict[str, object]) -> tuple[str, ...]:
     nodes: list[dict[str, object]] = cast(list[dict[str, object]], payload["nodes"])
     return tuple(str(node["id"]) for node in nodes)
+
+
+DIAMOND_LAYERS: int = 14
+
+
+def diamond_model_names() -> tuple[str, ...]:
+    """Return every model in the layered diamond project, in layer order."""
+
+    layer_names: tuple[tuple[str, str, str], ...] = tuple(
+        (f"orders_left_{layer}", f"orders_right_{layer}", f"orders_{layer}")
+        for layer in range(1, DIAMOND_LAYERS + 1)
+    )
+    return ("orders_0", *sum(layer_names, ()))
+
+
+def prepare_diamond_lineage_project(*, tmp_path: Path) -> Path:
+    """Write a project whose layers each fork into two models and join again."""
+
+    files: dict[str, str] = {
+        "sqlbuild_project.toml": 'name = "diamond_lineage"\nadapter = "duckdb"\n',
+        "models/orders_0.sql": "MODEL (materialized view);\n\nSELECT 1 AS order_id, 10 AS amount\n",
+    }
+    for layer in range(1, DIAMOND_LAYERS + 1):
+        previous: str = f"orders_{layer - 1}"
+        for side in ("left", "right"):
+            files[f"models/orders_{side}_{layer}.sql"] = (
+                f'MODEL (materialized view);\n\nSELECT order_id, amount FROM __ref("{previous}")\n'
+            )
+        files[f"models/orders_{layer}.sql"] = (
+            "MODEL (materialized view);\n\n"
+            "SELECT l.order_id, l.amount + r.amount AS amount\n"
+            f'FROM __ref("orders_left_{layer}") AS l\n'
+            f'JOIN __ref("orders_right_{layer}") AS r ON l.order_id = r.order_id\n'
+        )
+    return prepare_inline_project(
+        tmp_path=tmp_path, project_name="diamond_lineage", repo_files=files
+    )

tests/e2e/src/sqlbuild/cli/commands/main/lineage/test_lineage.py

@@ -1,26 +1,36 @@
 from __future__ import annotations
 
 import json
+import re
 import subprocess
 from pathlib import Path
 
 import pytest
 
 from tests.e2e.src.sqlbuild.cli.commands.main.lineage._test_types import (
     ColumnLineageCacheCliTestCase,
+    DiamondLineageTreeCliTestCase,
     LineageCacheCliTestCase,
     LineageCliTestCase,
     LineageErrorCliTestCase,
 )
 from tests.e2e.src.sqlbuild.cli.commands.main.lineage.helpers import (
+    DIAMOND_LAYERS,
+    diamond_model_names,
     lineage_node_ids,
+    prepare_diamond_lineage_project,
     prepare_lineage_cache_project,
 )
 from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import (
     prepare_waffle_shop,
     run_sqb,
 )
 
+_TREE_LINE: re.Pattern[str] = re.compile(r"^[│ ]*[├└]── .*$", re.MULTILINE)
+_EXPANDED_DIAMOND_NODE: re.Pattern[str] = re.compile(
+    r"^[│ ]*[├└]── (?![^\n]*\(already shown\))[^\n]*?(orders_(?:left_|right_)?\d+(?:\.amount)?)",
+    re.MULTILINE,
+)
 _LINEAGE_CACHE_RELATIVE_PATH: Path = Path("target/cache/lineage/v1/structural-graph.sqlite3")
 
 
@@ -439,3 +449,61 @@ def test_given_invalid_lineage_target_when_running_then_explains_the_error(
     assert result.returncode == 1, result.stdout + result.stderr
     for fragment in test_case.expected_fragments:
         assert fragment in result.stderr, result.stderr
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    (
+        DiamondLineageTreeCliTestCase(
+            description="downstream tree expands each shared model once",
+            command=("--no-color", "lineage", "orders_0", "--direction", "downstream"),
+            expected_expanded_names=diamond_model_names()[1:],
+            expected_max_lines=4 * len(diamond_model_names()),
+        ),
+        DiamondLineageTreeCliTestCase(
+            description="upstream tree expands each shared model once",
+            command=(
+                "--no-color",
+                "lineage",
+                f"orders_{DIAMOND_LAYERS}",
+                "--direction",
+                "upstream",
+            ),
+            expected_expanded_names=diamond_model_names()[:-1],
+            expected_max_lines=4 * len(diamond_model_names()),
+        ),
+        DiamondLineageTreeCliTestCase(
+            description="column trace expands each shared column once",
+            command=(
+                "--no-color",
+                "lineage",
+                f"orders_{DIAMOND_LAYERS}.amount",
+                "--direction",
+                "upstream",
+            ),
+            expected_expanded_names=(
+                f"orders_left_{DIAMOND_LAYERS}.amount",
+                f"orders_right_{DIAMOND_LAYERS}.amount",
+                f"orders_{DIAMOND_LAYERS - 1}.amount",
+            ),
+            expected_max_lines=4 * len(diamond_model_names()),
+        ),
+    ),
+    ids=lambda case: case.description,
+)
+def test_given_diamond_dependencies_when_rendering_lineage_tree_then_output_stays_linear(
+    test_case: DiamondLineageTreeCliTestCase,
+    tmp_path: Path,
+) -> None:
+    project_dir: Path = prepare_diamond_lineage_project(tmp_path=tmp_path)
+
+    result: subprocess.CompletedProcess[str] = run_sqb(
+        command=test_case.command, project_dir=project_dir
+    )
+
+    assert result.returncode == 0, result.stdout + result.stderr
+    tree_lines: list[str] = _TREE_LINE.findall(result.stdout)
+    expanded_nodes: list[str] = _EXPANDED_DIAMOND_NODE.findall(result.stdout)
+    assert len(tree_lines) <= test_case.expected_max_lines
+    assert len(expanded_nodes) == len(set(expanded_nodes))
+    assert set(test_case.expected_expanded_names) <= set(expanded_nodes)