From ec254ba8481bfeafa05d078de2b6525f4f795ba0 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Tue, 6 Oct 2026 13:08:08 +0100 Subject: [PATCH] fix: make sqb lineage use the shared project selector parser sqb lineage --select/--exclude now resolve through the same selector implementation as compile, plan and build, so path:models/... selectors, name globs and shared error codes behave identically across commands. The undocumented lineage-only path: form without the models/ root is now rejected with S012, matching every other command. --- .../sqlbuild/references/docs/cli/lineage.md | 4 +- .../commands/_helpers/lineage/selection.py | 336 +++--------------- src/sqlbuild/cli/commands/constants.py | 13 - src/sqlbuild/cli/commands/models.py | 20 -- .../planner/_helpers/graph/selectors.py | 178 +++++----- .../planner/main/selection/selector_match.py | 24 ++ .../planner/main/selection/selector_tokens.py | 27 ++ .../cli/commands/main/lineage/_test_types.py | 16 + .../cli/commands/main/lineage/test_lineage.py | 89 ++++- .../commands/_helpers/lineage/_test_types.py | 24 +- .../_helpers/lineage/test_selection.py | 207 ++++++++++- website/src/content/docs/docs/cli/lineage.mdx | 4 +- 12 files changed, 516 insertions(+), 426 deletions(-) create mode 100644 src/sqlbuild/compiler/planner/main/selection/selector_match.py create mode 100644 src/sqlbuild/compiler/planner/main/selection/selector_tokens.py diff --git a/src/sqlbuild/.agents/skills/sqlbuild/references/docs/cli/lineage.md b/src/sqlbuild/.agents/skills/sqlbuild/references/docs/cli/lineage.md index 86fd57c2b6..2014ce6a70 100644 --- a/src/sqlbuild/.agents/skills/sqlbuild/references/docs/cli/lineage.md +++ b/src/sqlbuild/.agents/skills/sqlbuild/references/docs/cli/lineage.md @@ -43,7 +43,7 @@ Exactly one of a positional target or `--select` is required. | `--format` | Output format: `tree` (default), `list`, or `json`. | | `--mode` | Column lineage analysis mode: `rich` (default) or `fast`. | | `--no-sql-analysis` | Disable compile-time SQL analysis (`--no-sql-validation` is an alias). | -| `--select`, `-s` | Select resources using standard selector syntax. | +| `--select`, `-s` | Select resources using the same selector syntax as `sqb compile`, `sqb plan` and `sqb build`, for example `path:models/marts`, `tag:finance` or `stg_*`. | | `--exclude` | Exclude resources from the selection. | ## Model lineage @@ -244,7 +244,7 @@ sqb lineage fact_orders --format json --direction upstream ## Depth limiting -`--depth` controls how many hops from the focus node(s) to include. In selector mode, `--depth` requires name, source, seed, or path-between selectors - it cannot be combined with tag or path selectors or comma-intersection selectors. +`--depth` controls how many hops from the focus node(s) to include. In selector mode without `--direction`, `--depth` requires name (including glob patterns), source, seed, or path-between selectors - it cannot be combined with tag or path selectors or comma-intersection selectors. ```bash # Only immediate parents diff --git a/src/sqlbuild/cli/commands/_helpers/lineage/selection.py b/src/sqlbuild/cli/commands/_helpers/lineage/selection.py index 9d00fcf380..db8086d4f7 100644 --- a/src/sqlbuild/cli/commands/_helpers/lineage/selection.py +++ b/src/sqlbuild/cli/commands/_helpers/lineage/selection.py @@ -10,13 +10,8 @@ COLUMN_TARGET_KIND, COLUMN_TARGET_SEPARATOR, DOWNSTREAM_DIRECTION, - PATH_BETWEEN_MARKER, - PATH_SEPARATOR, RICH_LINEAGE_STATUS_MODEL_THRESHOLD, - SELECTOR_EXPANSION_MARKER, SELECTOR_INTERSECTION_MARKER, - SELECTOR_KIND_SEPARATOR, - SUPPORTED_TYPED_SELECTOR_KINDS, UNLIMITED_DEPTH_VALUE, UPSTREAM_DIRECTION, ) @@ -27,8 +22,6 @@ LineageNode, LineageSelectionAnchors, LineageTarget, - ParsedLineagePathSelector, - ParsedLineageSelector, RelationLineageIndex, ) from sqlbuild.compiler.compile.models import ( @@ -47,6 +40,13 @@ ) from sqlbuild.compiler.lineage.types import ColumnLineageMode from sqlbuild.compiler.pipeline.models import ProjectGraph +from sqlbuild.compiler.planner.main.selection.selection import resolve_project_selectors +from sqlbuild.compiler.planner.main.selection.selector_match import match_project_selector +from sqlbuild.compiler.planner.main.selection.selector_parse import parse_project_selector +from sqlbuild.compiler.planner.main.selection.selector_tokens import ( + resolve_project_selector_tokens, +) +from sqlbuild.compiler.planner.models import ParsedSelector, PathSelector from sqlbuild.compiler.planner.types import SelectorKind from sqlbuild.presentation.main.maybe_status import maybe_status @@ -340,31 +340,26 @@ def select_selector_lineage( depth: int | None, direction: str | None = None, ) -> LineageGraph: - """Select lineage using selector semantics, expanding by direction when one is given.""" + """Select lineage using the shared project selector syntax, expanding by direction if given.""" - selected_keys: frozenset[CompiledObjectKey] = _resolve_selectors( + selected_keys: frozenset[CompiledObjectKey] = resolve_project_selectors( select=select, exclude=exclude, all_keys=graph.all_keys, - upstream=graph.upstream_deps, - downstream=graph.downstream_deps, + upstream_deps=graph.upstream_deps, + downstream_deps=graph.downstream_deps, tag_index=graph.tag_index, path_index=graph.path_index, ) if direction is not None: - excluded_keys: set[CompiledObjectKey] = set() - for raw_exclude in exclude: - for token in raw_exclude.split(): - excluded_keys.update( - _resolve_token( - token=token, - all_keys=graph.all_keys, - upstream=graph.upstream_deps, - downstream=graph.downstream_deps, - tag_index=graph.tag_index, - path_index=graph.path_index, - ) - ) + excluded_keys: frozenset[CompiledObjectKey] = resolve_project_selector_tokens( + selectors=exclude, + all_keys=graph.all_keys, + upstream_deps=graph.upstream_deps, + downstream_deps=graph.downstream_deps, + tag_index=graph.tag_index, + path_index=graph.path_index, + ) return build_lineage_graph( graph=graph, selected_keys=( @@ -377,11 +372,7 @@ def select_selector_lineage( ) anchors: LineageSelectionAnchors = _resolve_selector_anchors( select=select, - all_keys=graph.all_keys, - upstream_deps=graph.upstream_deps, - downstream_deps=graph.downstream_deps, - tag_index=graph.tag_index, - path_index=graph.path_index, + graph=graph, require_clear_anchors=depth is not None, ) if depth is not None: @@ -431,13 +422,11 @@ def build_lineage_graph( def _resolve_selector_anchors( *, select: tuple[str, ...], - all_keys: dict[str, CompiledObjectKey], - upstream_deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], - downstream_deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], - tag_index: dict[str, frozenset[CompiledObjectKey]], - path_index: dict[CompiledObjectKey, str], + graph: ProjectGraph | RelationLineageIndex, require_clear_anchors: bool, ) -> LineageSelectionAnchors: + """Find depth anchors for selectors already validated by the shared resolver.""" + upstream_anchors: set[CompiledObjectKey] = set() downstream_anchors: set[CompiledObjectKey] = set() retained: set[CompiledObjectKey] = set() @@ -449,43 +438,39 @@ def _resolve_selector_anchors( "--depth cannot be combined with comma-intersection selectors", code="C308", ) - parts: list[str] = token.split(SELECTOR_INTERSECTION_MARKER) - for part in parts: - parsed: ParsedLineageSelector | ParsedLineagePathSelector = _parse_selector(part) - if isinstance(parsed, ParsedLineagePathSelector): - start_key: CompiledObjectKey = _lookup_name( - name=parsed.start_name, all_keys=all_keys - ) - end_key: CompiledObjectKey = _lookup_name( - name=parsed.end_name, all_keys=all_keys - ) + for part in token.split(SELECTOR_INTERSECTION_MARKER): + parsed: ParsedSelector | PathSelector = parse_project_selector(part) + if isinstance(parsed, PathSelector): + start_key: CompiledObjectKey = graph.all_keys[parsed.start_name] + end_key: CompiledObjectKey = graph.all_keys[parsed.end_name] retained.update( - _find_path_keys(start=start_key, end=end_key, downstream=downstream_deps) + path_nodes(start=start_key, end=end_key, downstream=graph.downstream_deps) + or () ) upstream_anchors.add(start_key) downstream_anchors.add(end_key) continue + if parsed.kind in {SelectorKind.TAG, SelectorKind.PATH} and require_clear_anchors: + raise CliUserError( + "--depth requires name, source, seed, or path-between selectors", + code="C309", + ) + matched: frozenset[CompiledObjectKey] = match_project_selector( + parsed=parsed, + all_keys=graph.all_keys, + tag_index=graph.tag_index, + path_index=graph.path_index, + ) if parsed.kind in {SelectorKind.TAG, SelectorKind.PATH}: - if require_clear_anchors: - raise CliUserError( - "--depth requires name, source, seed, or path-between selectors", - code="C309", - ) - matched: frozenset[CompiledObjectKey] - if parsed.kind == SelectorKind.TAG: - matched = tag_index.get(parsed.value, frozenset()) - else: - matched = _match_path(folder=parsed.value, path_index=path_index) upstream_anchors.update(matched) downstream_anchors.update(matched) continue - key: CompiledObjectKey = _lookup_parsed_selector(parsed=parsed, all_keys=all_keys) if parsed.upstream: - upstream_anchors.add(key) + upstream_anchors.update(matched) if parsed.downstream: - downstream_anchors.add(key) + downstream_anchors.update(matched) if not parsed.upstream and not parsed.downstream: - retained.add(key) + retained.update(matched) return LineageSelectionAnchors( upstream=frozenset(upstream_anchors), downstream=frozenset(downstream_anchors), @@ -527,239 +512,6 @@ def _lookup_name(*, name: str, all_keys: dict[str, CompiledObjectKey]) -> Compil return key -def _resolve_selectors( - *, - select: tuple[str, ...], - exclude: tuple[str, ...], - all_keys: dict[str, CompiledObjectKey], - upstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], - downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], - tag_index: dict[str, frozenset[CompiledObjectKey]], - path_index: dict[CompiledObjectKey, str], -) -> frozenset[CompiledObjectKey]: - selected: set[CompiledObjectKey] = set() - for raw_select in select: - for token in raw_select.split(): - selected.update( - _resolve_token( - token=token, - all_keys=all_keys, - upstream=upstream, - downstream=downstream, - tag_index=tag_index, - path_index=path_index, - ) - ) - excluded: set[CompiledObjectKey] = set() - for raw_exclude in exclude: - for token in raw_exclude.split(): - excluded.update( - _resolve_token( - token=token, - all_keys=all_keys, - upstream=upstream, - downstream=downstream, - tag_index=tag_index, - path_index=path_index, - ) - ) - scoped: set[CompiledObjectKey] = selected - excluded - for upstream_key in transitive_closure_many( - starts=tuple(scoped), edges=upstream, include_starts=False - ): - if upstream_key.resource_type in { - CompiledResourceType.UDF, - CompiledResourceType.TABLE_FN, - }: - scoped.add(upstream_key) - return frozenset(scoped) - - -def _resolve_token( - *, - token: str, - all_keys: dict[str, CompiledObjectKey], - upstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], - downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], - tag_index: dict[str, frozenset[CompiledObjectKey]], - path_index: dict[CompiledObjectKey, str], -) -> frozenset[CompiledObjectKey]: - parts: list[str] = token.split(SELECTOR_INTERSECTION_MARKER) - resolved_parts: list[frozenset[CompiledObjectKey]] = [ - _resolve_single( - raw=part, - all_keys=all_keys, - upstream=upstream, - downstream=downstream, - tag_index=tag_index, - path_index=path_index, - ) - for part in parts - ] - result: frozenset[CompiledObjectKey] = resolved_parts[0] - for subsequent in resolved_parts[1:]: - result = result & subsequent - return result - - -def _resolve_single( - *, - raw: str, - all_keys: dict[str, CompiledObjectKey], - upstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], - downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], - tag_index: dict[str, frozenset[CompiledObjectKey]], - path_index: dict[CompiledObjectKey, str], -) -> frozenset[CompiledObjectKey]: - parsed: ParsedLineageSelector | ParsedLineagePathSelector = _parse_selector(raw) - if isinstance(parsed, ParsedLineagePathSelector): - start_key: CompiledObjectKey = _lookup_name(name=parsed.start_name, all_keys=all_keys) - end_key: CompiledObjectKey = _lookup_name(name=parsed.end_name, all_keys=all_keys) - result: set[CompiledObjectKey] = set( - _find_path_keys(start=start_key, end=end_key, downstream=downstream) - ) - if parsed.upstream: - result.update(transitive_closure(start=start_key, edges=upstream)) - if parsed.downstream: - result.update(transitive_closure(start=end_key, edges=downstream)) - return frozenset(result) - if parsed.kind == SelectorKind.TAG: - matched_keys: frozenset[CompiledObjectKey] = tag_index.get(parsed.value, frozenset()) - if not matched_keys: - raise CliUserError(f"no models found with tag '{parsed.value}'", code="C310") - return _apply_selector_expansion( - matched_keys=matched_keys, parsed=parsed, upstream=upstream, downstream=downstream - ) - if parsed.kind == SelectorKind.PATH: - matched_keys = _match_path(folder=parsed.value, path_index=path_index) - if not matched_keys: - raise CliUserError(f"no models found under path '{parsed.value}'", code="C311") - return _apply_selector_expansion( - matched_keys=matched_keys, parsed=parsed, upstream=upstream, downstream=downstream - ) - key: CompiledObjectKey = _lookup_parsed_selector(parsed=parsed, all_keys=all_keys) - return _apply_selector_expansion( - matched_keys=frozenset({key}), parsed=parsed, upstream=upstream, downstream=downstream - ) - - -def _apply_selector_expansion( - *, - matched_keys: frozenset[CompiledObjectKey], - parsed: ParsedLineageSelector, - upstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], - downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], -) -> frozenset[CompiledObjectKey]: - result: set[CompiledObjectKey] = set(matched_keys) - if parsed.upstream: - result.update( - transitive_closure_many(starts=matched_keys, edges=upstream, include_starts=False) - ) - if parsed.downstream: - result.update( - transitive_closure_many(starts=matched_keys, edges=downstream, include_starts=False) - ) - return frozenset(result) - - -def _parse_selector(raw: str) -> ParsedLineageSelector | ParsedLineagePathSelector: - stripped: str = raw.strip() - if not stripped: - raise CliUserError("empty selector", code="C312") - upstream: bool = stripped.startswith(SELECTOR_EXPANSION_MARKER) - downstream: bool = stripped.endswith(SELECTOR_EXPANSION_MARKER) - core: str = stripped.lstrip(SELECTOR_EXPANSION_MARKER).rstrip(SELECTOR_EXPANSION_MARKER) - if PATH_BETWEEN_MARKER in core: - start_name, end_name = (part.strip() for part in core.split(PATH_BETWEEN_MARKER, 1)) - if not start_name or not end_name: - raise CliUserError( - f"path selector '{stripped}' requires names on both sides of '~'", - code="C313", - ) - return ParsedLineagePathSelector( - start_name=start_name, - end_name=end_name, - upstream=upstream, - downstream=downstream, - ) - if not core: - raise CliUserError( - f"selector '{stripped}' has no name after removing '+' markers", - code="C314", - ) - if SELECTOR_KIND_SEPARATOR in core: - prefix, value = core.split(SELECTOR_KIND_SEPARATOR, 1) - if prefix not in SUPPORTED_TYPED_SELECTOR_KINDS: - raise CliUserError(f"unknown selector type '{prefix}' in '{stripped}'", code="C315") - if not value: - raise CliUserError(f"selector '{stripped}' has empty value after ':'", code="C316") - return ParsedLineageSelector( - kind=prefix, - value=value, - upstream=upstream, - downstream=downstream, - ) - if PATH_SEPARATOR in core: - return ParsedLineageSelector( - kind=SelectorKind.PATH, - value=core.strip(PATH_SEPARATOR), - upstream=upstream, - downstream=downstream, - ) - return ParsedLineageSelector( - kind=SelectorKind.NAME, - value=core, - upstream=upstream, - downstream=downstream, - ) - - -def _lookup_parsed_selector( - *, - parsed: ParsedLineageSelector, - all_keys: dict[str, CompiledObjectKey], -) -> CompiledObjectKey: - key: CompiledObjectKey | None = all_keys.get(parsed.value) - if key is None: - raise CliUserError(f"unknown lineage target '{parsed.value}'", code="C305") - if parsed.kind == SelectorKind.SOURCE and key.resource_type != CompiledResourceType.SOURCE: - raise CliUserError(f"unknown lineage source '{parsed.value}'", code="C317") - if parsed.kind == SelectorKind.SEED and key.resource_type != CompiledResourceType.SEED: - raise CliUserError(f"unknown lineage seed '{parsed.value}'", code="C318") - return key - - -def _find_path_keys( - *, - start: CompiledObjectKey, - end: CompiledObjectKey, - downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], -) -> frozenset[CompiledObjectKey]: - path_keys: frozenset[CompiledObjectKey] | None = path_nodes( - start=start, end=end, downstream=downstream - ) - if path_keys is None: - raise CliUserError( - f"'{end.resource_type}:{end.name}' is not downstream of " - f"'{start.resource_type}:{start.name}'", - code="C319", - ) - return path_keys - - -def _match_path( - *, - folder: str, - path_index: dict[CompiledObjectKey, str], -) -> frozenset[CompiledObjectKey]: - prefix: str = folder + "/" - return frozenset( - key - for key, model_folder in path_index.items() - if model_folder == folder or model_folder.startswith(prefix) - ) - - def _build_node( *, graph: ProjectGraph | RelationLineageIndex, key: CompiledObjectKey ) -> LineageNode: diff --git a/src/sqlbuild/cli/commands/constants.py b/src/sqlbuild/cli/commands/constants.py index 3935075e65..16a46bc828 100644 --- a/src/sqlbuild/cli/commands/constants.py +++ b/src/sqlbuild/cli/commands/constants.py @@ -4,7 +4,6 @@ from sqlbuild.cli.commands.types import CliCommand, CompileLineageMode, PlaygroundTemplate from sqlbuild.compiler.lineage.types import ColumnLineageMode -from sqlbuild.compiler.planner.types import SelectorKind from sqlbuild.spec.contracts.types import WarehouseGroup C0_CONTROL_CODE_LIMIT: int = 32 @@ -37,20 +36,8 @@ RULES_SKILLS_ACTION: str = "skills" SQL_FILE_SUFFIX: str = ".sql" SELECTOR_INTERSECTION_MARKER: str = "," -PATH_BETWEEN_MARKER: str = "~" -SELECTOR_KIND_SEPARATOR: str = ":" -PATH_SEPARATOR: str = "/" -SELECTOR_EXPANSION_MARKER: str = "+" COLUMN_TARGET_SEPARATOR: str = "." COLUMN_TARGET_KIND: str = "column" -SUPPORTED_TYPED_SELECTOR_KINDS: frozenset[str] = frozenset( - { - SelectorKind.SEED, - SelectorKind.SOURCE, - SelectorKind.TAG, - SelectorKind.PATH, - } -) PLAYGROUND_ORCHESTRATED_PROJECT_DIR: str = "waffle_shop" PLAYGROUND_ORCHESTRATED_TEMPLATES: frozenset[PlaygroundTemplate] = frozenset( {PlaygroundTemplate.DAGSTER, PlaygroundTemplate.RIVERS} diff --git a/src/sqlbuild/cli/commands/models.py b/src/sqlbuild/cli/commands/models.py index 11feb979d9..1733b95d8d 100644 --- a/src/sqlbuild/cli/commands/models.py +++ b/src/sqlbuild/cli/commands/models.py @@ -775,26 +775,6 @@ class LineageTarget: kind_matches: bool -@dataclass(frozen=True) -class ParsedLineageSelector: - """One parsed non-path lineage selector.""" - - kind: str - value: str - upstream: bool = False - downstream: bool = False - - -@dataclass(frozen=True) -class ParsedLineagePathSelector: - """One parsed path-between lineage selector.""" - - start_name: str - end_name: str - upstream: bool = False - downstream: bool = False - - @dataclass(frozen=True) class ColumnLineageTrace: """Selected column-level lineage trace.""" diff --git a/src/sqlbuild/compiler/planner/_helpers/graph/selectors.py b/src/sqlbuild/compiler/planner/_helpers/graph/selectors.py index a21235db77..b09b9a84c2 100644 --- a/src/sqlbuild/compiler/planner/_helpers/graph/selectors.py +++ b/src/sqlbuild/compiler/planner/_helpers/graph/selectors.py @@ -120,38 +120,28 @@ def resolve_selectors( if not select and not exclude: return frozenset(all_keys.values()) - selected: set[CompiledObjectKey] = set(all_keys.values()) if not select else set() - raw_select: str - for raw_select in select: - tokens: list[str] = raw_select.split() - token: str - for token in tokens: - resolved: frozenset[CompiledObjectKey] = _resolve_token( - token=token, - all_keys=all_keys, - upstream=upstream, - downstream=downstream, - tag_index=effective_tag_index, - path_index=effective_path_index, - ) - selected.update(resolved) - - excluded: set[CompiledObjectKey] = set() - raw_exclude: str - for raw_exclude in exclude: - tokens = raw_exclude.split() - for token in tokens: - resolved = _resolve_token( - token=token, - all_keys=all_keys, - upstream=upstream, - downstream=downstream, - tag_index=effective_tag_index, - path_index=effective_path_index, - ) - excluded.update(resolved) + selected: frozenset[CompiledObjectKey] = ( + resolve_selector_tokens( + selectors=select, + all_keys=all_keys, + upstream=upstream, + downstream=downstream, + tag_index=effective_tag_index, + path_index=effective_path_index, + ) + if select + else frozenset(all_keys.values()) + ) + excluded: frozenset[CompiledObjectKey] = resolve_selector_tokens( + selectors=exclude, + all_keys=all_keys, + upstream=upstream, + downstream=downstream, + tag_index=effective_tag_index, + path_index=effective_path_index, + ) - scoped: frozenset[CompiledObjectKey] = frozenset(selected - excluded) + scoped: frozenset[CompiledObjectKey] = selected - excluded return expand_required_build_resources( selected_keys=scoped, upstream=upstream, @@ -162,6 +152,60 @@ def resolve_selectors( ) +def resolve_selector_tokens( + *, + selectors: tuple[str, ...], + all_keys: dict[str, CompiledObjectKey], + upstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], + downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], + tag_index: dict[str, frozenset[CompiledObjectKey]], + path_index: dict[CompiledObjectKey, str], +) -> frozenset[CompiledObjectKey]: + """Union the keys matched by whitespace-separated selector tokens, without build expansion.""" + + resolved: set[CompiledObjectKey] = set() + raw_selector: str + for raw_selector in selectors: + token: str + for token in raw_selector.split(): + resolved.update( + _resolve_token( + token=token, + all_keys=all_keys, + upstream=upstream, + downstream=downstream, + tag_index=tag_index, + path_index=path_index, + ) + ) + return frozenset(resolved) + + +def match_selector_keys( + *, + parsed: ParsedSelector, + all_keys: dict[str, CompiledObjectKey], + tag_index: dict[str, frozenset[CompiledObjectKey]], + path_index: dict[CompiledObjectKey, str], +) -> frozenset[CompiledObjectKey]: + """Return the keys one parsed selector matches before `+` graph expansion.""" + + if parsed.kind == SelectorKind.TAG: + tagged_keys: frozenset[CompiledObjectKey] = tag_index.get(parsed.value, frozenset()) + if not tagged_keys: + raise PlannerInputError(f"no models found with tag '{parsed.value}'", code="S008") + return tagged_keys + + if parsed.kind == SelectorKind.PATH: + return _match_path(value=parsed.value, path_index=path_index) + + keys: frozenset[CompiledObjectKey] = _lookup_keys(parsed=parsed, all_keys=all_keys) + if not keys: + label: str = "pattern" if _is_name_pattern(parsed.value) else "name" + raise PlannerInputError(f"unknown selector {label} '{parsed.value}'", code="S007") + return keys + + def expand_required_build_resources( *, selected_keys: frozenset[CompiledObjectKey], @@ -272,27 +316,12 @@ def _resolve_single( result.update(expand_downstream(key=end_key, downstream=downstream)) return frozenset(result) - if parsed.kind == SelectorKind.TAG: - return _resolve_tag( - parsed=parsed, - tag_index=tag_index, - upstream=upstream, - downstream=downstream, - ) - - if parsed.kind == SelectorKind.PATH: - return _resolve_path( - parsed=parsed, - path_index=path_index, - upstream=upstream, - downstream=downstream, - ) - - keys: frozenset[CompiledObjectKey] = _lookup_keys(parsed=parsed, all_keys=all_keys) - if not keys: - label: str = "pattern" if _is_name_pattern(parsed.value) else "name" - raise PlannerInputError(f"unknown selector {label} '{parsed.value}'", code="S007") - + keys: frozenset[CompiledObjectKey] = match_selector_keys( + parsed=parsed, + all_keys=all_keys, + tag_index=tag_index, + path_index=path_index, + ) result: set[CompiledObjectKey] = set(keys) key: CompiledObjectKey if parsed.upstream: @@ -304,40 +333,14 @@ def _resolve_single( return frozenset(result) -def _resolve_tag( - *, - parsed: ParsedSelector, - tag_index: dict[str, frozenset[CompiledObjectKey]], - upstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], - downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], -) -> frozenset[CompiledObjectKey]: - """Resolve a tag selector to matching keys with optional graph expansion.""" - - tagged_keys: frozenset[CompiledObjectKey] = tag_index.get(parsed.value, frozenset()) - if not tagged_keys: - raise PlannerInputError(f"no models found with tag '{parsed.value}'", code="S008") - - result: set[CompiledObjectKey] = set(tagged_keys) - key: CompiledObjectKey - if parsed.upstream: - for key in tagged_keys: - result.update(expand_upstream(key=key, upstream=upstream)) - if parsed.downstream: - for key in tagged_keys: - result.update(expand_downstream(key=key, downstream=downstream)) - return frozenset(result) - - -def _resolve_path( +def _match_path( *, - parsed: ParsedSelector, + value: str, path_index: dict[CompiledObjectKey, str], - upstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], - downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], ) -> frozenset[CompiledObjectKey]: - """Resolve a path selector to matching keys with optional graph expansion.""" + """Match a `models/`-rooted path selector against model folders.""" - folder: str = _normalize_path_selector_value(parsed.value) + folder: str = _normalize_path_selector_value(value) selector_folder: str = _model_path_candidate(folder) matched_keys: frozenset[CompiledObjectKey] = frozenset( key @@ -354,16 +357,7 @@ def _resolve_path( ), code="S009", ) - - result: set[CompiledObjectKey] = set(matched_keys) - key: CompiledObjectKey - if parsed.upstream: - for key in matched_keys: - result.update(expand_upstream(key=key, upstream=upstream)) - if parsed.downstream: - for key in matched_keys: - result.update(expand_downstream(key=key, downstream=downstream)) - return frozenset(result) + return matched_keys def _normalize_path_selector_value(value: str) -> str: diff --git a/src/sqlbuild/compiler/planner/main/selection/selector_match.py b/src/sqlbuild/compiler/planner/main/selection/selector_match.py new file mode 100644 index 0000000000..88244d4e29 --- /dev/null +++ b/src/sqlbuild/compiler/planner/main/selection/selector_match.py @@ -0,0 +1,24 @@ +"""Public entrypoint matching one parsed selector before graph expansion.""" + +from __future__ import annotations + +from sqlbuild.compiler.compile.models import CompiledObjectKey +from sqlbuild.compiler.planner._helpers.graph.selectors import match_selector_keys +from sqlbuild.compiler.planner.models import ParsedSelector + + +def match_project_selector( + *, + parsed: ParsedSelector, + all_keys: dict[str, CompiledObjectKey], + tag_index: dict[str, frozenset[CompiledObjectKey]], + path_index: dict[CompiledObjectKey, str], +) -> frozenset[CompiledObjectKey]: + """Return the keys one parsed selector matches before `+` graph expansion.""" + + return match_selector_keys( + parsed=parsed, + all_keys=all_keys, + tag_index=tag_index, + path_index=path_index, + ) diff --git a/src/sqlbuild/compiler/planner/main/selection/selector_tokens.py b/src/sqlbuild/compiler/planner/main/selection/selector_tokens.py new file mode 100644 index 0000000000..a69d10ba0f --- /dev/null +++ b/src/sqlbuild/compiler/planner/main/selection/selector_tokens.py @@ -0,0 +1,27 @@ +"""Public selector-token resolution entrypoint without required build resources.""" + +from __future__ import annotations + +from sqlbuild.compiler.compile.models import CompiledObjectKey +from sqlbuild.compiler.planner._helpers.graph.selectors import resolve_selector_tokens + + +def resolve_project_selector_tokens( + *, + selectors: tuple[str, ...], + all_keys: dict[str, CompiledObjectKey], + upstream_deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], + downstream_deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], + tag_index: dict[str, frozenset[CompiledObjectKey]], + path_index: dict[CompiledObjectKey, str], +) -> frozenset[CompiledObjectKey]: + """Resolve selector strings to the keys they match, without adding required build resources.""" + + return resolve_selector_tokens( + selectors=selectors, + all_keys=all_keys, + upstream=upstream_deps, + downstream=downstream_deps, + tag_index=tag_index, + path_index=path_index, + ) 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 3f88729097..e67b6d85eb 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 @@ -39,3 +39,19 @@ class DiamondLineageTreeCliTestCase: command: tuple[str, ...] expected_expanded_names: tuple[str, ...] expected_max_lines: int + + +@dataclass(frozen=True) +class SharedSelectorCliTestCase: + description: str + selector: str + expected_node_ids: tuple[str, ...] + expected_selected_models: int + expected_selected_functions: int + + +@dataclass(frozen=True) +class SharedSelectorErrorCliTestCase: + description: str + selector: str + expected_fragment: str 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 27d9a44aab..fc11cb389c 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 @@ -13,6 +13,8 @@ LineageCacheCliTestCase, LineageCliTestCase, LineageErrorCliTestCase, + SharedSelectorCliTestCase, + SharedSelectorErrorCliTestCase, ) from tests.e2e.src.sqlbuild.cli.commands.main.lineage.helpers import ( DIAMOND_EDGE_COUNT, @@ -79,7 +81,7 @@ command=( "lineage", "--select", - "path:staging", + "path:models/staging", "--exclude", "tag:finance", "--direction", @@ -539,3 +541,88 @@ def test_given_diamond_dependencies_when_rendering_lineage_tree_then_output_stay 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) + + +@pytest.mark.parametrize( + "test_case", + ( + SharedSelectorCliTestCase( + description="models-rooted path", + selector="path:models/staging", + expected_node_ids=("model:stg_customers", "model:stg_orders", "model:stg_payments"), + expected_selected_models=3, + expected_selected_functions=0, + ), + SharedSelectorCliTestCase( + description="name glob", + selector="stg_*", + expected_node_ids=("model:stg_customers", "model:stg_orders", "model:stg_payments"), + expected_selected_models=3, + expected_selected_functions=0, + ), + SharedSelectorCliTestCase( + description="path and tag intersection", + selector="path:models/marts,tag:acceptance", + expected_node_ids=( + "model:daily_activity_rollup", + "model:hourly_activity_with_daily_context", + "udf:is_completed_order", + "udf:is_completed_order_py", + ), + expected_selected_models=2, + expected_selected_functions=2, + ), + ), + ids=lambda case: case.description, +) +def test_given_selector_when_running_lineage_and_compile_then_both_select_the_same_resources( + test_case: SharedSelectorCliTestCase, + tmp_path: Path, +) -> None: + project_dir: Path = prepare_waffle_shop(tmp_path) + + lineage: subprocess.CompletedProcess[str] = run_sqb( + command=("lineage", "--select", test_case.selector, "--format", "json"), + project_dir=project_dir, + ) + compiled: subprocess.CompletedProcess[str] = run_sqb( + command=("compile", "--select", test_case.selector, "--json"), + project_dir=project_dir, + ) + + assert lineage.returncode == 0, lineage.stdout + lineage.stderr + assert compiled.returncode == 0, compiled.stdout + compiled.stderr + summary: dict[str, int] = json.loads(compiled.stdout)["summary"] + assert lineage_node_ids(payload=json.loads(lineage.stdout)) == test_case.expected_node_ids + assert summary["selected_models"] == test_case.expected_selected_models + assert summary["selected_functions"] == test_case.expected_selected_functions + + +@pytest.mark.parametrize( + "test_case", + ( + SharedSelectorErrorCliTestCase( + description="path without the models root", + selector="path:staging", + expected_fragment="error[S012]", + ), + ), + ids=lambda case: case.description, +) +def test_given_invalid_selector_when_running_lineage_and_compile_then_both_report_it( + test_case: SharedSelectorErrorCliTestCase, + tmp_path: Path, +) -> None: + project_dir: Path = prepare_waffle_shop(tmp_path) + + results: tuple[subprocess.CompletedProcess[str], ...] = tuple( + run_sqb(command=command, project_dir=project_dir) + for command in ( + ("--no-color", "lineage", "--select", test_case.selector), + ("--no-color", "compile", "--select", test_case.selector), + ) + ) + + for result in results: + assert result.returncode == 1, result.stdout + result.stderr + assert test_case.expected_fragment in result.stdout + result.stderr diff --git a/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/_test_types.py b/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/_test_types.py index 8683957ff4..0a8e703a11 100644 --- a/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/_test_types.py +++ b/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/_test_types.py @@ -96,10 +96,24 @@ class LineagePathSelectorTestCase: class LineagePathSelectorErrorTestCase: description: str select: tuple[str, ...] - expected_error_code: str expected_error_fragment: str +@dataclass(frozen=True) +class SharedSelectorParityTestCase: + description: str + select: tuple[str, ...] + exclude: tuple[str, ...] + expected_node_ids: tuple[str, ...] + + +@dataclass(frozen=True) +class SharedSelectorErrorParityTestCase: + description: str + select: tuple[str, ...] + expected_code: str + + @dataclass(frozen=True) class NormalizeLineageTargetTestCase: description: str @@ -112,3 +126,11 @@ class NormalizeLineageTargetErrorTestCase: description: str target: str expected_code: str + + +@dataclass(frozen=True) +class SharedSelectorDepthTestCase: + description: str + select: tuple[str, ...] + direction: str | None + expected_node_ids: tuple[str, ...] diff --git a/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/test_selection.py b/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/test_selection.py index 5b776f78b1..c2b51307d9 100644 --- a/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/test_selection.py +++ b/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/test_selection.py @@ -12,6 +12,7 @@ ) from sqlbuild.cli.commands.exceptions import CliUserError from sqlbuild.cli.commands.models import ColumnLineageTrace, LineageGraph +from sqlbuild.compiler.compile.models import CompiledObjectKey from sqlbuild.compiler.compile.types import CompiledResourceType from sqlbuild.compiler.lineage.models import ( ColumnLineageEdge, @@ -20,6 +21,8 @@ ) from sqlbuild.compiler.lineage.types import ColumnLineageMode from sqlbuild.compiler.pipeline.models import ProjectGraph +from sqlbuild.compiler.planner.exceptions import PlannerInputError +from sqlbuild.compiler.planner.main.selection.selection import resolve_project_selectors from tests.unit.src.sqlbuild.cli.commands._helpers.lineage._test_types import ( ColumnLineageSelectionTestCase, LineagePathSelectorErrorTestCase, @@ -28,6 +31,9 @@ LineageSelectorDepthErrorTestCase, NormalizeLineageTargetErrorTestCase, NormalizeLineageTargetTestCase, + SharedSelectorDepthTestCase, + SharedSelectorErrorParityTestCase, + SharedSelectorParityTestCase, ) from tests.unit.src.sqlbuild.cli.commands._helpers.lineage.helpers import ( build_lineage_test_graph, @@ -314,7 +320,6 @@ def test_given_path_selector_when_selecting_lineage_then_returns_path_nodes( LineagePathSelectorErrorTestCase( description="rejects an end that is not downstream of the start", select=("daily_rollup~raw_orders",), - expected_error_code="C319", expected_error_fragment=( "'source:raw_orders' is not downstream of 'model:daily_rollup'" ), @@ -327,7 +332,7 @@ def test_given_unreachable_path_selector_when_selecting_lineage_then_raises_user ) -> None: graph: ProjectGraph = build_lineage_test_graph() - with pytest.raises(CliUserError) as error: + with pytest.raises(PlannerInputError) as error: select_selector_lineage( graph=graph, select=test_case.select, @@ -335,7 +340,6 @@ def test_given_unreachable_path_selector_when_selecting_lineage_then_raises_user depth=None, ) - assert error.value.code == test_case.expected_error_code assert test_case.expected_error_fragment in str(error.value) @@ -415,3 +419,200 @@ def test_given_mismatched_kind_prefix_when_normalizing_then_raises_unknown_targe _ = normalize_lineage_target(graph=graph, target=test_case.target) assert raised.value.code == test_case.expected_code + + +@pytest.mark.parametrize( + "test_case", + ( + SharedSelectorParityTestCase( + description="models-rooted path selects one folder", + select=("path:models/staging",), + exclude=(), + expected_node_ids=("model:stg_orders",), + ), + SharedSelectorParityTestCase( + description="models root selects every model", + select=("path:models",), + exclude=(), + expected_node_ids=("model:daily_rollup", "model:fact_orders", "model:stg_orders"), + ), + SharedSelectorParityTestCase( + description="bare models-rooted path is a path selector", + select=("models/marts",), + exclude=(), + expected_node_ids=("model:daily_rollup", "model:fact_orders"), + ), + SharedSelectorParityTestCase( + description="path selector expands upstream", + select=("+path:models/staging",), + exclude=(), + expected_node_ids=("model:stg_orders", "source:raw_orders"), + ), + SharedSelectorParityTestCase( + description="name glob matches every resource kind", + select=("*_orders",), + exclude=(), + expected_node_ids=("model:fact_orders", "model:stg_orders", "source:raw_orders"), + ), + SharedSelectorParityTestCase( + description="typed glob matches only that kind", + select=("source:raw_*",), + exclude=(), + expected_node_ids=("source:raw_orders",), + ), + SharedSelectorParityTestCase( + description="seed selector expands downstream", + select=("seed:waffle_types+",), + exclude=(), + expected_node_ids=("model:daily_rollup", "model:fact_orders", "seed:waffle_types"), + ), + SharedSelectorParityTestCase( + description="exclude removes a resolved upstream", + select=("+fact_orders",), + exclude=("source:raw_orders",), + expected_node_ids=("model:fact_orders", "model:stg_orders", "seed:waffle_types"), + ), + SharedSelectorParityTestCase( + description="comma intersects tag and path", + select=("tag:marts,path:models/marts",), + exclude=(), + expected_node_ids=("model:daily_rollup", "model:fact_orders"), + ), + SharedSelectorParityTestCase( + description="space-separated selectors are unioned", + select=("stg_orders daily_rollup",), + exclude=(), + expected_node_ids=("model:daily_rollup", "model:stg_orders"), + ), + SharedSelectorParityTestCase( + description="path-between keeps nodes on the path", + select=("raw_orders~fact_orders",), + exclude=(), + expected_node_ids=("model:fact_orders", "model:stg_orders", "source:raw_orders"), + ), + ), + ids=lambda case: case.description, +) +def test_given_selector_when_selecting_lineage_then_matches_shared_project_selection( + test_case: SharedSelectorParityTestCase, +) -> None: + graph: ProjectGraph = build_lineage_test_graph() + + lineage: LineageGraph = select_selector_lineage( + graph=graph, + select=test_case.select, + exclude=test_case.exclude, + depth=None, + ) + shared: frozenset[CompiledObjectKey] = resolve_project_selectors( + select=test_case.select, + exclude=test_case.exclude, + all_keys=graph.all_keys, + upstream_deps=graph.upstream_deps, + downstream_deps=graph.downstream_deps, + tag_index=graph.tag_index, + path_index=graph.path_index, + ) + + assert node_ids(lineage.nodes) == test_case.expected_node_ids + assert tuple(sorted(f"{key.resource_type}:{key.name}" for key in shared)) == ( + test_case.expected_node_ids + ) + + +@pytest.mark.parametrize( + "test_case", + ( + SharedSelectorErrorParityTestCase( + description="path without models root is rejected", + select=("path:staging",), + expected_code="S012", + ), + SharedSelectorErrorParityTestCase( + description="unknown models-rooted path is rejected", + select=("path:models/missing",), + expected_code="S009", + ), + SharedSelectorErrorParityTestCase( + description="unknown tag is rejected", + select=("tag:missing",), + expected_code="S008", + ), + SharedSelectorErrorParityTestCase( + description="unknown name is rejected", + select=("missing_model",), + expected_code="S007", + ), + SharedSelectorErrorParityTestCase( + description="unsupported selector kind is rejected", + select=("model:fact_orders",), + expected_code="S005", + ), + SharedSelectorErrorParityTestCase( + description="inner plus marker is rejected", + select=("stg+orders",), + expected_code="S002", + ), + ), + ids=lambda case: case.description, +) +def test_given_invalid_selector_when_selecting_lineage_then_raises_shared_selector_error( + test_case: SharedSelectorErrorParityTestCase, +) -> None: + graph: ProjectGraph = build_lineage_test_graph() + + with pytest.raises(PlannerInputError) as lineage_error: + _ = select_selector_lineage(graph=graph, select=test_case.select, exclude=(), depth=None) + with pytest.raises(PlannerInputError) as shared_error: + _ = resolve_project_selectors( + select=test_case.select, + exclude=(), + all_keys=graph.all_keys, + upstream_deps=graph.upstream_deps, + downstream_deps=graph.downstream_deps, + tag_index=graph.tag_index, + path_index=graph.path_index, + ) + + assert lineage_error.value.code == test_case.expected_code + assert shared_error.value.code == test_case.expected_code + assert str(lineage_error.value) == str(shared_error.value) + + +@pytest.mark.parametrize( + "test_case", + [ + SharedSelectorDepthTestCase( + description="glob anchors every match when trimming to depth", + select=("*_orders+",), + direction=None, + expected_node_ids=( + "model:daily_rollup", + "model:fact_orders", + "model:stg_orders", + "source:raw_orders", + ), + ), + SharedSelectorDepthTestCase( + description="models-rooted path expands by direction", + select=("path:models/staging",), + direction="downstream", + expected_node_ids=("model:fact_orders", "model:stg_orders"), + ), + ], + ids=lambda case: case.description, +) +def test_given_shared_selector_with_depth_when_selecting_lineage_then_trims_from_matches( + test_case: SharedSelectorDepthTestCase, +) -> None: + graph: ProjectGraph = build_lineage_test_graph() + + result: LineageGraph = select_selector_lineage( + graph=graph, + select=test_case.select, + exclude=(), + depth=1, + direction=test_case.direction, + ) + + assert node_ids(result.nodes) == test_case.expected_node_ids diff --git a/website/src/content/docs/docs/cli/lineage.mdx b/website/src/content/docs/docs/cli/lineage.mdx index 9979bf7cd8..f1e59e5c79 100644 --- a/website/src/content/docs/docs/cli/lineage.mdx +++ b/website/src/content/docs/docs/cli/lineage.mdx @@ -31,7 +31,7 @@ Exactly one of a positional target or `--select` is required. | `--format` | Output format: `tree` (default), `list`, or `json`. | | `--mode` | Column lineage analysis mode: `rich` (default) or `fast`. | | `--no-sql-analysis` | Disable compile-time SQL analysis (`--no-sql-validation` is an alias). | -| `--select`, `-s` | Select resources using standard selector syntax. | +| `--select`, `-s` | Select resources using the same selector syntax as `sqb compile`, `sqb plan` and `sqb build`, for example `path:models/marts`, `tag:finance` or `stg_*`. | | `--exclude` | Exclude resources from the selection. | ## Model lineage @@ -232,7 +232,7 @@ sqb lineage fact_orders --format json --direction upstream ## Depth limiting -`--depth` controls how many hops from the focus node(s) to include. In selector mode, `--depth` requires name, source, seed, or path-between selectors - it cannot be combined with tag or path selectors or comma-intersection selectors. +`--depth` controls how many hops from the focus node(s) to include. In selector mode without `--direction`, `--depth` requires name (including glob patterns), source, seed, or path-between selectors - it cannot be combined with tag or path selectors or comma-intersection selectors. ```bash # Only immediate parents