diff --git a/src/sqlbuild/compiler/lineage/_helpers/tree_render.py b/src/sqlbuild/compiler/lineage/_helpers/tree_render.py index e48eb0a19a..1873b21e45 100644 --- a/src/sqlbuild/compiler/lineage/_helpers/tree_render.py +++ b/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, diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/lineage/_test_types.py b/tests/e2e/src/sqlbuild/cli/commands/main/lineage/_test_types.py index 1823d74eb2..3f88729097 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/lineage/_test_types.py +++ b/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 diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/lineage/helpers.py b/tests/e2e/src/sqlbuild/cli/commands/main/lineage/helpers.py index fd181dba9e..afb6b5fce3 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/lineage/helpers.py +++ b/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 + ) diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/lineage/test_lineage.py b/tests/e2e/src/sqlbuild/cli/commands/main/lineage/test_lineage.py index 4a12f0a416..0602dba144 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/lineage/test_lineage.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/lineage/test_lineage.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import re import subprocess from pathlib import Path @@ -8,12 +9,16 @@ 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 ( @@ -21,6 +26,11 @@ 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)