From 720984a3c38c604dd9056fa53c765482154fc8d1 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:22:03 +0100 Subject: [PATCH 01/57] refactor(graph): share multi-start transitive closure Add compiler.graph transitive_closure_many (explicit include_starts) and replace the hand-rolled dependency-closure copies in check, lineage, compile diff, planner scopes/source loads/warehouse snapshot/audit scheduling, Python node selectors, virtual Python closure, build blocking, and Dagster selection. Behaviour is unchanged. Add characterization tests for build blocking and planner scope resolution, and unit tests for the new closure entrypoint. (cherry picked from commit 5e61a2d34fed7f30e3f8fc145cc7bb4d69c96061) --- .../cli/commands/_helpers/check/core.py | 15 ++-- .../commands/_helpers/lineage/selection.py | 60 ++++++------- .../compiler/graph/_helpers/algorithms.py | 21 ++++- .../graph/main/transitive_closure_many.py | 15 ++++ .../compiler/pipeline/_helpers/diff.py | 31 ++----- .../_helpers/graph/source_load_nodes.py | 21 +---- .../_helpers/output/audit_scheduling.py | 26 +----- .../planner/_helpers/planning/scopes.py | 26 ++---- .../planner/_helpers/warehouse/snapshot.py | 12 +-- .../python_nodes/_helpers/selectors.py | 70 ++++----------- .../executor/build/_helpers/blocking.py | 16 +--- .../dagster/_helpers/selection.py | 28 +++--- .../planner/_helpers/python_node_closure.py | 9 +- .../compiler/graph/main/_test_types.py | 9 ++ .../graph/main/test_graph_algorithms.py | 80 +++++++++++++++++ .../planner/_helpers/planning/_test_types.py | 9 ++ .../planner/_helpers/planning/test_scopes.py | 86 +++++++++++++++++++ .../executor/build/_helpers/_test_types.py | 9 ++ .../executor/build/_helpers/helpers.py | 15 ++++ .../executor/build/_helpers/test_blocking.py | 80 +++++++++++++++++ 20 files changed, 405 insertions(+), 233 deletions(-) create mode 100644 src/sqlbuild/compiler/graph/main/transitive_closure_many.py create mode 100644 tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/test_scopes.py create mode 100644 tests/unit/src/sqlbuild/executor/build/_helpers/test_blocking.py diff --git a/src/sqlbuild/cli/commands/_helpers/check/core.py b/src/sqlbuild/cli/commands/_helpers/check/core.py index 177b887708..5b09357aab 100644 --- a/src/sqlbuild/cli/commands/_helpers/check/core.py +++ b/src/sqlbuild/cli/commands/_helpers/check/core.py @@ -23,6 +23,7 @@ ) from sqlbuild.compiler.compile.models import CompiledRelationLocation from sqlbuild.compiler.discovery.models import DiscoveredCheckFunction, DiscoveredProjectInputs +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.pipeline.main.relation_targets import build_python_relation_targets from sqlbuild.compiler.pipeline.models import CompilePipelineResult from sqlbuild.compiler.planner.exceptions import PlannerInputError @@ -78,16 +79,10 @@ def check_dependency_closure( ) -> frozenset[str]: """Return transitive executable Python dependencies of selected checks.""" - dependencies: set[str] = set() - pending: list[str] = list(check_names) - while pending: - name: str = pending.pop() - for upstream_name in graph.upstream_deps.get(name, ()): - if upstream_name in dependencies: - continue - dependencies.add(upstream_name) - pending.append(upstream_name) - return frozenset(dependencies - check_names) + dependencies: frozenset[str] = transitive_closure_many( + starts=check_names, edges=graph.upstream_deps, include_starts=False + ) + return dependencies - check_names def _validate_check_selectors(*, select: tuple[str, ...]) -> None: diff --git a/src/sqlbuild/cli/commands/_helpers/lineage/selection.py b/src/sqlbuild/cli/commands/_helpers/lineage/selection.py index 5f93cd9f6b..371cacf6eb 100644 --- a/src/sqlbuild/cli/commands/_helpers/lineage/selection.py +++ b/src/sqlbuild/cli/commands/_helpers/lineage/selection.py @@ -34,6 +34,8 @@ CompiledProject, ) from sqlbuild.compiler.compile.types import CompiledResourceType +from sqlbuild.compiler.graph.main.transitive_closure import transitive_closure +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.lineage.main.columns import build_project_column_lineage from sqlbuild.compiler.lineage.models import ( ColumnLineageEdge, @@ -481,10 +483,7 @@ def _walk_bounded( max_depth: int | None, ) -> frozenset[CompiledObjectKey]: if max_depth is None: - result: set[CompiledObjectKey] = set() - for anchor in anchors: - result.update(_walk_all(key=anchor, deps=deps)) - return frozenset(result) + return transitive_closure_many(starts=anchors, edges=deps, include_starts=False) if max_depth == 0: return frozenset() visited: set[CompiledObjectKey] = set() @@ -501,23 +500,6 @@ def _walk_bounded( return frozenset(visited) -def _walk_all( - *, - key: CompiledObjectKey, - deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], -) -> frozenset[CompiledObjectKey]: - visited: set[CompiledObjectKey] = set() - stack: list[CompiledObjectKey] = [key] - while stack: - current: CompiledObjectKey = stack.pop() - for neighbor in deps.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - stack.append(neighbor) - return frozenset(visited) - - def _lookup_name(*, name: str, all_keys: dict[str, CompiledObjectKey]) -> CompiledObjectKey: key: CompiledObjectKey | None = all_keys.get(name) if key is None: @@ -562,13 +544,14 @@ def _resolve_selectors( ) ) scoped: set[CompiledObjectKey] = selected - excluded - for key in tuple(scoped): - for upstream_key in _walk_all(key=key, deps=upstream): - if upstream_key.resource_type in { - CompiledResourceType.UDF, - CompiledResourceType.TABLE_FN, - }: - scoped.add(upstream_key) + 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) @@ -616,9 +599,9 @@ def _resolve_single( _find_path_keys(start=start_key, end=end_key, downstream=downstream) ) if parsed.upstream: - result.update(_walk_all(key=start_key, deps=upstream)) + result.update(transitive_closure(start=start_key, edges=upstream)) if parsed.downstream: - result.update(_walk_all(key=end_key, deps=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()) @@ -648,11 +631,14 @@ def _apply_selector_expansion( downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], ) -> frozenset[CompiledObjectKey]: result: set[CompiledObjectKey] = set(matched_keys) - for key in matched_keys: - if parsed.upstream: - result.update(_walk_all(key=key, deps=upstream)) - if parsed.downstream: - result.update(_walk_all(key=key, deps=downstream)) + 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) @@ -729,7 +715,9 @@ def _find_path_keys( end: CompiledObjectKey, downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], ) -> frozenset[CompiledObjectKey]: - reachable_from_start: frozenset[CompiledObjectKey] = _walk_all(key=start, deps=downstream) + reachable_from_start: frozenset[CompiledObjectKey] = transitive_closure( + start=start, edges=downstream + ) if end not in reachable_from_start: raise CliUserError( f"'{end.resource_type}:{end.name}' is not downstream of " diff --git a/src/sqlbuild/compiler/graph/_helpers/algorithms.py b/src/sqlbuild/compiler/graph/_helpers/algorithms.py index e6bb069987..f0b611aaf9 100644 --- a/src/sqlbuild/compiler/graph/_helpers/algorithms.py +++ b/src/sqlbuild/compiler/graph/_helpers/algorithms.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterable, Mapping from typing import Any @@ -47,6 +47,25 @@ def transitive_closure_impl[K]( return frozenset(visited) +def transitive_closure_many_impl[K]( + *, starts: Iterable[K], edges: Mapping[K, Iterable[K]], include_starts: bool +) -> frozenset[K]: + """Return all graph keys reachable from any starting key.""" + + start_keys: tuple[K, ...] = tuple(starts) + visited: set[K] = set(start_keys) if include_starts else set() + frontier: list[K] = list(start_keys) + while frontier: + current: K = frontier.pop() + neighbor: K + for neighbor in edges.get(current, ()): + if neighbor in visited: + continue + visited.add(neighbor) + frontier.append(neighbor) + return frozenset(visited) + + def path_nodes_impl[K]( *, start: K, end: K, downstream: Mapping[K, tuple[K, ...]] ) -> frozenset[K] | None: diff --git a/src/sqlbuild/compiler/graph/main/transitive_closure_many.py b/src/sqlbuild/compiler/graph/main/transitive_closure_many.py new file mode 100644 index 0000000000..6b9ff93687 --- /dev/null +++ b/src/sqlbuild/compiler/graph/main/transitive_closure_many.py @@ -0,0 +1,15 @@ +"""Directed graph multi-start transitive closure entrypoint.""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping + +from sqlbuild.compiler.graph._helpers.algorithms import transitive_closure_many_impl + + +def transitive_closure_many[K]( + *, starts: Iterable[K], edges: Mapping[K, Iterable[K]], include_starts: bool +) -> frozenset[K]: + """Return keys reachable from any start, plus the starts when requested.""" + + return transitive_closure_many_impl(starts=starts, edges=edges, include_starts=include_starts) diff --git a/src/sqlbuild/compiler/pipeline/_helpers/diff.py b/src/sqlbuild/compiler/pipeline/_helpers/diff.py index 70f0e6a489..9edc690bd2 100644 --- a/src/sqlbuild/compiler/pipeline/_helpers/diff.py +++ b/src/sqlbuild/compiler/pipeline/_helpers/diff.py @@ -15,6 +15,7 @@ from sqlbuild.compiler.graph.main._build_lineage_upstream_deps import ( build_lineage_upstream_deps, ) +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.pipeline.constants import ( MODEL_PATH_ROOT, PATH_SELECTOR_PREFIX, @@ -222,28 +223,10 @@ def _expand_keys( downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], ) -> frozenset[CompiledObjectKey]: expanded: set[CompiledObjectKey] = set(keys) - key: CompiledObjectKey - for key in keys: - if upstream_requested: - expanded.update(_expand_graph(key=key, graph=upstream)) - if downstream_requested: - expanded.update(_expand_graph(key=key, graph=downstream)) + if upstream_requested: + expanded.update(transitive_closure_many(starts=keys, edges=upstream, include_starts=False)) + if downstream_requested: + expanded.update( + transitive_closure_many(starts=keys, edges=downstream, include_starts=False) + ) return frozenset(expanded) - - -def _expand_graph( - *, - key: CompiledObjectKey, - graph: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], -) -> frozenset[CompiledObjectKey]: - visited: set[CompiledObjectKey] = set() - stack: list[CompiledObjectKey] = [key] - while stack: - current: CompiledObjectKey = stack.pop() - neighbor: CompiledObjectKey - for neighbor in graph.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - stack.append(neighbor) - return frozenset(visited) diff --git a/src/sqlbuild/compiler/planner/_helpers/graph/source_load_nodes.py b/src/sqlbuild/compiler/planner/_helpers/graph/source_load_nodes.py index 3cb2578ba0..2e3218ca50 100644 --- a/src/sqlbuild/compiler/planner/_helpers/graph/source_load_nodes.py +++ b/src/sqlbuild/compiler/planner/_helpers/graph/source_load_nodes.py @@ -4,6 +4,7 @@ from sqlbuild.compiler.compile.models import CompiledObjectKey, CompiledProject from sqlbuild.compiler.compile.types import CompiledResourceType +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.planner._helpers.graph.loader_dag import ( build_intermediate_source_map, build_upstream_intermediate_source_map, @@ -22,7 +23,7 @@ def build_source_load_map( """Return source entries available to source-load and source-read planning.""" relevant_keys: frozenset[CompiledObjectKey] | None = ( - _upstream_closure(selected_keys=selected_keys, upstream_deps=upstream_deps) + transitive_closure_many(starts=selected_keys, edges=upstream_deps, include_starts=True) if upstream_deps is not None else None ) @@ -38,24 +39,6 @@ def build_source_load_map( return source_map -def _upstream_closure( - *, - selected_keys: frozenset[CompiledObjectKey], - upstream_deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], -) -> frozenset[CompiledObjectKey]: - """Return an adapter-neutral dependency closure for selected metadata reads.""" - - closure: set[CompiledObjectKey] = set(selected_keys) - pending: list[CompiledObjectKey] = list(selected_keys) - while pending: - key: CompiledObjectKey = pending.pop() - for upstream_key in upstream_deps.get(key, ()): - if upstream_key not in closure: - closure.add(upstream_key) - pending.append(upstream_key) - return frozenset(closure) - - def build_source_load_entries( *, execution_order: tuple[CompiledObjectKey, ...], diff --git a/src/sqlbuild/compiler/planner/_helpers/output/audit_scheduling.py b/src/sqlbuild/compiler/planner/_helpers/output/audit_scheduling.py index cfc3097ea0..9897711956 100644 --- a/src/sqlbuild/compiler/planner/_helpers/output/audit_scheduling.py +++ b/src/sqlbuild/compiler/planner/_helpers/output/audit_scheduling.py @@ -15,7 +15,7 @@ AttachedAuditTargetKind, CompiledResourceType, ) -from sqlbuild.compiler.planner._helpers.graph.core import expand_downstream +from sqlbuild.compiler.planner._helpers.graph.core import expand_downstream, expand_upstream from sqlbuild.compiler.planner.exceptions import PlannerInputError from sqlbuild.compiler.planner.types import MaterializationType from sqlbuild.compiler.references.types import SqlReferenceKind @@ -110,8 +110,8 @@ def _model_lifecycle_attachment_is_safe( attached_key: CompiledObjectKey = CompiledObjectKey( resource_type=CompiledResourceType.MODEL, name=audit.attached_target_name ) - attached_upstream: frozenset[CompiledObjectKey] = _expand_upstream( - key=attached_key, upstream_deps=upstream_deps + attached_upstream: frozenset[CompiledObjectKey] = expand_upstream( + key=attached_key, upstream=upstream_deps ) dep_key: CompiledObjectKey @@ -211,23 +211,3 @@ def _find_single_safe_latest_owner( return deepest.name return None - - -def _expand_upstream( - *, - key: CompiledObjectKey, - upstream_deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], -) -> frozenset[CompiledObjectKey]: - """Return all transitive upstream keys reachable from the given key.""" - - visited: set[CompiledObjectKey] = set() - stack: list[CompiledObjectKey] = [key] - while stack: - current: CompiledObjectKey = stack.pop() - neighbor: CompiledObjectKey - for neighbor in upstream_deps.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - stack.append(neighbor) - return frozenset(visited) diff --git a/src/sqlbuild/compiler/planner/_helpers/planning/scopes.py b/src/sqlbuild/compiler/planner/_helpers/planning/scopes.py index c4628c9172..333af898ce 100644 --- a/src/sqlbuild/compiler/planner/_helpers/planning/scopes.py +++ b/src/sqlbuild/compiler/planner/_helpers/planning/scopes.py @@ -5,6 +5,7 @@ from dataclasses import replace from sqlbuild.compiler.compile.models import CompiledObjectKey, CompiledProject +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.planner._helpers.graph.scope import build_planner_scope from sqlbuild.compiler.planner.models import ( PlannerPolicies, @@ -35,9 +36,10 @@ def resolve_planner_scopes( exclude=(), auto_load_sources=policies.auto_load_sources, ) - stale_warning_keys: frozenset[CompiledObjectKey] = _upstream_closure( - selected_keys=selected_scope.selected_keys, - upstream_deps=full_scope.upstream_deps, + stale_warning_keys: frozenset[CompiledObjectKey] = transitive_closure_many( + starts=selected_scope.selected_keys, + edges=full_scope.upstream_deps, + include_starts=True, ) stale_warning_scope: PlannerScope = replace( full_scope, @@ -60,21 +62,3 @@ def resolve_planner_scopes( stale_warning_scope=stale_warning_scope, inspection_scope=selected_scope, ) - - -def _upstream_closure( - *, - selected_keys: frozenset[CompiledObjectKey], - upstream_deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], -) -> frozenset[CompiledObjectKey]: - closure: set[CompiledObjectKey] = set(selected_keys) - pending: list[CompiledObjectKey] = list(selected_keys) - while pending: - key: CompiledObjectKey = pending.pop() - upstream_key: CompiledObjectKey - for upstream_key in upstream_deps.get(key, ()): - if upstream_key in closure: - continue - closure.add(upstream_key) - pending.append(upstream_key) - return frozenset(closure) diff --git a/src/sqlbuild/compiler/planner/_helpers/warehouse/snapshot.py b/src/sqlbuild/compiler/planner/_helpers/warehouse/snapshot.py index 45c5b2af52..ebc3905442 100644 --- a/src/sqlbuild/compiler/planner/_helpers/warehouse/snapshot.py +++ b/src/sqlbuild/compiler/planner/_helpers/warehouse/snapshot.py @@ -32,6 +32,7 @@ ) from sqlbuild.compiler.fingerprints.main.read import read_latest_fingerprints from sqlbuild.compiler.fingerprints.models import Fingerprint, FingerprintSet +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.planner._helpers.graph.buildability import check_buildability from sqlbuild.compiler.planner._helpers.graph.core import build_execution_upstream_deps from sqlbuild.compiler.planner._helpers.graph.loader_dag import ( @@ -291,16 +292,7 @@ def _relevant_state_keys( upstream_deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]] = ( build_execution_upstream_deps(project) ) - closure: set[CompiledObjectKey] = set(selected_keys) - pending: list[CompiledObjectKey] = list(selected_keys) - while pending: - key: CompiledObjectKey = pending.pop() - upstream_key: CompiledObjectKey - for upstream_key in upstream_deps.get(key, ()): - if upstream_key not in closure: - closure.add(upstream_key) - pending.append(upstream_key) - return frozenset(closure) + return transitive_closure_many(starts=selected_keys, edges=upstream_deps, include_starts=True) def _resolve_database( diff --git a/src/sqlbuild/compiler/python_nodes/_helpers/selectors.py b/src/sqlbuild/compiler/python_nodes/_helpers/selectors.py index dfe07fc7eb..007fb60e72 100644 --- a/src/sqlbuild/compiler/python_nodes/_helpers/selectors.py +++ b/src/sqlbuild/compiler/python_nodes/_helpers/selectors.py @@ -4,6 +4,7 @@ from fnmatch import fnmatchcase +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.planner.constants import PATH_SELECTOR_EXPLICIT_ROOT_ERROR from sqlbuild.compiler.planner.exceptions import PlannerInputError from sqlbuild.compiler.planner.main.selection.selector_parse import parse_project_selector @@ -83,15 +84,7 @@ def _resolve_single(*, raw: str, graph: PythonNodeGraph) -> frozenset[str]: if not node_names: raise PlannerInputError(f"unknown Python node selector '{parsed.value}'", code="S007") - result: set[str] = set(node_names) - node_name: str - if parsed.upstream: - for node_name in node_names: - result.update(_expand_upstream(name=node_name, graph=graph)) - if parsed.downstream: - for node_name in node_names: - result.update(_expand_downstream(name=node_name, graph=graph)) - return frozenset(result) + return _expand_selected(names=node_names, parsed=parsed, graph=graph) def _resolve_path(*, parsed: ParsedSelector, graph: PythonNodeGraph) -> frozenset[str]: @@ -105,15 +98,7 @@ def _resolve_path(*, parsed: ParsedSelector, graph: PythonNodeGraph) -> frozense if not matched_names: raise PlannerInputError(f"no Python nodes found under path '{folder}'", code="S009") - result: set[str] = set(matched_names) - node_name: str - if parsed.upstream: - for node_name in matched_names: - result.update(_expand_upstream(name=node_name, graph=graph)) - if parsed.downstream: - for node_name in matched_names: - result.update(_expand_downstream(name=node_name, graph=graph)) - return frozenset(result) + return _expand_selected(names=matched_names, parsed=parsed, graph=graph) def _path_matches(*, indexed_folder: str, selector_folder: str) -> bool: @@ -137,15 +122,7 @@ def _resolve_tag(*, parsed: ParsedSelector, graph: PythonNodeGraph) -> frozenset if not tagged_names: raise PlannerInputError(f"no Python nodes found with tag '{parsed.value}'", code="S008") - result: set[str] = set(tagged_names) - node_name: str - if parsed.upstream: - for node_name in tagged_names: - result.update(_expand_upstream(name=node_name, graph=graph)) - if parsed.downstream: - for node_name in tagged_names: - result.update(_expand_downstream(name=node_name, graph=graph)) - return frozenset(result) + return _expand_selected(names=tagged_names, parsed=parsed, graph=graph) def _lookup_node_names(*, parsed: ParsedSelector, graph: PythonNodeGraph) -> frozenset[str]: @@ -180,29 +157,16 @@ def _is_name_pattern(value: str) -> bool: return any(character in value for character in "*?[") -def _expand_upstream(*, name: str, graph: PythonNodeGraph) -> frozenset[str]: - visited: set[str] = set() - stack: list[str] = [name] - while stack: - current: str = stack.pop() - neighbor: str - for neighbor in graph.upstream_deps.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - stack.append(neighbor) - return frozenset(visited) - - -def _expand_downstream(*, name: str, graph: PythonNodeGraph) -> frozenset[str]: - visited: set[str] = set() - stack: list[str] = [name] - while stack: - current: str = stack.pop() - neighbor: str - for neighbor in graph.downstream_deps.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - stack.append(neighbor) - return frozenset(visited) +def _expand_selected( + *, names: frozenset[str], parsed: ParsedSelector, graph: PythonNodeGraph +) -> frozenset[str]: + result: set[str] = set(names) + if parsed.upstream: + result.update( + transitive_closure_many(starts=names, edges=graph.upstream_deps, include_starts=False) + ) + if parsed.downstream: + result.update( + transitive_closure_many(starts=names, edges=graph.downstream_deps, include_starts=False) + ) + return frozenset(result) diff --git a/src/sqlbuild/executor/build/_helpers/blocking.py b/src/sqlbuild/executor/build/_helpers/blocking.py index 56ecc74815..a26680457e 100644 --- a/src/sqlbuild/executor/build/_helpers/blocking.py +++ b/src/sqlbuild/executor/build/_helpers/blocking.py @@ -3,6 +3,7 @@ from __future__ import annotations from sqlbuild.compiler.compile.models import CompiledObjectKey +from sqlbuild.compiler.graph.main.transitive_closure import transitive_closure def downstream_blocked_keys( @@ -13,17 +14,4 @@ def downstream_blocked_keys( ) -> frozenset[CompiledObjectKey]: """Return all selected transitive downstream keys to block.""" - blocked: set[CompiledObjectKey] = set() - stack: list[CompiledObjectKey] = [failed_key] - visited: set[CompiledObjectKey] = set() - while stack: - current: CompiledObjectKey = stack.pop() - neighbor: CompiledObjectKey - for neighbor in downstream_deps.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - if neighbor in selected_keys: - blocked.add(neighbor) - stack.append(neighbor) - return frozenset(blocked) + return transitive_closure(start=failed_key, edges=downstream_deps) & selected_keys diff --git a/src/sqlbuild/integrations/dagster/_helpers/selection.py b/src/sqlbuild/integrations/dagster/_helpers/selection.py index 3d61e529dd..4d118eeca8 100644 --- a/src/sqlbuild/integrations/dagster/_helpers/selection.py +++ b/src/sqlbuild/integrations/dagster/_helpers/selection.py @@ -6,6 +6,7 @@ from collections.abc import Mapping, Sequence from typing import Any +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.planner.main.selection.selector_parse import parse_project_selector from sqlbuild.compiler.planner.models import ParsedSelector, PathSelector from sqlbuild.compiler.planner.types import SelectorKind @@ -118,16 +119,22 @@ def _resolve_atomic( end_id: str = _node_id_for_name(nodes_by_id=nodes_by_id, name=parsed.end_name) matched: set[str] = _shortest_path(start_id=start_id, end_id=end_id, downstream=downstream) if parsed.upstream: - matched.update(_expand(start_ids={start_id}, adjacency=upstream)) + matched.update( + transitive_closure_many(starts=(start_id,), edges=upstream, include_starts=True) + ) if parsed.downstream: - matched.update(_expand(start_ids={end_id}, adjacency=downstream)) + matched.update( + transitive_closure_many(starts=(end_id,), edges=downstream, include_starts=True) + ) return matched matched = _match_parsed_selector(parsed=parsed, dag=dag, nodes_by_id=nodes_by_id) if parsed.upstream: - matched.update(_expand(start_ids=matched, adjacency=upstream)) + matched.update(transitive_closure_many(starts=matched, edges=upstream, include_starts=True)) if parsed.downstream: - matched.update(_expand(start_ids=matched, adjacency=downstream)) + matched.update( + transitive_closure_many(starts=matched, edges=downstream, include_starts=True) + ) return matched @@ -199,19 +206,6 @@ def _dependency_indexes( return upstream, downstream -def _expand(*, start_ids: set[str], adjacency: Mapping[str, set[str]]) -> set[str]: - expanded: set[str] = set(start_ids) - pending: list[str] = list(start_ids) - while pending: - node_id: str = pending.pop() - for adjacent_id in adjacency.get(node_id, set()): - if adjacent_id in expanded: - continue - expanded.add(adjacent_id) - pending.append(adjacent_id) - return expanded - - def _node_id_for_name(*, nodes_by_id: Mapping[str, Mapping[str, Any]], name: str) -> str: matching_ids: list[str] = [ node_id for node_id, node in nodes_by_id.items() if str(node.get("name")) == name diff --git a/src/sqlbuild/virtual/planner/_helpers/python_node_closure.py b/src/sqlbuild/virtual/planner/_helpers/python_node_closure.py index 60ef0b0b18..8de231ee47 100644 --- a/src/sqlbuild/virtual/planner/_helpers/python_node_closure.py +++ b/src/sqlbuild/virtual/planner/_helpers/python_node_closure.py @@ -2,7 +2,7 @@ from __future__ import annotations -from sqlbuild.compiler.graph.main.transitive_closure import transitive_closure +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.planner.models import PlanOutput from sqlbuild.compiler.python_nodes.models import DiscoveredPythonNode, PythonNodeGraph from sqlbuild.compiler.python_nodes.types import PythonNodeKind @@ -49,10 +49,9 @@ def python_upstream_closure( ) -> frozenset[str]: """Return all Python upstreams for selected Python nodes.""" - names: set[str] = set() - for node_name in selected_python_names: - names.update(transitive_closure(start=node_name, edges=python_graph.upstream_deps)) - return frozenset(names) + return transitive_closure_many( + starts=selected_python_names, edges=python_graph.upstream_deps, include_starts=False + ) def python_downstream_closure( diff --git a/tests/unit/src/sqlbuild/compiler/graph/main/_test_types.py b/tests/unit/src/sqlbuild/compiler/graph/main/_test_types.py index 1f38e264b7..64a291db7e 100644 --- a/tests/unit/src/sqlbuild/compiler/graph/main/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/graph/main/_test_types.py @@ -17,6 +17,15 @@ class TransitiveClosureTestCase: expected_nodes: frozenset[str] +@dataclass(frozen=True) +class TransitiveClosureManyTestCase: + description: str + edges: dict[str, tuple[str, ...]] + starts: tuple[str, ...] + include_starts: bool + expected_nodes: frozenset[str] + + @dataclass(frozen=True) class PathNodesTestCase: description: str diff --git a/tests/unit/src/sqlbuild/compiler/graph/main/test_graph_algorithms.py b/tests/unit/src/sqlbuild/compiler/graph/main/test_graph_algorithms.py index f2b2db9a2c..c9198261b8 100644 --- a/tests/unit/src/sqlbuild/compiler/graph/main/test_graph_algorithms.py +++ b/tests/unit/src/sqlbuild/compiler/graph/main/test_graph_algorithms.py @@ -5,12 +5,23 @@ from sqlbuild.compiler.graph.main.invert_edges import invert_edges from sqlbuild.compiler.graph.main.path_nodes import path_nodes from sqlbuild.compiler.graph.main.transitive_closure import transitive_closure +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from tests.unit.src.sqlbuild.compiler.graph.main._test_types import ( InvertEdgesTestCase, PathNodesTestCase, + TransitiveClosureManyTestCase, TransitiveClosureTestCase, ) +_DIAMOND_UPSTREAM: dict[str, tuple[str, ...]] = { + "orders_report": ("orders_left", "orders_right"), + "orders_left": ("stg_orders",), + "orders_right": ("stg_orders",), + "stg_orders": (), + "inventory_report": ("stg_inventory",), + "stg_inventory": (), +} + @pytest.mark.parametrize( "test_case", @@ -66,6 +77,75 @@ def test_given_edges_when_finding_transitive_closure_then_returns_expected_nodes assert result == test_case.expected_nodes +@pytest.mark.parametrize( + "test_case", + [ + TransitiveClosureManyTestCase( + description="excludes a single start", + edges=_DIAMOND_UPSTREAM, + starts=("orders_report",), + include_starts=False, + expected_nodes=frozenset({"orders_left", "orders_right", "stg_orders"}), + ), + TransitiveClosureManyTestCase( + description="includes a single start", + edges=_DIAMOND_UPSTREAM, + starts=("orders_report",), + include_starts=True, + expected_nodes=frozenset( + {"orders_report", "orders_left", "orders_right", "stg_orders"} + ), + ), + TransitiveClosureManyTestCase( + description="unions disconnected starts", + edges=_DIAMOND_UPSTREAM, + starts=("orders_left", "inventory_report"), + include_starts=False, + expected_nodes=frozenset({"stg_orders", "stg_inventory"}), + ), + TransitiveClosureManyTestCase( + description="keeps an excluded start reachable from another start", + edges=_DIAMOND_UPSTREAM, + starts=("orders_report", "orders_left"), + include_starts=False, + expected_nodes=frozenset({"orders_left", "orders_right", "stg_orders"}), + ), + TransitiveClosureManyTestCase( + description="includes starts missing from the edge mapping", + edges=_DIAMOND_UPSTREAM, + starts=("customers",), + include_starts=True, + expected_nodes=frozenset({"customers"}), + ), + TransitiveClosureManyTestCase( + description="returns nothing without starts", + edges=_DIAMOND_UPSTREAM, + starts=(), + include_starts=True, + expected_nodes=frozenset(), + ), + TransitiveClosureManyTestCase( + description="terminates on cycles", + edges={"orders": ("customers",), "customers": ("orders",)}, + starts=("orders",), + include_starts=False, + expected_nodes=frozenset({"orders", "customers"}), + ), + ], + ids=lambda case: case.description, +) +def test_given_edges_when_finding_multi_start_closure_then_returns_expected_nodes( + test_case: TransitiveClosureManyTestCase, +) -> None: + result: frozenset[str] = transitive_closure_many( + starts=test_case.starts, + edges=test_case.edges, + include_starts=test_case.include_starts, + ) + + assert result == test_case.expected_nodes + + @pytest.mark.parametrize( "test_case", [ diff --git a/tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/_test_types.py b/tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/_test_types.py index 64a01b056a..cea5f69ec4 100644 --- a/tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/_test_types.py @@ -65,3 +65,12 @@ class ForeignSnapshotRelationTestCase: desired_days: int expected_direction: str expected_inspections: int + + +@dataclass(frozen=True) +class PlannerScopesTestCase: + description: str + model_deps: dict[str, tuple[str, ...]] + select: tuple[str, ...] + expected_selected_names: frozenset[str] + expected_stale_warning_names: tuple[str, ...] diff --git a/tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/test_scopes.py b/tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/test_scopes.py new file mode 100644 index 0000000000..e6f37aefd9 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/test_scopes.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import pytest + +from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.compiler.planner._helpers.planning.scopes import resolve_planner_scopes +from sqlbuild.compiler.planner.models import ( + PlannerPolicies, + PlannerScopeResolution, + PlannerSelection, +) +from tests.unit.src.sqlbuild.compiler.planner._helpers.helpers import build_test_project +from tests.unit.src.sqlbuild.compiler.planner._helpers.planning._test_types import ( + PlannerScopesTestCase, +) + +_MODEL_DEPS: dict[str, tuple[str, ...]] = { + "stg_orders": (), + "stg_customers": (), + "orders_enriched": ("stg_orders", "stg_customers"), + "orders_summary": ("orders_enriched",), + "orders_report": ("orders_summary",), + "inventory": (), +} + + +@pytest.mark.parametrize( + "test_case", + [ + PlannerScopesTestCase( + description="stale warning scope is the selection plus all upstreams", + model_deps=_MODEL_DEPS, + select=("orders_summary",), + expected_selected_names=frozenset({"orders_summary"}), + expected_stale_warning_names=( + "orders_enriched", + "orders_summary", + "stg_customers", + "stg_orders", + ), + ), + PlannerScopesTestCase( + description="stale warning scope unions upstreams of every selected model", + model_deps=_MODEL_DEPS, + select=("orders_enriched", "inventory"), + expected_selected_names=frozenset({"orders_enriched", "inventory"}), + expected_stale_warning_names=( + "inventory", + "orders_enriched", + "stg_customers", + "stg_orders", + ), + ), + PlannerScopesTestCase( + description="stale warning scope of a root model is the model itself", + model_deps=_MODEL_DEPS, + select=("stg_orders",), + expected_selected_names=frozenset({"stg_orders"}), + expected_stale_warning_names=("stg_orders",), + ), + ], + ids=lambda case: case.description, +) +def test_given_selection_when_resolving_scopes_then_stale_warning_scope_is_upstream_closure( + test_case: PlannerScopesTestCase, +) -> None: + project: CompiledProject = build_test_project(model_deps=test_case.model_deps) + + resolution: PlannerScopeResolution = resolve_planner_scopes( + project=project, + selection=PlannerSelection(select=test_case.select), + policies=PlannerPolicies(), + ) + + selected_names: frozenset[str] = frozenset( + key.name for key in resolution.selected_scope.selected_keys + ) + stale_scope_names: tuple[str, ...] = tuple(sorted(resolution.stale_warning_scope.all_keys)) + stale_order_names: tuple[str, ...] = tuple( + key.name for key in resolution.stale_warning_scope.execution_order + ) + assert selected_names == test_case.expected_selected_names + assert resolution.stale_warning_scope.selected_keys == resolution.selected_scope.selected_keys + assert stale_scope_names == test_case.expected_stale_warning_names + assert tuple(sorted(stale_order_names)) == test_case.expected_stale_warning_names + assert resolution.inspection_scope is resolution.selected_scope diff --git a/tests/unit/src/sqlbuild/executor/build/_helpers/_test_types.py b/tests/unit/src/sqlbuild/executor/build/_helpers/_test_types.py index efbd3a3413..aaedf2c0f3 100644 --- a/tests/unit/src/sqlbuild/executor/build/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/executor/build/_helpers/_test_types.py @@ -191,3 +191,12 @@ class BatchedRetentionReconciliationTestCase: expected_batch_request_ids: tuple[tuple[str, ...], ...] expected_single_inspections: tuple[str, ...] expected_statements: tuple[str, ...] + + +@dataclass(frozen=True) +class DownstreamBlockedKeysTestCase: + description: str + failed_name: str + downstream_deps: dict[str, tuple[str, ...]] + selected_names: frozenset[str] + expected_names: frozenset[str] diff --git a/tests/unit/src/sqlbuild/executor/build/_helpers/helpers.py b/tests/unit/src/sqlbuild/executor/build/_helpers/helpers.py index 13286a522e..da4a54d744 100644 --- a/tests/unit/src/sqlbuild/executor/build/_helpers/helpers.py +++ b/tests/unit/src/sqlbuild/executor/build/_helpers/helpers.py @@ -262,3 +262,18 @@ def batched_retention_request_ids(*, adapter: Mock) -> tuple[tuple[str, ...], .. for batch in adapter.inspect_retentions.call_args_list: batches.append(tuple(request.request_id for request in batch.kwargs["requests"])) return tuple(batches) + + +def model_key(name: str) -> CompiledObjectKey: + return CompiledObjectKey(resource_type=CompiledResourceType.MODEL, name=name) + + +def build_model_key_edges( + *, edges: dict[str, tuple[str, ...]] +) -> dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]]: + key_edges: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]] = {} + name: str + children: tuple[str, ...] + for name, children in edges.items(): + key_edges[model_key(name)] = tuple(model_key(child) for child in children) + return key_edges diff --git a/tests/unit/src/sqlbuild/executor/build/_helpers/test_blocking.py b/tests/unit/src/sqlbuild/executor/build/_helpers/test_blocking.py new file mode 100644 index 0000000000..e8bf977cb7 --- /dev/null +++ b/tests/unit/src/sqlbuild/executor/build/_helpers/test_blocking.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import pytest + +from sqlbuild.compiler.compile.models import CompiledObjectKey +from sqlbuild.executor.build._helpers.blocking import downstream_blocked_keys +from tests.unit.src.sqlbuild.executor.build._helpers._test_types import ( + DownstreamBlockedKeysTestCase, +) +from tests.unit.src.sqlbuild.executor.build._helpers.helpers import ( + build_model_key_edges, + model_key, +) + + +@pytest.mark.parametrize( + "test_case", + [ + DownstreamBlockedKeysTestCase( + description="blocks selected transitive downstreams but not the failed key", + failed_name="stg_orders", + downstream_deps={ + "stg_orders": ("orders_summary", "orders_audit"), + "orders_summary": ("orders_report",), + "orders_audit": (), + "orders_report": (), + "customers": ("orders_report",), + }, + selected_names=frozenset( + {"stg_orders", "orders_summary", "orders_audit", "orders_report", "customers"} + ), + expected_names=frozenset({"orders_summary", "orders_audit", "orders_report"}), + ), + DownstreamBlockedKeysTestCase( + description="walks through unselected keys to selected descendants", + failed_name="stg_orders", + downstream_deps={ + "stg_orders": ("orders_summary",), + "orders_summary": ("orders_report",), + "orders_report": (), + }, + selected_names=frozenset({"stg_orders", "orders_report"}), + expected_names=frozenset({"orders_report"}), + ), + DownstreamBlockedKeysTestCase( + description="blocks a diamond descendant once", + failed_name="stg_orders", + downstream_deps={ + "stg_orders": ("orders_left", "orders_right"), + "orders_left": ("orders_joined",), + "orders_right": ("orders_joined",), + "orders_joined": (), + }, + selected_names=frozenset({"orders_left", "orders_right", "orders_joined"}), + expected_names=frozenset({"orders_left", "orders_right", "orders_joined"}), + ), + DownstreamBlockedKeysTestCase( + description="returns nothing for a failed leaf", + failed_name="orders_report", + downstream_deps={"stg_orders": ("orders_report",), "orders_report": ()}, + selected_names=frozenset({"stg_orders", "orders_report"}), + expected_names=frozenset(), + ), + ], + ids=lambda case: case.description, +) +def test_given_failed_key_when_computing_blocked_keys_then_returns_selected_descendants( + test_case: DownstreamBlockedKeysTestCase, +) -> None: + downstream_deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]] = build_model_key_edges( + edges=test_case.downstream_deps + ) + + blocked: frozenset[CompiledObjectKey] = downstream_blocked_keys( + failed_key=model_key(test_case.failed_name), + downstream_deps=downstream_deps, + selected_keys=frozenset(model_key(name) for name in test_case.selected_names), + ) + + assert blocked == frozenset(model_key(name) for name in test_case.expected_names) From 5f974950a4fb1c58bda5e70039bd658011ba9580 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:42:31 +0100 Subject: [PATCH 02/57] fix(graph): bound closure depth by the shortest route The depth-limited transitive closure walked depth-first and marked nodes visited on first discovery, so a node first reached by a long route was never re-expanded from a shorter one and its descendants within the depth limit were dropped. Walk breadth-first so every node is expanded at its minimum depth, add max_depth to transitive_closure_many, and move the lineage depth-limited walk onto the shared owner. No production caller passed max_depth to the shared owner before; lineage already used its own breadth-first walk with identical semantics. (cherry picked from commit 0bf33f32dc74138dbbaa3b48bf4f91bf6d606059) --- .../commands/_helpers/lineage/selection.py | 51 ++++++++---------- .../compiler/graph/_helpers/algorithms.py | 37 ++++++------- .../graph/main/transitive_closure_many.py | 12 +++-- .../compiler/graph/main/_test_types.py | 1 + .../graph/main/test_graph_algorithms.py | 53 +++++++++++++++++++ 5 files changed, 101 insertions(+), 53 deletions(-) diff --git a/src/sqlbuild/cli/commands/_helpers/lineage/selection.py b/src/sqlbuild/cli/commands/_helpers/lineage/selection.py index 371cacf6eb..b38bc96934 100644 --- a/src/sqlbuild/cli/commands/_helpers/lineage/selection.py +++ b/src/sqlbuild/cli/commands/_helpers/lineage/selection.py @@ -126,9 +126,17 @@ def _expand( anchors: tuple[CompiledObjectKey, ...] = tuple(keys) expanded: set[CompiledObjectKey] = set() if direction in {UPSTREAM_DIRECTION, BOTH_DIRECTIONS}: - expanded.update(_walk_bounded(anchors=anchors, deps=graph.upstream_deps, max_depth=depth)) + expanded.update( + transitive_closure_many( + starts=anchors, edges=graph.upstream_deps, include_starts=False, max_depth=depth + ) + ) if direction in {DOWNSTREAM_DIRECTION, BOTH_DIRECTIONS}: - expanded.update(_walk_bounded(anchors=anchors, deps=graph.downstream_deps, max_depth=depth)) + expanded.update( + transitive_closure_many( + starts=anchors, edges=graph.downstream_deps, include_starts=False, max_depth=depth + ) + ) return frozenset(expanded) @@ -235,7 +243,7 @@ def _column_lineage_candidate_selection( deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]] = ( graph.downstream_deps if direction == DOWNSTREAM_DIRECTION else graph.upstream_deps ) - selected.update(_walk_bounded(anchors=(key,), deps=deps, max_depth=depth)) + selected.update(transitive_closure(start=key, edges=deps, max_depth=depth)) model_names: frozenset[str] = frozenset( selected_key.name for selected_key in selected @@ -244,7 +252,7 @@ def _column_lineage_candidate_selection( if depth is None: return _ColumnLineageCandidateSelection(model_names=model_names, truncated=False) extended: set[CompiledObjectKey] = {key} - extended.update(_walk_bounded(anchors=(key,), deps=deps, max_depth=depth + 1)) + extended.update(transitive_closure(start=key, edges=deps, max_depth=depth + 1)) return _ColumnLineageCandidateSelection( model_names=model_names, truncated=extended != selected, @@ -468,38 +476,21 @@ def _trim_selected_keys( retained.update(anchors.upstream) retained.update(anchors.downstream) retained.update( - _walk_bounded(anchors=anchors.upstream, deps=upstream_deps, max_depth=max_depth) + transitive_closure_many( + starts=anchors.upstream, edges=upstream_deps, include_starts=False, max_depth=max_depth + ) ) retained.update( - _walk_bounded(anchors=anchors.downstream, deps=downstream_deps, max_depth=max_depth) + transitive_closure_many( + starts=anchors.downstream, + edges=downstream_deps, + include_starts=False, + max_depth=max_depth, + ) ) return frozenset(selected_keys & retained) -def _walk_bounded( - *, - anchors: Iterable[CompiledObjectKey], - deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], - max_depth: int | None, -) -> frozenset[CompiledObjectKey]: - if max_depth is None: - return transitive_closure_many(starts=anchors, edges=deps, include_starts=False) - if max_depth == 0: - return frozenset() - visited: set[CompiledObjectKey] = set() - queue: list[tuple[CompiledObjectKey, int]] = [(anchor, 0) for anchor in anchors] - while queue: - current, current_depth = queue.pop(0) - if current_depth >= max_depth: - continue - for neighbor in deps.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - queue.append((neighbor, current_depth + 1)) - return frozenset(visited) - - def _lookup_name(*, name: str, all_keys: dict[str, CompiledObjectKey]) -> CompiledObjectKey: key: CompiledObjectKey | None = all_keys.get(name) if key is None: diff --git a/src/sqlbuild/compiler/graph/_helpers/algorithms.py b/src/sqlbuild/compiler/graph/_helpers/algorithms.py index f0b611aaf9..5f9fb1d266 100644 --- a/src/sqlbuild/compiler/graph/_helpers/algorithms.py +++ b/src/sqlbuild/compiler/graph/_helpers/algorithms.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections import deque from collections.abc import Callable, Iterable, Mapping from typing import Any @@ -30,39 +31,35 @@ def transitive_closure_impl[K]( ) -> frozenset[K]: """Return all graph keys reachable from a starting key.""" - visited: set[K] = set() - frontier: list[tuple[K, int]] = [(start, 0)] - while frontier: - current: K - depth: int - current, depth = frontier.pop() - if max_depth is not None and depth >= max_depth: - continue - neighbor: K - for neighbor in edges.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - frontier.append((neighbor, depth + 1)) - return frozenset(visited) + return transitive_closure_many_impl( + starts=(start,), edges=edges, include_starts=False, max_depth=max_depth + ) def transitive_closure_many_impl[K]( - *, starts: Iterable[K], edges: Mapping[K, Iterable[K]], include_starts: bool + *, + starts: Iterable[K], + edges: Mapping[K, Iterable[K]], + include_starts: bool, + max_depth: int | None = None, ) -> frozenset[K]: - """Return all graph keys reachable from any starting key.""" + """Return graph keys within an optional edge distance of any starting key.""" start_keys: tuple[K, ...] = tuple(starts) visited: set[K] = set(start_keys) if include_starts else set() - frontier: list[K] = list(start_keys) + frontier: deque[tuple[K, int]] = deque((start_key, 0) for start_key in start_keys) while frontier: - current: K = frontier.pop() + current: K + depth: int + current, depth = frontier.popleft() + if max_depth is not None and depth >= max_depth: + continue neighbor: K for neighbor in edges.get(current, ()): if neighbor in visited: continue visited.add(neighbor) - frontier.append(neighbor) + frontier.append((neighbor, depth + 1)) return frozenset(visited) diff --git a/src/sqlbuild/compiler/graph/main/transitive_closure_many.py b/src/sqlbuild/compiler/graph/main/transitive_closure_many.py index 6b9ff93687..fc4025dc0d 100644 --- a/src/sqlbuild/compiler/graph/main/transitive_closure_many.py +++ b/src/sqlbuild/compiler/graph/main/transitive_closure_many.py @@ -8,8 +8,14 @@ def transitive_closure_many[K]( - *, starts: Iterable[K], edges: Mapping[K, Iterable[K]], include_starts: bool + *, + starts: Iterable[K], + edges: Mapping[K, Iterable[K]], + include_starts: bool, + max_depth: int | None = None, ) -> frozenset[K]: - """Return keys reachable from any start, plus the starts when requested.""" + """Return keys within max_depth edges of any start, plus the starts when requested.""" - return transitive_closure_many_impl(starts=starts, edges=edges, include_starts=include_starts) + return transitive_closure_many_impl( + starts=starts, edges=edges, include_starts=include_starts, max_depth=max_depth + ) diff --git a/tests/unit/src/sqlbuild/compiler/graph/main/_test_types.py b/tests/unit/src/sqlbuild/compiler/graph/main/_test_types.py index 64a291db7e..b613351a59 100644 --- a/tests/unit/src/sqlbuild/compiler/graph/main/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/graph/main/_test_types.py @@ -24,6 +24,7 @@ class TransitiveClosureManyTestCase: starts: tuple[str, ...] include_starts: bool expected_nodes: frozenset[str] + max_depth: int | None = None @dataclass(frozen=True) diff --git a/tests/unit/src/sqlbuild/compiler/graph/main/test_graph_algorithms.py b/tests/unit/src/sqlbuild/compiler/graph/main/test_graph_algorithms.py index c9198261b8..232e5c32f5 100644 --- a/tests/unit/src/sqlbuild/compiler/graph/main/test_graph_algorithms.py +++ b/tests/unit/src/sqlbuild/compiler/graph/main/test_graph_algorithms.py @@ -22,6 +22,15 @@ "stg_inventory": (), } +_SHORTCUT_DOWNSTREAM: dict[str, tuple[str, ...]] = { + "orders": ("orders_left", "orders_right"), + "orders_right": ("orders_mid",), + "orders_mid": ("orders_joined",), + "orders_left": ("orders_joined",), + "orders_joined": ("orders_report",), + "orders_report": (), +} + @pytest.mark.parametrize( "test_case", @@ -62,6 +71,15 @@ def test_given_edges_when_inverting_then_returns_expected_edges( max_depth=1, expected_nodes=frozenset({"b"}), ), + TransitiveClosureTestCase( + description="bounds depth by the shortest route", + edges=_SHORTCUT_DOWNSTREAM, + start="orders", + max_depth=3, + expected_nodes=frozenset( + {"orders_left", "orders_right", "orders_mid", "orders_joined", "orders_report"} + ), + ), ], ids=lambda case: case.description, ) @@ -131,6 +149,40 @@ def test_given_edges_when_finding_transitive_closure_then_returns_expected_nodes include_starts=False, expected_nodes=frozenset({"orders", "customers"}), ), + TransitiveClosureManyTestCase( + description="returns nothing at zero depth without starts", + edges=_DIAMOND_UPSTREAM, + starts=("orders_report",), + include_starts=False, + max_depth=0, + expected_nodes=frozenset(), + ), + TransitiveClosureManyTestCase( + description="returns only starts at zero depth with starts", + edges=_DIAMOND_UPSTREAM, + starts=("orders_report",), + include_starts=True, + max_depth=0, + expected_nodes=frozenset({"orders_report"}), + ), + TransitiveClosureManyTestCase( + description="bounds depth from every start", + edges=_DIAMOND_UPSTREAM, + starts=("orders_report", "inventory_report"), + include_starts=False, + max_depth=1, + expected_nodes=frozenset({"orders_left", "orders_right", "stg_inventory"}), + ), + TransitiveClosureManyTestCase( + description="bounds multi-start depth by the shortest route", + edges=_SHORTCUT_DOWNSTREAM, + starts=("orders", "orders_right"), + include_starts=False, + max_depth=3, + expected_nodes=frozenset( + {"orders_left", "orders_right", "orders_mid", "orders_joined", "orders_report"} + ), + ), ], ids=lambda case: case.description, ) @@ -141,6 +193,7 @@ def test_given_edges_when_finding_multi_start_closure_then_returns_expected_node starts=test_case.starts, edges=test_case.edges, include_starts=test_case.include_starts, + max_depth=test_case.max_depth, ) assert result == test_case.expected_nodes From f762e8d090aa625bb1cdcab554fa5de0d38202e5 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:44:30 +0100 Subject: [PATCH 03/57] refactor(graph): reuse shared closure for model-name and path lineage Build the planner model-name closure and the virtual source-freshness incomplete-model set from transitive_closure_many, keeping the resource-type projection at each call site, and resolve lineage path-between selectors through the shared path_nodes owner. Add lineage path-selector tests covering path nodes, endpoint expansion, and the unreachable-endpoint error. run_despite_unchanged._upstream_source_names is left alone: it stops at the nearest source instead of walking through it, which is a pruned traversal rather than a projected closure. (cherry picked from commit 9ab50a13d2413ae8ee606ec376b27cdb55e3cb21) --- .../commands/_helpers/lineage/selection.py | 23 ++--- .../planner/_helpers/graph/model_closure.py | 20 ++--- .../virtual/planner/_helpers/planning.py | 23 ++--- .../commands/_helpers/lineage/_test_types.py | 15 ++++ .../_helpers/lineage/test_selection.py | 84 +++++++++++++++++++ 5 files changed, 118 insertions(+), 47 deletions(-) diff --git a/src/sqlbuild/cli/commands/_helpers/lineage/selection.py b/src/sqlbuild/cli/commands/_helpers/lineage/selection.py index b38bc96934..3a75fbb85b 100644 --- a/src/sqlbuild/cli/commands/_helpers/lineage/selection.py +++ b/src/sqlbuild/cli/commands/_helpers/lineage/selection.py @@ -34,6 +34,7 @@ CompiledProject, ) from sqlbuild.compiler.compile.types import CompiledResourceType +from sqlbuild.compiler.graph.main.path_nodes import path_nodes from sqlbuild.compiler.graph.main.transitive_closure import transitive_closure from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.lineage.main.columns import build_project_column_lineage @@ -706,30 +707,16 @@ def _find_path_keys( end: CompiledObjectKey, downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], ) -> frozenset[CompiledObjectKey]: - reachable_from_start: frozenset[CompiledObjectKey] = transitive_closure( - start=start, edges=downstream + path_keys: frozenset[CompiledObjectKey] | None = path_nodes( + start=start, end=end, downstream=downstream ) - if end not in reachable_from_start: + 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", ) - upstream: dict[CompiledObjectKey, list[CompiledObjectKey]] = {} - for key, dep_keys in downstream.items(): - for dep_key in dep_keys: - upstream.setdefault(dep_key, []).append(key) - upstream_from_end: set[CompiledObjectKey] = set() - stack: list[CompiledObjectKey] = [end] - while stack: - current: CompiledObjectKey = stack.pop() - if current in upstream_from_end: - continue - upstream_from_end.add(current) - for parent in upstream.get(current, ()): - if parent in reachable_from_start or parent == start: - stack.append(parent) - return frozenset(reachable_from_start & upstream_from_end | {start, end}) + return path_keys def _match_path( diff --git a/src/sqlbuild/compiler/planner/_helpers/graph/model_closure.py b/src/sqlbuild/compiler/planner/_helpers/graph/model_closure.py index 7a338b590a..39f0623cc6 100644 --- a/src/sqlbuild/compiler/planner/_helpers/graph/model_closure.py +++ b/src/sqlbuild/compiler/planner/_helpers/graph/model_closure.py @@ -4,6 +4,7 @@ from sqlbuild.compiler.compile.models import CompiledObjectKey from sqlbuild.compiler.compile.types import CompiledResourceType +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many def build_downstream_model_name_closure( @@ -31,17 +32,8 @@ def _build_model_name_closure( start_keys: tuple[CompiledObjectKey, ...], deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], ) -> frozenset[str]: - model_names: set[str] = set() - visited: set[CompiledObjectKey] = set() - stack: list[CompiledObjectKey] = list(start_keys) - while stack: - current: CompiledObjectKey = stack.pop() - if current in visited: - continue - visited.add(current) - if current.resource_type == CompiledResourceType.MODEL: - model_names.add(current.name) - neighbor: CompiledObjectKey - for neighbor in deps.get(current, ()): # pragma: no branch - stack.append(neighbor) - return frozenset(model_names) + return frozenset( + key.name + for key in transitive_closure_many(starts=start_keys, edges=deps, include_starts=True) + if key.resource_type == CompiledResourceType.MODEL + ) diff --git a/src/sqlbuild/virtual/planner/_helpers/planning.py b/src/sqlbuild/virtual/planner/_helpers/planning.py index 4ea8f8589f..27f96153ea 100644 --- a/src/sqlbuild/virtual/planner/_helpers/planning.py +++ b/src/sqlbuild/virtual/planner/_helpers/planning.py @@ -7,6 +7,7 @@ from sqlbuild.compiler.compile.models import CompiledObjectKey from sqlbuild.compiler.compile.types import CompiledResourceType +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.pipeline.models import ProjectGraph from sqlbuild.compiler.planner.exceptions import PlannerInputError from sqlbuild.compiler.planner.main.identity.seed_identity import build_seed_identity @@ -168,21 +169,13 @@ def build_source_freshness_incomplete_model_names( incomplete_source_keys: tuple[Any, ...] = tuple( source.key for source in graph.project.sources if source.name not in source_version_hashes ) - incomplete_model_names: set[str] = set() - source_key: Any - for source_key in incomplete_source_keys: - stack: list[Any] = list(graph.downstream_deps.get(source_key, ())) - visited: set[Any] = set() - while stack: - current: Any = stack.pop() - if current in visited: - continue - visited.add(current) - if current.resource_type == CompiledResourceType.MODEL: - incomplete_model_names.add(current.name) - downstream_key: Any - for downstream_key in graph.downstream_deps.get(current, ()): # pragma: no branch - stack.append(downstream_key) + incomplete_model_names: frozenset[str] = frozenset( + key.name + for key in transitive_closure_many( + starts=incomplete_source_keys, edges=graph.downstream_deps, include_starts=False + ) + if key.resource_type == CompiledResourceType.MODEL + ) return tuple(sorted(incomplete_model_names)) 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 00b008da24..6d016fad87 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 @@ -83,3 +83,18 @@ class LineageFingerprintAvailabilityTestCase: relative_path: str config: str expected_available: bool + + +@dataclass(frozen=True) +class LineagePathSelectorTestCase: + description: str + select: tuple[str, ...] + expected_node_ids: tuple[str, ...] + + +@dataclass(frozen=True) +class LineagePathSelectorErrorTestCase: + description: str + select: tuple[str, ...] + expected_error_code: str + expected_error_fragment: 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 78c628baeb..f23331dd46 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 @@ -21,6 +21,8 @@ from sqlbuild.compiler.pipeline.models import ProjectGraph from tests.unit.src.sqlbuild.cli.commands._helpers.lineage._test_types import ( ColumnLineageSelectionTestCase, + LineagePathSelectorErrorTestCase, + LineagePathSelectorTestCase, LineageSelectionTestCase, LineageSelectorDepthErrorTestCase, ) @@ -250,3 +252,85 @@ def test_given_unclear_selector_anchor_with_depth_when_selecting_then_raises_use ) assert test_case.expected_error_fragment in str(error.value) + + +@pytest.mark.parametrize( + "test_case", + [ + LineagePathSelectorTestCase( + description="keeps only nodes on paths between endpoints", + select=("raw_orders~daily_rollup",), + expected_node_ids=( + "model:daily_rollup", + "model:fact_orders", + "model:stg_orders", + "source:raw_orders", + ), + ), + LineagePathSelectorTestCase( + description="adds upstreams of the path start", + select=("+stg_orders~daily_rollup",), + expected_node_ids=( + "model:daily_rollup", + "model:fact_orders", + "model:stg_orders", + "source:raw_orders", + ), + ), + LineagePathSelectorTestCase( + description="adds downstreams of the path end", + select=("raw_orders~stg_orders+",), + expected_node_ids=( + "model:daily_rollup", + "model:fact_orders", + "model:stg_orders", + "source:raw_orders", + ), + ), + ], + ids=lambda case: case.description, +) +def test_given_path_selector_when_selecting_lineage_then_returns_path_nodes( + test_case: LineagePathSelectorTestCase, +) -> None: + graph: ProjectGraph = build_lineage_test_graph() + + result: LineageGraph = select_selector_lineage( + graph=graph, + select=test_case.select, + exclude=(), + depth=None, + ) + + assert node_ids(result.nodes) == test_case.expected_node_ids + + +@pytest.mark.parametrize( + "test_case", + [ + 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'" + ), + ) + ], + ids=lambda case: case.description, +) +def test_given_unreachable_path_selector_when_selecting_lineage_then_raises_user_error( + test_case: LineagePathSelectorErrorTestCase, +) -> None: + graph: ProjectGraph = build_lineage_test_graph() + + with pytest.raises(CliUserError) as error: + select_selector_lineage( + graph=graph, + select=test_case.select, + exclude=(), + depth=None, + ) + + assert error.value.code == test_case.expected_error_code + assert test_case.expected_error_fragment in str(error.value) From abe5c6ff0dd38673fa996e939e081c0cb2a0f87c Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:50:42 +0100 Subject: [PATCH 04/57] refactor(compiler): share authored value validators across packages Move the optional and required authored value validators out of the private discovery YAML primitives into compiler.authored_values, exposed as compiler-internal main entries, and replace the copies in audit argument parsing, model header parsing, schema column parsing, and dlt loader parsing. Error messages and defaults are preserved: the named boolean takes an explicit default so audit always_run keeps False while column nullable keeps None, and every caller keeps its own error class. Project config helpers keep their distinct path-free message contract and are unchanged. (cherry picked from commit b824326218be9b94f3fc63552cfdc567a4d7918e) --- .../auditing/_helpers/schema_audits.py | 43 +---- .../compiler/authored_values/__init__.py | 1 + .../authored_values/_helpers/__init__.py | 1 + .../authored_values/_helpers/fields.py | 138 ++++++++++++++ .../compiler/authored_values/main/__init__.py | 1 + .../authored_values/main/_optional_bool.py | 22 +++ .../authored_values/main/_optional_mapping.py | 22 +++ .../main/_optional_named_bool.py | 28 +++ .../main/_optional_named_string.py | 22 +++ .../main/_optional_non_empty_string.py | 22 +++ .../main/_optional_string_tuple.py | 22 +++ .../main/_require_non_empty_string.py | 22 +++ .../_helpers/attachment/model_config.py | 14 +- .../_helpers/integrations/loaders.py | 32 ++-- .../discovery/_helpers/sql/schema_columns.py | 39 +--- .../discovery/_helpers/yml/primitives.py | 115 +----------- .../compiler/discovery/_helpers/yml/schema.py | 12 +- .../discovery/_helpers/yml/sources.py | 16 +- .../authored_values/main/_test_types.py | 18 ++ .../compiler/authored_values/main/helpers.py | 112 +++++++++++ .../main/test_authored_values.py | 174 ++++++++++++++++++ 21 files changed, 655 insertions(+), 221 deletions(-) create mode 100644 src/sqlbuild/compiler/authored_values/__init__.py create mode 100644 src/sqlbuild/compiler/authored_values/_helpers/__init__.py create mode 100644 src/sqlbuild/compiler/authored_values/_helpers/fields.py create mode 100644 src/sqlbuild/compiler/authored_values/main/__init__.py create mode 100644 src/sqlbuild/compiler/authored_values/main/_optional_bool.py create mode 100644 src/sqlbuild/compiler/authored_values/main/_optional_mapping.py create mode 100644 src/sqlbuild/compiler/authored_values/main/_optional_named_bool.py create mode 100644 src/sqlbuild/compiler/authored_values/main/_optional_named_string.py create mode 100644 src/sqlbuild/compiler/authored_values/main/_optional_non_empty_string.py create mode 100644 src/sqlbuild/compiler/authored_values/main/_optional_string_tuple.py create mode 100644 src/sqlbuild/compiler/authored_values/main/_require_non_empty_string.py create mode 100644 tests/unit/src/sqlbuild/compiler/authored_values/main/_test_types.py create mode 100644 tests/unit/src/sqlbuild/compiler/authored_values/main/helpers.py create mode 100644 tests/unit/src/sqlbuild/compiler/authored_values/main/test_authored_values.py diff --git a/src/sqlbuild/compiler/auditing/_helpers/schema_audits.py b/src/sqlbuild/compiler/auditing/_helpers/schema_audits.py index 2d39fc3357..94c5395bca 100644 --- a/src/sqlbuild/compiler/auditing/_helpers/schema_audits.py +++ b/src/sqlbuild/compiler/auditing/_helpers/schema_audits.py @@ -14,6 +14,8 @@ from sqlbuild.compiler.auditing.exceptions import MeasurementAuditError from sqlbuild.compiler.auditing.models import MeasurementThresholdBound, MeasurementThresholds from sqlbuild.compiler.auditing.types import AuditSeverity, ThresholdOperator +from sqlbuild.compiler.authored_values.main._optional_named_bool import optional_named_bool +from sqlbuild.compiler.authored_values.main._optional_named_string import optional_named_string from sqlbuild.compiler.resource_names.main._validate_resource_identity import ( validate_resource_identity, ) @@ -62,7 +64,7 @@ def parse_audit_instance_impl( argument_mapping: dict[str, object] = cast(dict[str, object], raw_arguments) option_label: str = f"{label} audit '{definition_name}'" - name: str | None = _optional_named_string( + name: str | None = optional_named_string( raw_value=argument_mapping.get("name"), file_path=file_path, label=option_label, @@ -75,14 +77,14 @@ def parse_audit_instance_impl( kind=f"{label} audit instance", path=file_path, ) - description: str | None = _optional_named_string( + description: str | None = optional_named_string( raw_value=argument_mapping.get("description"), file_path=file_path, label=option_label, key="description", error_class=error_class, ) - raw_severity: str | None = _optional_named_string( + raw_severity: str | None = optional_named_string( raw_value=argument_mapping.get("severity"), file_path=file_path, label=option_label, @@ -98,19 +100,20 @@ def parse_audit_instance_impl( raise error_class( f"{file_path} {option_label} 'severity' must be one of: {allowed}" ) from error - run_scope: str | None = _optional_named_string( + run_scope: str | None = optional_named_string( raw_value=argument_mapping.get("run_scope"), file_path=file_path, label=option_label, key="run_scope", error_class=error_class, ) - always_run: bool = _optional_named_bool( + always_run: bool = optional_named_bool( raw_value=argument_mapping.get("always_run"), file_path=file_path, label=option_label, key="always_run", error_class=error_class, + default=False, ) thresholds: MeasurementThresholds | None = parse_measurement_thresholds( raw_value=argument_mapping.get("thresholds"), @@ -244,33 +247,3 @@ def _parse_threshold_bound(raw_value: object | None) -> MeasurementThresholdBoun if isinstance(raw_limit, bool) or not isinstance(raw_limit, (int, float)): raise MeasurementAuditError(f"{operator.value} threshold requires one numeric value") return MeasurementThresholdBound(operator=operator, limit=float(raw_limit)) - - -def _optional_named_string( - *, - raw_value: object | None, - file_path: Path, - label: str, - key: str, - error_class: type[Exception], -) -> str | None: - if raw_value is None: - return None - if not isinstance(raw_value, str) or not raw_value.strip(): - raise error_class(f"{file_path} {label} '{key}' must be a non-empty string") - return raw_value - - -def _optional_named_bool( - *, - raw_value: object | None, - file_path: Path, - label: str, - key: str, - error_class: type[Exception], -) -> bool: - if raw_value is None: - return False - if not isinstance(raw_value, bool): - raise error_class(f"{file_path} {label} '{key}' must be a boolean") - return raw_value diff --git a/src/sqlbuild/compiler/authored_values/__init__.py b/src/sqlbuild/compiler/authored_values/__init__.py new file mode 100644 index 0000000000..4c2eb9cf6a --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/__init__.py @@ -0,0 +1 @@ +"""Validation of optional and required authored configuration values.""" diff --git a/src/sqlbuild/compiler/authored_values/_helpers/__init__.py b/src/sqlbuild/compiler/authored_values/_helpers/__init__.py new file mode 100644 index 0000000000..9c587e407e --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/_helpers/__init__.py @@ -0,0 +1 @@ +"""Authored configuration value validation implementations.""" diff --git a/src/sqlbuild/compiler/authored_values/_helpers/fields.py b/src/sqlbuild/compiler/authored_values/_helpers/fields.py new file mode 100644 index 0000000000..a67aad1518 --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/_helpers/fields.py @@ -0,0 +1,138 @@ +"""Authored configuration value validation implementations.""" + +from __future__ import annotations + +from pathlib import Path +from typing import cast + + +def require_non_empty_string_impl( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> str: + """Extract a required non-empty string from a mapping.""" + + raw_value: object | None = entry.get(key) + if not isinstance(raw_value, str) or not raw_value.strip(): + raise error_class(f"{file_path} {label} must define non-empty string '{key}'") + return raw_value + + +def optional_non_empty_string_impl( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> str | None: + """Extract an optional non-empty string from a mapping.""" + + return optional_named_string_impl( + raw_value=entry.get(key), + file_path=file_path, + label=label, + key=key, + error_class=error_class, + ) + + +def optional_named_string_impl( + *, + raw_value: object | None, + file_path: Path, + label: str, + key: str, + error_class: type[Exception], +) -> str | None: + """Validate and return an optional named string value.""" + + if raw_value is None: + return None + if not isinstance(raw_value, str) or not raw_value.strip(): + raise error_class(f"{file_path} {label} '{key}' must be a non-empty string") + return raw_value + + +def optional_named_bool_impl[D: (bool, None)]( + *, + raw_value: object | None, + file_path: Path, + label: str, + key: str, + error_class: type[Exception], + default: D, +) -> bool | D: + """Validate and return an optional named boolean value.""" + + if raw_value is None: + return default + if not isinstance(raw_value, bool): + raise error_class(f"{file_path} {label} '{key}' must be a boolean") + return raw_value + + +def optional_bool_impl( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> bool | None: + """Extract an optional boolean from a mapping.""" + + return optional_named_bool_impl( + raw_value=entry.get(key), + file_path=file_path, + label=label, + key=key, + error_class=error_class, + default=None, + ) + + +def optional_mapping_impl( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> dict[str, object]: + """Extract an optional mapping from a mapping, defaulting to empty.""" + + raw_value: object | None = entry.get(key) + if raw_value is None: + return {} + if not isinstance(raw_value, dict): + raise error_class(f"{file_path} {label} '{key}' must be a mapping") + return cast(dict[str, object], raw_value) + + +def optional_string_tuple_impl( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> tuple[str, ...]: + """Extract an optional list of strings from a mapping.""" + + raw_value: object | None = entry.get(key) + if raw_value is None: + return () + if not isinstance(raw_value, list): + raise error_class(f"{file_path} {label} '{key}' must be a list") + items: list[str] = [] + item: object + for item in raw_value: + if not isinstance(item, str): + raise error_class(f"{file_path} {label} '{key}' entries must be strings") + items.append(item) + return tuple(items) diff --git a/src/sqlbuild/compiler/authored_values/main/__init__.py b/src/sqlbuild/compiler/authored_values/main/__init__.py new file mode 100644 index 0000000000..1f01d5721c --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/__init__.py @@ -0,0 +1 @@ +"""Authored configuration value public entries.""" diff --git a/src/sqlbuild/compiler/authored_values/main/_optional_bool.py b/src/sqlbuild/compiler/authored_values/main/_optional_bool.py new file mode 100644 index 0000000000..2c2f3f1372 --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/_optional_bool.py @@ -0,0 +1,22 @@ +"""Public optional bool entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values._helpers.fields import optional_bool_impl + + +def optional_bool( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> bool | None: + """Extract an optional boolean from a mapping.""" + + return optional_bool_impl( + entry=entry, key=key, file_path=file_path, label=label, error_class=error_class + ) diff --git a/src/sqlbuild/compiler/authored_values/main/_optional_mapping.py b/src/sqlbuild/compiler/authored_values/main/_optional_mapping.py new file mode 100644 index 0000000000..c44bd6e5be --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/_optional_mapping.py @@ -0,0 +1,22 @@ +"""Public optional mapping entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values._helpers.fields import optional_mapping_impl + + +def optional_mapping( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> dict[str, object]: + """Extract an optional mapping from a mapping, defaulting to empty.""" + + return optional_mapping_impl( + entry=entry, key=key, file_path=file_path, label=label, error_class=error_class + ) diff --git a/src/sqlbuild/compiler/authored_values/main/_optional_named_bool.py b/src/sqlbuild/compiler/authored_values/main/_optional_named_bool.py new file mode 100644 index 0000000000..8e57a6da90 --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/_optional_named_bool.py @@ -0,0 +1,28 @@ +"""Public optional named bool entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values._helpers.fields import optional_named_bool_impl + + +def optional_named_bool[D: (bool, None)]( + *, + raw_value: object | None, + file_path: Path, + label: str, + key: str, + error_class: type[Exception], + default: D, +) -> bool | D: + """Validate an optional named boolean, returning default when absent.""" + + return optional_named_bool_impl( + raw_value=raw_value, + file_path=file_path, + label=label, + key=key, + error_class=error_class, + default=default, + ) diff --git a/src/sqlbuild/compiler/authored_values/main/_optional_named_string.py b/src/sqlbuild/compiler/authored_values/main/_optional_named_string.py new file mode 100644 index 0000000000..2756d0745a --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/_optional_named_string.py @@ -0,0 +1,22 @@ +"""Public optional named string entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values._helpers.fields import optional_named_string_impl + + +def optional_named_string( + *, + raw_value: object | None, + file_path: Path, + label: str, + key: str, + error_class: type[Exception], +) -> str | None: + """Validate and return an optional named string value.""" + + return optional_named_string_impl( + raw_value=raw_value, file_path=file_path, label=label, key=key, error_class=error_class + ) diff --git a/src/sqlbuild/compiler/authored_values/main/_optional_non_empty_string.py b/src/sqlbuild/compiler/authored_values/main/_optional_non_empty_string.py new file mode 100644 index 0000000000..28f1af346f --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/_optional_non_empty_string.py @@ -0,0 +1,22 @@ +"""Public optional non empty string entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values._helpers.fields import optional_non_empty_string_impl + + +def optional_non_empty_string( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> str | None: + """Extract an optional non-empty string from a mapping.""" + + return optional_non_empty_string_impl( + entry=entry, key=key, file_path=file_path, label=label, error_class=error_class + ) diff --git a/src/sqlbuild/compiler/authored_values/main/_optional_string_tuple.py b/src/sqlbuild/compiler/authored_values/main/_optional_string_tuple.py new file mode 100644 index 0000000000..716bf965c4 --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/_optional_string_tuple.py @@ -0,0 +1,22 @@ +"""Public optional string tuple entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values._helpers.fields import optional_string_tuple_impl + + +def optional_string_tuple( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> tuple[str, ...]: + """Extract an optional list of strings from a mapping.""" + + return optional_string_tuple_impl( + entry=entry, key=key, file_path=file_path, label=label, error_class=error_class + ) diff --git a/src/sqlbuild/compiler/authored_values/main/_require_non_empty_string.py b/src/sqlbuild/compiler/authored_values/main/_require_non_empty_string.py new file mode 100644 index 0000000000..bd7f087b9a --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/_require_non_empty_string.py @@ -0,0 +1,22 @@ +"""Public require non empty string entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values._helpers.fields import require_non_empty_string_impl + + +def require_non_empty_string( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> str: + """Extract a required non-empty string from a mapping.""" + + return require_non_empty_string_impl( + entry=entry, key=key, file_path=file_path, label=label, error_class=error_class + ) diff --git a/src/sqlbuild/compiler/compile/_helpers/attachment/model_config.py b/src/sqlbuild/compiler/compile/_helpers/attachment/model_config.py index 67e92c85ed..2cd8790d7b 100644 --- a/src/sqlbuild/compiler/compile/_helpers/attachment/model_config.py +++ b/src/sqlbuild/compiler/compile/_helpers/attachment/model_config.py @@ -7,6 +7,7 @@ from typing import cast from sqlbuild.compiler.auditing.main._parse_audit_instance import parse_audit_instance +from sqlbuild.compiler.authored_values.main._optional_named_string import optional_named_string from sqlbuild.compiler.compile._helpers.audit_factories.core import ( merge_validated_model_audits, parse_model_header_audit_factories, @@ -265,11 +266,12 @@ def build_model_header_schema_entry( ): return None - model_description: str | None = _optional_model_header_string( + model_description: str | None = optional_named_string( raw_value=raw_description, file_path=file_path, label="model", key="description", + error_class=CompileInputError, ) description: str | None = model_description or model_schema_description local_columns: tuple[SchemaColumn, ...] = _parse_model_header_columns( @@ -514,16 +516,6 @@ def _parse_model_header_audits( ) -def _optional_model_header_string( - *, raw_value: object | None, file_path: Path, label: str, key: str -) -> str | None: - if raw_value is None: - return None - if not isinstance(raw_value, str) or not raw_value.strip(): - raise CompileInputError(f"{file_path} {label} '{key}' must be a non-empty string") - return raw_value - - def _resolve_model_schema( *, values: dict[str, object], diff --git a/src/sqlbuild/compiler/discovery/_helpers/integrations/loaders.py b/src/sqlbuild/compiler/discovery/_helpers/integrations/loaders.py index 2175e6b6ee..81ced7e510 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/integrations/loaders.py +++ b/src/sqlbuild/compiler/discovery/_helpers/integrations/loaders.py @@ -5,8 +5,11 @@ from pathlib import Path from typing import cast -from sqlbuild.compiler.discovery._helpers.yml.primitives import ( +from sqlbuild.compiler.authored_values.main._optional_mapping import optional_mapping +from sqlbuild.compiler.authored_values.main._optional_non_empty_string import ( optional_non_empty_string, +) +from sqlbuild.compiler.authored_values.main._require_non_empty_string import ( require_non_empty_string, ) from sqlbuild.compiler.discovery.constants import ( @@ -191,8 +194,12 @@ def _parse_dlt_source_group( if not isinstance(raw_config, dict): raise SourceParseError(f"{file_path} dlt source config must be a mapping") group_config: dict[str, object] = cast(dict[str, object], raw_config) - destination_config: dict[str, object] = _optional_mapping_value( - entry=group, key="destination", file_path=file_path, label="dlt source" + destination_config: dict[str, object] = optional_mapping( + entry=group, + key="destination", + file_path=file_path, + label="dlt source", + error_class=SourceParseError, ) _validate_dlt_destination_config(config=destination_config, file_path=file_path) group_schema: str | None = optional_non_empty_string( @@ -274,8 +281,12 @@ def _parse_dlt_resource_entry( write_disposition=write_disposition, primary_key=resource.get("primary_key"), merge_key=resource.get("merge_key"), - incremental=_optional_mapping_value( - entry=resource, key="incremental", file_path=file_path, label=f"dlt resource '{name}'" + incremental=optional_mapping( + entry=resource, + key="incremental", + file_path=file_path, + label=f"dlt resource '{name}'", + error_class=SourceParseError, ), ) return SourceEntry( @@ -372,14 +383,3 @@ def _raw_dlt_resource_config(*, source_type: str, resource: dict[str, object]) - if source_type == DLT_SOURCE_TYPE_REST_API: return {"endpoint": resource["endpoint"]} return {key: value for key, value in resource.items() if key not in excluded} - - -def _optional_mapping_value( - *, entry: dict[str, object], key: str, file_path: Path, label: str -) -> dict[str, object]: - value: object | None = entry.get(key) - if value is None: - return {} - if not isinstance(value, dict): - raise SourceParseError(f"{file_path} {label} '{key}' must be a mapping") - return cast(dict[str, object], value) diff --git a/src/sqlbuild/compiler/discovery/_helpers/sql/schema_columns.py b/src/sqlbuild/compiler/discovery/_helpers/sql/schema_columns.py index 35661dcfcf..1515de5667 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/sql/schema_columns.py +++ b/src/sqlbuild/compiler/discovery/_helpers/sql/schema_columns.py @@ -7,6 +7,8 @@ from typing import cast from sqlbuild.compiler.auditing.main._parse_audit_instance import parse_audit_instance +from sqlbuild.compiler.authored_values.main._optional_named_bool import optional_named_bool +from sqlbuild.compiler.authored_values.main._optional_named_string import optional_named_string from sqlbuild.compiler.compile.constants import NOT_NULL_AUDIT_NAME from sqlbuild.compiler.compile.exceptions import CompileInputError from sqlbuild.compiler.discovery.exceptions import DeclarationParseError @@ -67,12 +69,13 @@ def parse_schema_columns( f"{file_path} {label} column '{raw_column_name}' has unknown metadata keys: " f"{', '.join(sorted(unknown_keys))}" ) - nullable: bool | None = _optional_bool( + nullable: bool | None = optional_named_bool( raw_value=column_metadata.get("nullable"), file_path=file_path, label=f"{label} column '{raw_column_name}'", key="nullable", error_class=error_class, + default=None, ) column_location: SourceLocation | None = locations.get(raw_column_name) audits: tuple[SchemaAuditInstance, ...] = tuple( @@ -101,7 +104,7 @@ def parse_schema_columns( parsed_columns.append( SchemaColumn( name=raw_column_name, - type=_optional_string( + type=optional_named_string( raw_value=column_metadata.get("type"), file_path=file_path, label=f"{label} column '{raw_column_name}'", @@ -109,7 +112,7 @@ def parse_schema_columns( error_class=error_class, ), nullable=nullable, - description=_optional_string( + description=optional_named_string( raw_value=column_metadata.get("description"), file_path=file_path, label=f"{label} column '{raw_column_name}'", @@ -139,33 +142,3 @@ def _parse_audits( ) for raw_audit in raw_audits ) - - -def _optional_string( - *, - raw_value: object | None, - file_path: Path, - label: str, - key: str, - error_class: _SchemaColumnParseError, -) -> str | None: - if raw_value is None: - return None - if not isinstance(raw_value, str) or not raw_value.strip(): - raise error_class(f"{file_path} {label} '{key}' must be a non-empty string") - return raw_value - - -def _optional_bool( - *, - raw_value: object | None, - file_path: Path, - label: str, - key: str, - error_class: _SchemaColumnParseError, -) -> bool | None: - if raw_value is None: - return None - if not isinstance(raw_value, bool): - raise error_class(f"{file_path} {label} '{key}' must be a boolean") - return raw_value diff --git a/src/sqlbuild/compiler/discovery/_helpers/yml/primitives.py b/src/sqlbuild/compiler/discovery/_helpers/yml/primitives.py index a41b03f36b..69c8e8dc48 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/yml/primitives.py +++ b/src/sqlbuild/compiler/discovery/_helpers/yml/primitives.py @@ -1,9 +1,8 @@ -"""Shared YAML field-extraction primitives for discovery parsers.""" +"""Shared YAML audit-list parsing for discovery parsers.""" from __future__ import annotations from pathlib import Path -from typing import cast from sqlbuild.compiler.auditing.main._parse_audit_instance import parse_audit_instance from sqlbuild.spec.contracts.models import SchemaAuditInstance @@ -34,115 +33,3 @@ def parse_audit_instances( ) ) return tuple(parsed_audits) - - -def require_non_empty_string( - *, - entry: dict[str, object], - key: str, - file_path: Path, - label: str, - error_class: type[Exception], -) -> str: - """Extract a required non-empty string from a YAML mapping.""" - - raw_value: object | None = entry.get(key) - if not isinstance(raw_value, str) or not raw_value.strip(): - raise error_class(f"{file_path} {label} must define non-empty string '{key}'") - return raw_value - - -def optional_non_empty_string( - *, - entry: dict[str, object], - key: str, - file_path: Path, - label: str, - error_class: type[Exception], -) -> str | None: - """Extract an optional non-empty string from a YAML mapping.""" - - return optional_named_string( - raw_value=entry.get(key), - file_path=file_path, - label=label, - key=key, - error_class=error_class, - ) - - -def optional_named_string( - *, - raw_value: object | None, - file_path: Path, - label: str, - key: str, - error_class: type[Exception], -) -> str | None: - """Validate and return an optional named string value.""" - - if raw_value is None: - return None - if not isinstance(raw_value, str) or not raw_value.strip(): - raise error_class(f"{file_path} {label} '{key}' must be a non-empty string") - return raw_value - - -def optional_bool( - *, - entry: dict[str, object], - key: str, - file_path: Path, - label: str, - error_class: type[Exception], -) -> bool | None: - """Extract an optional boolean from a YAML mapping.""" - - raw_value: object | None = entry.get(key) - if raw_value is None: - return None - if not isinstance(raw_value, bool): - raise error_class(f"{file_path} {label} '{key}' must be a boolean") - return raw_value - - -def optional_mapping( - *, - entry: dict[str, object], - key: str, - file_path: Path, - label: str, - error_class: type[Exception], -) -> dict[str, object]: - """Extract an optional mapping from a YAML mapping.""" - - raw_value: object | None = entry.get(key) - if raw_value is None: - return {} - if not isinstance(raw_value, dict): - raise error_class(f"{file_path} {label} '{key}' must be a mapping") - return cast(dict[str, object], raw_value) - - -def optional_string_tuple( - *, - entry: dict[str, object], - key: str, - file_path: Path, - label: str, - error_class: type[Exception], -) -> tuple[str, ...]: - """Extract an optional list of strings from a YAML mapping.""" - - raw_value: object | None = entry.get(key) - if raw_value is None: - return () - if not isinstance(raw_value, list): - raise error_class(f"{file_path} {label} '{key}' must be a list") - items: list[str] = [] - item: object - for item in raw_value: - if not isinstance(item, str): - raise error_class(f"{file_path} {label} '{key}' entries must be strings") - items.append(item) - return tuple(items) diff --git a/src/sqlbuild/compiler/discovery/_helpers/yml/schema.py b/src/sqlbuild/compiler/discovery/_helpers/yml/schema.py index 45447c5501..c41d9573d9 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/yml/schema.py +++ b/src/sqlbuild/compiler/discovery/_helpers/yml/schema.py @@ -8,14 +8,16 @@ import yaml from yaml import YAMLError -from sqlbuild.compiler.discovery._helpers.yml.primitives import ( - optional_bool, - optional_mapping, +from sqlbuild.compiler.authored_values.main._optional_bool import optional_bool +from sqlbuild.compiler.authored_values.main._optional_mapping import optional_mapping +from sqlbuild.compiler.authored_values.main._optional_non_empty_string import ( optional_non_empty_string, - optional_string_tuple, - parse_audit_instances, +) +from sqlbuild.compiler.authored_values.main._optional_string_tuple import optional_string_tuple +from sqlbuild.compiler.authored_values.main._require_non_empty_string import ( require_non_empty_string, ) +from sqlbuild.compiler.discovery._helpers.yml.primitives import parse_audit_instances from sqlbuild.compiler.discovery.constants import NOT_NULL_AUDIT_NAME, SEEDS_DIRECTORY_NAME from sqlbuild.compiler.discovery.exceptions import SchemaParseError from sqlbuild.spec.contracts.models import ( diff --git a/src/sqlbuild/compiler/discovery/_helpers/yml/sources.py b/src/sqlbuild/compiler/discovery/_helpers/yml/sources.py index ebfd0fe8cb..28bdab662f 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/yml/sources.py +++ b/src/sqlbuild/compiler/discovery/_helpers/yml/sources.py @@ -8,18 +8,20 @@ import yaml from yaml import YAMLError +from sqlbuild.compiler.authored_values.main._optional_bool import optional_bool +from sqlbuild.compiler.authored_values.main._optional_mapping import optional_mapping +from sqlbuild.compiler.authored_values.main._optional_non_empty_string import ( + optional_non_empty_string, +) +from sqlbuild.compiler.authored_values.main._require_non_empty_string import ( + require_non_empty_string, +) from sqlbuild.compiler.discovery._helpers.integrations.loaders import ( integration_loader_name, parse_dlt_sources, parse_source_integration_loader, ) -from sqlbuild.compiler.discovery._helpers.yml.primitives import ( - optional_bool, - optional_mapping, - optional_non_empty_string, - parse_audit_instances, - require_non_empty_string, -) +from sqlbuild.compiler.discovery._helpers.yml.primitives import parse_audit_instances from sqlbuild.compiler.discovery.constants import ( NOT_NULL_AUDIT_NAME, SOURCE_AGE_POLICY_CONFIG_KEY, diff --git a/tests/unit/src/sqlbuild/compiler/authored_values/main/_test_types.py b/tests/unit/src/sqlbuild/compiler/authored_values/main/_test_types.py new file mode 100644 index 0000000000..0f86900ab7 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/authored_values/main/_test_types.py @@ -0,0 +1,18 @@ +from collections.abc import Callable +from dataclasses import dataclass + + +@dataclass(frozen=True) +class AuthoredValueTestCase: + description: str + parse: Callable[[object | None], object] + raw_value: object | None + expected_value: object + + +@dataclass(frozen=True) +class AuthoredValueErrorTestCase: + description: str + parse: Callable[[object | None], object] + raw_value: object | None + expected_message: str diff --git a/tests/unit/src/sqlbuild/compiler/authored_values/main/helpers.py b/tests/unit/src/sqlbuild/compiler/authored_values/main/helpers.py new file mode 100644 index 0000000000..2b3a1b94bc --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/authored_values/main/helpers.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values.main._optional_bool import optional_bool +from sqlbuild.compiler.authored_values.main._optional_mapping import optional_mapping +from sqlbuild.compiler.authored_values.main._optional_named_bool import optional_named_bool +from sqlbuild.compiler.authored_values.main._optional_named_string import ( + optional_named_string, +) +from sqlbuild.compiler.authored_values.main._optional_non_empty_string import ( + optional_non_empty_string, +) +from sqlbuild.compiler.authored_values.main._optional_string_tuple import ( + optional_string_tuple, +) +from sqlbuild.compiler.authored_values.main._require_non_empty_string import ( + require_non_empty_string, +) + +AUTHORED_VALUES_FILE_PATH: Path = Path("models/orders.yml") +AUTHORED_VALUES_LABEL: str = "model 'orders'" + + +class AuthoredValueTestError(Exception): + """Error class injected into authored value parsers under test.""" + + +def _entry(raw_value: object | None) -> dict[str, object]: + return {"setting": raw_value} + + +def parse_required_string(raw_value: object | None) -> object: + return require_non_empty_string( + entry=_entry(raw_value), + key="setting", + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + error_class=AuthoredValueTestError, + ) + + +def parse_optional_string(raw_value: object | None) -> object: + return optional_non_empty_string( + entry=_entry(raw_value), + key="setting", + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + error_class=AuthoredValueTestError, + ) + + +def parse_named_string(raw_value: object | None) -> object: + return optional_named_string( + raw_value=raw_value, + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + key="setting", + error_class=AuthoredValueTestError, + ) + + +def parse_optional_bool(raw_value: object | None) -> object: + return optional_bool( + entry=_entry(raw_value), + key="setting", + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + error_class=AuthoredValueTestError, + ) + + +def parse_named_bool_default_false(raw_value: object | None) -> object: + return optional_named_bool( + raw_value=raw_value, + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + key="setting", + error_class=AuthoredValueTestError, + default=False, + ) + + +def parse_named_bool_default_none(raw_value: object | None) -> object: + return optional_named_bool( + raw_value=raw_value, + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + key="setting", + error_class=AuthoredValueTestError, + default=None, + ) + + +def parse_optional_mapping(raw_value: object | None) -> object: + return optional_mapping( + entry=_entry(raw_value), + key="setting", + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + error_class=AuthoredValueTestError, + ) + + +def parse_optional_string_tuple(raw_value: object | None) -> object: + return optional_string_tuple( + entry=_entry(raw_value), + key="setting", + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + error_class=AuthoredValueTestError, + ) diff --git a/tests/unit/src/sqlbuild/compiler/authored_values/main/test_authored_values.py b/tests/unit/src/sqlbuild/compiler/authored_values/main/test_authored_values.py new file mode 100644 index 0000000000..518c6ae372 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/authored_values/main/test_authored_values.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import pytest + +from tests.unit.src.sqlbuild.compiler.authored_values.main._test_types import ( + AuthoredValueErrorTestCase, + AuthoredValueTestCase, +) +from tests.unit.src.sqlbuild.compiler.authored_values.main.helpers import ( + AuthoredValueTestError, + parse_named_bool_default_false, + parse_named_bool_default_none, + parse_named_string, + parse_optional_bool, + parse_optional_mapping, + parse_optional_string, + parse_optional_string_tuple, + parse_required_string, +) + + +@pytest.mark.parametrize( + "test_case", + [ + AuthoredValueTestCase( + description="required string keeps surrounding whitespace", + parse=parse_required_string, + raw_value=" orders ", + expected_value=" orders ", + ), + AuthoredValueTestCase( + description="optional string defaults to none", + parse=parse_optional_string, + raw_value=None, + expected_value=None, + ), + AuthoredValueTestCase( + description="optional string returns value", + parse=parse_optional_string, + raw_value="orders", + expected_value="orders", + ), + AuthoredValueTestCase( + description="named string defaults to none", + parse=parse_named_string, + raw_value=None, + expected_value=None, + ), + AuthoredValueTestCase( + description="optional bool defaults to none", + parse=parse_optional_bool, + raw_value=None, + expected_value=None, + ), + AuthoredValueTestCase( + description="optional bool returns false", + parse=parse_optional_bool, + raw_value=False, + expected_value=False, + ), + AuthoredValueTestCase( + description="named bool uses false default", + parse=parse_named_bool_default_false, + raw_value=None, + expected_value=False, + ), + AuthoredValueTestCase( + description="named bool uses none default", + parse=parse_named_bool_default_none, + raw_value=None, + expected_value=None, + ), + AuthoredValueTestCase( + description="named bool returns true over default", + parse=parse_named_bool_default_false, + raw_value=True, + expected_value=True, + ), + AuthoredValueTestCase( + description="optional mapping defaults to empty", + parse=parse_optional_mapping, + raw_value=None, + expected_value={}, + ), + AuthoredValueTestCase( + description="optional mapping returns value", + parse=parse_optional_mapping, + raw_value={"region": "west"}, + expected_value={"region": "west"}, + ), + AuthoredValueTestCase( + description="optional string list defaults to empty", + parse=parse_optional_string_tuple, + raw_value=None, + expected_value=(), + ), + AuthoredValueTestCase( + description="optional string list returns tuple", + parse=parse_optional_string_tuple, + raw_value=["orders", "customers"], + expected_value=("orders", "customers"), + ), + ], + ids=lambda case: case.description, +) +def test_given_authored_value_when_parsing_then_returns_expected_value( + test_case: AuthoredValueTestCase, +) -> None: + assert test_case.parse(test_case.raw_value) == test_case.expected_value + + +@pytest.mark.parametrize( + "test_case", + [ + AuthoredValueErrorTestCase( + description="required string rejects missing value", + parse=parse_required_string, + raw_value=None, + expected_message=( + "models/orders.yml model 'orders' must define non-empty string 'setting'" + ), + ), + AuthoredValueErrorTestCase( + description="optional string rejects blank value", + parse=parse_optional_string, + raw_value=" ", + expected_message="models/orders.yml model 'orders' 'setting' must be a non-empty string", + ), + AuthoredValueErrorTestCase( + description="named string rejects non-string value", + parse=parse_named_string, + raw_value=3, + expected_message="models/orders.yml model 'orders' 'setting' must be a non-empty string", + ), + AuthoredValueErrorTestCase( + description="optional bool rejects string value", + parse=parse_optional_bool, + raw_value="yes", + expected_message="models/orders.yml model 'orders' 'setting' must be a boolean", + ), + AuthoredValueErrorTestCase( + description="named bool rejects integer value", + parse=parse_named_bool_default_false, + raw_value=1, + expected_message="models/orders.yml model 'orders' 'setting' must be a boolean", + ), + AuthoredValueErrorTestCase( + description="optional mapping rejects list value", + parse=parse_optional_mapping, + raw_value=["region"], + expected_message="models/orders.yml model 'orders' 'setting' must be a mapping", + ), + AuthoredValueErrorTestCase( + description="optional string list rejects mapping value", + parse=parse_optional_string_tuple, + raw_value={"region": "west"}, + expected_message="models/orders.yml model 'orders' 'setting' must be a list", + ), + AuthoredValueErrorTestCase( + description="optional string list rejects non-string entries", + parse=parse_optional_string_tuple, + raw_value=["orders", 3], + expected_message="models/orders.yml model 'orders' 'setting' entries must be strings", + ), + ], + ids=lambda case: case.description, +) +def test_given_invalid_authored_value_when_parsing_then_raises_injected_error( + test_case: AuthoredValueErrorTestCase, +) -> None: + with pytest.raises(AuthoredValueTestError) as error: + test_case.parse(test_case.raw_value) + + assert str(error.value) == test_case.expected_message From bd9226c488b56aa8958ebd91d53d043e622eb194 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:54:22 +0100 Subject: [PATCH 05/57] refactor(hooks): share one lifecycle hook serializer Compile output, the manifest, and the model version identity each copied the same SqlHookEntry/PythonHookEntry serializer. Replace them with compiler.discovery serialize_hook_entries, which takes an explicit SQL hook field allowlist and per-name Python hook fields. Every caller still uses the full field list, so output and identities are unchanged. (cherry picked from commit 9b21e5fe093aa196180d8e53af9b810489903f20) --- .../cli/commands/_helpers/compile/output.py | 25 +---- .../discovery/_helpers/sql/hook_entries.py | 53 +++++++++ src/sqlbuild/compiler/discovery/constants.py | 9 ++ .../discovery/main/serialize_hook_entries.py | 20 ++++ .../compiler/manifest/_helpers/model_nodes.py | 38 ++----- .../version_identity_model_metadata.py | 43 ++------ .../compiler/discovery/main/_test_types.py | 10 ++ .../main/test_serialize_hook_entries.py | 104 ++++++++++++++++++ 8 files changed, 216 insertions(+), 86 deletions(-) create mode 100644 src/sqlbuild/compiler/discovery/_helpers/sql/hook_entries.py create mode 100644 src/sqlbuild/compiler/discovery/main/serialize_hook_entries.py create mode 100644 tests/unit/src/sqlbuild/compiler/discovery/main/_test_types.py create mode 100644 tests/unit/src/sqlbuild/compiler/discovery/main/test_serialize_hook_entries.py diff --git a/src/sqlbuild/cli/commands/_helpers/compile/output.py b/src/sqlbuild/cli/commands/_helpers/compile/output.py index 7edd8c5456..75944bcc6d 100644 --- a/src/sqlbuild/cli/commands/_helpers/compile/output.py +++ b/src/sqlbuild/cli/commands/_helpers/compile/output.py @@ -27,7 +27,8 @@ RelatedLocation, ) from sqlbuild.compiler.compile.types import CompiledResourceType, DiagnosticSeverity -from sqlbuild.compiler.discovery.models import PythonHookEntry, SqlHookEntry +from sqlbuild.compiler.discovery.constants import SQL_HOOK_OUTPUT_FIELDS +from sqlbuild.compiler.discovery.main.serialize_hook_entries import serialize_hook_entries from sqlbuild.compiler.lineage.models import ProjectColumnLineage from sqlbuild.compiler.pipeline.models import ProjectGraph from sqlbuild.compiler.python_nodes.main.hook_identities import build_hook_identities @@ -280,27 +281,7 @@ def _model_resource( def _hook_resources(*, hooks: object) -> list[dict[str, object]]: - if not isinstance(hooks, (list, tuple)): - return [] - resources: list[dict[str, object]] = [] - hook: object - for hook in hooks: - if isinstance(hook, SqlHookEntry): - resource: dict[str, object] = {"type": "sql", "statement": hook.statement} - if hook.name is not None: - resource["name"] = hook.name - if hook.relative_path is not None: - resource["relative_path"] = hook.relative_path.as_posix() - if hook.definition_sql is not None: - resource["definition_sql"] = hook.definition_sql - if hook.kwargs is not None: - resource["kwargs"] = hook.kwargs - if hook.description is not None: - resource["description"] = hook.description - resources.append(resource) - elif isinstance(hook, PythonHookEntry): - resources.append({"type": "python", "name": hook.name, "kwargs": hook.kwargs}) - return resources + return serialize_hook_entries(value=hooks, sql_fields=SQL_HOOK_OUTPUT_FIELDS) def _hook_definitions(project: CompiledProject) -> list[dict[str, object]]: diff --git a/src/sqlbuild/compiler/discovery/_helpers/sql/hook_entries.py b/src/sqlbuild/compiler/discovery/_helpers/sql/hook_entries.py new file mode 100644 index 0000000000..e92f5e5db8 --- /dev/null +++ b/src/sqlbuild/compiler/discovery/_helpers/sql/hook_entries.py @@ -0,0 +1,53 @@ +"""Serialization of resolved model lifecycle hook entries.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from sqlbuild.compiler.discovery.models import PythonHookEntry, SqlHookEntry + + +def serialize_hook_entries_impl( + *, + value: object, + sql_fields: tuple[str, ...], + python_hook_fields: Mapping[str, Mapping[str, object]], +) -> list[dict[str, object]]: + """Serialize resolved hook entries, keeping only the requested SQL hook fields.""" + + if not isinstance(value, list | tuple): + return [] + hooks: list[dict[str, object]] = [] + entry: object + for entry in value: + if isinstance(entry, SqlHookEntry): + hooks.append(_serialize_sql_hook(entry=entry, sql_fields=sql_fields)) + elif isinstance(entry, PythonHookEntry): + python_hook: dict[str, object] = { + "type": "python", + "name": entry.name, + "kwargs": entry.kwargs, + } + python_hook.update(python_hook_fields.get(entry.name, {})) + hooks.append(python_hook) + return hooks + + +def _serialize_sql_hook(*, entry: SqlHookEntry, sql_fields: tuple[str, ...]) -> dict[str, object]: + field_values: dict[str, object | None] = { + "statement": entry.statement, + "name": entry.name, + "relative_path": ( + entry.relative_path.as_posix() if entry.relative_path is not None else None + ), + "definition_sql": entry.definition_sql, + "kwargs": entry.kwargs, + "description": entry.description, + } + hook: dict[str, object] = {"type": "sql"} + field_name: str + for field_name in sql_fields: + field_value: object | None = field_values[field_name] + if field_value is not None: + hook[field_name] = field_value + return hook diff --git a/src/sqlbuild/compiler/discovery/constants.py b/src/sqlbuild/compiler/discovery/constants.py index c91382938c..1955fc953a 100644 --- a/src/sqlbuild/compiler/discovery/constants.py +++ b/src/sqlbuild/compiler/discovery/constants.py @@ -62,3 +62,12 @@ ) NOT_NULL_AUDIT_NAME: str = "not_null" + +SQL_HOOK_OUTPUT_FIELDS: tuple[str, ...] = ( + "statement", + "name", + "relative_path", + "definition_sql", + "kwargs", + "description", +) diff --git a/src/sqlbuild/compiler/discovery/main/serialize_hook_entries.py b/src/sqlbuild/compiler/discovery/main/serialize_hook_entries.py new file mode 100644 index 0000000000..dcf654f9c1 --- /dev/null +++ b/src/sqlbuild/compiler/discovery/main/serialize_hook_entries.py @@ -0,0 +1,20 @@ +"""Model lifecycle hook serialization entrypoint.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from sqlbuild.compiler.discovery._helpers.sql.hook_entries import serialize_hook_entries_impl + + +def serialize_hook_entries( + *, + value: object, + sql_fields: tuple[str, ...], + python_hook_fields: Mapping[str, Mapping[str, object]] | None = None, +) -> list[dict[str, object]]: + """Serialize resolved hook entries with an explicit SQL hook field allowlist.""" + + return serialize_hook_entries_impl( + value=value, sql_fields=sql_fields, python_hook_fields=python_hook_fields or {} + ) diff --git a/src/sqlbuild/compiler/manifest/_helpers/model_nodes.py b/src/sqlbuild/compiler/manifest/_helpers/model_nodes.py index 2ded851ef4..6acd78d8cd 100644 --- a/src/sqlbuild/compiler/manifest/_helpers/model_nodes.py +++ b/src/sqlbuild/compiler/manifest/_helpers/model_nodes.py @@ -5,7 +5,8 @@ from pathlib import Path from sqlbuild.compiler.compile.models import CompiledModel -from sqlbuild.compiler.discovery.models import PythonHookEntry, SqlHookEntry +from sqlbuild.compiler.discovery.constants import SQL_HOOK_OUTPUT_FIELDS +from sqlbuild.compiler.discovery.main.serialize_hook_entries import serialize_hook_entries from sqlbuild.compiler.fingerprints.main.compute_query_hash import compute_query_hash from sqlbuild.compiler.manifest._helpers.shared import ( build_columns_dict, @@ -151,33 +152,8 @@ def _extract_meta( def _serialize_hooks( *, value: object, python_hook_metadata: dict[str, dict[str, object]] ) -> list[dict[str, object]]: - if not isinstance(value, list | tuple): - return [] - hooks: list[dict[str, object]] = [] - entry: object - for entry in value: - if isinstance(entry, SqlHookEntry): - hook: dict[str, object] = { - "type": "sql", - "statement": entry.statement, - } - if entry.name is not None: - hook["name"] = entry.name - if entry.relative_path is not None: - hook["relative_path"] = entry.relative_path.as_posix() - if entry.definition_sql is not None: - hook["definition_sql"] = entry.definition_sql - if entry.kwargs is not None: - hook["kwargs"] = entry.kwargs - if entry.description is not None: - hook["description"] = entry.description - hooks.append(hook) - elif isinstance(entry, PythonHookEntry): - python_hook: dict[str, object] = { - "type": "python", - "name": entry.name, - "kwargs": entry.kwargs, - } - python_hook.update(python_hook_metadata.get(entry.name, {})) - hooks.append(python_hook) - return hooks + return serialize_hook_entries( + value=value, + sql_fields=SQL_HOOK_OUTPUT_FIELDS, + python_hook_fields=python_hook_metadata, + ) diff --git a/src/sqlbuild/compiler/planner/main/identity/version_identity_model_metadata.py b/src/sqlbuild/compiler/planner/main/identity/version_identity_model_metadata.py index 2f8de0cdd3..b0bedb2544 100644 --- a/src/sqlbuild/compiler/planner/main/identity/version_identity_model_metadata.py +++ b/src/sqlbuild/compiler/planner/main/identity/version_identity_model_metadata.py @@ -6,7 +6,8 @@ from sqlbuild.compiler.compile.models import CompiledModel from sqlbuild.compiler.compile.types import CompiledResourceType -from sqlbuild.compiler.discovery.models import PythonHookEntry, SqlHookEntry +from sqlbuild.compiler.discovery.constants import SQL_HOOK_OUTPUT_FIELDS +from sqlbuild.compiler.discovery.main.serialize_hook_entries import serialize_hook_entries from sqlbuild.compiler.planner._helpers.identity.model_metadata import contract_output_signature from sqlbuild.compiler.planner.constants import ( MODEL_CUSTOM_CONFIG_KEY, @@ -77,35 +78,11 @@ def _model_execution_signature( def _hook_execution_signature( *, value: object, hook_version_hashes: dict[str, str] ) -> list[dict[str, object]]: - if not isinstance(value, list | tuple): - return [] - hooks: list[dict[str, object]] = [] - entry: object - for entry in value: - if isinstance(entry, SqlHookEntry): - hook: dict[str, object] = { - "type": "sql", - "statement": entry.statement, - } - if entry.name is not None: - hook["name"] = entry.name - if entry.relative_path is not None: - hook["relative_path"] = entry.relative_path.as_posix() - if entry.definition_sql is not None: - hook["definition_sql"] = entry.definition_sql - if entry.kwargs is not None: - hook["kwargs"] = entry.kwargs - if entry.description is not None: - hook["description"] = entry.description - hooks.append(hook) - elif isinstance(entry, PythonHookEntry): - python_hook: dict[str, object] = { - "type": "python", - "name": entry.name, - "kwargs": entry.kwargs, - } - version_hash: str | None = hook_version_hashes.get(entry.name) - if version_hash is not None: - python_hook["version_hash"] = version_hash - hooks.append(python_hook) - return hooks + return serialize_hook_entries( + value=value, + sql_fields=SQL_HOOK_OUTPUT_FIELDS, + python_hook_fields={ + hook_name: {"version_hash": version_hash} + for hook_name, version_hash in hook_version_hashes.items() + }, + ) diff --git a/tests/unit/src/sqlbuild/compiler/discovery/main/_test_types.py b/tests/unit/src/sqlbuild/compiler/discovery/main/_test_types.py new file mode 100644 index 0000000000..cd0f8f07c7 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/discovery/main/_test_types.py @@ -0,0 +1,10 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class SerializeHookEntriesTestCase: + description: str + value: object + sql_fields: tuple[str, ...] + python_hook_fields: dict[str, dict[str, object]] + expected_hooks: list[dict[str, object]] diff --git a/tests/unit/src/sqlbuild/compiler/discovery/main/test_serialize_hook_entries.py b/tests/unit/src/sqlbuild/compiler/discovery/main/test_serialize_hook_entries.py new file mode 100644 index 0000000000..2df0c8979e --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/discovery/main/test_serialize_hook_entries.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from sqlbuild.compiler.discovery.constants import SQL_HOOK_OUTPUT_FIELDS +from sqlbuild.compiler.discovery.main.serialize_hook_entries import serialize_hook_entries +from sqlbuild.compiler.discovery.models import PythonHookEntry, SqlHookEntry +from tests.unit.src.sqlbuild.compiler.discovery.main._test_types import ( + SerializeHookEntriesTestCase, +) + +_NAMED_SQL_HOOK: SqlHookEntry = SqlHookEntry( + statement="INSERT INTO order_audit SELECT 'orders'", + name="record_order_access", + relative_path=Path("hooks/sql/record_order_access.sql"), + definition_sql="INSERT INTO order_audit SELECT {{ table_name }}", + kwargs={"table_name": "'orders'"}, + description="Record order table access", +) + + +@pytest.mark.parametrize( + "test_case", + [ + SerializeHookEntriesTestCase( + description="serializes every output field of a named SQL hook in order", + value=[_NAMED_SQL_HOOK], + sql_fields=SQL_HOOK_OUTPUT_FIELDS, + python_hook_fields={}, + expected_hooks=[ + { + "type": "sql", + "statement": "INSERT INTO order_audit SELECT 'orders'", + "name": "record_order_access", + "relative_path": "hooks/sql/record_order_access.sql", + "definition_sql": "INSERT INTO order_audit SELECT {{ table_name }}", + "kwargs": {"table_name": "'orders'"}, + "description": "Record order table access", + } + ], + ), + SerializeHookEntriesTestCase( + description="omits unset fields of an inline SQL hook", + value=(SqlHookEntry(statement="ANALYZE orders"),), + sql_fields=SQL_HOOK_OUTPUT_FIELDS, + python_hook_fields={}, + expected_hooks=[{"type": "sql", "statement": "ANALYZE orders"}], + ), + SerializeHookEntriesTestCase( + description="keeps only allowlisted SQL hook fields", + value=[_NAMED_SQL_HOOK], + sql_fields=("statement", "kwargs"), + python_hook_fields={}, + expected_hooks=[ + { + "type": "sql", + "statement": "INSERT INTO order_audit SELECT 'orders'", + "kwargs": {"table_name": "'orders'"}, + } + ], + ), + SerializeHookEntriesTestCase( + description="adds per-name fields to Python hooks and keeps hook order", + value=[ + PythonHookEntry(name="notify_orders", kwargs={"channel": "orders"}), + SqlHookEntry(statement="ANALYZE orders"), + PythonHookEntry(name="refresh_customers", kwargs={}), + ], + sql_fields=SQL_HOOK_OUTPUT_FIELDS, + python_hook_fields={"notify_orders": {"version_hash": "notify-hash"}}, + expected_hooks=[ + { + "type": "python", + "name": "notify_orders", + "kwargs": {"channel": "orders"}, + "version_hash": "notify-hash", + }, + {"type": "sql", "statement": "ANALYZE orders"}, + {"type": "python", "name": "refresh_customers", "kwargs": {}}, + ], + ), + SerializeHookEntriesTestCase( + description="returns no hooks for a non-sequence value", + value="ANALYZE orders", + sql_fields=SQL_HOOK_OUTPUT_FIELDS, + python_hook_fields={}, + expected_hooks=[], + ), + ], + ids=lambda case: case.description, +) +def test_given_hook_entries_when_serializing_then_returns_allowlisted_payloads( + test_case: SerializeHookEntriesTestCase, +) -> None: + hooks: list[dict[str, object]] = serialize_hook_entries( + value=test_case.value, + sql_fields=test_case.sql_fields, + python_hook_fields=test_case.python_hook_fields, + ) + + assert hooks == test_case.expected_hooks + assert [list(hook) for hook in hooks] == [list(hook) for hook in test_case.expected_hooks] From 24fee76110037868944b9b48ee870af929934b50 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:59:15 +0100 Subject: [PATCH 06/57] fix(identity): exclude hook documentation from model version identity The model version identity hashed every serialized SQL hook field, including the hook resource description and its file location, so editing a hook description or moving its file created new model versions and rebuilds without any execution change. The identity now uses an explicit SQL hook field allowlist (statement, name, definition_sql, kwargs); Python hooks keep name, kwargs, and version hash. Compile output and the manifest still include description and relative_path. Upgrading changes the version identity of every model that has a SQL hook with a description or a hook resource file, so the first plan or build after upgrading produces one round of new model versions for those models. Models without SQL hooks are unaffected. (cherry picked from commit 08472bb53ddd6c7d96a4d19874d3ef7ed14a24e8) --- src/sqlbuild/compiler/discovery/constants.py | 1 + .../version_identity_model_metadata.py | 4 +- .../cli/commands/main/build/_test_types.py | 11 ++ .../main/build/test_sql_hook_identity.py | 111 ++++++++++++++ .../planner/main/identity/_test_types.py | 16 ++ .../compiler/planner/main/identity/helpers.py | 18 +++ .../identity/test_version_identity_hooks.py | 145 ++++++++++++++++++ 7 files changed, 304 insertions(+), 2 deletions(-) create mode 100644 tests/e2e/src/sqlbuild/cli/commands/main/build/test_sql_hook_identity.py create mode 100644 tests/unit/src/sqlbuild/compiler/planner/main/identity/test_version_identity_hooks.py diff --git a/src/sqlbuild/compiler/discovery/constants.py b/src/sqlbuild/compiler/discovery/constants.py index 1955fc953a..34d287fe32 100644 --- a/src/sqlbuild/compiler/discovery/constants.py +++ b/src/sqlbuild/compiler/discovery/constants.py @@ -71,3 +71,4 @@ "kwargs", "description", ) +SQL_HOOK_IDENTITY_FIELDS: tuple[str, ...] = ("statement", "name", "definition_sql", "kwargs") diff --git a/src/sqlbuild/compiler/planner/main/identity/version_identity_model_metadata.py b/src/sqlbuild/compiler/planner/main/identity/version_identity_model_metadata.py index b0bedb2544..ad8445d44f 100644 --- a/src/sqlbuild/compiler/planner/main/identity/version_identity_model_metadata.py +++ b/src/sqlbuild/compiler/planner/main/identity/version_identity_model_metadata.py @@ -6,7 +6,7 @@ from sqlbuild.compiler.compile.models import CompiledModel from sqlbuild.compiler.compile.types import CompiledResourceType -from sqlbuild.compiler.discovery.constants import SQL_HOOK_OUTPUT_FIELDS +from sqlbuild.compiler.discovery.constants import SQL_HOOK_IDENTITY_FIELDS from sqlbuild.compiler.discovery.main.serialize_hook_entries import serialize_hook_entries from sqlbuild.compiler.planner._helpers.identity.model_metadata import contract_output_signature from sqlbuild.compiler.planner.constants import ( @@ -80,7 +80,7 @@ def _hook_execution_signature( ) -> list[dict[str, object]]: return serialize_hook_entries( value=value, - sql_fields=SQL_HOOK_OUTPUT_FIELDS, + sql_fields=SQL_HOOK_IDENTITY_FIELDS, python_hook_fields={ hook_name: {"version_hash": version_hash} for hook_name, version_hash in hook_version_hashes.items() diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/build/_test_types.py b/tests/e2e/src/sqlbuild/cli/commands/main/build/_test_types.py index 7687fc89aa..fc1ee9f4f3 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/build/_test_types.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/build/_test_types.py @@ -1085,3 +1085,14 @@ class WindowLimitedReplayE2ETestCase: changed_rows_sql: str expected_changed_rows: int expected_total_rows: int + + +@dataclass(frozen=True) +class SqlHookIdentityBuildE2ETestCase: + """Test case for model version identity after a SQL hook resource edit.""" + + description: str + removed_hook_path: str + added_hook_path: str + added_hook_contents: str + expected_reason: str diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/build/test_sql_hook_identity.py b/tests/e2e/src/sqlbuild/cli/commands/main/build/test_sql_hook_identity.py new file mode 100644 index 0000000000..86f73a6061 --- /dev/null +++ b/tests/e2e/src/sqlbuild/cli/commands/main/build/test_sql_hook_identity.py @@ -0,0 +1,111 @@ +"""E2E tests for model version identity across SQL hook resource edits.""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path +from textwrap import dedent + +import pytest + +from tests.e2e.src.sqlbuild.cli.commands.main.build._test_types import ( + SqlHookIdentityBuildE2ETestCase, +) +from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import prepare_inline_project, run_sqb + +_HOOK_PATH: str = "hooks/sql/record_order_access.sql" +_HOOK_CONTENTS: str = ( + dedent( + """ + HOOK (description "Record order table access"); + + CREATE OR REPLACE TABLE main.order_access AS SELECT @label AS label + """ + ).strip() + + "\n" +) + + +@pytest.mark.parametrize( + "test_case", + [ + SqlHookIdentityBuildE2ETestCase( + description="hook description edit keeps the model unchanged", + removed_hook_path=_HOOK_PATH, + added_hook_path=_HOOK_PATH, + added_hook_contents=_HOOK_CONTENTS.replace( + "Record order table access", "Audit order reads" + ), + expected_reason="no_change", + ), + SqlHookIdentityBuildE2ETestCase( + description="hook file move keeps the model unchanged", + removed_hook_path=_HOOK_PATH, + added_hook_path="hooks/sql/audit/record_order_access.sql", + added_hook_contents=_HOOK_CONTENTS, + expected_reason="no_change", + ), + SqlHookIdentityBuildE2ETestCase( + description="hook SQL edit changes the model", + removed_hook_path=_HOOK_PATH, + added_hook_path=_HOOK_PATH, + added_hook_contents=_HOOK_CONTENTS.replace("@label AS label", "@label AS access"), + expected_reason="config_changed", + ), + ], + ids=lambda case: case.description, +) +def test_given_built_model_with_sql_hook_when_editing_hook_then_plans_expected_reason( + test_case: SqlHookIdentityBuildE2ETestCase, + tmp_path: Path, +) -> None: + project_dir: Path = prepare_inline_project( + tmp_path=tmp_path, + project_name="sql_hook_identity_project", + repo_files={ + "sqlbuild_project.toml": dedent( + """ + name = "sql_hook_identity_project" + adapter = "duckdb" + + [connection] + database = "sql_hook_identity_project.duckdb" + """ + ).strip() + + "\n", + _HOOK_PATH: _HOOK_CONTENTS, + "models/orders.sql": dedent( + """ + MODEL ( + materialized table, + post_hooks [sql("record_order_access", label: 1)] + ); + + SELECT 1 AS order_id + """ + ).strip() + + "\n", + }, + ) + first_build: subprocess.CompletedProcess[str] = run_sqb( + command=("--no-color", "build"), + project_dir=project_dir, + ) + assert first_build.returncode == 0, first_build.stdout + first_build.stderr + + (project_dir / test_case.removed_hook_path).unlink() + added_hook: Path = project_dir / test_case.added_hook_path + added_hook.parent.mkdir(parents=True, exist_ok=True) + added_hook.write_text(test_case.added_hook_contents) + plan: subprocess.CompletedProcess[str] = run_sqb( + command=("plan", "--json"), + project_dir=project_dir, + ) + + assert plan.returncode == 0, plan.stdout + plan.stderr + payload: dict[str, object] = json.loads(plan.stdout) + reasons_by_name: dict[str, str] = { + str(entry["name"]): str(entry["reason"]) for entry in payload["models"] + } + assert reasons_by_name["orders"] == test_case.expected_reason diff --git a/tests/unit/src/sqlbuild/compiler/planner/main/identity/_test_types.py b/tests/unit/src/sqlbuild/compiler/planner/main/identity/_test_types.py index a7acd351aa..53a0099164 100644 --- a/tests/unit/src/sqlbuild/compiler/planner/main/identity/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/planner/main/identity/_test_types.py @@ -30,3 +30,19 @@ class CursorRoleIdentityTestCase: original_config: dict[str, object] changed_config: dict[str, object] expected_equal: bool + + +@dataclass(frozen=True) +class HookIdentityTestCase: + description: str + original_hooks: dict[str, object] + changed_hooks: dict[str, object] + expected_changed: bool + + +@dataclass(frozen=True) +class HookIdentityPayloadTestCase: + description: str + hooks: dict[str, object] + expected_fragments: tuple[str, ...] + forbidden_fragments: tuple[str, ...] diff --git a/tests/unit/src/sqlbuild/compiler/planner/main/identity/helpers.py b/tests/unit/src/sqlbuild/compiler/planner/main/identity/helpers.py index 69a18b815c..65fecd6cca 100644 --- a/tests/unit/src/sqlbuild/compiler/planner/main/identity/helpers.py +++ b/tests/unit/src/sqlbuild/compiler/planner/main/identity/helpers.py @@ -14,6 +14,9 @@ from sqlbuild.compiler.planner._helpers.identity.direct import ( build_direct_model_version_identities, ) +from sqlbuild.compiler.planner.main.identity.version_identity_model_metadata import ( + build_model_version_identity_metadata_json, +) from sqlbuild.compiler.planner.models import DirectModelVersionIdentities, PlannerScope from tests.unit.src.sqlbuild.compiler.planner._helpers.helpers import build_compiled_function @@ -88,3 +91,18 @@ def _build_model( qualified_name=key.name, ), ) + + +def build_hooked_model_identity_json(*, hook_config: dict[str, object]) -> str: + model: CompiledModel = replace( + _build_model( + key=CompiledObjectKey(resource_type=CompiledResourceType.MODEL, name="orders"), + query_sql="SELECT 1 AS order_id", + deps=(), + ), + config=CompileModelConfig(values=hook_config), + ) + return build_model_version_identity_metadata_json( + model=model, + hook_version_hashes={"notify_orders": "notify-hash-v1"}, + ) diff --git a/tests/unit/src/sqlbuild/compiler/planner/main/identity/test_version_identity_hooks.py b/tests/unit/src/sqlbuild/compiler/planner/main/identity/test_version_identity_hooks.py new file mode 100644 index 0000000000..b84f7ff474 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/planner/main/identity/test_version_identity_hooks.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path + +import pytest + +from sqlbuild.compiler.discovery.models import PythonHookEntry, SqlHookEntry +from tests.unit.src.sqlbuild.compiler.planner.main.identity._test_types import ( + HookIdentityPayloadTestCase, + HookIdentityTestCase, +) +from tests.unit.src.sqlbuild.compiler.planner.main.identity.helpers import ( + build_hooked_model_identity_json, +) + +_RECORD_ACCESS: SqlHookEntry = SqlHookEntry( + statement="INSERT INTO order_audit SELECT 'orders'", + name="record_order_access", + relative_path=Path("hooks/sql/record_order_access.sql"), + definition_sql="INSERT INTO order_audit SELECT {{ table_name }}", + kwargs={"table_name": "'orders'"}, + description="Record order table access", +) +_ANALYZE: SqlHookEntry = SqlHookEntry(statement="ANALYZE orders") +_NOTIFY: PythonHookEntry = PythonHookEntry(name="notify_orders", kwargs={"channel": "orders"}) + + +@pytest.mark.parametrize( + "test_case", + [ + HookIdentityTestCase( + description="ignores a hook description change", + original_hooks={"pre_hooks": [_RECORD_ACCESS]}, + changed_hooks={"pre_hooks": [replace(_RECORD_ACCESS, description="Audit order reads")]}, + expected_changed=False, + ), + HookIdentityTestCase( + description="ignores a hook file location change", + original_hooks={"pre_hooks": [_RECORD_ACCESS]}, + changed_hooks={ + "pre_hooks": [ + replace( + _RECORD_ACCESS, + relative_path=Path("hooks/sql/audit/record_order_access.sql"), + ) + ] + }, + expected_changed=False, + ), + HookIdentityTestCase( + description="changes when hook SQL changes", + original_hooks={"pre_hooks": [_RECORD_ACCESS]}, + changed_hooks={ + "pre_hooks": [ + replace(_RECORD_ACCESS, statement="INSERT INTO order_audit SELECT 'all'") + ] + }, + expected_changed=True, + ), + HookIdentityTestCase( + description="changes when hook definition changes", + original_hooks={"pre_hooks": [_RECORD_ACCESS]}, + changed_hooks={ + "pre_hooks": [ + replace( + _RECORD_ACCESS, + definition_sql="INSERT INTO order_audit SELECT {{ table_name }} -- v2", + ) + ] + }, + expected_changed=True, + ), + HookIdentityTestCase( + description="changes when hook arguments change", + original_hooks={"pre_hooks": [_RECORD_ACCESS]}, + changed_hooks={ + "pre_hooks": [replace(_RECORD_ACCESS, kwargs={"table_name": "'customers'"})] + }, + expected_changed=True, + ), + HookIdentityTestCase( + description="changes when hook name changes", + original_hooks={"pre_hooks": [_RECORD_ACCESS]}, + changed_hooks={"pre_hooks": [replace(_RECORD_ACCESS, name="record_access")]}, + expected_changed=True, + ), + HookIdentityTestCase( + description="changes when a hook moves between phases", + original_hooks={"pre_hooks": [_RECORD_ACCESS]}, + changed_hooks={"post_hooks": [_RECORD_ACCESS]}, + expected_changed=True, + ), + HookIdentityTestCase( + description="changes when hook order changes", + original_hooks={"post_hooks": [_RECORD_ACCESS, _ANALYZE]}, + changed_hooks={"post_hooks": [_ANALYZE, _RECORD_ACCESS]}, + expected_changed=True, + ), + HookIdentityTestCase( + description="changes when Python hook arguments change", + original_hooks={"post_hooks": [_NOTIFY]}, + changed_hooks={"post_hooks": [replace(_NOTIFY, kwargs={"channel": "customers"})]}, + expected_changed=True, + ), + ], + ids=lambda case: case.description, +) +def test_given_hook_change_when_building_model_identity_then_tracks_only_execution_fields( + test_case: HookIdentityTestCase, +) -> None: + original_json: str = build_hooked_model_identity_json(hook_config=test_case.original_hooks) + changed_json: str = build_hooked_model_identity_json(hook_config=test_case.changed_hooks) + + assert (original_json != changed_json) is test_case.expected_changed + + +@pytest.mark.parametrize( + "test_case", + [ + HookIdentityPayloadTestCase( + description="keeps execution fields and omits documentation fields", + hooks={"pre_hooks": [_RECORD_ACCESS], "post_hooks": [_NOTIFY]}, + expected_fragments=( + '"definition_sql":"INSERT INTO order_audit SELECT {{ table_name }}"', + '"name":"record_order_access"', + '"version_hash":"notify-hash-v1"', + ), + forbidden_fragments=( + "Record order table access", + "hooks/sql/record_order_access.sql", + '"description"', + '"relative_path"', + ), + ) + ], + ids=lambda case: case.description, +) +def test_given_documented_hooks_when_building_model_identity_then_omits_doc_fields( + test_case: HookIdentityPayloadTestCase, +) -> None: + identity_json: str = build_hooked_model_identity_json(hook_config=test_case.hooks) + + assert all(fragment in identity_json for fragment in test_case.expected_fragments) + assert not any(fragment in identity_json for fragment in test_case.forbidden_fragments) From aab80c95a3e91e93f38768cd9a2fa8ad302d2c99 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:02:27 +0100 Subject: [PATCH 07/57] refactor(auditing): share authored audit list parsing Replace discovery's parse_audit_instances primitive, the model header audit parser, and the schema column audit parser with one auditing entry, parse_audit_instances. Messages are unchanged. null_as_empty keeps the existing split: MODEL headers and column metadata treat a null audits value as no audits, while YAML model, source, and column entries keep rejecting an explicit null with "audits must be a list". (cherry picked from commit f413e445dc7cfa2e132b4ba9ca812e199a96404f) --- .../auditing/_helpers/schema_audits.py | 25 ++++++ .../auditing/main/_parse_audit_instances.py | 27 ++++++ .../_helpers/attachment/model_config.py | 24 +----- .../discovery/_helpers/sql/schema_columns.py | 23 +---- .../discovery/_helpers/yml/primitives.py | 35 -------- .../compiler/discovery/_helpers/yml/schema.py | 11 ++- .../discovery/_helpers/yml/sources.py | 11 ++- .../compiler/auditing/main/_test_types.py | 16 ++++ .../main/test_parse_audit_instances.py | 83 +++++++++++++++++++ .../discovery/_helpers/test_yml_sources.py | 9 ++ 10 files changed, 183 insertions(+), 81 deletions(-) create mode 100644 src/sqlbuild/compiler/auditing/main/_parse_audit_instances.py delete mode 100644 src/sqlbuild/compiler/discovery/_helpers/yml/primitives.py create mode 100644 tests/unit/src/sqlbuild/compiler/auditing/main/test_parse_audit_instances.py diff --git a/src/sqlbuild/compiler/auditing/_helpers/schema_audits.py b/src/sqlbuild/compiler/auditing/_helpers/schema_audits.py index 94c5395bca..b838e2837f 100644 --- a/src/sqlbuild/compiler/auditing/_helpers/schema_audits.py +++ b/src/sqlbuild/compiler/auditing/_helpers/schema_audits.py @@ -22,6 +22,31 @@ from sqlbuild.spec.contracts.models import SchemaAuditInstance +def parse_audit_instances_impl( + *, + raw_audits: object | None, + file_path: Path, + label: str, + error_class: type[Exception], + null_as_empty: bool, +) -> tuple[SchemaAuditInstance, ...]: + """Parse an authored audit list, treating null as empty only when requested.""" + + if raw_audits is None and null_as_empty: + return () + if not isinstance(raw_audits, list): + raise error_class(f"{file_path} {label} audits must be a list") + return tuple( + parse_audit_instance_impl( + raw_audit=raw_audit, + file_path=file_path, + label=label, + error_class=error_class, + ) + for raw_audit in raw_audits + ) + + def parse_audit_instance_impl( *, raw_audit: object, diff --git a/src/sqlbuild/compiler/auditing/main/_parse_audit_instances.py b/src/sqlbuild/compiler/auditing/main/_parse_audit_instances.py new file mode 100644 index 0000000000..4b451e5a09 --- /dev/null +++ b/src/sqlbuild/compiler/auditing/main/_parse_audit_instances.py @@ -0,0 +1,27 @@ +"""Schema-attached audit list parsing entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.auditing._helpers.schema_audits import parse_audit_instances_impl +from sqlbuild.spec.contracts.models import SchemaAuditInstance + + +def parse_audit_instances( + *, + raw_audits: object | None, + file_path: Path, + label: str, + error_class: type[Exception], + null_as_empty: bool, +) -> tuple[SchemaAuditInstance, ...]: + """Parse an authored audit list, treating null as empty only when requested.""" + + return parse_audit_instances_impl( + raw_audits=raw_audits, + file_path=file_path, + label=label, + error_class=error_class, + null_as_empty=null_as_empty, + ) diff --git a/src/sqlbuild/compiler/compile/_helpers/attachment/model_config.py b/src/sqlbuild/compiler/compile/_helpers/attachment/model_config.py index 2cd8790d7b..ea7c56e8c0 100644 --- a/src/sqlbuild/compiler/compile/_helpers/attachment/model_config.py +++ b/src/sqlbuild/compiler/compile/_helpers/attachment/model_config.py @@ -6,7 +6,7 @@ from pathlib import Path from typing import cast -from sqlbuild.compiler.auditing.main._parse_audit_instance import parse_audit_instance +from sqlbuild.compiler.auditing.main._parse_audit_instances import parse_audit_instances from sqlbuild.compiler.authored_values.main._optional_named_string import optional_named_string from sqlbuild.compiler.compile._helpers.audit_factories.core import ( merge_validated_model_audits, @@ -291,10 +291,12 @@ def build_model_header_schema_entry( model_name=model_name, file_path=file_path, ) - audits: tuple[SchemaAuditInstance, ...] = _parse_model_header_audits( + audits: tuple[SchemaAuditInstance, ...] = parse_audit_instances( raw_audits=raw_audits, file_path=file_path, label="model", + error_class=CompileInputError, + null_as_empty=True, ) generated_audits: tuple[SchemaAuditInstance, ...] = parse_model_header_audit_factories( raw_audit_factories=raw_audit_factories, @@ -498,24 +500,6 @@ def _validate_model_schema_audit_augmentation( ) -def _parse_model_header_audits( - *, raw_audits: object | None, file_path: Path, label: str -) -> tuple[SchemaAuditInstance, ...]: - if raw_audits is None: - return () - if not isinstance(raw_audits, list): - raise CompileInputError(f"{file_path} {label} audits must be a list") - return tuple( - parse_audit_instance( - raw_audit=raw_audit, - file_path=file_path, - label=label, - error_class=CompileInputError, - ) - for raw_audit in raw_audits - ) - - def _resolve_model_schema( *, values: dict[str, object], diff --git a/src/sqlbuild/compiler/discovery/_helpers/sql/schema_columns.py b/src/sqlbuild/compiler/discovery/_helpers/sql/schema_columns.py index 1515de5667..e42eb709bc 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/sql/schema_columns.py +++ b/src/sqlbuild/compiler/discovery/_helpers/sql/schema_columns.py @@ -6,7 +6,7 @@ from pathlib import Path from typing import cast -from sqlbuild.compiler.auditing.main._parse_audit_instance import parse_audit_instance +from sqlbuild.compiler.auditing.main._parse_audit_instances import parse_audit_instances from sqlbuild.compiler.authored_values.main._optional_named_bool import optional_named_bool from sqlbuild.compiler.authored_values.main._optional_named_string import optional_named_string from sqlbuild.compiler.compile.constants import NOT_NULL_AUDIT_NAME @@ -80,11 +80,12 @@ def parse_schema_columns( column_location: SourceLocation | None = locations.get(raw_column_name) audits: tuple[SchemaAuditInstance, ...] = tuple( replace(audit, location=column_location) - for audit in _parse_audits( + for audit in parse_audit_instances( raw_audits=column_metadata.get("audits"), file_path=file_path, label=f"{label} column '{raw_column_name}'", error_class=error_class, + null_as_empty=True, ) ) if nullable is True and any( @@ -124,21 +125,3 @@ def parse_schema_columns( ) ) return tuple(parsed_columns) - - -def _parse_audits( - *, raw_audits: object | None, file_path: Path, label: str, error_class: _SchemaColumnParseError -) -> tuple[SchemaAuditInstance, ...]: - if raw_audits is None: - return () - if not isinstance(raw_audits, list): - raise error_class(f"{file_path} {label} audits must be a list") - return tuple( - parse_audit_instance( - raw_audit=raw_audit, - file_path=file_path, - label=label, - error_class=error_class, - ) - for raw_audit in raw_audits - ) diff --git a/src/sqlbuild/compiler/discovery/_helpers/yml/primitives.py b/src/sqlbuild/compiler/discovery/_helpers/yml/primitives.py deleted file mode 100644 index 69c8e8dc48..0000000000 --- a/src/sqlbuild/compiler/discovery/_helpers/yml/primitives.py +++ /dev/null @@ -1,35 +0,0 @@ -"""Shared YAML audit-list parsing for discovery parsers.""" - -from __future__ import annotations - -from pathlib import Path - -from sqlbuild.compiler.auditing.main._parse_audit_instance import parse_audit_instance -from sqlbuild.spec.contracts.models import SchemaAuditInstance - - -def parse_audit_instances( - *, - entry: dict[str, object], - file_path: Path, - label: str, - error_class: type[Exception], -) -> tuple[SchemaAuditInstance, ...]: - """Parse a list of schema-attached audit instances from a YAML entry.""" - - raw_audits: object = entry.get("audits", []) - if not isinstance(raw_audits, list): - raise error_class(f"{file_path} {label} audits must be a list") - - parsed_audits: list[SchemaAuditInstance] = [] - raw_audit: object - for raw_audit in raw_audits: - parsed_audits.append( - parse_audit_instance( - raw_audit=raw_audit, - file_path=file_path, - label=label, - error_class=error_class, - ) - ) - return tuple(parsed_audits) diff --git a/src/sqlbuild/compiler/discovery/_helpers/yml/schema.py b/src/sqlbuild/compiler/discovery/_helpers/yml/schema.py index c41d9573d9..ccbd79ac00 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/yml/schema.py +++ b/src/sqlbuild/compiler/discovery/_helpers/yml/schema.py @@ -8,6 +8,7 @@ import yaml from yaml import YAMLError +from sqlbuild.compiler.auditing.main._parse_audit_instances import parse_audit_instances from sqlbuild.compiler.authored_values.main._optional_bool import optional_bool from sqlbuild.compiler.authored_values.main._optional_mapping import optional_mapping from sqlbuild.compiler.authored_values.main._optional_non_empty_string import ( @@ -17,7 +18,6 @@ from sqlbuild.compiler.authored_values.main._require_non_empty_string import ( require_non_empty_string, ) -from sqlbuild.compiler.discovery._helpers.yml.primitives import parse_audit_instances from sqlbuild.compiler.discovery.constants import NOT_NULL_AUDIT_NAME, SEEDS_DIRECTORY_NAME from sqlbuild.compiler.discovery.exceptions import SchemaParseError from sqlbuild.spec.contracts.models import ( @@ -160,7 +160,11 @@ def _parse_model_entry(*, entry: dict[str, object], file_path: Path) -> SchemaMo ), columns=_parse_columns(entry=entry, file_path=file_path, label="model"), audits=parse_audit_instances( - entry=entry, file_path=file_path, label="model", error_class=SchemaParseError + raw_audits=entry.get("audits", []), + file_path=file_path, + label="model", + error_class=SchemaParseError, + null_as_empty=False, ), tags=optional_string_tuple( entry=entry, @@ -330,10 +334,11 @@ def _parse_columns( error_class=SchemaParseError, ) audits: tuple[SchemaAuditInstance, ...] = parse_audit_instances( - entry=column, + raw_audits=column.get("audits", []), file_path=file_path, label=column_label, error_class=SchemaParseError, + null_as_empty=False, ) _validate_nullable_audits( file_path=file_path, diff --git a/src/sqlbuild/compiler/discovery/_helpers/yml/sources.py b/src/sqlbuild/compiler/discovery/_helpers/yml/sources.py index 28bdab662f..35ac84ddbc 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/yml/sources.py +++ b/src/sqlbuild/compiler/discovery/_helpers/yml/sources.py @@ -8,6 +8,7 @@ import yaml from yaml import YAMLError +from sqlbuild.compiler.auditing.main._parse_audit_instances import parse_audit_instances from sqlbuild.compiler.authored_values.main._optional_bool import optional_bool from sqlbuild.compiler.authored_values.main._optional_mapping import optional_mapping from sqlbuild.compiler.authored_values.main._optional_non_empty_string import ( @@ -21,7 +22,6 @@ parse_dlt_sources, parse_source_integration_loader, ) -from sqlbuild.compiler.discovery._helpers.yml.primitives import parse_audit_instances from sqlbuild.compiler.discovery.constants import ( NOT_NULL_AUDIT_NAME, SOURCE_AGE_POLICY_CONFIG_KEY, @@ -210,7 +210,11 @@ def _parse_source_entry(*, entry: dict[str, object], file_path: Path) -> SourceE ), columns=columns, audits=parse_audit_instances( - entry=entry, file_path=file_path, label="source", error_class=SourceParseError + raw_audits=entry.get("audits", []), + file_path=file_path, + label="source", + error_class=SourceParseError, + null_as_empty=False, ), ) _validate_source_entry(entry=source_entry, file_path=file_path) @@ -591,10 +595,11 @@ def _parse_columns(*, entry: dict[str, object], file_path: Path) -> tuple[Source error_class=SourceParseError, ) audits: tuple[SchemaAuditInstance, ...] = parse_audit_instances( - entry=column, + raw_audits=column.get("audits", []), file_path=file_path, label=column_label, error_class=SourceParseError, + null_as_empty=False, ) _validate_nullable_audits( file_path=file_path, diff --git a/tests/unit/src/sqlbuild/compiler/auditing/main/_test_types.py b/tests/unit/src/sqlbuild/compiler/auditing/main/_test_types.py index 6c6bb9abaf..ae1f3dd4ca 100644 --- a/tests/unit/src/sqlbuild/compiler/auditing/main/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/auditing/main/_test_types.py @@ -109,3 +109,19 @@ class InvalidMeasuredValueTestCase: description: str measured_value: float expected_error_message: str + + +@dataclass(frozen=True) +class ParseAuditInstancesTestCase: + description: str + raw_audits: object | None + null_as_empty: bool + expected_definition_names: tuple[str, ...] + + +@dataclass(frozen=True) +class ParseAuditInstancesErrorTestCase: + description: str + raw_audits: object | None + null_as_empty: bool + expected_message: str diff --git a/tests/unit/src/sqlbuild/compiler/auditing/main/test_parse_audit_instances.py b/tests/unit/src/sqlbuild/compiler/auditing/main/test_parse_audit_instances.py new file mode 100644 index 0000000000..d1c526fa55 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/auditing/main/test_parse_audit_instances.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from sqlbuild.compiler.auditing.main._parse_audit_instances import parse_audit_instances +from sqlbuild.spec.contracts.models import SchemaAuditInstance +from tests.unit.src.sqlbuild.compiler.auditing.main._test_types import ( + ParseAuditInstancesErrorTestCase, + ParseAuditInstancesTestCase, +) + + +@pytest.mark.parametrize( + "test_case", + [ + ParseAuditInstancesTestCase( + description="parses string and mapping audits in order", + raw_audits=["not_null", {"accepted_values": {"values": ["open", "closed"]}}], + null_as_empty=False, + expected_definition_names=("not_null", "accepted_values"), + ), + ParseAuditInstancesTestCase( + description="treats null as empty when requested", + raw_audits=None, + null_as_empty=True, + expected_definition_names=(), + ), + ParseAuditInstancesTestCase( + description="accepts an empty list", + raw_audits=[], + null_as_empty=False, + expected_definition_names=(), + ), + ], + ids=lambda case: case.description, +) +def test_given_audit_list_when_parsing_then_returns_audit_instances( + test_case: ParseAuditInstancesTestCase, +) -> None: + result: tuple[SchemaAuditInstance, ...] = parse_audit_instances( + raw_audits=test_case.raw_audits, + file_path=Path("models/orders.yml"), + label="model orders", + error_class=ValueError, + null_as_empty=test_case.null_as_empty, + ) + + assert tuple(audit.definition_name for audit in result) == (test_case.expected_definition_names) + + +@pytest.mark.parametrize( + "test_case", + [ + ParseAuditInstancesErrorTestCase( + description="rejects null when null is not empty", + raw_audits=None, + null_as_empty=False, + expected_message="models/orders.yml model orders audits must be a list", + ), + ParseAuditInstancesErrorTestCase( + description="rejects a mapping", + raw_audits={"not_null": {}}, + null_as_empty=True, + expected_message="models/orders.yml model orders audits must be a list", + ), + ], + ids=lambda case: case.description, +) +def test_given_invalid_audit_list_when_parsing_then_raises_injected_error( + test_case: ParseAuditInstancesErrorTestCase, +) -> None: + with pytest.raises(ValueError) as error: + parse_audit_instances( + raw_audits=test_case.raw_audits, + file_path=Path("models/orders.yml"), + label="model orders", + error_class=ValueError, + null_as_empty=test_case.null_as_empty, + ) + + assert str(error.value) == test_case.expected_message diff --git a/tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_yml_sources.py b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_yml_sources.py index 2adc2a2267..fec22682e2 100644 --- a/tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_yml_sources.py +++ b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_yml_sources.py @@ -1119,6 +1119,15 @@ def test_given_dlt_sources_yaml_when_parsing_then_expands_managed_sources( """, expected_error_fragment="source audits must be a list", ), + ParseSourcesYamlErrorTestCase( + description="raises when source audits is null", + contents=""" + sources: + - name: raw_orders + audits: + """, + expected_error_fragment="source audits must be a list", + ), ParseSourcesYamlErrorTestCase( description="raises when source columns is not a list", contents=""" From c3b9123bbfc37e94cd72c606055cefdf3fc0f5b5 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:07:47 +0100 Subject: [PATCH 08/57] test(state): characterize shared state backend reads on DuckDB and Postgres Cover function version upsert/get, physical relation lookup and listing order, virtual environment listing order, and active versus expired lock listings against both state backends before consolidating their SQL. (cherry picked from commit 1fae3905de8c823d642de7caf2959ce4086c3172) --- .../virtual/state/classes/_test_types.py | 11 ++ .../sqlbuild/virtual/state/classes/helpers.py | 154 +++++++++++++++++- .../state/classes/postgres/_test_types.py | 10 ++ .../state/classes/postgres/test_postgres.py | 56 +++++++ .../virtual/state/classes/test_duckdb.py | 56 +++++++ 5 files changed, 286 insertions(+), 1 deletion(-) diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/_test_types.py b/tests/integration/src/sqlbuild/virtual/state/classes/_test_types.py index d52b2139a0..d542eda509 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/_test_types.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/_test_types.py @@ -2,6 +2,7 @@ from dataclasses import dataclass +from sqlbuild.virtual.state.types import VirtualEnvironmentStatus from tests.integration.src.sqlbuild.virtual.state.classes.helpers import ( ConditionalPublicationPayload, ) @@ -252,3 +253,13 @@ class MicrobatchStateRoundTripTestCase: description: str expected_event_count: int + + +@dataclass(frozen=True) +class StateReadContractTestCase: + description: str + schema: str + sqlbuild_version: str + expected_environments: tuple[tuple[str, VirtualEnvironmentStatus], ...] + expected_active_lock_keys: tuple[str, ...] + expected_expired_lock_keys: tuple[str, ...] diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py b/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py index 5cc4b6c320..499d57a473 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py @@ -1,11 +1,15 @@ from __future__ import annotations from dataclasses import replace +from datetime import datetime, timedelta from pathlib import Path from typing import Any, NamedTuple from sqlbuild.virtual.state.classes.duckdb import DuckDbStateBackend +from sqlbuild.virtual.state.classes.state_backend import StateBackend from sqlbuild.virtual.state.models import ( + FunctionVersionRecord, + PhysicalRelationRecord, VirtualEnvironmentCheckpointFunctionRefRecord, VirtualEnvironmentCheckpointModelRefRecord, VirtualEnvironmentCheckpointRecord, @@ -13,7 +17,11 @@ VirtualEnvironmentNodeRefRecord, VirtualEnvironmentRecord, ) -from sqlbuild.virtual.state.types import VirtualEnvironmentStatus +from sqlbuild.virtual.state.types import ( + ModelVersionStatus, + PhysicalArtifactType, + VirtualEnvironmentStatus, +) _ENVIRONMENT: str = "dev" _CHECKPOINT_ID: str = "checkpoint-1" @@ -135,3 +143,147 @@ def open_duckdb_state_backend(*, db_path: Path) -> tuple[DuckDbStateBackend, Any def fetch_all(connection: Any, sql: str) -> list[tuple[Any, ...]]: return connection.execute(sql).fetchall() + + +STATE_READ_CONTRACT_FUNCTION: FunctionVersionRecord = FunctionVersionRecord( + function_name="is_large_order", + version_hash="function-v1", + language="sql", + returns="BOOLEAN", + arguments_json_b64="W3sibmFtZSI6ImFtb3VudCIsInR5cGUiOiJJTlRFR0VSIn1d", + return_columns_json_b64="W10=", + packages_json_b64="W10=", + runtime_version=None, + entry_point=None, + body_sql_b64="YW1vdW50ID4gOQ==", + definition_text_b64="YW1vdW50ID4gOQ==", + status=ModelVersionStatus.READY, +) +STATE_READ_CONTRACT_RELATION_V1: PhysicalRelationRecord = PhysicalRelationRecord( + artifact_type=PhysicalArtifactType.MODEL, + artifact_name="orders", + version_hash="orders-v1", + database_name=None, + schema_name="sqb_orders", + relation_name="orders__v1", + relation_type="table", +) +STATE_READ_CONTRACT_RELATION_V2: PhysicalRelationRecord = replace( + STATE_READ_CONTRACT_RELATION_V1, version_hash="orders-v2", relation_name="orders__v2" +) + + +class StateReadContractObservation(NamedTuple): + """Values read back after writing the shared state read contract fixture.""" + + function_before_upsert: FunctionVersionRecord | None + function_after_upsert: FunctionVersionRecord | None + function_after_second_upsert: FunctionVersionRecord | None + relation_v1: PhysicalRelationRecord | None + relation_missing: PhysicalRelationRecord | None + relations: tuple[PhysicalRelationRecord, ...] + environments: tuple[tuple[str, VirtualEnvironmentStatus], ...] + active_lock_keys: tuple[str, ...] + expired_lock_keys: tuple[str, ...] + + +def exercise_state_read_contract( + *, backend: StateBackend, connection: Any, schema: str +) -> StateReadContractObservation: + function_before_upsert: FunctionVersionRecord | None = backend.get_function_version( + connection=connection, + schema=schema, + function_name=STATE_READ_CONTRACT_FUNCTION.function_name, + version_hash=STATE_READ_CONTRACT_FUNCTION.version_hash, + ) + backend.upsert_function_version( + connection=connection, schema=schema, record=STATE_READ_CONTRACT_FUNCTION + ) + function_after_upsert: FunctionVersionRecord | None = backend.get_function_version( + connection=connection, + schema=schema, + function_name=STATE_READ_CONTRACT_FUNCTION.function_name, + version_hash=STATE_READ_CONTRACT_FUNCTION.version_hash, + ) + backend.upsert_function_version( + connection=connection, + schema=schema, + record=replace(STATE_READ_CONTRACT_FUNCTION, status=ModelVersionStatus.FAILED), + ) + function_after_second_upsert: FunctionVersionRecord | None = backend.get_function_version( + connection=connection, + schema=schema, + function_name=STATE_READ_CONTRACT_FUNCTION.function_name, + version_hash=STATE_READ_CONTRACT_FUNCTION.version_hash, + ) + backend.upsert_physical_relation( + connection=connection, schema=schema, record=STATE_READ_CONTRACT_RELATION_V1 + ) + backend.upsert_physical_relation( + connection=connection, schema=schema, record=STATE_READ_CONTRACT_RELATION_V2 + ) + relation_v1: PhysicalRelationRecord | None = backend.get_physical_relation_for_artifact( + connection=connection, + schema=schema, + artifact_type=PhysicalArtifactType.MODEL, + artifact_name="orders", + version_hash="orders-v1", + ) + relation_missing: PhysicalRelationRecord | None = backend.get_physical_relation_for_artifact( + connection=connection, + schema=schema, + artifact_type=PhysicalArtifactType.MODEL, + artifact_name="orders", + version_hash="orders-v3", + ) + relations: tuple[PhysicalRelationRecord, ...] = backend.list_physical_relations_for_artifact( + connection=connection, + schema=schema, + artifact_type=PhysicalArtifactType.MODEL, + artifact_name="orders", + ) + backend.upsert_virtual_environment( + connection=connection, + schema=schema, + record=VirtualEnvironmentRecord("dev", VirtualEnvironmentStatus.ACTIVE), + ) + backend.upsert_virtual_environment( + connection=connection, + schema=schema, + record=VirtualEnvironmentRecord("prod", VirtualEnvironmentStatus.FINALIZED), + ) + environments: tuple[tuple[str, VirtualEnvironmentStatus], ...] = tuple( + (record.virtual_environment_name, record.status) + for record in backend.list_virtual_environments(connection=connection, schema=schema) + ) + backend.acquire_lock( + connection=connection, + schema=schema, + lock_key="lock_expired", + owner_id="run-1", + expires_at=datetime.now() - timedelta(days=2), + ) + backend.acquire_lock( + connection=connection, + schema=schema, + lock_key="lock_active", + owner_id="run-1", + expires_at=datetime.now() + timedelta(days=2), + ) + active_lock_keys: tuple[str, ...] = tuple( + lock.lock_key for lock in backend.list_active_locks(connection=connection, schema=schema) + ) + expired_lock_keys: tuple[str, ...] = tuple( + lock.lock_key for lock in backend.list_expired_locks(connection=connection, schema=schema) + ) + return StateReadContractObservation( + function_before_upsert=function_before_upsert, + function_after_upsert=function_after_upsert, + function_after_second_upsert=function_after_second_upsert, + relation_v1=relation_v1, + relation_missing=relation_missing, + relations=relations, + environments=environments, + active_lock_keys=active_lock_keys, + expired_lock_keys=expired_lock_keys, + ) diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/postgres/_test_types.py b/tests/integration/src/sqlbuild/virtual/state/classes/postgres/_test_types.py index d93e797c00..c70a492d09 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/postgres/_test_types.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/postgres/_test_types.py @@ -2,6 +2,7 @@ from dataclasses import dataclass +from sqlbuild.virtual.state.types import VirtualEnvironmentStatus from tests.integration.src.sqlbuild.virtual.state.classes.helpers import ( ConditionalPublicationPayload, ) @@ -209,3 +210,12 @@ class PostgresRetiredMicrobatchRecordTestCase: description: str retired_record_type: str expected_event_count: int + + +@dataclass(frozen=True) +class PostgresStateReadContractTestCase: + description: str + sqlbuild_version: str + expected_environments: tuple[tuple[str, VirtualEnvironmentStatus], ...] + expected_active_lock_keys: tuple[str, ...] + expected_expired_lock_keys: tuple[str, ...] diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/postgres/test_postgres.py b/tests/integration/src/sqlbuild/virtual/state/classes/postgres/test_postgres.py index e0fd36e5bf..48a35f740f 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/postgres/test_postgres.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/postgres/test_postgres.py @@ -2,6 +2,7 @@ from collections.abc import Callable from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import replace from datetime import UTC, datetime, timedelta from typing import Any @@ -75,7 +76,12 @@ REF_ENVIRONMENT_PAYLOAD, REF_NODE_TYPE_PAYLOAD, SEED_EXTRA_PAYLOAD, + STATE_READ_CONTRACT_FUNCTION, + STATE_READ_CONTRACT_RELATION_V1, + STATE_READ_CONTRACT_RELATION_V2, VALID_PAYLOAD, + StateReadContractObservation, + exercise_state_read_contract, ) from tests.integration.src.sqlbuild.virtual.state.classes.postgres._test_types import ( PostgresAtomicFinalizedVirtualPublishTestCase, @@ -99,6 +105,7 @@ PostgresStateBackendTableCreationTestCase, PostgresStateBackendTransactionRollbackTestCase, PostgresStateBackendValidationTestCase, + PostgresStateReadContractTestCase, ) from tests.integration.src.sqlbuild.virtual.state.classes.postgres.helpers import ( fetch_all, @@ -2499,3 +2506,52 @@ def _insert_then_fail(**kwargs: Any) -> None: ) == checkpoint_refs ) + + +@pytest.mark.parametrize( + "test_case", + [ + PostgresStateReadContractTestCase( + description="round trips function versions relations environments and locks", + sqlbuild_version="0.0.test", + expected_environments=( + ("prod", VirtualEnvironmentStatus.FINALIZED), + ("dev", VirtualEnvironmentStatus.ACTIVE), + ), + expected_active_lock_keys=("lock_active",), + expected_expired_lock_keys=("lock_expired",), + ) + ], + ids=lambda case: case.description, +) +def test_given_postgres_state_backend_when_reading_state_then_matches_read_contract( + test_case: PostgresStateReadContractTestCase, + postgres_state_backend: PostgresStateBackend, + postgres_state_connection: Any, + postgres_state_schema: str, +) -> None: + postgres_state_backend.initialize( + connection=postgres_state_connection, + schema=postgres_state_schema, + sqlbuild_version=test_case.sqlbuild_version, + ) + observation: StateReadContractObservation = exercise_state_read_contract( + backend=postgres_state_backend, + connection=postgres_state_connection, + schema=postgres_state_schema, + ) + + assert observation.function_before_upsert is None + assert observation.function_after_upsert == STATE_READ_CONTRACT_FUNCTION + assert observation.function_after_second_upsert == replace( + STATE_READ_CONTRACT_FUNCTION, status=ModelVersionStatus.FAILED + ) + assert observation.relation_v1 == STATE_READ_CONTRACT_RELATION_V1 + assert observation.relation_missing is None + assert observation.relations == ( + STATE_READ_CONTRACT_RELATION_V2, + STATE_READ_CONTRACT_RELATION_V1, + ) + assert observation.environments == test_case.expected_environments + assert observation.active_lock_keys == test_case.expected_active_lock_keys + assert observation.expired_lock_keys == test_case.expected_expired_lock_keys diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/test_duckdb.py b/tests/integration/src/sqlbuild/virtual/state/classes/test_duckdb.py index 177c6687e9..adc2e9d10a 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/test_duckdb.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/test_duckdb.py @@ -97,6 +97,7 @@ DuckDbStateBackendTransactionRollbackTestCase, DuckDbStateBackendValidationTestCase, MicrobatchStateRoundTripTestCase, + StateReadContractTestCase, ) from tests.integration.src.sqlbuild.virtual.state.classes.helpers import ( ACTIVE_CHECKPOINT_PAYLOAD, @@ -114,7 +115,12 @@ REF_ENVIRONMENT_PAYLOAD, REF_NODE_TYPE_PAYLOAD, SEED_EXTRA_PAYLOAD, + STATE_READ_CONTRACT_FUNCTION, + STATE_READ_CONTRACT_RELATION_V1, + STATE_READ_CONTRACT_RELATION_V2, VALID_PAYLOAD, + StateReadContractObservation, + exercise_state_read_contract, fetch_all, open_duckdb_state_backend, ) @@ -2802,3 +2808,53 @@ def test_given_duckdb_state_backend_when_two_connections_acquire_same_lock_then_ finally: backend.close(first_connection) second_backend.close(second_connection) + + +@pytest.mark.parametrize( + "test_case", + [ + StateReadContractTestCase( + description="round trips function versions relations environments and locks", + schema="sqlbuild_state", + sqlbuild_version="0.0.test", + expected_environments=( + ("prod", VirtualEnvironmentStatus.FINALIZED), + ("dev", VirtualEnvironmentStatus.ACTIVE), + ), + expected_active_lock_keys=("lock_active",), + expected_expired_lock_keys=("lock_expired",), + ) + ], + ids=lambda case: case.description, +) +def test_given_duckdb_state_backend_when_reading_state_then_matches_read_contract( + test_case: StateReadContractTestCase, + tmp_path: Path, +) -> None: + backend, connection = open_duckdb_state_backend(db_path=tmp_path / "state.duckdb") + try: + backend.initialize( + connection=connection, + schema=test_case.schema, + sqlbuild_version=test_case.sqlbuild_version, + ) + observation: StateReadContractObservation = exercise_state_read_contract( + backend=backend, connection=connection, schema=test_case.schema + ) + finally: + backend.close(connection) + + assert observation.function_before_upsert is None + assert observation.function_after_upsert == STATE_READ_CONTRACT_FUNCTION + assert observation.function_after_second_upsert == replace( + STATE_READ_CONTRACT_FUNCTION, status=ModelVersionStatus.FAILED + ) + assert observation.relation_v1 == STATE_READ_CONTRACT_RELATION_V1 + assert observation.relation_missing is None + assert observation.relations == ( + STATE_READ_CONTRACT_RELATION_V2, + STATE_READ_CONTRACT_RELATION_V1, + ) + assert observation.environments == test_case.expected_environments + assert observation.active_lock_keys == test_case.expected_active_lock_keys + assert observation.expired_lock_keys == test_case.expected_expired_lock_keys From 4c7af875f9bf51d62b1f835419a2b9e69e1b2920 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:12:18 +0100 Subject: [PATCH 09/57] refactor(state): share relational state backend getters Add SqlStateBackend between StateBackend and the DuckDB and Postgres state backends. It owns the read-only getters whose SQL, parameters, and row mapping were identical in both backends (model, function, seed and Python node versions, physical relations and ancestry, virtual environments, node/model/function/seed/Python refs, source freshness, checkpoints and checkpoint refs, and state operations). Backends supply only the placeholder style and the fetch calls: DuckDB runs on the connection, Postgres on a short-lived cursor, as before. Table layout, stored formats, and query results are unchanged. The Postgres file is now under the line limit, so its FFR601 exception is removed. (cherry picked from commit 85c893469ed1c3add44fae4d0038bdf2f2bd8aef) --- fensu.toml | 5 - .../state/classes/_sql_state_backend.py | 544 ++++++++++++++++++ src/sqlbuild/virtual/state/classes/duckdb.py | 465 +-------------- .../virtual/state/classes/postgres.py | 511 +--------------- 4 files changed, 590 insertions(+), 935 deletions(-) create mode 100644 src/sqlbuild/virtual/state/classes/_sql_state_backend.py diff --git a/fensu.toml b/fensu.toml index 6d449018d8..a5409ee41b 100644 --- a/fensu.toml +++ b/fensu.toml @@ -93,11 +93,6 @@ rule = "FFR601" path = "src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py" reason = "The SQL Server adapter intentionally keeps overrideable backend behavior on one class." -[[targets.python.rule_exceptions]] -rule = "FFR601" -path = "src/sqlbuild/virtual/state/classes/postgres.py" -reason = "The Postgres state backend intentionally keeps its transactional storage contract on one class." - [[targets.python.rule_exceptions]] rule = "FFS120" path = "src/sqlbuild/integrations/dbt/_helpers/profile/render.py" diff --git a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py new file mode 100644 index 0000000000..3bb25b1426 --- /dev/null +++ b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py @@ -0,0 +1,544 @@ +"""Shared SQL for relational virtual-state backends.""" + +from __future__ import annotations + +from abc import abstractmethod +from collections.abc import Sequence +from typing import Any, ClassVar + +from sqlbuild.virtual.state._helpers.state_storage.datetime import ( + from_naive_utc_wall_clock, +) +from sqlbuild.virtual.state.classes.state_backend import StateBackend +from sqlbuild.virtual.state.constants import ( + FUNCTION_VERSION_TABLE, + MODEL_VERSION_TABLE, + PHYSICAL_RELATION_ANCESTRY_TABLE, + PHYSICAL_RELATION_TABLE, + PYTHON_NODE_VERSION_TABLE, + SEED_VERSION_TABLE, + SOURCE_FRESHNESS_OBSERVATION_TABLE, + STATE_OPERATION_TABLE, + VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, + VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, + VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, + VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE, + VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, + VIRTUAL_ENVIRONMENT_TABLE, +) +from sqlbuild.virtual.state.models import ( + FunctionVersionRecord, + ModelVersionRecord, + PhysicalRelationAncestryRecord, + PhysicalRelationRecord, + PythonNodeVersionRecord, + SeedVersionRecord, + SourceFreshnessRecord, + StateOperationRecord, + VirtualEnvironmentCheckpointFunctionRefRecord, + VirtualEnvironmentCheckpointModelRefRecord, + VirtualEnvironmentCheckpointRecord, + VirtualEnvironmentCheckpointSeedRefRecord, + VirtualEnvironmentFunctionRefRecord, + VirtualEnvironmentModelRefRecord, + VirtualEnvironmentNodeRefRecord, + VirtualEnvironmentPythonNodeRefRecord, + VirtualEnvironmentRecord, + VirtualEnvironmentSeedRefRecord, +) +from sqlbuild.virtual.state.types import ( + ModelVersionStatus, + PhysicalArtifactType, + StateOperationStatus, + StateOperationType, + VirtualEnvironmentStatus, +) + + +class SqlStateBackend(StateBackend): + """Relational state backend owning SQL shared by DuckDB and Postgres.""" + + _placeholder: ClassVar[str] + + @abstractmethod + def _fetch_one( + self, *, connection: Any, sql: str, params: Sequence[object] | None = None + ) -> tuple[Any, ...] | None: + """Run one read statement and return its first row.""" + ... + + @abstractmethod + def _fetch_all( + self, *, connection: Any, sql: str, params: Sequence[object] | None = None + ) -> list[tuple[Any, ...]]: + """Run one read statement and return all rows.""" + ... + + def _quote_identifier(self, identifier: str) -> str: + return '"' + identifier.replace('"', '""') + '"' + + def _qualified_name(self, *, schema: str, table: str) -> str: + return f"{self._quote_identifier(schema)}.{self._quote_identifier(table)}" + + def get_model_version( + self, *, connection: Any, schema: str, model_name: str, version_hash: str + ) -> ModelVersionRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT model_name, version_hash, definition_identity_hash, " + "identity_metadata_hash, definition_text_b64, identity_metadata_json_b64, " + "compiled_sql_b64, status " + f"FROM {self._qualified_name(schema=schema, table=MODEL_VERSION_TABLE)} " + f"WHERE model_name = {p} AND version_hash = {p}", + params=[model_name, version_hash], + ) + if row is None: + return None + return ModelVersionRecord( + model_name=row[0], + version_hash=row[1], + definition_identity_hash=row[2], + identity_metadata_hash=row[3], + definition_text_b64=row[4], + identity_metadata_json_b64=row[5], + compiled_sql_b64=row[6], + status=ModelVersionStatus(row[7]), + ) + + def get_function_version( + self, *, connection: Any, schema: str, function_name: str, version_hash: str + ) -> FunctionVersionRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT function_name, version_hash, language, returns, arguments_json_b64, " + "return_columns_json_b64, packages_json_b64, runtime_version, entry_point, " + "body_sql_b64, definition_text_b64, status " + f"FROM {self._qualified_name(schema=schema, table=FUNCTION_VERSION_TABLE)} " + f"WHERE function_name = {p} AND version_hash = {p}", + params=[function_name, version_hash], + ) + if row is None: + return None + return FunctionVersionRecord( + function_name=row[0], + version_hash=row[1], + language=row[2], + returns=row[3], + arguments_json_b64=row[4], + return_columns_json_b64=row[5], + packages_json_b64=row[6], + runtime_version=row[7], + entry_point=row[8], + body_sql_b64=row[9], + definition_text_b64=row[10], + status=ModelVersionStatus(row[11]), + ) + + def get_seed_version( + self, *, connection: Any, schema: str, seed_name: str, version_hash: str + ) -> SeedVersionRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT seed_name, version_hash, identity_metadata_hash, " + "identity_metadata_json_b64, status " + f"FROM {self._qualified_name(schema=schema, table=SEED_VERSION_TABLE)} " + f"WHERE seed_name = {p} AND version_hash = {p}", + params=[seed_name, version_hash], + ) + if row is None: + return None + return SeedVersionRecord( + seed_name=row[0], + version_hash=row[1], + identity_metadata_hash=row[2], + identity_metadata_json_b64=row[3], + status=ModelVersionStatus(row[4]), + ) + + def get_python_node_version( + self, + *, + connection: Any, + schema: str, + node_type: str, + node_name: str, + version_hash: str, + ) -> PythonNodeVersionRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT node_type, node_name, version_hash, definition_hash, " + "identity_metadata_hash, definition_json_b64, identity_metadata_json_b64, status " + f"FROM {self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " + f"WHERE node_type = {p} AND node_name = {p} AND version_hash = {p}", + params=[node_type, node_name, version_hash], + ) + if row is None: + return None + return PythonNodeVersionRecord( + node_type=row[0], + node_name=row[1], + version_hash=row[2], + definition_hash=row[3], + identity_metadata_hash=row[4], + definition_json_b64=row[5], + identity_metadata_json_b64=row[6], + status=ModelVersionStatus(row[7]), + ) + + def get_physical_relation_for_artifact( + self, + *, + connection: Any, + schema: str, + artifact_type: PhysicalArtifactType, + artifact_name: str, + version_hash: str, + ) -> PhysicalRelationRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT artifact_type, artifact_name, version_hash, database_name, schema_name, " + "relation_name, relation_type " + f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " + f"WHERE artifact_type = {p} AND artifact_name = {p} AND version_hash = {p}", + params=[artifact_type.value, artifact_name, version_hash], + ) + if row is None: + return None + return PhysicalRelationRecord( + artifact_type=PhysicalArtifactType(row[0]), + artifact_name=row[1], + version_hash=row[2], + database_name=row[3], + schema_name=row[4], + relation_name=row[5], + relation_type=row[6], + ) + + def list_physical_relations_for_artifact( + self, + *, + connection: Any, + schema: str, + artifact_type: PhysicalArtifactType, + artifact_name: str, + ) -> tuple[PhysicalRelationRecord, ...]: + p: str = self._placeholder + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT artifact_type, artifact_name, version_hash, database_name, schema_name, " + "relation_name, relation_type " + f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " + f"WHERE artifact_type = {p} AND artifact_name = {p} " + "ORDER BY updated_at DESC, version_hash DESC", + params=[artifact_type.value, artifact_name], + ) + return tuple( + PhysicalRelationRecord( + artifact_type=PhysicalArtifactType(row[0]), + artifact_name=row[1], + version_hash=row[2], + database_name=row[3], + schema_name=row[4], + relation_name=row[5], + relation_type=row[6], + ) + for row in rows + ) + + def get_physical_relation_ancestry( + self, *, connection: Any, schema: str, model_name: str, version_hash: str + ) -> PhysicalRelationAncestryRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT model_name, version_hash, parent_model_name, parent_version_hash, " + "seed_strategy " + f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_ANCESTRY_TABLE)} " + f"WHERE model_name = {p} AND version_hash = {p}", + params=[model_name, version_hash], + ) + if row is None: + return None + return PhysicalRelationAncestryRecord( + model_name=row[0], + version_hash=row[1], + parent_model_name=row[2], + parent_version_hash=row[3], + seed_strategy=row[4], + ) + + def get_virtual_environment( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> VirtualEnvironmentRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT virtual_environment_name, status, baseline_virtual_environment_name, " + "finalized_at " + f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " + f"WHERE virtual_environment_name = {p}", + params=[virtual_environment_name], + ) + if row is None: + return None + return VirtualEnvironmentRecord( + virtual_environment_name=row[0], + status=VirtualEnvironmentStatus(row[1]), + baseline_virtual_environment_name=row[2], + finalized_at=row[3], + ) + + def get_virtual_environment_node_refs( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + node_type: str, + ) -> tuple[VirtualEnvironmentNodeRefRecord, ...]: + p: str = self._placeholder + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT virtual_environment_name, node_type, node_name, version_hash " + f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " + f"WHERE virtual_environment_name = {p} AND node_type = {p} ORDER BY node_name", + params=[virtual_environment_name, node_type], + ) + return tuple( + VirtualEnvironmentNodeRefRecord( + virtual_environment_name=row[0], + node_type=row[1], + node_name=row[2], + version_hash=row[3], + ) + for row in rows + ) + + def get_virtual_environment_model_refs( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> tuple[VirtualEnvironmentModelRefRecord, ...]: + refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = self.get_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + node_type="model", + ) + return tuple( + VirtualEnvironmentModelRefRecord( + virtual_environment_name=ref.virtual_environment_name, + model_name=ref.node_name, + version_hash=ref.version_hash, + ) + for ref in refs + ) + + def get_virtual_environment_function_refs( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> tuple[VirtualEnvironmentFunctionRefRecord, ...]: + refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + *self.get_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + node_type="udf", + ), + *self.get_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + node_type="table_fn", + ), + ) + return tuple( + VirtualEnvironmentFunctionRefRecord( + virtual_environment_name=ref.virtual_environment_name, + node_type=ref.node_type, + function_name=ref.node_name, + version_hash=ref.version_hash, + ) + for ref in sorted(refs, key=lambda item: (item.node_type, item.node_name)) + ) + + def get_virtual_environment_seed_refs( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> tuple[VirtualEnvironmentSeedRefRecord, ...]: + refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = self.get_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + node_type="seed", + ) + return tuple( + VirtualEnvironmentSeedRefRecord( + virtual_environment_name=ref.virtual_environment_name, + seed_name=ref.node_name, + version_hash=ref.version_hash, + ) + for ref in refs + ) + + def get_virtual_environment_python_node_refs( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> tuple[VirtualEnvironmentPythonNodeRefRecord, ...]: + p: str = self._placeholder + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT virtual_environment_name, node_type, node_name, version_hash " + f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " + f"WHERE virtual_environment_name = {p} " + "AND node_type IN ('task', 'loader', 'asset', 'check', 'hook') " + "ORDER BY node_type, node_name", + params=[virtual_environment_name], + ) + return tuple( + VirtualEnvironmentPythonNodeRefRecord( + virtual_environment_name=row[0], + node_type=row[1], + node_name=row[2], + version_hash=row[3], + ) + for row in rows + ) + + def get_virtual_environment_source_freshness( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> tuple[SourceFreshnessRecord, ...]: + p: str = self._placeholder + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT virtual_environment_name, source_name, strategy, value_kind, " + "data_version, data_version_hash, observed_at " + f"FROM {self._qualified_name(schema=schema, table=SOURCE_FRESHNESS_OBSERVATION_TABLE)} " + f"WHERE virtual_environment_name = {p} ORDER BY source_name", + params=[virtual_environment_name], + ) + return tuple( + SourceFreshnessRecord( + virtual_environment_name=row[0], + source_name=row[1], + strategy=row[2], + value_kind=row[3], + data_version=row[4], + data_version_hash=row[5], + observed_at=from_naive_utc_wall_clock(row[6]), + ) + for row in rows + ) + + def list_virtual_environment_checkpoints( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> tuple[VirtualEnvironmentCheckpointRecord, ...]: + p: str = self._placeholder + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT checkpoint_id, virtual_environment_name, created_at " + "FROM " + f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE)} " + f"WHERE virtual_environment_name = {p} ORDER BY created_at DESC, checkpoint_id DESC", + params=[virtual_environment_name], + ) + return tuple( + VirtualEnvironmentCheckpointRecord( + checkpoint_id=row[0], + virtual_environment_name=row[1], + created_at=row[2], + ) + for row in rows + ) + + def get_virtual_environment_checkpoint_model_refs( + self, *, connection: Any, schema: str, checkpoint_id: str + ) -> tuple[VirtualEnvironmentCheckpointModelRefRecord, ...]: + p: str = self._placeholder + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT checkpoint_id, model_name, version_hash FROM " + + self._qualified_name( + schema=schema, + table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, + ) + + " " + f"WHERE checkpoint_id = {p} ORDER BY model_name", + params=[checkpoint_id], + ) + return tuple( + VirtualEnvironmentCheckpointModelRefRecord( + checkpoint_id=row[0], + model_name=row[1], + version_hash=row[2], + ) + for row in rows + ) + + def get_virtual_environment_checkpoint_function_refs( + self, *, connection: Any, schema: str, checkpoint_id: str + ) -> tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...]: + p: str = self._placeholder + checkpoint_function_ref_table: str = self._qualified_name( + schema=schema, + table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, + ) + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql=f"SELECT checkpoint_id, function_name, version_hash " + f"FROM {checkpoint_function_ref_table} " + f"WHERE checkpoint_id = {p} ORDER BY function_name", + params=[checkpoint_id], + ) + return tuple( + VirtualEnvironmentCheckpointFunctionRefRecord( + checkpoint_id=row[0], + function_name=row[1], + version_hash=row[2], + ) + for row in rows + ) + + def get_virtual_environment_checkpoint_seed_refs( + self, *, connection: Any, schema: str, checkpoint_id: str + ) -> tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...]: + p: str = self._placeholder + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT checkpoint_id, seed_name, version_hash FROM " + + self._qualified_name( + schema=schema, + table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, + ) + + " " + f"WHERE checkpoint_id = {p} ORDER BY seed_name", + params=[checkpoint_id], + ) + return tuple( + VirtualEnvironmentCheckpointSeedRefRecord( + checkpoint_id=row[0], + seed_name=row[1], + version_hash=row[2], + ) + for row in rows + ) + + def get_state_operation( + self, *, connection: Any, schema: str, operation_id: str + ) -> StateOperationRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT operation_id, operation_type, status, virtual_environment_name " + f"FROM {self._qualified_name(schema=schema, table=STATE_OPERATION_TABLE)} " + f"WHERE operation_id = {p}", + params=[operation_id], + ) + if row is None: + return None + return StateOperationRecord( + operation_id=row[0], + operation_type=StateOperationType(row[1]), + status=StateOperationStatus(row[2]), + virtual_environment_name=row[3], + ) diff --git a/src/sqlbuild/virtual/state/classes/duckdb.py b/src/sqlbuild/virtual/state/classes/duckdb.py index 9c1405d714..56fb0e97da 100644 --- a/src/sqlbuild/virtual/state/classes/duckdb.py +++ b/src/sqlbuild/virtual/state/classes/duckdb.py @@ -2,8 +2,9 @@ from __future__ import annotations +from collections.abc import Sequence from datetime import datetime -from typing import Any +from typing import Any, ClassVar from sqlbuild.adapter.contract.classes.observed_connection import ObservedConnection from sqlbuild.compiler.compile.types import CompiledResourceType @@ -16,7 +17,6 @@ ) from sqlbuild.microbatches.models import MicrobatchEvent, MicrobatchScope, MicrobatchWriteResult from sqlbuild.virtual.state._helpers.state_storage.datetime import ( - from_naive_utc_wall_clock, to_naive_utc_wall_clock, ) from sqlbuild.virtual.state._helpers.state_storage.events import backup_id, event_id @@ -31,7 +31,7 @@ from sqlbuild.virtual.state.classes._duckdb_conditional_publish import ( DuckDbConditionalPublishMixin, ) -from sqlbuild.virtual.state.classes.state_backend import StateBackend +from sqlbuild.virtual.state.classes._sql_state_backend import SqlStateBackend from sqlbuild.virtual.state.constants import ( CURRENT_STATE_SCHEMA_VERSION, DUCKDB_DATETIME_TYPE_TOKEN, @@ -96,20 +96,18 @@ VirtualEnvironmentSeedRefRecord, ) from sqlbuild.virtual.state.types import ( - ModelVersionStatus, - PhysicalArtifactType, StateColumnType, StateMigrationAction, StateMigrationStatus, - StateOperationStatus, - StateOperationType, VirtualEnvironmentStatus, ) -class DuckDbStateBackend(DuckDbConditionalPublishMixin, StateBackend): +class DuckDbStateBackend(DuckDbConditionalPublishMixin, SqlStateBackend): """DuckDB implementation for virtual state.""" + _placeholder: ClassVar[str] = "?" + def connect(self, config: dict[str, object]) -> Any: import duckdb @@ -121,6 +119,20 @@ def connect(self, config: dict[str, object]) -> Any: def close(self, connection: Any) -> None: connection.close() + def _fetch_one( + self, *, connection: Any, sql: str, params: Sequence[object] | None = None + ) -> tuple[Any, ...] | None: + if params is None: + return connection.execute(sql).fetchone() + return connection.execute(sql, params).fetchone() + + def _fetch_all( + self, *, connection: Any, sql: str, params: Sequence[object] | None = None + ) -> list[tuple[Any, ...]]: + if params is None: + return connection.execute(sql).fetchall() + return connection.execute(sql, params).fetchall() + def initialize(self, *, connection: Any, schema: str, sqlbuild_version: str) -> None: connection.execute("BEGIN") try: @@ -323,30 +335,6 @@ def upsert_model_version( connection.execute("ROLLBACK") raise - def get_model_version( - self, *, connection: Any, schema: str, model_name: str, version_hash: str - ) -> ModelVersionRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT model_name, version_hash, definition_identity_hash, " - "identity_metadata_hash, definition_text_b64, identity_metadata_json_b64, " - "compiled_sql_b64, status " - f"FROM {self._qualified_name(schema=schema, table=MODEL_VERSION_TABLE)} " - "WHERE model_name = ? AND version_hash = ?", - [model_name, version_hash], - ).fetchone() - if row is None: - return None - return ModelVersionRecord( - model_name=row[0], - version_hash=row[1], - definition_identity_hash=row[2], - identity_metadata_hash=row[3], - definition_text_b64=row[4], - identity_metadata_json_b64=row[5], - compiled_sql_b64=row[6], - status=ModelVersionStatus(row[7]), - ) - def upsert_function_version( self, *, connection: Any, schema: str, record: FunctionVersionRecord ) -> None: @@ -392,34 +380,6 @@ def upsert_function_version( connection.execute("ROLLBACK") raise - def get_function_version( - self, *, connection: Any, schema: str, function_name: str, version_hash: str - ) -> FunctionVersionRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT function_name, version_hash, language, returns, arguments_json_b64, " - "return_columns_json_b64, packages_json_b64, runtime_version, entry_point, " - "body_sql_b64, definition_text_b64, status " - f"FROM {self._qualified_name(schema=schema, table=FUNCTION_VERSION_TABLE)} " - "WHERE function_name = ? AND version_hash = ?", - [function_name, version_hash], - ).fetchone() - if row is None: - return None - return FunctionVersionRecord( - function_name=row[0], - version_hash=row[1], - language=row[2], - returns=row[3], - arguments_json_b64=row[4], - return_columns_json_b64=row[5], - packages_json_b64=row[6], - runtime_version=row[7], - entry_point=row[8], - body_sql_b64=row[9], - definition_text_b64=row[10], - status=ModelVersionStatus(row[11]), - ) - def upsert_seed_version( self, *, connection: Any, schema: str, record: SeedVersionRecord ) -> None: @@ -456,26 +416,6 @@ def upsert_seed_version( connection.execute("ROLLBACK") raise - def get_seed_version( - self, *, connection: Any, schema: str, seed_name: str, version_hash: str - ) -> SeedVersionRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT seed_name, version_hash, identity_metadata_hash, " - "identity_metadata_json_b64, status " - f"FROM {self._qualified_name(schema=schema, table=SEED_VERSION_TABLE)} " - "WHERE seed_name = ? AND version_hash = ?", - [seed_name, version_hash], - ).fetchone() - if row is None: - return None - return SeedVersionRecord( - seed_name=row[0], - version_hash=row[1], - identity_metadata_hash=row[2], - identity_metadata_json_b64=row[3], - status=ModelVersionStatus(row[4]), - ) - def upsert_python_node_version( self, *, connection: Any, schema: str, record: PythonNodeVersionRecord ) -> None: @@ -519,35 +459,6 @@ def upsert_python_node_version( connection.execute("ROLLBACK") raise - def get_python_node_version( - self, - *, - connection: Any, - schema: str, - node_type: str, - node_name: str, - version_hash: str, - ) -> PythonNodeVersionRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT node_type, node_name, version_hash, definition_hash, " - "identity_metadata_hash, definition_json_b64, identity_metadata_json_b64, status " - f"FROM {self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " - "WHERE node_type = ? AND node_name = ? AND version_hash = ?", - [node_type, node_name, version_hash], - ).fetchone() - if row is None: - return None - return PythonNodeVersionRecord( - node_type=row[0], - node_name=row[1], - version_hash=row[2], - definition_hash=row[3], - identity_metadata_hash=row[4], - definition_json_b64=row[5], - identity_metadata_json_b64=row[6], - status=ModelVersionStatus(row[7]), - ) - def insert_node_result( self, *, @@ -710,63 +621,6 @@ def upsert_physical_relation( connection.execute("ROLLBACK") raise - def get_physical_relation_for_artifact( - self, - *, - connection: Any, - schema: str, - artifact_type: PhysicalArtifactType, - artifact_name: str, - version_hash: str, - ) -> PhysicalRelationRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT artifact_type, artifact_name, version_hash, database_name, schema_name, " - "relation_name, relation_type " - f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "WHERE artifact_type = ? AND artifact_name = ? AND version_hash = ?", - [artifact_type.value, artifact_name, version_hash], - ).fetchone() - if row is None: - return None - return PhysicalRelationRecord( - artifact_type=PhysicalArtifactType(row[0]), - artifact_name=row[1], - version_hash=row[2], - database_name=row[3], - schema_name=row[4], - relation_name=row[5], - relation_type=row[6], - ) - - def list_physical_relations_for_artifact( - self, - *, - connection: Any, - schema: str, - artifact_type: PhysicalArtifactType, - artifact_name: str, - ) -> tuple[PhysicalRelationRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT artifact_type, artifact_name, version_hash, database_name, schema_name, " - "relation_name, relation_type " - f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "WHERE artifact_type = ? AND artifact_name = ? " - "ORDER BY updated_at DESC, version_hash DESC", - [artifact_type.value, artifact_name], - ).fetchall() - return tuple( - PhysicalRelationRecord( - artifact_type=PhysicalArtifactType(row[0]), - artifact_name=row[1], - version_hash=row[2], - database_name=row[3], - schema_name=row[4], - relation_name=row[5], - relation_type=row[6], - ) - for row in rows - ) - def upsert_physical_relation_ancestry( self, *, connection: Any, schema: str, record: PhysicalRelationAncestryRecord ) -> None: @@ -805,26 +659,6 @@ def upsert_physical_relation_ancestry( connection.execute("ROLLBACK") raise - def get_physical_relation_ancestry( - self, *, connection: Any, schema: str, model_name: str, version_hash: str - ) -> PhysicalRelationAncestryRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT model_name, version_hash, parent_model_name, parent_version_hash, " - "seed_strategy " - f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_ANCESTRY_TABLE)} " - "WHERE model_name = ? AND version_hash = ?", - [model_name, version_hash], - ).fetchone() - if row is None: - return None - return PhysicalRelationAncestryRecord( - model_name=row[0], - version_hash=row[1], - parent_model_name=row[2], - parent_version_hash=row[3], - seed_strategy=row[4], - ) - def upsert_virtual_environment( self, *, connection: Any, schema: str, record: VirtualEnvironmentRecord ) -> None: @@ -838,25 +672,6 @@ def upsert_virtual_environment( connection.execute("ROLLBACK") raise - def get_virtual_environment( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> VirtualEnvironmentRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT virtual_environment_name, status, baseline_virtual_environment_name, " - "finalized_at " - f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "WHERE virtual_environment_name = ?", - [virtual_environment_name], - ).fetchone() - if row is None: - return None - return VirtualEnvironmentRecord( - virtual_environment_name=row[0], - status=VirtualEnvironmentStatus(row[1]), - baseline_virtual_environment_name=row[2], - finalized_at=row[3], - ) - def list_virtual_environments( self, *, connection: Any, schema: str ) -> tuple[VirtualEnvironmentRetentionRecord, ...]: @@ -963,30 +778,6 @@ def upsert_virtual_environment_and_replace_node_ref_groups( connection.execute("ROLLBACK") raise - def get_virtual_environment_node_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - node_type: str, - ) -> tuple[VirtualEnvironmentNodeRefRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT virtual_environment_name, node_type, node_name, version_hash " - f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "WHERE virtual_environment_name = ? AND node_type = ? ORDER BY node_name", - [virtual_environment_name, node_type], - ).fetchall() - return tuple( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=row[0], - node_type=row[1], - node_name=row[2], - version_hash=row[3], - ) - for row in rows - ) - def upsert_virtual_environment_node_ref( self, *, @@ -1028,24 +819,6 @@ def replace_virtual_environment_model_refs( ), ) - def get_virtual_environment_model_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentModelRefRecord, ...]: - refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = self.get_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="model", - ) - return tuple( - VirtualEnvironmentModelRefRecord( - virtual_environment_name=ref.virtual_environment_name, - model_name=ref.node_name, - version_hash=ref.version_hash, - ) - for ref in refs - ) - def replace_virtual_environment_function_refs( self, *, @@ -1082,33 +855,6 @@ def replace_virtual_environment_function_refs( refs_by_node_type=refs_by_node_type, ) - def get_virtual_environment_function_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentFunctionRefRecord, ...]: - refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( - *self.get_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="udf", - ), - *self.get_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="table_fn", - ), - ) - return tuple( - VirtualEnvironmentFunctionRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type=ref.node_type, - function_name=ref.node_name, - version_hash=ref.version_hash, - ) - for ref in sorted(refs, key=lambda item: (item.node_type, item.node_name)) - ) - def replace_virtual_environment_seed_refs( self, *, @@ -1133,24 +879,6 @@ def replace_virtual_environment_seed_refs( ), ) - def get_virtual_environment_seed_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentSeedRefRecord, ...]: - refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = self.get_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="seed", - ) - return tuple( - VirtualEnvironmentSeedRefRecord( - virtual_environment_name=ref.virtual_environment_name, - seed_name=ref.node_name, - version_hash=ref.version_hash, - ) - for ref in refs - ) - def upsert_virtual_environment_python_node_ref( self, *, @@ -1169,27 +897,6 @@ def upsert_virtual_environment_python_node_ref( ), ) - def get_virtual_environment_python_node_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentPythonNodeRefRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT virtual_environment_name, node_type, node_name, version_hash " - f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "WHERE virtual_environment_name = ? " - "AND node_type IN ('task', 'loader', 'asset', 'check', 'hook') " - "ORDER BY node_type, node_name", - [virtual_environment_name], - ).fetchall() - return tuple( - VirtualEnvironmentPythonNodeRefRecord( - virtual_environment_name=row[0], - node_type=row[1], - node_name=row[2], - version_hash=row[3], - ) - for row in rows - ) - def count_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: row: tuple[Any, ...] = connection.execute( "SELECT COUNT(*) " @@ -1296,29 +1003,6 @@ def replace_virtual_environment_source_freshness( connection.execute("ROLLBACK") raise - def get_virtual_environment_source_freshness( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[SourceFreshnessRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT virtual_environment_name, source_name, strategy, value_kind, " - "data_version, data_version_hash, observed_at " - f"FROM {self._qualified_name(schema=schema, table=SOURCE_FRESHNESS_OBSERVATION_TABLE)} " - "WHERE virtual_environment_name = ? ORDER BY source_name", - [virtual_environment_name], - ).fetchall() - return tuple( - SourceFreshnessRecord( - virtual_environment_name=row[0], - source_name=row[1], - strategy=row[2], - value_kind=row[3], - data_version=row[4], - data_version_hash=row[5], - observed_at=from_naive_utc_wall_clock(row[6]), - ) - for row in rows - ) - def create_virtual_environment_checkpoint( self, *, @@ -1344,91 +1028,6 @@ def create_virtual_environment_checkpoint( connection.execute("ROLLBACK") raise - def list_virtual_environment_checkpoints( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentCheckpointRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT checkpoint_id, virtual_environment_name, created_at " - "FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE)} " - "WHERE virtual_environment_name = ? ORDER BY created_at DESC, checkpoint_id DESC", - [virtual_environment_name], - ).fetchall() - return tuple( - VirtualEnvironmentCheckpointRecord( - checkpoint_id=row[0], - virtual_environment_name=row[1], - created_at=row[2], - ) - for row in rows - ) - - def get_virtual_environment_checkpoint_model_refs( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> tuple[VirtualEnvironmentCheckpointModelRefRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT checkpoint_id, model_name, version_hash FROM " - + self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, - ) - + " " - "WHERE checkpoint_id = ? ORDER BY model_name", - [checkpoint_id], - ).fetchall() - return tuple( - VirtualEnvironmentCheckpointModelRefRecord( - checkpoint_id=row[0], - model_name=row[1], - version_hash=row[2], - ) - for row in rows - ) - - def get_virtual_environment_checkpoint_function_refs( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...]: - checkpoint_function_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, - ) - rows: list[tuple[Any, ...]] = connection.execute( - f"SELECT checkpoint_id, function_name, version_hash " - f"FROM {checkpoint_function_ref_table} " - "WHERE checkpoint_id = ? ORDER BY function_name", - [checkpoint_id], - ).fetchall() - return tuple( - VirtualEnvironmentCheckpointFunctionRefRecord( - checkpoint_id=row[0], - function_name=row[1], - version_hash=row[2], - ) - for row in rows - ) - - def get_virtual_environment_checkpoint_seed_refs( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT checkpoint_id, seed_name, version_hash FROM " - + self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, - ) - + " " - "WHERE checkpoint_id = ? ORDER BY seed_name", - [checkpoint_id], - ).fetchall() - return tuple( - VirtualEnvironmentCheckpointSeedRefRecord( - checkpoint_id=row[0], - seed_name=row[1], - version_hash=row[2], - ) - for row in rows - ) - def delete_virtual_environment_checkpoint( self, *, connection: Any, schema: str, checkpoint_id: str ) -> None: @@ -1505,24 +1104,6 @@ def upsert_state_operation( connection.execute("ROLLBACK") raise - def get_state_operation( - self, *, connection: Any, schema: str, operation_id: str - ) -> StateOperationRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT operation_id, operation_type, status, virtual_environment_name " - f"FROM {self._qualified_name(schema=schema, table=STATE_OPERATION_TABLE)} " - "WHERE operation_id = ?", - [operation_id], - ).fetchone() - if row is None: - return None - return StateOperationRecord( - operation_id=row[0], - operation_type=StateOperationType(row[1]), - status=StateOperationStatus(row[2]), - virtual_environment_name=row[3], - ) - def create_state_operation_event( self, *, connection: Any, schema: str, record: StateOperationEventRecord ) -> None: @@ -1821,12 +1402,6 @@ def _parse_materialized(self, value: object) -> bool | None: return None return str(value).lower() == STATE_BOOLEAN_TRUE - def _quote_identifier(self, identifier: str) -> str: - return '"' + identifier.replace('"', '""') + '"' - - def _qualified_name(self, *, schema: str, table: str) -> str: - return f"{self._quote_identifier(schema)}.{self._quote_identifier(table)}" - def _validate_source_freshness_records( self, *, diff --git a/src/sqlbuild/virtual/state/classes/postgres.py b/src/sqlbuild/virtual/state/classes/postgres.py index 48e910d476..2495bb685f 100644 --- a/src/sqlbuild/virtual/state/classes/postgres.py +++ b/src/sqlbuild/virtual/state/classes/postgres.py @@ -2,8 +2,9 @@ from __future__ import annotations +from collections.abc import Sequence from datetime import datetime -from typing import Any +from typing import Any, ClassVar from sqlbuild.adapter.contract.classes.observed_connection import ObservedConnection from sqlbuild.compiler.compile.types import CompiledResourceType @@ -22,7 +23,6 @@ ) from sqlbuild.microbatches.models import MicrobatchEvent, MicrobatchScope, MicrobatchWriteResult from sqlbuild.virtual.state._helpers.state_storage.datetime import ( - from_naive_utc_wall_clock, to_naive_utc_wall_clock, ) from sqlbuild.virtual.state._helpers.state_storage.events import backup_id, event_id @@ -30,7 +30,7 @@ build_validation_result, validate_conditional_virtual_environment_publication, ) -from sqlbuild.virtual.state.classes.state_backend import StateBackend +from sqlbuild.virtual.state.classes._sql_state_backend import SqlStateBackend from sqlbuild.virtual.state.constants import ( CURRENT_STATE_SCHEMA_VERSION, FUNCTION_VERSION_TABLE, @@ -95,20 +95,18 @@ VirtualEnvironmentSeedRefRecord, ) from sqlbuild.virtual.state.types import ( - ModelVersionStatus, - PhysicalArtifactType, StateColumnType, StateMigrationAction, StateMigrationStatus, - StateOperationStatus, - StateOperationType, VirtualEnvironmentStatus, ) -class PostgresStateBackend(StateBackend): +class PostgresStateBackend(SqlStateBackend): """Postgres implementation for virtual state.""" + _placeholder: ClassVar[str] = "%s" + def connect(self, config: dict[str, object]) -> Any: try: import psycopg @@ -136,6 +134,26 @@ def connect(self, config: dict[str, object]) -> Any: def close(self, connection: Any) -> None: connection.close() + def _fetch_one( + self, *, connection: Any, sql: str, params: Sequence[object] | None = None + ) -> tuple[Any, ...] | None: + with connection.cursor() as cursor: + if params is None: + cursor.execute(sql) + else: + cursor.execute(sql, params) + return cursor.fetchone() + + def _fetch_all( + self, *, connection: Any, sql: str, params: Sequence[object] | None = None + ) -> list[tuple[Any, ...]]: + with connection.cursor() as cursor: + if params is None: + cursor.execute(sql) + else: + cursor.execute(sql, params) + return cursor.fetchall() + def initialize(self, *, connection: Any, schema: str, sqlbuild_version: str) -> None: with connection.cursor() as cursor: cursor.execute("BEGIN") @@ -350,32 +368,6 @@ def upsert_model_version( cursor.execute("ROLLBACK") raise - def get_model_version( - self, *, connection: Any, schema: str, model_name: str, version_hash: str - ) -> ModelVersionRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT model_name, version_hash, definition_identity_hash, " - "identity_metadata_hash, definition_text_b64, identity_metadata_json_b64, " - "compiled_sql_b64, status " - f"FROM {self._qualified_name(schema=schema, table=MODEL_VERSION_TABLE)} " - "WHERE model_name = %s AND version_hash = %s", - [model_name, version_hash], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return ModelVersionRecord( - model_name=row[0], - version_hash=row[1], - definition_identity_hash=row[2], - identity_metadata_hash=row[3], - definition_text_b64=row[4], - identity_metadata_json_b64=row[5], - compiled_sql_b64=row[6], - status=ModelVersionStatus(row[7]), - ) - def upsert_function_version( self, *, connection: Any, schema: str, record: FunctionVersionRecord ) -> None: @@ -424,36 +416,6 @@ def upsert_function_version( cursor.execute("ROLLBACK") raise - def get_function_version( - self, *, connection: Any, schema: str, function_name: str, version_hash: str - ) -> FunctionVersionRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT function_name, version_hash, language, returns, arguments_json_b64, " - "return_columns_json_b64, packages_json_b64, runtime_version, entry_point, " - "body_sql_b64, definition_text_b64, status " - f"FROM {self._qualified_name(schema=schema, table=FUNCTION_VERSION_TABLE)} " - "WHERE function_name = %s AND version_hash = %s", - [function_name, version_hash], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return FunctionVersionRecord( - function_name=row[0], - version_hash=row[1], - language=row[2], - returns=row[3], - arguments_json_b64=row[4], - return_columns_json_b64=row[5], - packages_json_b64=row[6], - runtime_version=row[7], - entry_point=row[8], - body_sql_b64=row[9], - definition_text_b64=row[10], - status=ModelVersionStatus(row[11]), - ) - def upsert_seed_version( self, *, connection: Any, schema: str, record: SeedVersionRecord ) -> None: @@ -492,28 +454,6 @@ def upsert_seed_version( cursor.execute("ROLLBACK") raise - def get_seed_version( - self, *, connection: Any, schema: str, seed_name: str, version_hash: str - ) -> SeedVersionRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT seed_name, version_hash, identity_metadata_hash, " - "identity_metadata_json_b64, status " - f"FROM {self._qualified_name(schema=schema, table=SEED_VERSION_TABLE)} " - "WHERE seed_name = %s AND version_hash = %s", - [seed_name, version_hash], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return SeedVersionRecord( - seed_name=row[0], - version_hash=row[1], - identity_metadata_hash=row[2], - identity_metadata_json_b64=row[3], - status=ModelVersionStatus(row[4]), - ) - def upsert_python_node_version( self, *, connection: Any, schema: str, record: PythonNodeVersionRecord ) -> None: @@ -558,37 +498,6 @@ def upsert_python_node_version( cursor.execute("ROLLBACK") raise - def get_python_node_version( - self, - *, - connection: Any, - schema: str, - node_type: str, - node_name: str, - version_hash: str, - ) -> PythonNodeVersionRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT node_type, node_name, version_hash, definition_hash, " - "identity_metadata_hash, definition_json_b64, identity_metadata_json_b64, status " - f"FROM {self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " - "WHERE node_type = %s AND node_name = %s AND version_hash = %s", - [node_type, node_name, version_hash], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return PythonNodeVersionRecord( - node_type=row[0], - node_name=row[1], - version_hash=row[2], - definition_hash=row[3], - identity_metadata_hash=row[4], - definition_json_b64=row[5], - identity_metadata_json_b64=row[6], - status=ModelVersionStatus(row[7]), - ) - def insert_node_result( self, *, @@ -813,67 +722,6 @@ def upsert_physical_relation( cursor.execute("ROLLBACK") raise - def get_physical_relation_for_artifact( - self, - *, - connection: Any, - schema: str, - artifact_type: PhysicalArtifactType, - artifact_name: str, - version_hash: str, - ) -> PhysicalRelationRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT artifact_type, artifact_name, version_hash, database_name, schema_name, " - "relation_name, relation_type " - f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "WHERE artifact_type = %s AND artifact_name = %s AND version_hash = %s", - [artifact_type.value, artifact_name, version_hash], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return PhysicalRelationRecord( - artifact_type=PhysicalArtifactType(row[0]), - artifact_name=row[1], - version_hash=row[2], - database_name=row[3], - schema_name=row[4], - relation_name=row[5], - relation_type=row[6], - ) - - def list_physical_relations_for_artifact( - self, - *, - connection: Any, - schema: str, - artifact_type: PhysicalArtifactType, - artifact_name: str, - ) -> tuple[PhysicalRelationRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT artifact_type, artifact_name, version_hash, database_name, schema_name, " - "relation_name, relation_type " - f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "WHERE artifact_type = %s AND artifact_name = %s " - "ORDER BY updated_at DESC, version_hash DESC", - [artifact_type.value, artifact_name], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - PhysicalRelationRecord( - artifact_type=PhysicalArtifactType(row[0]), - artifact_name=row[1], - version_hash=row[2], - database_name=row[3], - schema_name=row[4], - relation_name=row[5], - relation_type=row[6], - ) - for row in rows - ) - def upsert_physical_relation_ancestry( self, *, connection: Any, schema: str, record: PhysicalRelationAncestryRecord ) -> None: @@ -916,29 +764,6 @@ def upsert_physical_relation_ancestry( cursor.execute("ROLLBACK") raise - def get_physical_relation_ancestry( - self, *, connection: Any, schema: str, model_name: str, version_hash: str - ) -> PhysicalRelationAncestryRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT model_name, version_hash, parent_model_name, parent_version_hash, " - "seed_strategy " - "FROM " - f"{self._qualified_name(schema=schema, table=PHYSICAL_RELATION_ANCESTRY_TABLE)} " - "WHERE model_name = %s AND version_hash = %s", - [model_name, version_hash], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return PhysicalRelationAncestryRecord( - model_name=row[0], - version_hash=row[1], - parent_model_name=row[2], - parent_version_hash=row[3], - seed_strategy=row[4], - ) - def upsert_virtual_environment( self, *, connection: Any, schema: str, record: VirtualEnvironmentRecord ) -> None: @@ -951,27 +776,6 @@ def upsert_virtual_environment( cursor.execute("ROLLBACK") raise - def get_virtual_environment( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> VirtualEnvironmentRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT virtual_environment_name, status, baseline_virtual_environment_name, " - "finalized_at " - f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "WHERE virtual_environment_name = %s", - [virtual_environment_name], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return VirtualEnvironmentRecord( - virtual_environment_name=row[0], - status=VirtualEnvironmentStatus(row[1]), - baseline_virtual_environment_name=row[2], - finalized_at=row[3], - ) - def list_virtual_environments( self, *, connection: Any, schema: str ) -> tuple[VirtualEnvironmentRetentionRecord, ...]: @@ -1147,33 +951,6 @@ def upsert_virtual_environment_and_replace_node_ref_groups_if_locks_owned( cursor.execute("ROLLBACK") raise - def get_virtual_environment_node_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - node_type: str, - ) -> tuple[VirtualEnvironmentNodeRefRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT virtual_environment_name, node_type, node_name, version_hash " - "FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "WHERE virtual_environment_name = %s AND node_type = %s ORDER BY node_name", - [virtual_environment_name, node_type], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=row[0], - node_type=row[1], - node_name=row[2], - version_hash=row[3], - ) - for row in rows - ) - def upsert_virtual_environment_node_ref( self, *, @@ -1217,24 +994,6 @@ def replace_virtual_environment_model_refs( ), ) - def get_virtual_environment_model_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentModelRefRecord, ...]: - refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = self.get_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="model", - ) - return tuple( - VirtualEnvironmentModelRefRecord( - virtual_environment_name=ref.virtual_environment_name, - model_name=ref.node_name, - version_hash=ref.version_hash, - ) - for ref in refs - ) - def replace_virtual_environment_function_refs( self, *, @@ -1271,33 +1030,6 @@ def replace_virtual_environment_function_refs( refs_by_node_type=refs_by_node_type, ) - def get_virtual_environment_function_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentFunctionRefRecord, ...]: - refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( - *self.get_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="udf", - ), - *self.get_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="table_fn", - ), - ) - return tuple( - VirtualEnvironmentFunctionRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type=ref.node_type, - function_name=ref.node_name, - version_hash=ref.version_hash, - ) - for ref in sorted(refs, key=lambda item: (item.node_type, item.node_name)) - ) - def replace_virtual_environment_seed_refs( self, *, @@ -1322,24 +1054,6 @@ def replace_virtual_environment_seed_refs( ), ) - def get_virtual_environment_seed_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentSeedRefRecord, ...]: - refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = self.get_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="seed", - ) - return tuple( - VirtualEnvironmentSeedRefRecord( - virtual_environment_name=ref.virtual_environment_name, - seed_name=ref.node_name, - version_hash=ref.version_hash, - ) - for ref in refs - ) - def upsert_virtual_environment_python_node_ref( self, *, @@ -1358,30 +1072,6 @@ def upsert_virtual_environment_python_node_ref( ), ) - def get_virtual_environment_python_node_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentPythonNodeRefRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT virtual_environment_name, node_type, node_name, version_hash " - "FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "WHERE virtual_environment_name = %s " - "AND node_type IN ('task', 'loader', 'asset', 'check', 'hook') " - "ORDER BY node_type, node_name", - [virtual_environment_name], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - VirtualEnvironmentPythonNodeRefRecord( - virtual_environment_name=row[0], - node_type=row[1], - node_name=row[2], - version_hash=row[3], - ) - for row in rows - ) - def count_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: with connection.cursor() as cursor: cursor.execute( @@ -1472,32 +1162,6 @@ def replace_virtual_environment_source_freshness( cursor.execute("ROLLBACK") raise - def get_virtual_environment_source_freshness( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[SourceFreshnessRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT virtual_environment_name, source_name, strategy, value_kind, " - "data_version, data_version_hash, observed_at " - "FROM " - f"{self._qualified_name(schema=schema, table=SOURCE_FRESHNESS_OBSERVATION_TABLE)} " - "WHERE virtual_environment_name = %s ORDER BY source_name", - [virtual_environment_name], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - SourceFreshnessRecord( - virtual_environment_name=row[0], - source_name=row[1], - strategy=row[2], - value_kind=row[3], - data_version=row[4], - data_version_hash=row[5], - observed_at=from_naive_utc_wall_clock(row[6]), - ) - for row in rows - ) - def create_virtual_environment_checkpoint( self, *, @@ -1574,103 +1238,6 @@ def _insert_virtual_environment_checkpoint_rows( [seed_ref.checkpoint_id, seed_ref.seed_name, seed_ref.version_hash], ) - def list_virtual_environment_checkpoints( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentCheckpointRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT checkpoint_id, virtual_environment_name, created_at " - "FROM " - + self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE, - ) - + " " - "WHERE virtual_environment_name = %s ORDER BY created_at DESC, checkpoint_id DESC", - [virtual_environment_name], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - VirtualEnvironmentCheckpointRecord( - checkpoint_id=row[0], - virtual_environment_name=row[1], - created_at=row[2], - ) - for row in rows - ) - - def get_virtual_environment_checkpoint_model_refs( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> tuple[VirtualEnvironmentCheckpointModelRefRecord, ...]: - with connection.cursor() as cursor: - checkpoint_model_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, - ) - cursor.execute( - f"SELECT checkpoint_id, model_name, version_hash " - f"FROM {checkpoint_model_ref_table} " - "WHERE checkpoint_id = %s ORDER BY model_name", - [checkpoint_id], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - VirtualEnvironmentCheckpointModelRefRecord( - checkpoint_id=row[0], - model_name=row[1], - version_hash=row[2], - ) - for row in rows - ) - - def get_virtual_environment_checkpoint_function_refs( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...]: - with connection.cursor() as cursor: - checkpoint_function_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, - ) - cursor.execute( - f"SELECT checkpoint_id, function_name, version_hash " - f"FROM {checkpoint_function_ref_table} " - "WHERE checkpoint_id = %s ORDER BY function_name", - [checkpoint_id], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - VirtualEnvironmentCheckpointFunctionRefRecord( - checkpoint_id=row[0], - function_name=row[1], - version_hash=row[2], - ) - for row in rows - ) - - def get_virtual_environment_checkpoint_seed_refs( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...]: - with connection.cursor() as cursor: - checkpoint_seed_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, - ) - cursor.execute( - f"SELECT checkpoint_id, seed_name, version_hash " - f"FROM {checkpoint_seed_ref_table} " - "WHERE checkpoint_id = %s ORDER BY seed_name", - [checkpoint_id], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - VirtualEnvironmentCheckpointSeedRefRecord( - checkpoint_id=row[0], - seed_name=row[1], - version_hash=row[2], - ) - for row in rows - ) - def delete_virtual_environment_checkpoint( self, *, connection: Any, schema: str, checkpoint_id: str ) -> None: @@ -1754,26 +1321,6 @@ def upsert_state_operation( cursor.execute("ROLLBACK") raise - def get_state_operation( - self, *, connection: Any, schema: str, operation_id: str - ) -> StateOperationRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT operation_id, operation_type, status, virtual_environment_name " - f"FROM {self._qualified_name(schema=schema, table=STATE_OPERATION_TABLE)} " - "WHERE operation_id = %s", - [operation_id], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return StateOperationRecord( - operation_id=row[0], - operation_type=StateOperationType(row[1]), - status=StateOperationStatus(row[2]), - virtual_environment_name=row[3], - ) - def create_state_operation_event( self, *, connection: Any, schema: str, record: StateOperationEventRecord ) -> None: @@ -2090,12 +1637,6 @@ def _parse_materialized(self, value: object) -> bool | None: return None return str(value).lower() == STATE_BOOLEAN_TRUE - def _quote_identifier(self, identifier: str) -> str: - return '"' + identifier.replace('"', '""') + '"' - - def _qualified_name(self, *, schema: str, table: str) -> str: - return f"{self._quote_identifier(schema)}.{self._quote_identifier(table)}" - def _validate_source_freshness_records( self, *, From 4e6bb330bf54a479966df974d8da55af738b25dd Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:14:45 +0100 Subject: [PATCH 10/57] refactor(state): share relational state backend upserts Move the model, function, seed, Python node, physical relation, ancestry, state operation, and virtual environment upserts into SqlStateBackend on one created_at-preserving replace helper. Each upsert still runs the same lookup, DELETE, and INSERT for the same key and columns inside one BEGIN/COMMIT with ROLLBACK on failure; backends supply the transaction executor (the DuckDB connection or a Postgres cursor) and statement calls. Stored rows are unchanged: the same columns and values, created_at kept from the existing row, updated_at set to CURRENT_TIMESTAMP. (cherry picked from commit b68614cc2816a0045b3afd8715d15694fc08699d) --- .../classes/_duckdb_conditional_publish.py | 2 +- .../state/classes/_sql_state_backend.py | 232 ++++++++++ src/sqlbuild/virtual/state/classes/duckdb.py | 373 +--------------- .../virtual/state/classes/postgres.py | 399 ++---------------- 4 files changed, 282 insertions(+), 724 deletions(-) diff --git a/src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py b/src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py index 565604db6d..d49c5cddd9 100644 --- a/src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py +++ b/src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py @@ -65,7 +65,7 @@ def upsert_virtual_environment_and_replace_node_ref_groups_if_locks_owned( connection.execute("ROLLBACK") return False backend._upsert_virtual_environment_record( - connection=connection, schema=schema, record=record + executor=connection, schema=schema, record=record ) backend._replace_virtual_environment_node_ref_groups( connection=connection, diff --git a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py index 3bb25b1426..885a3fe500 100644 --- a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py +++ b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py @@ -4,6 +4,8 @@ from abc import abstractmethod from collections.abc import Sequence +from contextlib import AbstractContextManager +from datetime import datetime from typing import Any, ClassVar from sqlbuild.virtual.state._helpers.state_storage.datetime import ( @@ -74,6 +76,236 @@ def _fetch_all( """Run one read statement and return all rows.""" ... + @abstractmethod + def _write_transaction(self, *, connection: Any) -> AbstractContextManager[Any]: + """Open a write transaction and yield the executor that runs its statements.""" + ... + + @abstractmethod + def _execute_in(self, *, executor: Any, sql: str, params: Sequence[object]) -> None: + """Run one statement on a transaction executor.""" + ... + + @abstractmethod + def _fetch_one_in( + self, *, executor: Any, sql: str, params: Sequence[object] + ) -> tuple[Any, ...] | None: + """Run one read statement on a transaction executor and return its first row.""" + ... + + def upsert_model_version( + self, *, connection: Any, schema: str, record: ModelVersionRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=MODEL_VERSION_TABLE, + key_values={"model_name": record.model_name, "version_hash": record.version_hash}, + row_values={ + "definition_identity_hash": record.definition_identity_hash, + "identity_metadata_hash": record.identity_metadata_hash, + "definition_text_b64": record.definition_text_b64, + "identity_metadata_json_b64": record.identity_metadata_json_b64, + "compiled_sql_b64": record.compiled_sql_b64, + "status": record.status.value, + }, + ) + + def upsert_function_version( + self, *, connection: Any, schema: str, record: FunctionVersionRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=FUNCTION_VERSION_TABLE, + key_values={ + "function_name": record.function_name, + "version_hash": record.version_hash, + }, + row_values={ + "language": record.language, + "returns": record.returns, + "arguments_json_b64": record.arguments_json_b64, + "return_columns_json_b64": record.return_columns_json_b64, + "packages_json_b64": record.packages_json_b64, + "runtime_version": record.runtime_version, + "entry_point": record.entry_point, + "body_sql_b64": record.body_sql_b64, + "definition_text_b64": record.definition_text_b64, + "status": record.status.value, + }, + ) + + def upsert_seed_version( + self, *, connection: Any, schema: str, record: SeedVersionRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=SEED_VERSION_TABLE, + key_values={"seed_name": record.seed_name, "version_hash": record.version_hash}, + row_values={ + "identity_metadata_hash": record.identity_metadata_hash, + "identity_metadata_json_b64": record.identity_metadata_json_b64, + "status": record.status.value, + }, + ) + + def upsert_python_node_version( + self, *, connection: Any, schema: str, record: PythonNodeVersionRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=PYTHON_NODE_VERSION_TABLE, + key_values={ + "node_type": record.node_type, + "node_name": record.node_name, + "version_hash": record.version_hash, + }, + row_values={ + "definition_hash": record.definition_hash, + "identity_metadata_hash": record.identity_metadata_hash, + "definition_json_b64": record.definition_json_b64, + "identity_metadata_json_b64": record.identity_metadata_json_b64, + "status": record.status.value, + }, + ) + + def upsert_physical_relation( + self, *, connection: Any, schema: str, record: PhysicalRelationRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=PHYSICAL_RELATION_TABLE, + key_values={ + "artifact_type": record.artifact_type.value, + "artifact_name": record.artifact_name, + "version_hash": record.version_hash, + }, + row_values={ + "database_name": record.database_name, + "schema_name": record.schema_name, + "relation_name": record.relation_name, + "relation_type": record.relation_type, + }, + ) + + def upsert_physical_relation_ancestry( + self, *, connection: Any, schema: str, record: PhysicalRelationAncestryRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=PHYSICAL_RELATION_ANCESTRY_TABLE, + key_values={"model_name": record.model_name, "version_hash": record.version_hash}, + row_values={ + "parent_model_name": record.parent_model_name, + "parent_version_hash": record.parent_version_hash, + "seed_strategy": record.seed_strategy, + }, + ) + + def upsert_virtual_environment( + self, *, connection: Any, schema: str, record: VirtualEnvironmentRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._upsert_virtual_environment_record(executor=executor, schema=schema, record=record) + + def upsert_state_operation( + self, *, connection: Any, schema: str, record: StateOperationRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=STATE_OPERATION_TABLE, + key_values={"operation_id": record.operation_id}, + row_values={ + "operation_type": record.operation_type.value, + "status": record.status.value, + "virtual_environment_name": record.virtual_environment_name, + }, + ) + + def _upsert_virtual_environment_record( + self, *, executor: Any, schema: str, record: VirtualEnvironmentRecord + ) -> None: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=VIRTUAL_ENVIRONMENT_TABLE, + key_values={"virtual_environment_name": record.virtual_environment_name}, + row_values={ + "status": record.status.value, + "baseline_virtual_environment_name": record.baseline_virtual_environment_name, + "finalized_at": record.finalized_at, + }, + ) + + def _replace_row_preserving_created_at( + self, + *, + executor: Any, + schema: str, + table_name: str, + key_values: dict[str, object], + row_values: dict[str, object], + ) -> None: + p: str = self._placeholder + table: str = self._qualified_name(schema=schema, table=table_name) + where_sql: str = " AND ".join(f"{column} = {p}" for column in key_values) + key_params: list[object] = list(key_values.values()) + existing_created_at: datetime | None = self._created_at_for_key( + executor=executor, + schema=schema, + table_name=table_name, + where_sql=where_sql, + params=key_params, + ) + self._execute_in( + executor=executor, sql=f"DELETE FROM {table} WHERE {where_sql}", params=key_params + ) + columns: tuple[str, ...] = (*key_values, *row_values) + value_placeholders: str = ", ".join(p for _ in columns) + self._execute_in( + executor=executor, + sql=( + f"INSERT INTO {table} ({', '.join(columns)}, created_at, updated_at) " + f"VALUES ({value_placeholders}, COALESCE({p}, CURRENT_TIMESTAMP), " + "CURRENT_TIMESTAMP)" + ), + params=[*key_params, *row_values.values(), existing_created_at], + ) + + def _created_at_for_key( + self, + *, + executor: Any, + schema: str, + table_name: str, + where_sql: str, + params: list[object], + ) -> datetime | None: + row: tuple[Any, ...] | None = self._fetch_one_in( + executor=executor, + sql=( + f"SELECT created_at FROM {self._qualified_name(schema=schema, table=table_name)} " + f"WHERE {where_sql}" + ), + params=params, + ) + if row is None: + return None + return row[0] + def _quote_identifier(self, identifier: str) -> str: return '"' + identifier.replace('"', '""') + '"' diff --git a/src/sqlbuild/virtual/state/classes/duckdb.py b/src/sqlbuild/virtual/state/classes/duckdb.py index 56fb0e97da..c5a30d4a92 100644 --- a/src/sqlbuild/virtual/state/classes/duckdb.py +++ b/src/sqlbuild/virtual/state/classes/duckdb.py @@ -2,7 +2,8 @@ from __future__ import annotations -from collections.abc import Sequence +from collections.abc import Iterator, Sequence +from contextlib import contextmanager from datetime import datetime from typing import Any, ClassVar @@ -37,22 +38,16 @@ DUCKDB_DATETIME_TYPE_TOKEN, DUCKDB_INTEGER_TYPE_TOKEN, DUCKDB_TIMESTAMP_TYPE_TOKEN, - FUNCTION_VERSION_TABLE, LOCK_TABLE, MICROBATCH_EVENT_TABLE, - MODEL_VERSION_TABLE, NODE_RESULTS_TABLE, NON_UNIQUE_STATE_INDEXES, - PHYSICAL_RELATION_ANCESTRY_TABLE, - PHYSICAL_RELATION_TABLE, PYTHON_NODE_VERSION_TABLE, RECONCILE_EVENT_TABLE, - SEED_VERSION_TABLE, SOURCE_FRESHNESS_OBSERVATION_TABLE, STATE_BOOLEAN_TRUE, STATE_MIGRATION_EVENTS_TABLE, STATE_OPERATION_EVENT_TABLE, - STATE_OPERATION_TABLE, STATE_TABLE_COLUMNS, STATE_TABLE_INDEXES, STATE_TABLES, @@ -70,18 +65,11 @@ StateSchemaInvalidError, ) from sqlbuild.virtual.state.models import ( - FunctionVersionRecord, - ModelVersionRecord, - PhysicalRelationAncestryRecord, - PhysicalRelationRecord, - PythonNodeVersionRecord, ReconcileEventRecord, - SeedVersionRecord, SourceFreshnessRecord, StateBackupRecord, StateLockRecord, StateOperationEventRecord, - StateOperationRecord, StateSchemaValidationResult, VirtualEnvironmentCheckpointFunctionRefRecord, VirtualEnvironmentCheckpointModelRefRecord, @@ -133,6 +121,24 @@ def _fetch_all( return connection.execute(sql).fetchall() return connection.execute(sql, params).fetchall() + @contextmanager + def _write_transaction(self, *, connection: Any) -> Iterator[Any]: + connection.execute("BEGIN") + try: + yield connection + connection.execute("COMMIT") + except BaseException: + connection.execute("ROLLBACK") + raise + + def _execute_in(self, *, executor: Any, sql: str, params: Sequence[object]) -> None: + executor.execute(sql, params) + + def _fetch_one_in( + self, *, executor: Any, sql: str, params: Sequence[object] + ) -> tuple[Any, ...] | None: + return executor.execute(sql, params).fetchone() + def initialize(self, *, connection: Any, schema: str, sqlbuild_version: str) -> None: connection.execute("BEGIN") try: @@ -293,172 +299,6 @@ def reset(self, *, connection: Any, schema: str) -> None: connection.execute("ROLLBACK") raise - def upsert_model_version( - self, *, connection: Any, schema: str, record: ModelVersionRecord - ) -> None: - connection.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=MODEL_VERSION_TABLE, - where_sql="model_name = ? AND version_hash = ?", - params=[record.model_name, record.version_hash], - ) - connection.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=MODEL_VERSION_TABLE)} " - "WHERE model_name = ? AND version_hash = ?", - [record.model_name, record.version_hash], - ) - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=MODEL_VERSION_TABLE)} " - "(model_name, version_hash, definition_identity_hash, " - "identity_metadata_hash, definition_text_b64, identity_metadata_json_b64, " - "compiled_sql_b64, status, " - "created_at, updated_at) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, " - "COALESCE(?, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.model_name, - record.version_hash, - record.definition_identity_hash, - record.identity_metadata_hash, - record.definition_text_b64, - record.identity_metadata_json_b64, - record.compiled_sql_b64, - record.status.value, - existing_created_at, - ], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def upsert_function_version( - self, *, connection: Any, schema: str, record: FunctionVersionRecord - ) -> None: - connection.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=FUNCTION_VERSION_TABLE, - where_sql="function_name = ? AND version_hash = ?", - params=[record.function_name, record.version_hash], - ) - connection.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=FUNCTION_VERSION_TABLE)} " - "WHERE function_name = ? AND version_hash = ?", - [record.function_name, record.version_hash], - ) - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=FUNCTION_VERSION_TABLE)} " - "(function_name, version_hash, language, returns, arguments_json_b64, " - "return_columns_json_b64, packages_json_b64, runtime_version, entry_point, " - "body_sql_b64, definition_text_b64, status, created_at, updated_at) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, " - "COALESCE(?, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.function_name, - record.version_hash, - record.language, - record.returns, - record.arguments_json_b64, - record.return_columns_json_b64, - record.packages_json_b64, - record.runtime_version, - record.entry_point, - record.body_sql_b64, - record.definition_text_b64, - record.status.value, - existing_created_at, - ], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def upsert_seed_version( - self, *, connection: Any, schema: str, record: SeedVersionRecord - ) -> None: - connection.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=SEED_VERSION_TABLE, - where_sql="seed_name = ? AND version_hash = ?", - params=[record.seed_name, record.version_hash], - ) - connection.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=SEED_VERSION_TABLE)} " - "WHERE seed_name = ? AND version_hash = ?", - [record.seed_name, record.version_hash], - ) - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=SEED_VERSION_TABLE)} " - "(seed_name, version_hash, identity_metadata_hash, " - "identity_metadata_json_b64, status, created_at, updated_at) " - "VALUES (?, ?, ?, ?, ?, COALESCE(?, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.seed_name, - record.version_hash, - record.identity_metadata_hash, - record.identity_metadata_json_b64, - record.status.value, - existing_created_at, - ], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def upsert_python_node_version( - self, *, connection: Any, schema: str, record: PythonNodeVersionRecord - ) -> None: - connection.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=PYTHON_NODE_VERSION_TABLE, - where_sql="node_type = ? AND node_name = ? AND version_hash = ?", - params=[record.node_type, record.node_name, record.version_hash], - ) - connection.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " - "WHERE node_type = ? AND node_name = ? AND version_hash = ?", - [record.node_type, record.node_name, record.version_hash], - ) - connection.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " - "(node_type, node_name, version_hash, definition_hash, " - "identity_metadata_hash, definition_json_b64, identity_metadata_json_b64, " - "status, created_at, updated_at) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, COALESCE(?, CURRENT_TIMESTAMP), " - "CURRENT_TIMESTAMP)", - [ - record.node_type, - record.node_name, - record.version_hash, - record.definition_hash, - record.identity_metadata_hash, - record.definition_json_b64, - record.identity_metadata_json_b64, - record.status.value, - existing_created_at, - ], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - def insert_node_result( self, *, @@ -583,95 +423,6 @@ def read_microbatch_model_history( scope=scope, ) - def upsert_physical_relation( - self, *, connection: Any, schema: str, record: PhysicalRelationRecord - ) -> None: - connection.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=PHYSICAL_RELATION_TABLE, - where_sql="artifact_type = ? AND artifact_name = ? AND version_hash = ?", - params=[record.artifact_type.value, record.artifact_name, record.version_hash], - ) - connection.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "WHERE artifact_type = ? AND artifact_name = ? AND version_hash = ?", - [record.artifact_type.value, record.artifact_name, record.version_hash], - ) - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "(artifact_type, artifact_name, version_hash, database_name, schema_name, " - "relation_name, relation_type, created_at, updated_at) " - "VALUES (?, ?, ?, ?, ?, ?, ?, COALESCE(?, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.artifact_type.value, - record.artifact_name, - record.version_hash, - record.database_name, - record.schema_name, - record.relation_name, - record.relation_type, - existing_created_at, - ], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def upsert_physical_relation_ancestry( - self, *, connection: Any, schema: str, record: PhysicalRelationAncestryRecord - ) -> None: - connection.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=PHYSICAL_RELATION_ANCESTRY_TABLE, - where_sql="model_name = ? AND version_hash = ?", - params=[record.model_name, record.version_hash], - ) - connection.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=PHYSICAL_RELATION_ANCESTRY_TABLE)} " - "WHERE model_name = ? AND version_hash = ?", - [record.model_name, record.version_hash], - ) - connection.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=PHYSICAL_RELATION_ANCESTRY_TABLE)} " - "(model_name, version_hash, parent_model_name, parent_version_hash, " - "seed_strategy, created_at, updated_at) " - "VALUES (?, ?, ?, ?, ?, COALESCE(?, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.model_name, - record.version_hash, - record.parent_model_name, - record.parent_version_hash, - record.seed_strategy, - existing_created_at, - ], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def upsert_virtual_environment( - self, *, connection: Any, schema: str, record: VirtualEnvironmentRecord - ) -> None: - connection.execute("BEGIN") - try: - self._upsert_virtual_environment_record( - connection=connection, schema=schema, record=record - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - def list_virtual_environments( self, *, connection: Any, schema: str ) -> tuple[VirtualEnvironmentRetentionRecord, ...]: @@ -765,7 +516,7 @@ def upsert_virtual_environment_and_replace_node_ref_groups( connection.execute("BEGIN") try: self._upsert_virtual_environment_record( - connection=connection, schema=schema, record=record + executor=connection, schema=schema, record=record ) self._replace_virtual_environment_node_ref_groups( connection=connection, @@ -1069,41 +820,6 @@ def delete_virtual_environment_checkpoint( connection.execute("ROLLBACK") raise - def upsert_state_operation( - self, *, connection: Any, schema: str, record: StateOperationRecord - ) -> None: - connection.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=STATE_OPERATION_TABLE, - where_sql="operation_id = ?", - params=[record.operation_id], - ) - connection.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=STATE_OPERATION_TABLE)} " - "WHERE operation_id = ?", - [record.operation_id], - ) - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=STATE_OPERATION_TABLE)} " - "(operation_id, operation_type, status, virtual_environment_name, " - "created_at, updated_at) " - "VALUES (?, ?, ?, ?, COALESCE(?, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.operation_id, - record.operation_type.value, - record.status.value, - record.virtual_environment_name, - existing_created_at, - ], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - def create_state_operation_event( self, *, connection: Any, schema: str, record: StateOperationEventRecord ) -> None: @@ -1296,24 +1012,6 @@ def _record_event( [event_id(), action.value, backup_id_value, status.value, message], ) - def _created_at_for_key( - self, - *, - connection: Any, - schema: str, - table_name: str, - where_sql: str, - params: list[object], - ) -> datetime | None: - row: tuple[Any, ...] | None = connection.execute( - f"SELECT created_at FROM {self._qualified_name(schema=schema, table=table_name)} " - f"WHERE {where_sql}", - params, - ).fetchone() - if row is None: - return None - return row[0] - def _create_additional_state_tables(self, *, connection: Any, schema: str) -> None: table_name: str columns: dict[str, StateColumnType] @@ -1445,35 +1143,6 @@ def _validate_node_ref_replacement( ) seen_node_names.add(ref.node_name) - def _upsert_virtual_environment_record( - self, *, connection: Any, schema: str, record: VirtualEnvironmentRecord - ) -> None: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=VIRTUAL_ENVIRONMENT_TABLE, - where_sql="virtual_environment_name = ?", - params=[record.virtual_environment_name], - ) - connection.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "WHERE virtual_environment_name = ?", - [record.virtual_environment_name], - ) - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "(virtual_environment_name, status, baseline_virtual_environment_name, " - "created_at, updated_at, finalized_at) " - "VALUES (?, ?, ?, COALESCE(?, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP, ?)", - [ - record.virtual_environment_name, - record.status.value, - record.baseline_virtual_environment_name, - existing_created_at, - record.finalized_at, - ], - ) - def _replace_virtual_environment_node_ref_groups( self, *, diff --git a/src/sqlbuild/virtual/state/classes/postgres.py b/src/sqlbuild/virtual/state/classes/postgres.py index 2495bb685f..95c440933e 100644 --- a/src/sqlbuild/virtual/state/classes/postgres.py +++ b/src/sqlbuild/virtual/state/classes/postgres.py @@ -2,7 +2,8 @@ from __future__ import annotations -from collections.abc import Sequence +from collections.abc import Iterator, Sequence +from contextlib import contextmanager from datetime import datetime from typing import Any, ClassVar @@ -33,24 +34,18 @@ from sqlbuild.virtual.state.classes._sql_state_backend import SqlStateBackend from sqlbuild.virtual.state.constants import ( CURRENT_STATE_SCHEMA_VERSION, - FUNCTION_VERSION_TABLE, LOCK_TABLE, MICROBATCH_EVENT_TABLE, - MODEL_VERSION_TABLE, NODE_RESULTS_TABLE, NON_UNIQUE_STATE_INDEXES, - PHYSICAL_RELATION_ANCESTRY_TABLE, - PHYSICAL_RELATION_TABLE, POSTGRES_INTEGER_TYPES, POSTGRES_TEXT_TYPES, PYTHON_NODE_VERSION_TABLE, RECONCILE_EVENT_TABLE, - SEED_VERSION_TABLE, SOURCE_FRESHNESS_OBSERVATION_TABLE, STATE_BOOLEAN_TRUE, STATE_MIGRATION_EVENTS_TABLE, STATE_OPERATION_EVENT_TABLE, - STATE_OPERATION_TABLE, STATE_TABLE_COLUMNS, STATE_TABLE_INDEXES, STATE_TABLES, @@ -68,19 +63,12 @@ StateSchemaInvalidError, ) from sqlbuild.virtual.state.models import ( - FunctionVersionRecord, - ModelVersionRecord, - PhysicalRelationAncestryRecord, - PhysicalRelationRecord, - PythonNodeVersionRecord, ReconcileEventRecord, - SeedVersionRecord, SourceFreshnessRecord, StateBackupRecord, StateLockLease, StateLockRecord, StateOperationEventRecord, - StateOperationRecord, StateSchemaValidationResult, VirtualEnvironmentCheckpointFunctionRefRecord, VirtualEnvironmentCheckpointModelRefRecord, @@ -154,6 +142,26 @@ def _fetch_all( cursor.execute(sql, params) return cursor.fetchall() + @contextmanager + def _write_transaction(self, *, connection: Any) -> Iterator[Any]: + with connection.cursor() as cursor: + cursor.execute("BEGIN") + try: + yield cursor + cursor.execute("COMMIT") + except BaseException: + cursor.execute("ROLLBACK") + raise + + def _execute_in(self, *, executor: Any, sql: str, params: Sequence[object]) -> None: + executor.execute(sql, params) + + def _fetch_one_in( + self, *, executor: Any, sql: str, params: Sequence[object] + ) -> tuple[Any, ...] | None: + executor.execute(sql, params) + return executor.fetchone() + def initialize(self, *, connection: Any, schema: str, sqlbuild_version: str) -> None: with connection.cursor() as cursor: cursor.execute("BEGIN") @@ -325,179 +333,6 @@ def reset(self, *, connection: Any, schema: str) -> None: cursor.execute("ROLLBACK") raise - def upsert_model_version( - self, *, connection: Any, schema: str, record: ModelVersionRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=MODEL_VERSION_TABLE, - where_sql="model_name = %s AND version_hash = %s", - params=[record.model_name, record.version_hash], - ) - cursor.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=MODEL_VERSION_TABLE)} " - "WHERE model_name = %s AND version_hash = %s", - [record.model_name, record.version_hash], - ) - cursor.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=MODEL_VERSION_TABLE)} " - "(model_name, version_hash, definition_identity_hash, " - "identity_metadata_hash, definition_text_b64, identity_metadata_json_b64, " - "compiled_sql_b64, status, " - "created_at, updated_at) " - "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, " - "COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.model_name, - record.version_hash, - record.definition_identity_hash, - record.identity_metadata_hash, - record.definition_text_b64, - record.identity_metadata_json_b64, - record.compiled_sql_b64, - record.status.value, - existing_created_at, - ], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def upsert_function_version( - self, *, connection: Any, schema: str, record: FunctionVersionRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=FUNCTION_VERSION_TABLE, - where_sql="function_name = %s AND version_hash = %s", - params=[record.function_name, record.version_hash], - ) - cursor.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=FUNCTION_VERSION_TABLE)} " - "WHERE function_name = %s AND version_hash = %s", - [record.function_name, record.version_hash], - ) - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=FUNCTION_VERSION_TABLE)} " - "(function_name, version_hash, language, returns, arguments_json_b64, " - "return_columns_json_b64, packages_json_b64, runtime_version, entry_point, " - "body_sql_b64, definition_text_b64, status, created_at, updated_at) " - "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, " - "COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.function_name, - record.version_hash, - record.language, - record.returns, - record.arguments_json_b64, - record.return_columns_json_b64, - record.packages_json_b64, - record.runtime_version, - record.entry_point, - record.body_sql_b64, - record.definition_text_b64, - record.status.value, - existing_created_at, - ], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def upsert_seed_version( - self, *, connection: Any, schema: str, record: SeedVersionRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=SEED_VERSION_TABLE, - where_sql="seed_name = %s AND version_hash = %s", - params=[record.seed_name, record.version_hash], - ) - cursor.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=SEED_VERSION_TABLE)} " - "WHERE seed_name = %s AND version_hash = %s", - [record.seed_name, record.version_hash], - ) - cursor.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=SEED_VERSION_TABLE)} " - "(seed_name, version_hash, identity_metadata_hash, " - "identity_metadata_json_b64, status, created_at, updated_at) " - "VALUES (%s, %s, %s, %s, %s, " - "COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.seed_name, - record.version_hash, - record.identity_metadata_hash, - record.identity_metadata_json_b64, - record.status.value, - existing_created_at, - ], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def upsert_python_node_version( - self, *, connection: Any, schema: str, record: PythonNodeVersionRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=PYTHON_NODE_VERSION_TABLE, - where_sql="node_type = %s AND node_name = %s AND version_hash = %s", - params=[record.node_type, record.node_name, record.version_hash], - ) - cursor.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " - "WHERE node_type = %s AND node_name = %s AND version_hash = %s", - [record.node_type, record.node_name, record.version_hash], - ) - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " - "(node_type, node_name, version_hash, definition_hash, " - "identity_metadata_hash, definition_json_b64, identity_metadata_json_b64, " - "status, created_at, updated_at) " - "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, " - "COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.node_type, - record.node_name, - record.version_hash, - record.definition_hash, - record.identity_metadata_hash, - record.definition_json_b64, - record.identity_metadata_json_b64, - record.status.value, - existing_created_at, - ], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - def insert_node_result( self, *, @@ -680,102 +515,6 @@ def read_microbatch_model_history( rows: list[tuple[Any, ...]] = cursor.fetchall() return MicrobatchEventCodec.from_rows(tuple(row) for row in rows) - def upsert_physical_relation( - self, *, connection: Any, schema: str, record: PhysicalRelationRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=PHYSICAL_RELATION_TABLE, - where_sql="artifact_type = %s AND artifact_name = %s AND version_hash = %s", - params=[record.artifact_type.value, record.artifact_name, record.version_hash], - ) - cursor.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "WHERE artifact_type = %s AND artifact_name = %s AND version_hash = %s", - [record.artifact_type.value, record.artifact_name, record.version_hash], - ) - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "(artifact_type, artifact_name, version_hash, database_name, schema_name, " - "relation_name, relation_type, created_at, updated_at) " - "VALUES (%s, %s, %s, %s, %s, %s, %s, " - "COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.artifact_type.value, - record.artifact_name, - record.version_hash, - record.database_name, - record.schema_name, - record.relation_name, - record.relation_type, - existing_created_at, - ], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def upsert_physical_relation_ancestry( - self, *, connection: Any, schema: str, record: PhysicalRelationAncestryRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=PHYSICAL_RELATION_ANCESTRY_TABLE, - where_sql="model_name = %s AND version_hash = %s", - params=[record.model_name, record.version_hash], - ) - cursor.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=PHYSICAL_RELATION_ANCESTRY_TABLE)}" - " " - "WHERE model_name = %s AND version_hash = %s", - [record.model_name, record.version_hash], - ) - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=PHYSICAL_RELATION_ANCESTRY_TABLE)}" - " " - "(model_name, version_hash, parent_model_name, parent_version_hash, " - "seed_strategy, created_at, updated_at) " - "VALUES (%s, %s, %s, %s, %s, " - "COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.model_name, - record.version_hash, - record.parent_model_name, - record.parent_version_hash, - record.seed_strategy, - existing_created_at, - ], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def upsert_virtual_environment( - self, *, connection: Any, schema: str, record: VirtualEnvironmentRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - self._upsert_virtual_environment_record(cursor=cursor, schema=schema, record=record) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - def list_virtual_environments( self, *, connection: Any, schema: str ) -> tuple[VirtualEnvironmentRetentionRecord, ...]: @@ -881,7 +620,9 @@ def upsert_virtual_environment_and_replace_node_ref_groups( with connection.cursor() as cursor: cursor.execute("BEGIN") try: - self._upsert_virtual_environment_record(cursor=cursor, schema=schema, record=record) + self._upsert_virtual_environment_record( + executor=cursor, schema=schema, record=record + ) self._replace_virtual_environment_node_ref_groups( cursor=cursor, schema=schema, @@ -929,7 +670,9 @@ def upsert_virtual_environment_and_replace_node_ref_groups_if_locks_owned( if cursor.fetchone() is None: cursor.execute("ROLLBACK") return False - self._upsert_virtual_environment_record(cursor=cursor, schema=schema, record=record) + self._upsert_virtual_environment_record( + executor=cursor, schema=schema, record=record + ) self._replace_virtual_environment_node_ref_groups( cursor=cursor, schema=schema, @@ -1283,44 +1026,6 @@ def delete_virtual_environment_checkpoint( cursor.execute("ROLLBACK") raise - def upsert_state_operation( - self, *, connection: Any, schema: str, record: StateOperationRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=STATE_OPERATION_TABLE, - where_sql="operation_id = %s", - params=[record.operation_id], - ) - cursor.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=STATE_OPERATION_TABLE)} " - "WHERE operation_id = %s", - [record.operation_id], - ) - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=STATE_OPERATION_TABLE)} " - "(operation_id, operation_type, status, virtual_environment_name, " - "created_at, updated_at) " - "VALUES (%s, %s, %s, %s, COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.operation_id, - record.operation_type.value, - record.status.value, - record.virtual_environment_name, - existing_created_at, - ], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - def create_state_operation_event( self, *, connection: Any, schema: str, record: StateOperationEventRecord ) -> None: @@ -1569,25 +1274,6 @@ def _create_state_indexes(self, *, cursor: Any, schema: str) -> None: f"ON {self._qualified_name(schema=schema, table=table_name)} ({column_sql})" ) - def _created_at_for_key( - self, - *, - cursor: Any, - schema: str, - table_name: str, - where_sql: str, - params: list[object], - ) -> datetime | None: - cursor.execute( - f"SELECT created_at FROM {self._qualified_name(schema=schema, table=table_name)} " - f"WHERE {where_sql}", - params, - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return row[0] - def _state_column_sql_type(self, column_type: StateColumnType) -> str: match column_type: case StateColumnType.INTEGER: @@ -1680,35 +1366,6 @@ def _validate_node_ref_replacement( ) seen_node_names.add(ref.node_name) - def _upsert_virtual_environment_record( - self, *, cursor: Any, schema: str, record: VirtualEnvironmentRecord - ) -> None: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=VIRTUAL_ENVIRONMENT_TABLE, - where_sql="virtual_environment_name = %s", - params=[record.virtual_environment_name], - ) - cursor.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "WHERE virtual_environment_name = %s", - [record.virtual_environment_name], - ) - cursor.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "(virtual_environment_name, status, baseline_virtual_environment_name, " - "created_at, updated_at, finalized_at) " - "VALUES (%s, %s, %s, COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP, %s)", - [ - record.virtual_environment_name, - record.status.value, - record.baseline_virtual_environment_name, - existing_created_at, - record.finalized_at, - ], - ) - def _replace_virtual_environment_node_ref_groups( self, *, From a60574a929666a02c61e39c99e89a113c8ecbb48 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:16:14 +0100 Subject: [PATCH 11/57] refactor(state): share relational state backend listings and lock reads Move the virtual environment listing, active and expired lock listings, and lock renewal into SqlStateBackend. Their SQL, ordering, and row mapping were identical in both backends. acquire_lock, release_lock, and delete_lock stay per backend: DuckDB tolerates a failed ROLLBACK after a conflicting insert, and Postgres commits explicitly after deleting a lock. (cherry picked from commit a5f25a846b3320dcd5f970e601fc0ac46168f050) --- .../state/classes/_sql_state_backend.py | 63 ++++++++++++++++++ src/sqlbuild/virtual/state/classes/duckdb.py | 58 ----------------- .../virtual/state/classes/postgres.py | 65 ------------------- 3 files changed, 63 insertions(+), 123 deletions(-) diff --git a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py index 885a3fe500..364ff428ba 100644 --- a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py +++ b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py @@ -14,6 +14,7 @@ from sqlbuild.virtual.state.classes.state_backend import StateBackend from sqlbuild.virtual.state.constants import ( FUNCTION_VERSION_TABLE, + LOCK_TABLE, MODEL_VERSION_TABLE, PHYSICAL_RELATION_ANCESTRY_TABLE, PHYSICAL_RELATION_TABLE, @@ -36,6 +37,7 @@ PythonNodeVersionRecord, SeedVersionRecord, SourceFreshnessRecord, + StateLockRecord, StateOperationRecord, VirtualEnvironmentCheckpointFunctionRefRecord, VirtualEnvironmentCheckpointModelRefRecord, @@ -46,6 +48,7 @@ VirtualEnvironmentNodeRefRecord, VirtualEnvironmentPythonNodeRefRecord, VirtualEnvironmentRecord, + VirtualEnvironmentRetentionRecord, VirtualEnvironmentSeedRefRecord, ) from sqlbuild.virtual.state.types import ( @@ -250,6 +253,66 @@ def _upsert_virtual_environment_record( }, ) + def list_virtual_environments( + self, *, connection: Any, schema: str + ) -> tuple[VirtualEnvironmentRetentionRecord, ...]: + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT virtual_environment_name, status, updated_at " + f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " + "ORDER BY updated_at DESC, virtual_environment_name DESC", + ) + return tuple( + VirtualEnvironmentRetentionRecord( + virtual_environment_name=row[0], + status=VirtualEnvironmentStatus(row[1]), + updated_at=row[2], + ) + for row in rows + ) + + def renew_lock( + self, + *, + connection: Any, + schema: str, + lock_key: str, + owner_id: str, + expires_at: datetime, + ) -> bool: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql=f"UPDATE {self._qualified_name(schema=schema, table=LOCK_TABLE)} " + f"SET expires_at = {p}, updated_at = CURRENT_TIMESTAMP " + f"WHERE lock_key = {p} AND owner_id = {p} AND expires_at > CURRENT_TIMESTAMP " + "RETURNING lock_key", + params=[expires_at, lock_key, owner_id], + ) + return row is not None + + def list_active_locks(self, *, connection: Any, schema: str) -> tuple[StateLockRecord, ...]: + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT lock_key, owner_id, expires_at FROM " + f"{self._qualified_name(schema=schema, table=LOCK_TABLE)} " + "WHERE expires_at > CURRENT_TIMESTAMP ORDER BY lock_key", + ) + return tuple( + StateLockRecord(lock_key=row[0], owner_id=row[1], expires_at=row[2]) for row in rows + ) + + def list_expired_locks(self, *, connection: Any, schema: str) -> tuple[StateLockRecord, ...]: + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT lock_key, owner_id, expires_at FROM " + f"{self._qualified_name(schema=schema, table=LOCK_TABLE)} " + "WHERE expires_at <= CURRENT_TIMESTAMP ORDER BY lock_key", + ) + return tuple( + StateLockRecord(lock_key=row[0], owner_id=row[1], expires_at=row[2]) for row in rows + ) + def _replace_row_preserving_created_at( self, *, diff --git a/src/sqlbuild/virtual/state/classes/duckdb.py b/src/sqlbuild/virtual/state/classes/duckdb.py index c5a30d4a92..3a069e651f 100644 --- a/src/sqlbuild/virtual/state/classes/duckdb.py +++ b/src/sqlbuild/virtual/state/classes/duckdb.py @@ -68,7 +68,6 @@ ReconcileEventRecord, SourceFreshnessRecord, StateBackupRecord, - StateLockRecord, StateOperationEventRecord, StateSchemaValidationResult, VirtualEnvironmentCheckpointFunctionRefRecord, @@ -80,14 +79,12 @@ VirtualEnvironmentNodeRefRecord, VirtualEnvironmentPythonNodeRefRecord, VirtualEnvironmentRecord, - VirtualEnvironmentRetentionRecord, VirtualEnvironmentSeedRefRecord, ) from sqlbuild.virtual.state.types import ( StateColumnType, StateMigrationAction, StateMigrationStatus, - VirtualEnvironmentStatus, ) @@ -423,23 +420,6 @@ def read_microbatch_model_history( scope=scope, ) - def list_virtual_environments( - self, *, connection: Any, schema: str - ) -> tuple[VirtualEnvironmentRetentionRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT virtual_environment_name, status, updated_at " - f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "ORDER BY updated_at DESC, virtual_environment_name DESC" - ).fetchall() - return tuple( - VirtualEnvironmentRetentionRecord( - virtual_environment_name=row[0], - status=VirtualEnvironmentStatus(row[1]), - updated_at=row[2], - ) - for row in rows - ) - def delete_virtual_environment( self, *, connection: Any, schema: str, virtual_environment_name: str ) -> None: @@ -906,44 +886,6 @@ def release_lock(self, *, connection: Any, schema: str, lock_key: str, owner_id: connection.execute("ROLLBACK") raise - def renew_lock( - self, - *, - connection: Any, - schema: str, - lock_key: str, - owner_id: str, - expires_at: datetime, - ) -> bool: - row: tuple[Any, ...] | None = connection.execute( - f"UPDATE {self._qualified_name(schema=schema, table=LOCK_TABLE)} " - "SET expires_at = ?, updated_at = CURRENT_TIMESTAMP " - "WHERE lock_key = ? AND owner_id = ? AND expires_at > CURRENT_TIMESTAMP " - "RETURNING lock_key", - [expires_at, lock_key, owner_id], - ).fetchone() - return row is not None - - def list_active_locks(self, *, connection: Any, schema: str) -> tuple[StateLockRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT lock_key, owner_id, expires_at FROM " - f"{self._qualified_name(schema=schema, table=LOCK_TABLE)} " - "WHERE expires_at > CURRENT_TIMESTAMP ORDER BY lock_key" - ).fetchall() - return tuple( - StateLockRecord(lock_key=row[0], owner_id=row[1], expires_at=row[2]) for row in rows - ) - - def list_expired_locks(self, *, connection: Any, schema: str) -> tuple[StateLockRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT lock_key, owner_id, expires_at FROM " - f"{self._qualified_name(schema=schema, table=LOCK_TABLE)} " - "WHERE expires_at <= CURRENT_TIMESTAMP ORDER BY lock_key" - ).fetchall() - return tuple( - StateLockRecord(lock_key=row[0], owner_id=row[1], expires_at=row[2]) for row in rows - ) - def delete_lock(self, *, connection: Any, schema: str, lock_key: str) -> None: connection.execute( "DELETE FROM " diff --git a/src/sqlbuild/virtual/state/classes/postgres.py b/src/sqlbuild/virtual/state/classes/postgres.py index 95c440933e..01e5f816d0 100644 --- a/src/sqlbuild/virtual/state/classes/postgres.py +++ b/src/sqlbuild/virtual/state/classes/postgres.py @@ -67,7 +67,6 @@ SourceFreshnessRecord, StateBackupRecord, StateLockLease, - StateLockRecord, StateOperationEventRecord, StateSchemaValidationResult, VirtualEnvironmentCheckpointFunctionRefRecord, @@ -79,14 +78,12 @@ VirtualEnvironmentNodeRefRecord, VirtualEnvironmentPythonNodeRefRecord, VirtualEnvironmentRecord, - VirtualEnvironmentRetentionRecord, VirtualEnvironmentSeedRefRecord, ) from sqlbuild.virtual.state.types import ( StateColumnType, StateMigrationAction, StateMigrationStatus, - VirtualEnvironmentStatus, ) @@ -515,25 +512,6 @@ def read_microbatch_model_history( rows: list[tuple[Any, ...]] = cursor.fetchall() return MicrobatchEventCodec.from_rows(tuple(row) for row in rows) - def list_virtual_environments( - self, *, connection: Any, schema: str - ) -> tuple[VirtualEnvironmentRetentionRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT virtual_environment_name, status, updated_at " - f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "ORDER BY updated_at DESC, virtual_environment_name DESC" - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - VirtualEnvironmentRetentionRecord( - virtual_environment_name=row[0], - status=VirtualEnvironmentStatus(row[1]), - updated_at=row[2], - ) - for row in rows - ) - def delete_virtual_environment( self, *, connection: Any, schema: str, virtual_environment_name: str ) -> None: @@ -1114,49 +1092,6 @@ def release_lock(self, *, connection: Any, schema: str, lock_key: str, owner_id: cursor.execute("ROLLBACK") raise - def renew_lock( - self, - *, - connection: Any, - schema: str, - lock_key: str, - owner_id: str, - expires_at: datetime, - ) -> bool: - with connection.cursor() as cursor: - cursor.execute( - f"UPDATE {self._qualified_name(schema=schema, table=LOCK_TABLE)} " - "SET expires_at = %s, updated_at = CURRENT_TIMESTAMP " - "WHERE lock_key = %s AND owner_id = %s AND expires_at > CURRENT_TIMESTAMP " - "RETURNING lock_key", - [expires_at, lock_key, owner_id], - ) - return cursor.fetchone() is not None - - def list_active_locks(self, *, connection: Any, schema: str) -> tuple[StateLockRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT lock_key, owner_id, expires_at FROM " - f"{self._qualified_name(schema=schema, table=LOCK_TABLE)} " - "WHERE expires_at > CURRENT_TIMESTAMP ORDER BY lock_key" - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - StateLockRecord(lock_key=row[0], owner_id=row[1], expires_at=row[2]) for row in rows - ) - - def list_expired_locks(self, *, connection: Any, schema: str) -> tuple[StateLockRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT lock_key, owner_id, expires_at FROM " - f"{self._qualified_name(schema=schema, table=LOCK_TABLE)} " - "WHERE expires_at <= CURRENT_TIMESTAMP ORDER BY lock_key" - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - StateLockRecord(lock_key=row[0], owner_id=row[1], expires_at=row[2]) for row in rows - ) - def delete_lock(self, *, connection: Any, schema: str, lock_key: str) -> None: with connection.cursor() as cursor: cursor.execute( From 4ec076abcdbd9407e0c35332526cde46c183b88d Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:27:08 +0100 Subject: [PATCH 12/57] fix(dagster): expand +x+ selectors from the matched nodes only The Dagster selection resolver expanded downstream from the set it had already expanded upstream, so +x+ also selected every descendant of x's ancestors (for example siblings of x). Expand both directions from the originally matched nodes, matching the planner, Python node, lineage, and diff selectors. Existing +x+ Dagster selections may now select fewer assets: nodes that are descendants of x's ancestors but not of x itself are no longer included. (cherry picked from commit bfc4ddfb8e9be0746fca69f14105ac4936cbd1c1) --- .../dagster/_helpers/selection.py | 5 +- .../sqlbuild/integrations/dagster/helpers.py | 44 +++++++++++++++ .../integrations/dagster/test_assets.py | 55 +++++++++++++++++++ 3 files changed, 102 insertions(+), 2 deletions(-) diff --git a/src/sqlbuild/integrations/dagster/_helpers/selection.py b/src/sqlbuild/integrations/dagster/_helpers/selection.py index 4d118eeca8..8aa32e977a 100644 --- a/src/sqlbuild/integrations/dagster/_helpers/selection.py +++ b/src/sqlbuild/integrations/dagster/_helpers/selection.py @@ -129,11 +129,12 @@ def _resolve_atomic( return matched matched = _match_parsed_selector(parsed=parsed, dag=dag, nodes_by_id=nodes_by_id) + anchors: frozenset[str] = frozenset(matched) if parsed.upstream: - matched.update(transitive_closure_many(starts=matched, edges=upstream, include_starts=True)) + matched.update(transitive_closure_many(starts=anchors, edges=upstream, include_starts=True)) if parsed.downstream: matched.update( - transitive_closure_many(starts=matched, edges=downstream, include_starts=True) + transitive_closure_many(starts=anchors, edges=downstream, include_starts=True) ) return matched diff --git a/tests/unit/src/sqlbuild/integrations/dagster/helpers.py b/tests/unit/src/sqlbuild/integrations/dagster/helpers.py index d66bc16284..c89331807f 100644 --- a/tests/unit/src/sqlbuild/integrations/dagster/helpers.py +++ b/tests/unit/src/sqlbuild/integrations/dagster/helpers.py @@ -568,3 +568,47 @@ def assert_json_output_file_behavior(*, command: tuple[str, ...]) -> None: assert "--json" not in command json_output_index: int = command.index("--json-output") + 1 assert not Path(command[json_output_index]).exists() + + +def _selection_model_node(name: str) -> dict[str, Any]: + return { + "id": f"model:{name}", + "kind": "model", + "name": name, + "asset_key": ["analytics", name], + "path": f"models/{name}.sql", + "materialization_type": "table", + } + + +def _selection_edge(from_name: str, to_name: str) -> dict[str, str]: + return {"from_id": f"model:{from_name}", "to_id": f"model:{to_name}"} + + +def build_dagster_selection_graph_dag() -> Mapping[str, Any]: + return { + "version": 1, + "project_name": "dagster_selection_project", + "nodes": [ + _selection_model_node("stg_orders"), + _selection_model_node("orders"), + _selection_model_node("customers"), + _selection_model_node("raw_events"), + _selection_model_node("events_left"), + _selection_model_node("events_right"), + _selection_model_node("events_joined"), + _selection_model_node("events_report"), + _selection_model_node("events_side"), + ], + "edges": [ + _selection_edge("stg_orders", "orders"), + _selection_edge("stg_orders", "customers"), + _selection_edge("raw_events", "events_left"), + _selection_edge("raw_events", "events_right"), + _selection_edge("raw_events", "events_side"), + _selection_edge("events_left", "events_joined"), + _selection_edge("events_right", "events_joined"), + _selection_edge("events_joined", "events_report"), + ], + "checks": [], + } diff --git a/tests/unit/src/sqlbuild/integrations/dagster/test_assets.py b/tests/unit/src/sqlbuild/integrations/dagster/test_assets.py index e3379076b0..6577ceb989 100644 --- a/tests/unit/src/sqlbuild/integrations/dagster/test_assets.py +++ b/tests/unit/src/sqlbuild/integrations/dagster/test_assets.py @@ -32,6 +32,7 @@ DagsterSourceAssetKeyTestCase, ) from tests.unit.src.sqlbuild.integrations.dagster.helpers import ( + build_dagster_selection_graph_dag, build_dagster_test_dag, build_python_augmented_dagster_test_dag, ) @@ -226,6 +227,60 @@ def assets_def() -> dg.MaterializeResult: ) +@pytest.mark.parametrize( + "test_case", + ( + DagsterAssetSelectionTestCase( + description="both-direction expansion keeps only ancestors and descendants", + select="+orders+", + exclude=None, + expected_asset_keys=(("analytics", "orders"), ("analytics", "stg_orders")), + ), + DagsterAssetSelectionTestCase( + description="downstream expansion includes every child", + select="stg_orders+", + exclude=None, + expected_asset_keys=( + ("analytics", "customers"), + ("analytics", "orders"), + ("analytics", "stg_orders"), + ), + ), + DagsterAssetSelectionTestCase( + description="both-direction expansion of a middle node skips siblings", + select="+events_left+", + exclude=None, + expected_asset_keys=( + ("analytics", "events_joined"), + ("analytics", "events_left"), + ("analytics", "events_report"), + ("analytics", "raw_events"), + ), + ), + ), + ids=lambda case: case.description, +) +def test_given_branching_dag_when_building_asset_selection_then_expands_from_matched_nodes( + test_case: DagsterAssetSelectionTestCase, +) -> None: + dag: Mapping[str, Any] = build_dagster_selection_graph_dag() + + @sqlbuild_assets(dag=dag) + def assets_def() -> dg.MaterializeResult: + return dg.MaterializeResult() + + selection: Any = build_sqlbuild_asset_selection( + sqlbuild_assets=[assets_def], + dag=dag, + sqlbuild_select=test_case.select, + sqlbuild_exclude=test_case.exclude, + ) + + assert tuple(sorted(tuple(key.path) for key in selection.resolve([assets_def]))) == tuple( + sorted(test_case.expected_asset_keys) + ) + + @pytest.mark.parametrize( "test_case", [ From 56b2380a8c376ad4ddf967ab5d653994283fb3c8 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:29:51 +0100 Subject: [PATCH 13/57] fix(dagster): keep every path node for a~b selectors The Dagster resolver kept only one breadth-first shortest path between a and b, and which branch won a tie depended on set iteration order. Resolve path selectors through the shared path_nodes helper so a~b keeps every node on every directed path between a and b, deterministically, matching the planner and lineage selectors. Behaviour change: as in the planner, x~x and a~b where b is not downstream of a now raise DagsterDagInputError ("no SQLBuild DAG path exists between ..."). Previously x~x selected only x. Add a waffle shop e2e test that checks Dagster selection against sqb plan for +x+ and path selectors. (cherry picked from commit 0c8ec48a3fbcaace15916da0de8344313db96603) --- .../dagster/_helpers/selection.py | 28 +++---- .../integrations/dagster/_test_types.py | 9 ++ .../dagster/test_selection_parity.py | 84 +++++++++++++++++++ .../integrations/dagster/_test_types.py | 7 ++ .../integrations/dagster/test_assets.py | 68 +++++++++++++++ 5 files changed, 181 insertions(+), 15 deletions(-) create mode 100644 tests/e2e/src/sqlbuild/integrations/dagster/test_selection_parity.py diff --git a/src/sqlbuild/integrations/dagster/_helpers/selection.py b/src/sqlbuild/integrations/dagster/_helpers/selection.py index 8aa32e977a..05ae1aa30f 100644 --- a/src/sqlbuild/integrations/dagster/_helpers/selection.py +++ b/src/sqlbuild/integrations/dagster/_helpers/selection.py @@ -2,10 +2,10 @@ from __future__ import annotations -from collections import deque from collections.abc import Mapping, Sequence from typing import Any +from sqlbuild.compiler.graph.main.path_nodes import path_nodes from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.planner.main.selection.selector_parse import parse_project_selector from sqlbuild.compiler.planner.models import ParsedSelector, PathSelector @@ -117,7 +117,7 @@ def _resolve_atomic( if isinstance(parsed, PathSelector): start_id: str = _node_id_for_name(nodes_by_id=nodes_by_id, name=parsed.start_name) end_id: str = _node_id_for_name(nodes_by_id=nodes_by_id, name=parsed.end_name) - matched: set[str] = _shortest_path(start_id=start_id, end_id=end_id, downstream=downstream) + matched: set[str] = _path_node_ids(start_id=start_id, end_id=end_id, downstream=downstream) if parsed.upstream: matched.update( transitive_closure_many(starts=(start_id,), edges=upstream, include_starts=True) @@ -218,16 +218,14 @@ def _node_id_for_name(*, nodes_by_id: Mapping[str, Mapping[str, Any]], name: str return matching_ids[0] -def _shortest_path(*, start_id: str, end_id: str, downstream: Mapping[str, set[str]]) -> set[str]: - pending: deque[tuple[str, tuple[str, ...]]] = deque([(start_id, (start_id,))]) - visited: set[str] = {start_id} - while pending: - node_id, path = pending.popleft() - if node_id == end_id: - return set(path) - for adjacent_id in downstream.get(node_id, set()): - if adjacent_id in visited: - continue - visited.add(adjacent_id) - pending.append((adjacent_id, (*path, adjacent_id))) - raise DagsterDagInputError(f"no SQLBuild DAG path exists between {start_id!r} and {end_id!r}") +def _path_node_ids(*, start_id: str, end_id: str, downstream: Mapping[str, set[str]]) -> set[str]: + on_path: frozenset[str] | None = path_nodes( + start=start_id, + end=end_id, + downstream={node_id: tuple(sorted(targets)) for node_id, targets in downstream.items()}, + ) + if on_path is None: + raise DagsterDagInputError( + f"no SQLBuild DAG path exists between {start_id!r} and {end_id!r}" + ) + return set(on_path) diff --git a/tests/e2e/src/sqlbuild/integrations/dagster/_test_types.py b/tests/e2e/src/sqlbuild/integrations/dagster/_test_types.py index f5edc37879..713f303186 100644 --- a/tests/e2e/src/sqlbuild/integrations/dagster/_test_types.py +++ b/tests/e2e/src/sqlbuild/integrations/dagster/_test_types.py @@ -126,3 +126,12 @@ class DagsterComputeLogE2ETestCase: expected_stdout_fragments: tuple[str, ...] unexpected_structured_log_fragment: str expected_materialization_fragment: str + + +@dataclass(frozen=True) +class DagsterSelectorParityE2ETestCase: + """Test case comparing Dagster asset selection with the SQLBuild planner.""" + + description: str + select: str + expected_model_names: tuple[str, ...] diff --git a/tests/e2e/src/sqlbuild/integrations/dagster/test_selection_parity.py b/tests/e2e/src/sqlbuild/integrations/dagster/test_selection_parity.py new file mode 100644 index 0000000000..2ce58984c4 --- /dev/null +++ b/tests/e2e/src/sqlbuild/integrations/dagster/test_selection_parity.py @@ -0,0 +1,84 @@ +"""E2E parity between Dagster asset selection and SQLBuild planner selection.""" + +import json +import subprocess +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import dagster as dg +import pytest + +from sqlbuild.integrations.dagster import build_sqlbuild_asset_selection, sqlbuild_assets +from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import prepare_waffle_shop, run_sqb +from tests.e2e.src.sqlbuild.integrations.dagster._test_types import ( + DagsterSelectorParityE2ETestCase, +) + + +@pytest.mark.parametrize( + "test_case", + ( + DagsterSelectorParityE2ETestCase( + description="both-direction expansion matches the planner", + select="+daily_revenue+", + expected_model_names=("daily_revenue", "stg_orders", "stg_payments"), + ), + DagsterSelectorParityE2ETestCase( + description="path selector keeps every path like the planner", + select="+fact_orders~hourly_activity_with_daily_context", + expected_model_names=( + "daily_activity_rollup", + "fact_orders", + "hourly_activity_with_daily_context", + "hourly_order_activity", + "stg_orders", + "stg_payments", + ), + ), + ), + ids=lambda case: case.description, +) +def test_given_waffle_shop_selector_when_resolving_in_dagster_then_models_match_planner( + test_case: DagsterSelectorParityE2ETestCase, + tmp_path: Path, +) -> None: + project_dir: Path = prepare_waffle_shop(tmp_path) + dag_result: subprocess.CompletedProcess[str] = run_sqb( + command=("dag", "--json"), project_dir=project_dir + ) + assert dag_result.returncode == 0, dag_result.stdout + dag_result.stderr + dag: Mapping[str, Any] = json.loads(dag_result.stdout) + plan_result: subprocess.CompletedProcess[str] = run_sqb( + command=("plan", "--select", test_case.select, "--json"), project_dir=project_dir + ) + assert plan_result.returncode == 0, plan_result.stdout + plan_result.stderr + planned_model_names: tuple[str, ...] = tuple( + sorted(str(entry["name"]) for entry in json.loads(plan_result.stdout)["models"]) + ) + + @sqlbuild_assets(dag=dag) + def waffle_shop_assets() -> dg.MaterializeResult: + return dg.MaterializeResult() + + selection: Any = build_sqlbuild_asset_selection( + sqlbuild_assets=[waffle_shop_assets], + dag=dag, + sqlbuild_select=test_case.select, + ) + kind_and_name_by_asset_key: dict[tuple[str, ...], tuple[str, str]] = { + tuple(node["asset_key"]): (str(node["kind"]), str(node["name"])) for node in dag["nodes"] + } + selected_kind_and_names: set[tuple[str, str]] = { + kind_and_name_by_asset_key[tuple(key.path)] + for key in selection.resolve([waffle_shop_assets]) + } + model_kind_and_names: set[tuple[str, str]] = { + ("model", name) for _, name in kind_and_name_by_asset_key.values() + } + dagster_model_names: tuple[str, ...] = tuple( + sorted(name for _, name in selected_kind_and_names & model_kind_and_names) + ) + + assert planned_model_names == test_case.expected_model_names + assert dagster_model_names == test_case.expected_model_names diff --git a/tests/unit/src/sqlbuild/integrations/dagster/_test_types.py b/tests/unit/src/sqlbuild/integrations/dagster/_test_types.py index 11b90ee991..043c02b7af 100644 --- a/tests/unit/src/sqlbuild/integrations/dagster/_test_types.py +++ b/tests/unit/src/sqlbuild/integrations/dagster/_test_types.py @@ -289,3 +289,10 @@ class DagsterProjectPrepareFailureTestCase: class DagsterProjectDecoratorTestCase: description: str expected_asset_keys: tuple[tuple[str, ...], ...] + + +@dataclass(frozen=True) +class DagsterAssetSelectionErrorTestCase: + description: str + select: str + expected_error_fragment: str diff --git a/tests/unit/src/sqlbuild/integrations/dagster/test_assets.py b/tests/unit/src/sqlbuild/integrations/dagster/test_assets.py index 6577ceb989..a6afd656e0 100644 --- a/tests/unit/src/sqlbuild/integrations/dagster/test_assets.py +++ b/tests/unit/src/sqlbuild/integrations/dagster/test_assets.py @@ -21,6 +21,7 @@ from sqlbuild.integrations.dagster.models import SqlBuildProject from tests.unit.src.sqlbuild.integrations.dagster._test_types import ( DagsterAssetCheckFilterTestCase, + DagsterAssetSelectionErrorTestCase, DagsterAssetSelectionTestCase, DagsterAssetSpecTestCase, DagsterConflictingInputTestCase, @@ -246,6 +247,40 @@ def assets_def() -> dg.MaterializeResult: ("analytics", "stg_orders"), ), ), + DagsterAssetSelectionTestCase( + description="path selector keeps every node on every path", + select="raw_events~events_report", + exclude=None, + expected_asset_keys=( + ("analytics", "events_joined"), + ("analytics", "events_left"), + ("analytics", "events_report"), + ("analytics", "events_right"), + ("analytics", "raw_events"), + ), + ), + DagsterAssetSelectionTestCase( + description="path selector keeps both equal-length diamond branches", + select="raw_events~events_joined", + exclude=None, + expected_asset_keys=( + ("analytics", "events_joined"), + ("analytics", "events_left"), + ("analytics", "events_right"), + ("analytics", "raw_events"), + ), + ), + DagsterAssetSelectionTestCase( + description="path selector expands upstream of start and downstream of end", + select="+events_left~events_joined+", + exclude=None, + expected_asset_keys=( + ("analytics", "events_joined"), + ("analytics", "events_left"), + ("analytics", "events_report"), + ("analytics", "raw_events"), + ), + ), DagsterAssetSelectionTestCase( description="both-direction expansion of a middle node skips siblings", select="+events_left+", @@ -584,3 +619,36 @@ def test_given_sqlbuild_dag_when_building_check_specs_then_can_exclude_scenarios assert check_names == test_case.expected_check_names assert not set(test_case.unexpected_check_names).intersection(check_names) + + +@pytest.mark.parametrize( + "test_case", + [ + DagsterAssetSelectionErrorTestCase( + description="rejects a path whose end is not downstream of its start", + select="events_report~raw_events", + expected_error_fragment="no SQLBuild DAG path exists between", + ), + DagsterAssetSelectionErrorTestCase( + description="rejects a path from a node to itself like the planner", + select="orders~orders", + expected_error_fragment="no SQLBuild DAG path exists between", + ), + ], + ids=lambda case: case.description, +) +def test_given_unreachable_path_selector_when_building_asset_selection_then_raises_input_error( + test_case: DagsterAssetSelectionErrorTestCase, +) -> None: + dag: Mapping[str, Any] = build_dagster_selection_graph_dag() + + @sqlbuild_assets(dag=dag) + def assets_def() -> dg.MaterializeResult: + return dg.MaterializeResult() + + with pytest.raises(DagsterDagInputError, match=test_case.expected_error_fragment): + build_sqlbuild_asset_selection( + sqlbuild_assets=[assets_def], + dag=dag, + sqlbuild_select=test_case.select, + ) From 72266697a9182cc0468370f0466a5f63dc8149c5 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:34:19 +0100 Subject: [PATCH 14/57] test(state): characterize ref replacement on DuckDB and Postgres state backends Cover single-type node ref replacement, function ref replacement across udf and table_fn, grouped ref replacement, and environment publish with refs on both backends before consolidating the remaining shared state SQL. (cherry picked from commit b5b64939a3284919c87edc6bb0a2765573e587ce) --- .../virtual/state/classes/_test_types.py | 9 ++ .../sqlbuild/virtual/state/classes/helpers.py | 137 ++++++++++++++++++ .../state/classes/postgres/_test_types.py | 8 + .../state/classes/postgres/test_postgres.py | 35 +++++ .../virtual/state/classes/test_duckdb.py | 36 +++++ 5 files changed, 225 insertions(+) diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/_test_types.py b/tests/integration/src/sqlbuild/virtual/state/classes/_test_types.py index d542eda509..f1ae320a32 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/_test_types.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/_test_types.py @@ -5,6 +5,7 @@ from sqlbuild.virtual.state.types import VirtualEnvironmentStatus from tests.integration.src.sqlbuild.virtual.state.classes.helpers import ( ConditionalPublicationPayload, + StateRefContractObservation, ) @@ -263,3 +264,11 @@ class StateReadContractTestCase: expected_environments: tuple[tuple[str, VirtualEnvironmentStatus], ...] expected_active_lock_keys: tuple[str, ...] expected_expired_lock_keys: tuple[str, ...] + + +@dataclass(frozen=True) +class DuckDbStateBackendRefContractTestCase: + description: str + expected_observation: StateRefContractObservation + schema: str + sqlbuild_version: str diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py b/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py index 499d57a473..1ea7658a61 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py @@ -14,6 +14,7 @@ VirtualEnvironmentCheckpointModelRefRecord, VirtualEnvironmentCheckpointRecord, VirtualEnvironmentCheckpointSeedRefRecord, + VirtualEnvironmentFunctionRefRecord, VirtualEnvironmentNodeRefRecord, VirtualEnvironmentRecord, ) @@ -287,3 +288,139 @@ def exercise_state_read_contract( active_lock_keys=active_lock_keys, expired_lock_keys=expired_lock_keys, ) + + +class StateRefContractObservation(NamedTuple): + """Values read back after writing the shared state ref contract fixture.""" + + node_refs_after_first_replace: tuple[VirtualEnvironmentNodeRefRecord, ...] + node_refs_after_second_replace: tuple[VirtualEnvironmentNodeRefRecord, ...] + untouched_seed_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] + function_refs: tuple[VirtualEnvironmentFunctionRefRecord, ...] + grouped_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] + grouped_seed_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] + published_environment: VirtualEnvironmentRecord | None + published_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] + + +def exercise_state_ref_contract( + *, backend: StateBackend, connection: Any, schema: str +) -> StateRefContractObservation: + backend.replace_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name="dev", + node_type="seed", + refs=(VirtualEnvironmentNodeRefRecord("dev", "seed", "countries", "seed-v1"),), + ) + backend.replace_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name="dev", + node_type="model", + refs=( + VirtualEnvironmentNodeRefRecord("dev", "model", "orders", "orders-v1"), + VirtualEnvironmentNodeRefRecord("dev", "model", "customers", "customers-v1"), + ), + ) + node_refs_after_first_replace: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="dev", node_type="model" + ) + ) + backend.replace_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name="dev", + node_type="model", + refs=(VirtualEnvironmentNodeRefRecord("dev", "model", "orders", "orders-v2"),), + ) + node_refs_after_second_replace: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="dev", node_type="model" + ) + ) + untouched_seed_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="dev", node_type="seed" + ) + ) + backend.replace_virtual_environment_function_refs( + connection=connection, + schema=schema, + virtual_environment_name="dev", + refs=( + VirtualEnvironmentFunctionRefRecord("dev", "table_fn", "customer_orders", "fn-v1"), + VirtualEnvironmentFunctionRefRecord("dev", "udf", "normalize_email", "fn-v2"), + ), + ) + function_refs: tuple[VirtualEnvironmentFunctionRefRecord, ...] = ( + backend.get_virtual_environment_function_refs( + connection=connection, schema=schema, virtual_environment_name="dev" + ) + ) + backend.replace_virtual_environment_node_ref_groups( + connection=connection, + schema=schema, + virtual_environment_name="qa", + refs_by_node_type={ + "model": (VirtualEnvironmentNodeRefRecord("qa", "model", "orders", "orders-v3"),), + "seed": (VirtualEnvironmentNodeRefRecord("qa", "seed", "countries", "seed-v2"),), + }, + ) + grouped_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="qa", node_type="model" + ) + ) + grouped_seed_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="qa", node_type="seed" + ) + ) + backend.upsert_virtual_environment_and_replace_node_ref_groups( + connection=connection, + schema=schema, + record=VirtualEnvironmentRecord("prod", VirtualEnvironmentStatus.ACTIVE), + refs_by_node_type={ + "model": (VirtualEnvironmentNodeRefRecord("prod", "model", "orders", "orders-v1"),) + }, + ) + published_environment: VirtualEnvironmentRecord | None = backend.get_virtual_environment( + connection=connection, schema=schema, virtual_environment_name="prod" + ) + published_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="prod", node_type="model" + ) + ) + return StateRefContractObservation( + node_refs_after_first_replace=node_refs_after_first_replace, + node_refs_after_second_replace=node_refs_after_second_replace, + untouched_seed_refs=untouched_seed_refs, + function_refs=function_refs, + grouped_model_refs=grouped_model_refs, + grouped_seed_refs=grouped_seed_refs, + published_environment=published_environment, + published_model_refs=published_model_refs, + ) + + +EXPECTED_STATE_REF_CONTRACT_OBSERVATION: StateRefContractObservation = StateRefContractObservation( + node_refs_after_first_replace=( + VirtualEnvironmentNodeRefRecord("dev", "model", "customers", "customers-v1"), + VirtualEnvironmentNodeRefRecord("dev", "model", "orders", "orders-v1"), + ), + node_refs_after_second_replace=( + VirtualEnvironmentNodeRefRecord("dev", "model", "orders", "orders-v2"), + ), + untouched_seed_refs=(VirtualEnvironmentNodeRefRecord("dev", "seed", "countries", "seed-v1"),), + function_refs=( + VirtualEnvironmentFunctionRefRecord("dev", "table_fn", "customer_orders", "fn-v1"), + VirtualEnvironmentFunctionRefRecord("dev", "udf", "normalize_email", "fn-v2"), + ), + grouped_model_refs=(VirtualEnvironmentNodeRefRecord("qa", "model", "orders", "orders-v3"),), + grouped_seed_refs=(VirtualEnvironmentNodeRefRecord("qa", "seed", "countries", "seed-v2"),), + published_environment=VirtualEnvironmentRecord("prod", VirtualEnvironmentStatus.ACTIVE), + published_model_refs=(VirtualEnvironmentNodeRefRecord("prod", "model", "orders", "orders-v1"),), +) diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/postgres/_test_types.py b/tests/integration/src/sqlbuild/virtual/state/classes/postgres/_test_types.py index c70a492d09..f993dd7f92 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/postgres/_test_types.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/postgres/_test_types.py @@ -5,6 +5,7 @@ from sqlbuild.virtual.state.types import VirtualEnvironmentStatus from tests.integration.src.sqlbuild.virtual.state.classes.helpers import ( ConditionalPublicationPayload, + StateRefContractObservation, ) @@ -219,3 +220,10 @@ class PostgresStateReadContractTestCase: expected_environments: tuple[tuple[str, VirtualEnvironmentStatus], ...] expected_active_lock_keys: tuple[str, ...] expected_expired_lock_keys: tuple[str, ...] + + +@dataclass(frozen=True) +class PostgresStateBackendRefContractTestCase: + description: str + expected_observation: StateRefContractObservation + sqlbuild_version: str diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/postgres/test_postgres.py b/tests/integration/src/sqlbuild/virtual/state/classes/postgres/test_postgres.py index 48a35f740f..2c97970d33 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/postgres/test_postgres.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/postgres/test_postgres.py @@ -67,6 +67,7 @@ CHECKPOINT_ENVIRONMENT_PAYLOAD, CHECKPOINT_ID_PAYLOAD, DETACHED_CHECKPOINT_PAYLOAD, + EXPECTED_STATE_REF_CONTRACT_OBSERVATION, FAILED_CHECKPOINT_PAYLOAD, FINALIZING_CHECKPOINT_PAYLOAD, FUNCTION_OMISSION_PAYLOAD, @@ -81,7 +82,9 @@ STATE_READ_CONTRACT_RELATION_V2, VALID_PAYLOAD, StateReadContractObservation, + StateRefContractObservation, exercise_state_read_contract, + exercise_state_ref_contract, ) from tests.integration.src.sqlbuild.virtual.state.classes.postgres._test_types import ( PostgresAtomicFinalizedVirtualPublishTestCase, @@ -100,6 +103,7 @@ PostgresStateBackendNodeResultTestCase, PostgresStateBackendOperationEventTestCase, PostgresStateBackendPythonNodeIdentityTestCase, + PostgresStateBackendRefContractTestCase, PostgresStateBackendSeedRefTestCase, PostgresStateBackendSourceFreshnessTestCase, PostgresStateBackendTableCreationTestCase, @@ -2555,3 +2559,34 @@ def test_given_postgres_state_backend_when_reading_state_then_matches_read_contr assert observation.environments == test_case.expected_environments assert observation.active_lock_keys == test_case.expected_active_lock_keys assert observation.expired_lock_keys == test_case.expected_expired_lock_keys + + +@pytest.mark.parametrize( + "test_case", + [ + PostgresStateBackendRefContractTestCase( + description="replaces node function and grouped refs per environment and type", + expected_observation=EXPECTED_STATE_REF_CONTRACT_OBSERVATION, + sqlbuild_version="0.0.test", + ) + ], + ids=lambda case: case.description, +) +def test_given_postgres_state_backend_when_replacing_refs_then_matches_ref_contract( + test_case: PostgresStateBackendRefContractTestCase, + postgres_state_backend: PostgresStateBackend, + postgres_state_connection: Any, + postgres_state_schema: str, +) -> None: + postgres_state_backend.initialize( + connection=postgres_state_connection, + schema=postgres_state_schema, + sqlbuild_version=test_case.sqlbuild_version, + ) + observation: StateRefContractObservation = exercise_state_ref_contract( + backend=postgres_state_backend, + connection=postgres_state_connection, + schema=postgres_state_schema, + ) + + assert observation == test_case.expected_observation diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/test_duckdb.py b/tests/integration/src/sqlbuild/virtual/state/classes/test_duckdb.py index adc2e9d10a..66e92b8089 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/test_duckdb.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/test_duckdb.py @@ -90,6 +90,7 @@ DuckDbStateBackendNodeResultTestCase, DuckDbStateBackendOperationEventTestCase, DuckDbStateBackendPythonNodeIdentityTestCase, + DuckDbStateBackendRefContractTestCase, DuckDbStateBackendRollbackTestCase, DuckDbStateBackendSeedRefTestCase, DuckDbStateBackendSourceFreshnessTestCase, @@ -106,6 +107,7 @@ CHECKPOINT_ENVIRONMENT_PAYLOAD, CHECKPOINT_ID_PAYLOAD, DETACHED_CHECKPOINT_PAYLOAD, + EXPECTED_STATE_REF_CONTRACT_OBSERVATION, FAILED_CHECKPOINT_PAYLOAD, FINALIZING_CHECKPOINT_PAYLOAD, FUNCTION_OMISSION_PAYLOAD, @@ -120,7 +122,9 @@ STATE_READ_CONTRACT_RELATION_V2, VALID_PAYLOAD, StateReadContractObservation, + StateRefContractObservation, exercise_state_read_contract, + exercise_state_ref_contract, fetch_all, open_duckdb_state_backend, ) @@ -2858,3 +2862,35 @@ def test_given_duckdb_state_backend_when_reading_state_then_matches_read_contrac assert observation.environments == test_case.expected_environments assert observation.active_lock_keys == test_case.expected_active_lock_keys assert observation.expired_lock_keys == test_case.expected_expired_lock_keys + + +@pytest.mark.parametrize( + "test_case", + [ + DuckDbStateBackendRefContractTestCase( + description="replaces node function and grouped refs per environment and type", + expected_observation=EXPECTED_STATE_REF_CONTRACT_OBSERVATION, + schema="sqlbuild_state", + sqlbuild_version="0.0.test", + ) + ], + ids=lambda case: case.description, +) +def test_given_duckdb_state_backend_when_replacing_refs_then_matches_ref_contract( + test_case: DuckDbStateBackendRefContractTestCase, + tmp_path: Path, +) -> None: + backend, connection = open_duckdb_state_backend(db_path=tmp_path / "state.duckdb") + try: + backend.initialize( + connection=connection, + schema=test_case.schema, + sqlbuild_version=test_case.sqlbuild_version, + ) + observation: StateRefContractObservation = exercise_state_ref_contract( + backend=backend, connection=connection, schema=test_case.schema + ) + finally: + backend.close(connection) + + assert observation == test_case.expected_observation From 1f6bb2994e4901591143566adf45a27df0f361a4 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:36:11 +0100 Subject: [PATCH 15/57] refactor(state): share virtual environment ref replacement wrappers Move the node, model, seed, function, and grouped ref replacement methods, the combined environment-and-refs publish, the Python node ref upsert, and ref validation into SqlStateBackend. They were identical in both backends apart from the executor passed to the node-ref group helper. _replace_virtual_environment_node_ref_groups stays per backend behind an abstract hook: DuckDB stages refs in a temporary table and deletes missing rows before an upsert, while Postgres deletes the environment's rows for the node type and reinserts them with ON CONFLICT. (cherry picked from commit 142f22fd33bfa25c15a37fbdd8d67249988f1ebd) --- .../classes/_duckdb_conditional_publish.py | 2 +- .../state/classes/_sql_state_backend.py | 188 ++++++++++++++++ src/sqlbuild/virtual/state/classes/duckdb.py | 206 +----------------- .../virtual/state/classes/postgres.py | 201 +---------------- 4 files changed, 200 insertions(+), 397 deletions(-) diff --git a/src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py b/src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py index d49c5cddd9..6a9f220ad6 100644 --- a/src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py +++ b/src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py @@ -68,7 +68,7 @@ def upsert_virtual_environment_and_replace_node_ref_groups_if_locks_owned( executor=connection, schema=schema, record=record ) backend._replace_virtual_environment_node_ref_groups( - connection=connection, + executor=connection, schema=schema, virtual_environment_name=record.virtual_environment_name, refs_by_node_type=refs_by_node_type, diff --git a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py index 364ff428ba..560ae7fca9 100644 --- a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py +++ b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py @@ -8,6 +8,7 @@ from datetime import datetime from typing import Any, ClassVar +from sqlbuild.compiler.compile.types import CompiledResourceType from sqlbuild.virtual.state._helpers.state_storage.datetime import ( from_naive_utc_wall_clock, ) @@ -29,6 +30,7 @@ VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, VIRTUAL_ENVIRONMENT_TABLE, ) +from sqlbuild.virtual.state.exceptions import StateBackendConfigError from sqlbuild.virtual.state.models import ( FunctionVersionRecord, ModelVersionRecord, @@ -313,6 +315,192 @@ def list_expired_locks(self, *, connection: Any, schema: str) -> tuple[StateLock StateLockRecord(lock_key=row[0], owner_id=row[1], expires_at=row[2]) for row in rows ) + @abstractmethod + def _replace_virtual_environment_node_ref_groups( + self, + *, + executor: Any, + schema: str, + virtual_environment_name: str, + refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], + ) -> None: + """Replace node refs per node type inside an open write transaction.""" + ... + + def replace_virtual_environment_node_refs( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + node_type: str, + refs: tuple[VirtualEnvironmentNodeRefRecord, ...], + ) -> None: + self.replace_virtual_environment_node_ref_groups( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + refs_by_node_type={node_type: refs}, + ) + + def replace_virtual_environment_node_ref_groups( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_virtual_environment_node_ref_groups( + executor=executor, + schema=schema, + virtual_environment_name=virtual_environment_name, + refs_by_node_type=refs_by_node_type, + ) + + def upsert_virtual_environment_and_replace_node_ref_groups( + self, + *, + connection: Any, + schema: str, + record: VirtualEnvironmentRecord, + refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._upsert_virtual_environment_record(executor=executor, schema=schema, record=record) + self._replace_virtual_environment_node_ref_groups( + executor=executor, + schema=schema, + virtual_environment_name=record.virtual_environment_name, + refs_by_node_type=refs_by_node_type, + ) + + def replace_virtual_environment_model_refs( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + refs: tuple[VirtualEnvironmentModelRefRecord, ...], + ) -> None: + self.replace_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + node_type="model", + refs=tuple( + VirtualEnvironmentNodeRefRecord( + virtual_environment_name=ref.virtual_environment_name, + node_type="model", + node_name=ref.model_name, + version_hash=ref.version_hash, + ) + for ref in refs + ), + ) + + def replace_virtual_environment_function_refs( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + refs: tuple[VirtualEnvironmentFunctionRefRecord, ...], + ) -> None: + ref: VirtualEnvironmentFunctionRefRecord + for ref in refs: + if ref.node_type not in { + CompiledResourceType.UDF, + CompiledResourceType.TABLE_FN, + }: + raise StateBackendConfigError("Function ref node_type must be 'udf' or 'table_fn'") + refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]] = {} + for node_type in ("udf", "table_fn"): + node_refs: list[VirtualEnvironmentNodeRefRecord] = [] + for ref in refs: + if ref.node_type == node_type: + node_refs.append( + VirtualEnvironmentNodeRefRecord( + virtual_environment_name=ref.virtual_environment_name, + node_type=ref.node_type, + node_name=ref.function_name, + version_hash=ref.version_hash, + ) + ) + refs_by_node_type[node_type] = tuple(node_refs) + self.replace_virtual_environment_node_ref_groups( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + refs_by_node_type=refs_by_node_type, + ) + + def replace_virtual_environment_seed_refs( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + refs: tuple[VirtualEnvironmentSeedRefRecord, ...], + ) -> None: + self.replace_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + node_type="seed", + refs=tuple( + VirtualEnvironmentNodeRefRecord( + virtual_environment_name=ref.virtual_environment_name, + node_type="seed", + node_name=ref.seed_name, + version_hash=ref.version_hash, + ) + for ref in refs + ), + ) + + def upsert_virtual_environment_python_node_ref( + self, + *, + connection: Any, + schema: str, + ref: VirtualEnvironmentPythonNodeRefRecord, + ) -> None: + self.upsert_virtual_environment_node_ref( + connection=connection, + schema=schema, + ref=VirtualEnvironmentNodeRefRecord( + virtual_environment_name=ref.virtual_environment_name, + node_type=ref.node_type, + node_name=ref.node_name, + version_hash=ref.version_hash, + ), + ) + + def _validate_node_ref_replacement( + self, + *, + virtual_environment_name: str, + node_type: str, + refs: tuple[VirtualEnvironmentNodeRefRecord, ...], + ) -> None: + seen_node_names: set[str] = set() + ref: VirtualEnvironmentNodeRefRecord + for ref in refs: + if ref.virtual_environment_name != virtual_environment_name: + raise StateBackendConfigError( + "Node ref virtual_environment_name must match replacement " + "virtual_environment_name" + ) + if ref.node_type != node_type: + raise StateBackendConfigError("Node ref node_type must match replacement node_type") + if ref.node_name in seen_node_names: + raise StateBackendConfigError( + f"Duplicate node ref for node type '{node_type}' and name '{ref.node_name}'" + ) + seen_node_names.add(ref.node_name) + def _replace_row_preserving_created_at( self, *, diff --git a/src/sqlbuild/virtual/state/classes/duckdb.py b/src/sqlbuild/virtual/state/classes/duckdb.py index 3a069e651f..1d4e1d8ce6 100644 --- a/src/sqlbuild/virtual/state/classes/duckdb.py +++ b/src/sqlbuild/virtual/state/classes/duckdb.py @@ -8,7 +8,6 @@ from typing import Any, ClassVar from sqlbuild.adapter.contract.classes.observed_connection import ObservedConnection -from sqlbuild.compiler.compile.types import CompiledResourceType from sqlbuild.executor.node_results.main.decode_json import decode_node_result_json from sqlbuild.executor.node_results.main.encode_json import encode_node_result_json from sqlbuild.executor.node_results.models import ( @@ -74,12 +73,7 @@ VirtualEnvironmentCheckpointModelRefRecord, VirtualEnvironmentCheckpointRecord, VirtualEnvironmentCheckpointSeedRefRecord, - VirtualEnvironmentFunctionRefRecord, - VirtualEnvironmentModelRefRecord, VirtualEnvironmentNodeRefRecord, - VirtualEnvironmentPythonNodeRefRecord, - VirtualEnvironmentRecord, - VirtualEnvironmentSeedRefRecord, ) from sqlbuild.virtual.state.types import ( StateColumnType, @@ -448,67 +442,6 @@ def delete_virtual_environment( connection.execute("ROLLBACK") raise - def replace_virtual_environment_node_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - node_type: str, - refs: tuple[VirtualEnvironmentNodeRefRecord, ...], - ) -> None: - self.replace_virtual_environment_node_ref_groups( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - refs_by_node_type={node_type: refs}, - ) - - def replace_virtual_environment_node_ref_groups( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], - ) -> None: - connection.execute("BEGIN") - try: - self._replace_virtual_environment_node_ref_groups( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def upsert_virtual_environment_and_replace_node_ref_groups( - self, - *, - connection: Any, - schema: str, - record: VirtualEnvironmentRecord, - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], - ) -> None: - connection.execute("BEGIN") - try: - self._upsert_virtual_environment_record( - executor=connection, schema=schema, record=record - ) - self._replace_virtual_environment_node_ref_groups( - connection=connection, - schema=schema, - virtual_environment_name=record.virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - def upsert_virtual_environment_node_ref( self, *, @@ -526,108 +459,6 @@ def upsert_virtual_environment_node_ref( [ref.virtual_environment_name, ref.node_type, ref.node_name, ref.version_hash], ) - def replace_virtual_environment_model_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs: tuple[VirtualEnvironmentModelRefRecord, ...], - ) -> None: - self.replace_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="model", - refs=tuple( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type="model", - node_name=ref.model_name, - version_hash=ref.version_hash, - ) - for ref in refs - ), - ) - - def replace_virtual_environment_function_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs: tuple[VirtualEnvironmentFunctionRefRecord, ...], - ) -> None: - ref: VirtualEnvironmentFunctionRefRecord - for ref in refs: - if ref.node_type not in { - CompiledResourceType.UDF, - CompiledResourceType.TABLE_FN, - }: - raise StateBackendConfigError("Function ref node_type must be 'udf' or 'table_fn'") - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]] = {} - for node_type in ("udf", "table_fn"): - node_refs: list[VirtualEnvironmentNodeRefRecord] = [] - for ref in refs: - if ref.node_type == node_type: - node_refs.append( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type=ref.node_type, - node_name=ref.function_name, - version_hash=ref.version_hash, - ) - ) - refs_by_node_type[node_type] = tuple(node_refs) - self.replace_virtual_environment_node_ref_groups( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - - def replace_virtual_environment_seed_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs: tuple[VirtualEnvironmentSeedRefRecord, ...], - ) -> None: - self.replace_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="seed", - refs=tuple( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type="seed", - node_name=ref.seed_name, - version_hash=ref.version_hash, - ) - for ref in refs - ), - ) - - def upsert_virtual_environment_python_node_ref( - self, - *, - connection: Any, - schema: str, - ref: VirtualEnvironmentPythonNodeRefRecord, - ) -> None: - self.upsert_virtual_environment_node_ref( - connection=connection, - schema=schema, - ref=VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type=ref.node_type, - node_name=ref.node_name, - version_hash=ref.version_hash, - ), - ) - def count_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: row: tuple[Any, ...] = connection.execute( "SELECT COUNT(*) " @@ -1062,33 +893,10 @@ def _validate_source_freshness_records( ) seen_source_names.add(record.source_name) - def _validate_node_ref_replacement( - self, - *, - virtual_environment_name: str, - node_type: str, - refs: tuple[VirtualEnvironmentNodeRefRecord, ...], - ) -> None: - seen_node_names: set[str] = set() - ref: VirtualEnvironmentNodeRefRecord - for ref in refs: - if ref.virtual_environment_name != virtual_environment_name: - raise StateBackendConfigError( - "Node ref virtual_environment_name must match replacement " - "virtual_environment_name" - ) - if ref.node_type != node_type: - raise StateBackendConfigError("Node ref node_type must match replacement node_type") - if ref.node_name in seen_node_names: - raise StateBackendConfigError( - f"Duplicate node ref for node type '{node_type}' and name '{ref.node_name}'" - ) - seen_node_names.add(ref.node_name) - def _replace_virtual_environment_node_ref_groups( self, *, - connection: Any, + executor: Any, schema: str, virtual_environment_name: str, refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], @@ -1102,8 +910,8 @@ def _replace_virtual_environment_node_ref_groups( node_type=node_type, refs=refs, ) - connection.execute(f"DROP TABLE IF EXISTS {temp_table_name}") - connection.execute( + executor.execute(f"DROP TABLE IF EXISTS {temp_table_name}") + executor.execute( f"CREATE TEMP TABLE {temp_table_name} (" "virtual_environment_name TEXT NOT NULL, " "node_type TEXT NOT NULL, " @@ -1113,7 +921,7 @@ def _replace_virtual_environment_node_ref_groups( ) ref: VirtualEnvironmentNodeRefRecord for ref in refs: - connection.execute( + executor.execute( f"INSERT INTO {temp_table_name} " "(virtual_environment_name, node_type, node_name, version_hash) " "VALUES (?, ?, ?, ?)", @@ -1124,7 +932,7 @@ def _replace_virtual_environment_node_ref_groups( ref.version_hash, ], ) - connection.execute( + executor.execute( "DELETE FROM " f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " "WHERE virtual_environment_name = ? AND node_type = ? " @@ -1135,7 +943,7 @@ def _replace_virtual_environment_node_ref_groups( ".node_name)", [virtual_environment_name, node_type], ) - connection.execute( + executor.execute( "INSERT INTO " f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " "(virtual_environment_name, node_type, node_name, version_hash, updated_at) " @@ -1144,7 +952,7 @@ def _replace_virtual_environment_node_ref_groups( "ON CONFLICT (virtual_environment_name, node_type, node_name) " "DO UPDATE SET version_hash = excluded.version_hash, updated_at = now()" ) - connection.execute(f"DROP TABLE IF EXISTS {temp_table_name}") + executor.execute(f"DROP TABLE IF EXISTS {temp_table_name}") def _backup_schema_name(self, *, schema: str, backup_id_value: str) -> str: return f"{schema}__backup_{backup_id_value}" diff --git a/src/sqlbuild/virtual/state/classes/postgres.py b/src/sqlbuild/virtual/state/classes/postgres.py index 01e5f816d0..4e41cfacd0 100644 --- a/src/sqlbuild/virtual/state/classes/postgres.py +++ b/src/sqlbuild/virtual/state/classes/postgres.py @@ -8,7 +8,6 @@ from typing import Any, ClassVar from sqlbuild.adapter.contract.classes.observed_connection import ObservedConnection -from sqlbuild.compiler.compile.types import CompiledResourceType from sqlbuild.executor.node_results.main.decode_json import decode_node_result_json from sqlbuild.executor.node_results.main.encode_json import encode_node_result_json from sqlbuild.executor.node_results.models import ( @@ -73,12 +72,8 @@ VirtualEnvironmentCheckpointModelRefRecord, VirtualEnvironmentCheckpointRecord, VirtualEnvironmentCheckpointSeedRefRecord, - VirtualEnvironmentFunctionRefRecord, - VirtualEnvironmentModelRefRecord, VirtualEnvironmentNodeRefRecord, - VirtualEnvironmentPythonNodeRefRecord, VirtualEnvironmentRecord, - VirtualEnvironmentSeedRefRecord, ) from sqlbuild.virtual.state.types import ( StateColumnType, @@ -549,69 +544,6 @@ def delete_virtual_environment( cursor.execute("ROLLBACK") raise - def replace_virtual_environment_node_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - node_type: str, - refs: tuple[VirtualEnvironmentNodeRefRecord, ...], - ) -> None: - self.replace_virtual_environment_node_ref_groups( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - refs_by_node_type={node_type: refs}, - ) - - def replace_virtual_environment_node_ref_groups( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - self._replace_virtual_environment_node_ref_groups( - cursor=cursor, - schema=schema, - virtual_environment_name=virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def upsert_virtual_environment_and_replace_node_ref_groups( - self, - *, - connection: Any, - schema: str, - record: VirtualEnvironmentRecord, - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - self._upsert_virtual_environment_record( - executor=cursor, schema=schema, record=record - ) - self._replace_virtual_environment_node_ref_groups( - cursor=cursor, - schema=schema, - virtual_environment_name=record.virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - def upsert_virtual_environment_and_replace_node_ref_groups_if_locks_owned( self, *, @@ -652,7 +584,7 @@ def upsert_virtual_environment_and_replace_node_ref_groups_if_locks_owned( executor=cursor, schema=schema, record=record ) self._replace_virtual_environment_node_ref_groups( - cursor=cursor, + executor=cursor, schema=schema, virtual_environment_name=record.virtual_environment_name, refs_by_node_type=refs_by_node_type, @@ -691,108 +623,6 @@ def upsert_virtual_environment_node_ref( [ref.virtual_environment_name, ref.node_type, ref.node_name, ref.version_hash], ) - def replace_virtual_environment_model_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs: tuple[VirtualEnvironmentModelRefRecord, ...], - ) -> None: - self.replace_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="model", - refs=tuple( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type="model", - node_name=ref.model_name, - version_hash=ref.version_hash, - ) - for ref in refs - ), - ) - - def replace_virtual_environment_function_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs: tuple[VirtualEnvironmentFunctionRefRecord, ...], - ) -> None: - ref: VirtualEnvironmentFunctionRefRecord - for ref in refs: - if ref.node_type not in { - CompiledResourceType.UDF, - CompiledResourceType.TABLE_FN, - }: - raise StateBackendConfigError("Function ref node_type must be 'udf' or 'table_fn'") - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]] = {} - for node_type in ("udf", "table_fn"): - node_refs: list[VirtualEnvironmentNodeRefRecord] = [] - for ref in refs: - if ref.node_type == node_type: - node_refs.append( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type=ref.node_type, - node_name=ref.function_name, - version_hash=ref.version_hash, - ) - ) - refs_by_node_type[node_type] = tuple(node_refs) - self.replace_virtual_environment_node_ref_groups( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - - def replace_virtual_environment_seed_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs: tuple[VirtualEnvironmentSeedRefRecord, ...], - ) -> None: - self.replace_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="seed", - refs=tuple( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type="seed", - node_name=ref.seed_name, - version_hash=ref.version_hash, - ) - for ref in refs - ), - ) - - def upsert_virtual_environment_python_node_ref( - self, - *, - connection: Any, - schema: str, - ref: VirtualEnvironmentPythonNodeRefRecord, - ) -> None: - self.upsert_virtual_environment_node_ref( - connection=connection, - schema=schema, - ref=VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type=ref.node_type, - node_name=ref.node_name, - version_hash=ref.version_hash, - ), - ) - def count_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: with connection.cursor() as cursor: cursor.execute( @@ -1278,33 +1108,10 @@ def _validate_source_freshness_records( ) seen_source_names.add(record.source_name) - def _validate_node_ref_replacement( - self, - *, - virtual_environment_name: str, - node_type: str, - refs: tuple[VirtualEnvironmentNodeRefRecord, ...], - ) -> None: - seen_node_names: set[str] = set() - ref: VirtualEnvironmentNodeRefRecord - for ref in refs: - if ref.virtual_environment_name != virtual_environment_name: - raise StateBackendConfigError( - "Node ref virtual_environment_name must match replacement " - "virtual_environment_name" - ) - if ref.node_type != node_type: - raise StateBackendConfigError("Node ref node_type must match replacement node_type") - if ref.node_name in seen_node_names: - raise StateBackendConfigError( - f"Duplicate node ref for node type '{node_type}' and name '{ref.node_name}'" - ) - seen_node_names.add(ref.node_name) - def _replace_virtual_environment_node_ref_groups( self, *, - cursor: Any, + executor: Any, schema: str, virtual_environment_name: str, refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], @@ -1317,7 +1124,7 @@ def _replace_virtual_environment_node_ref_groups( node_type=node_type, refs=refs, ) - cursor.execute( + executor.execute( "DELETE FROM " f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " "WHERE virtual_environment_name = %s AND node_type = %s", @@ -1325,7 +1132,7 @@ def _replace_virtual_environment_node_ref_groups( ) ref: VirtualEnvironmentNodeRefRecord for ref in refs: - cursor.execute( + executor.execute( "INSERT INTO " + self._qualified_name( schema=schema, From 835da6c0903de4f19427c51d03be7ba4ad01ed7f Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:38:14 +0100 Subject: [PATCH 16/57] refactor(state): share checkpoint writes and conditional publication Move checkpoint creation and deletion, checkpoint row inserts, and the lease-fenced conditional environment publication into SqlStateBackend, and delete the DuckDB conditional-publish mixin that held the DuckDB copies. Statement order, BEGIN/COMMIT/ROLLBACK placement (including the ROLLBACK and False return when a lease is not held), and the inserted rows are unchanged. The lease fence stays per backend behind _lease_is_owned: DuckDB bumps updated_at with UPDATE ... RETURNING, while Postgres takes a row lock with SELECT ... FOR UPDATE. (cherry picked from commit 656e12058b102e875342b3098069f5b0b760f3f7) --- .../classes/_duckdb_conditional_publish.py | 140 ------------ .../state/classes/_sql_state_backend.py | 182 ++++++++++++++- src/sqlbuild/virtual/state/classes/duckdb.py | 102 ++------- .../virtual/state/classes/postgres.py | 213 ++---------------- 4 files changed, 222 insertions(+), 415 deletions(-) delete mode 100644 src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py diff --git a/src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py b/src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py deleted file mode 100644 index 6a9f220ad6..0000000000 --- a/src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py +++ /dev/null @@ -1,140 +0,0 @@ -"""Conditional DuckDB virtual-environment publication.""" - -from __future__ import annotations - -from typing import Any, cast - -from sqlbuild.virtual.state._helpers.state_storage.validation import ( - validate_conditional_virtual_environment_publication, -) -from sqlbuild.virtual.state.constants import ( - LOCK_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE, -) -from sqlbuild.virtual.state.models import ( - StateLockLease, - VirtualEnvironmentCheckpointFunctionRefRecord, - VirtualEnvironmentCheckpointModelRefRecord, - VirtualEnvironmentCheckpointRecord, - VirtualEnvironmentCheckpointSeedRefRecord, - VirtualEnvironmentNodeRefRecord, - VirtualEnvironmentRecord, -) - - -class DuckDbConditionalPublishMixin: - """Publish virtual refs in the same transaction that validates lease ownership.""" - - def upsert_virtual_environment_and_replace_node_ref_groups_if_locks_owned( - self, - *, - connection: Any, - schema: str, - record: VirtualEnvironmentRecord, - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], - leases: tuple[StateLockLease, ...], - checkpoint: VirtualEnvironmentCheckpointRecord | None = None, - checkpoint_refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...] = (), - checkpoint_function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...] = (), - checkpoint_seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...] = (), - ) -> bool: - backend: Any = cast(Any, self) - validate_conditional_virtual_environment_publication( - record=record, - refs_by_node_type=refs_by_node_type, - checkpoint=checkpoint, - checkpoint_refs=checkpoint_refs, - checkpoint_function_refs=checkpoint_function_refs, - checkpoint_seed_refs=checkpoint_seed_refs, - ) - connection.execute("BEGIN") - try: - lock_table: str = backend._qualified_name(schema=schema, table=LOCK_TABLE) - lease: StateLockLease - for lease in leases: - owned: tuple[Any, ...] | None = connection.execute( - f"UPDATE {lock_table} SET updated_at = CURRENT_TIMESTAMP " - "WHERE lock_key = ? AND owner_id = ? AND expires_at > CURRENT_TIMESTAMP " - "RETURNING lock_key", - [lease.lock_key, lease.owner_id], - ).fetchone() - if owned is None: - connection.execute("ROLLBACK") - return False - backend._upsert_virtual_environment_record( - executor=connection, schema=schema, record=record - ) - backend._replace_virtual_environment_node_ref_groups( - executor=connection, - schema=schema, - virtual_environment_name=record.virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - if checkpoint is not None: - self._insert_virtual_environment_checkpoint_rows( - connection=connection, - schema=schema, - checkpoint=checkpoint, - refs=checkpoint_refs, - function_refs=checkpoint_function_refs, - seed_refs=checkpoint_seed_refs, - ) - connection.execute("COMMIT") - return True - except BaseException: - connection.execute("ROLLBACK") - raise - - def _insert_virtual_environment_checkpoint_rows( - self, - *, - connection: Any, - schema: str, - checkpoint: VirtualEnvironmentCheckpointRecord, - refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...], - function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...], - seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...], - ) -> None: - backend: Any = cast(Any, self) - checkpoint_table: str = backend._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE - ) - model_ref_table: str = backend._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE - ) - function_ref_table: str = backend._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE - ) - seed_ref_table: str = backend._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE - ) - connection.execute( - f"INSERT INTO {checkpoint_table} " - "(checkpoint_id, virtual_environment_name, created_at) " - "VALUES (?, ?, CURRENT_TIMESTAMP)", - [checkpoint.checkpoint_id, checkpoint.virtual_environment_name], - ) - ref: VirtualEnvironmentCheckpointModelRefRecord - for ref in refs: - connection.execute( - f"INSERT INTO {model_ref_table} " - "(checkpoint_id, model_name, version_hash) VALUES (?, ?, ?)", - [ref.checkpoint_id, ref.model_name, ref.version_hash], - ) - function_ref: VirtualEnvironmentCheckpointFunctionRefRecord - for function_ref in function_refs: - connection.execute( - f"INSERT INTO {function_ref_table} " - "(checkpoint_id, function_name, version_hash) VALUES (?, ?, ?)", - [function_ref.checkpoint_id, function_ref.function_name, function_ref.version_hash], - ) - seed_ref: VirtualEnvironmentCheckpointSeedRefRecord - for seed_ref in seed_refs: - connection.execute( - f"INSERT INTO {seed_ref_table} " - "(checkpoint_id, seed_name, version_hash) VALUES (?, ?, ?)", - [seed_ref.checkpoint_id, seed_ref.seed_name, seed_ref.version_hash], - ) diff --git a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py index 560ae7fca9..f8b732414c 100644 --- a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py +++ b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py @@ -12,6 +12,9 @@ from sqlbuild.virtual.state._helpers.state_storage.datetime import ( from_naive_utc_wall_clock, ) +from sqlbuild.virtual.state._helpers.state_storage.validation import ( + validate_conditional_virtual_environment_publication, +) from sqlbuild.virtual.state.classes.state_backend import StateBackend from sqlbuild.virtual.state.constants import ( FUNCTION_VERSION_TABLE, @@ -39,6 +42,7 @@ PythonNodeVersionRecord, SeedVersionRecord, SourceFreshnessRecord, + StateLockLease, StateLockRecord, StateOperationRecord, VirtualEnvironmentCheckpointFunctionRefRecord, @@ -87,8 +91,20 @@ def _write_transaction(self, *, connection: Any) -> AbstractContextManager[Any]: ... @abstractmethod - def _execute_in(self, *, executor: Any, sql: str, params: Sequence[object]) -> None: - """Run one statement on a transaction executor.""" + def _statement_executor(self, *, connection: Any) -> AbstractContextManager[Any]: + """Open an executor for explicitly delimited statements on a connection.""" + ... + + @abstractmethod + def _execute_in( + self, *, executor: Any, sql: str, params: Sequence[object] | None = None + ) -> None: + """Run one statement on an executor.""" + ... + + @abstractmethod + def _lease_is_owned(self, *, executor: Any, schema: str, lease: StateLockLease) -> bool: + """Fence one lock lease inside an open transaction and report whether it is held.""" ... @abstractmethod @@ -501,6 +517,168 @@ def _validate_node_ref_replacement( ) seen_node_names.add(ref.node_name) + def upsert_virtual_environment_and_replace_node_ref_groups_if_locks_owned( + self, + *, + connection: Any, + schema: str, + record: VirtualEnvironmentRecord, + refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], + leases: tuple[StateLockLease, ...], + checkpoint: VirtualEnvironmentCheckpointRecord | None = None, + checkpoint_refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...] = (), + checkpoint_function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...] = (), + checkpoint_seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...] = (), + ) -> bool: + validate_conditional_virtual_environment_publication( + record=record, + refs_by_node_type=refs_by_node_type, + checkpoint=checkpoint, + checkpoint_refs=checkpoint_refs, + checkpoint_function_refs=checkpoint_function_refs, + checkpoint_seed_refs=checkpoint_seed_refs, + ) + with self._statement_executor(connection=connection) as executor: + self._execute_in(executor=executor, sql="BEGIN") + try: + lease: StateLockLease + for lease in leases: + if not self._lease_is_owned(executor=executor, schema=schema, lease=lease): + self._execute_in(executor=executor, sql="ROLLBACK") + return False + self._upsert_virtual_environment_record( + executor=executor, schema=schema, record=record + ) + self._replace_virtual_environment_node_ref_groups( + executor=executor, + schema=schema, + virtual_environment_name=record.virtual_environment_name, + refs_by_node_type=refs_by_node_type, + ) + if checkpoint is not None: + self._insert_virtual_environment_checkpoint_rows( + executor=executor, + schema=schema, + checkpoint=checkpoint, + refs=checkpoint_refs, + function_refs=checkpoint_function_refs, + seed_refs=checkpoint_seed_refs, + ) + self._execute_in(executor=executor, sql="COMMIT") + return True + except BaseException: + self._execute_in(executor=executor, sql="ROLLBACK") + raise + + def create_virtual_environment_checkpoint( + self, + *, + connection: Any, + schema: str, + checkpoint: VirtualEnvironmentCheckpointRecord, + refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...], + function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...] = (), + seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...] = (), + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._insert_virtual_environment_checkpoint_rows( + executor=executor, + schema=schema, + checkpoint=checkpoint, + refs=refs, + function_refs=function_refs, + seed_refs=seed_refs, + ) + + def delete_virtual_environment_checkpoint( + self, *, connection: Any, schema: str, checkpoint_id: str + ) -> None: + p: str = self._placeholder + table_names: tuple[str, ...] = ( + VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, + VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, + VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, + VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE, + ) + with self._write_transaction(connection=connection) as executor: + table_name: str + for table_name in table_names: + self._execute_in( + executor=executor, + sql=( + f"DELETE FROM {self._qualified_name(schema=schema, table=table_name)} " + f"WHERE checkpoint_id = {p}" + ), + params=[checkpoint_id], + ) + + def _insert_virtual_environment_checkpoint_rows( + self, + *, + executor: Any, + schema: str, + checkpoint: VirtualEnvironmentCheckpointRecord, + refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...], + function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...], + seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...], + ) -> None: + p: str = self._placeholder + checkpoint_table: str = self._qualified_name( + schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE + ) + model_ref_table: str = self._qualified_name( + schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE + ) + function_ref_table: str = self._qualified_name( + schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE + ) + seed_ref_table: str = self._qualified_name( + schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE + ) + self._execute_in( + executor=executor, + sql=( + f"INSERT INTO {checkpoint_table} " + "(checkpoint_id, virtual_environment_name, created_at) " + f"VALUES ({p}, {p}, CURRENT_TIMESTAMP)" + ), + params=[checkpoint.checkpoint_id, checkpoint.virtual_environment_name], + ) + ref: VirtualEnvironmentCheckpointModelRefRecord + for ref in refs: + self._execute_in( + executor=executor, + sql=( + f"INSERT INTO {model_ref_table} " + f"(checkpoint_id, model_name, version_hash) VALUES ({p}, {p}, {p})" + ), + params=[ref.checkpoint_id, ref.model_name, ref.version_hash], + ) + function_ref: VirtualEnvironmentCheckpointFunctionRefRecord + for function_ref in function_refs: + self._execute_in( + executor=executor, + sql=( + f"INSERT INTO {function_ref_table} " + f"(checkpoint_id, function_name, version_hash) VALUES ({p}, {p}, {p})" + ), + params=[ + function_ref.checkpoint_id, + function_ref.function_name, + function_ref.version_hash, + ], + ) + seed_ref: VirtualEnvironmentCheckpointSeedRefRecord + for seed_ref in seed_refs: + self._execute_in( + executor=executor, + sql=( + f"INSERT INTO {seed_ref_table} " + f"(checkpoint_id, seed_name, version_hash) VALUES ({p}, {p}, {p})" + ), + params=[seed_ref.checkpoint_id, seed_ref.seed_name, seed_ref.version_hash], + ) + def _replace_row_preserving_created_at( self, *, diff --git a/src/sqlbuild/virtual/state/classes/duckdb.py b/src/sqlbuild/virtual/state/classes/duckdb.py index 1d4e1d8ce6..53bd79d671 100644 --- a/src/sqlbuild/virtual/state/classes/duckdb.py +++ b/src/sqlbuild/virtual/state/classes/duckdb.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Iterator, Sequence -from contextlib import contextmanager +from contextlib import AbstractContextManager, contextmanager, nullcontext from datetime import datetime from typing import Any, ClassVar @@ -28,9 +28,6 @@ read_duckdb_microbatch_scope_history, ) from sqlbuild.virtual.state._helpers.state_storage.validation import build_validation_result -from sqlbuild.virtual.state.classes._duckdb_conditional_publish import ( - DuckDbConditionalPublishMixin, -) from sqlbuild.virtual.state.classes._sql_state_backend import SqlStateBackend from sqlbuild.virtual.state.constants import ( CURRENT_STATE_SCHEMA_VERSION, @@ -51,10 +48,6 @@ STATE_TABLE_INDEXES, STATE_TABLES, STATE_VERSION_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE, VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, VIRTUAL_ENVIRONMENT_TABLE, ) @@ -67,12 +60,9 @@ ReconcileEventRecord, SourceFreshnessRecord, StateBackupRecord, + StateLockLease, StateOperationEventRecord, StateSchemaValidationResult, - VirtualEnvironmentCheckpointFunctionRefRecord, - VirtualEnvironmentCheckpointModelRefRecord, - VirtualEnvironmentCheckpointRecord, - VirtualEnvironmentCheckpointSeedRefRecord, VirtualEnvironmentNodeRefRecord, ) from sqlbuild.virtual.state.types import ( @@ -82,7 +72,7 @@ ) -class DuckDbStateBackend(DuckDbConditionalPublishMixin, SqlStateBackend): +class DuckDbStateBackend(SqlStateBackend): """DuckDB implementation for virtual state.""" _placeholder: ClassVar[str] = "?" @@ -122,9 +112,27 @@ def _write_transaction(self, *, connection: Any) -> Iterator[Any]: connection.execute("ROLLBACK") raise - def _execute_in(self, *, executor: Any, sql: str, params: Sequence[object]) -> None: + def _statement_executor(self, *, connection: Any) -> AbstractContextManager[Any]: + return nullcontext(connection) + + def _execute_in( + self, *, executor: Any, sql: str, params: Sequence[object] | None = None + ) -> None: + if params is None: + executor.execute(sql) + return executor.execute(sql, params) + def _lease_is_owned(self, *, executor: Any, schema: str, lease: StateLockLease) -> bool: + owned: tuple[Any, ...] | None = executor.execute( + f"UPDATE {self._qualified_name(schema=schema, table=LOCK_TABLE)} " + "SET updated_at = CURRENT_TIMESTAMP " + "WHERE lock_key = ? AND owner_id = ? AND expires_at > CURRENT_TIMESTAMP " + "RETURNING lock_key", + [lease.lock_key, lease.owner_id], + ).fetchone() + return owned is not None + def _fetch_one_in( self, *, executor: Any, sql: str, params: Sequence[object] ) -> tuple[Any, ...] | None: @@ -565,72 +573,6 @@ def replace_virtual_environment_source_freshness( connection.execute("ROLLBACK") raise - def create_virtual_environment_checkpoint( - self, - *, - connection: Any, - schema: str, - checkpoint: VirtualEnvironmentCheckpointRecord, - refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...], - function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...] = (), - seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...] = (), - ) -> None: - connection.execute("BEGIN") - try: - self._insert_virtual_environment_checkpoint_rows( - connection=connection, - schema=schema, - checkpoint=checkpoint, - refs=refs, - function_refs=function_refs, - seed_refs=seed_refs, - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def delete_virtual_environment_checkpoint( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> None: - connection.execute("BEGIN") - try: - checkpoint_function_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, - ) - checkpoint_seed_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, - ) - checkpoint_model_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, - ) - connection.execute( - f"DELETE FROM {checkpoint_seed_ref_table} WHERE checkpoint_id = ?", - [checkpoint_id], - ) - connection.execute( - f"DELETE FROM {checkpoint_function_ref_table} WHERE checkpoint_id = ?", - [checkpoint_id], - ) - connection.execute( - f"DELETE FROM {checkpoint_model_ref_table} WHERE checkpoint_id = ?", - [checkpoint_id], - ) - connection.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE)}" - " " - "WHERE checkpoint_id = ?", - [checkpoint_id], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - def create_state_operation_event( self, *, connection: Any, schema: str, record: StateOperationEventRecord ) -> None: diff --git a/src/sqlbuild/virtual/state/classes/postgres.py b/src/sqlbuild/virtual/state/classes/postgres.py index 4e41cfacd0..e3e7bd9bbf 100644 --- a/src/sqlbuild/virtual/state/classes/postgres.py +++ b/src/sqlbuild/virtual/state/classes/postgres.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Iterator, Sequence -from contextlib import contextmanager +from contextlib import AbstractContextManager, contextmanager from datetime import datetime from typing import Any, ClassVar @@ -28,7 +28,6 @@ from sqlbuild.virtual.state._helpers.state_storage.events import backup_id, event_id from sqlbuild.virtual.state._helpers.state_storage.validation import ( build_validation_result, - validate_conditional_virtual_environment_publication, ) from sqlbuild.virtual.state.classes._sql_state_backend import SqlStateBackend from sqlbuild.virtual.state.constants import ( @@ -49,10 +48,6 @@ STATE_TABLE_INDEXES, STATE_TABLES, STATE_VERSION_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE, VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, VIRTUAL_ENVIRONMENT_TABLE, ) @@ -68,12 +63,7 @@ StateLockLease, StateOperationEventRecord, StateSchemaValidationResult, - VirtualEnvironmentCheckpointFunctionRefRecord, - VirtualEnvironmentCheckpointModelRefRecord, - VirtualEnvironmentCheckpointRecord, - VirtualEnvironmentCheckpointSeedRefRecord, VirtualEnvironmentNodeRefRecord, - VirtualEnvironmentRecord, ) from sqlbuild.virtual.state.types import ( StateColumnType, @@ -145,9 +135,27 @@ def _write_transaction(self, *, connection: Any) -> Iterator[Any]: cursor.execute("ROLLBACK") raise - def _execute_in(self, *, executor: Any, sql: str, params: Sequence[object]) -> None: + def _statement_executor(self, *, connection: Any) -> AbstractContextManager[Any]: + return connection.cursor() + + def _execute_in( + self, *, executor: Any, sql: str, params: Sequence[object] | None = None + ) -> None: + if params is None: + executor.execute(sql) + return executor.execute(sql, params) + def _lease_is_owned(self, *, executor: Any, schema: str, lease: StateLockLease) -> bool: + lock_table: str = self._qualified_name(schema=schema, table=LOCK_TABLE) + executor.execute( + f"SELECT lock_key FROM {lock_table} " + "WHERE lock_key = %s AND owner_id = %s " + "AND expires_at > CURRENT_TIMESTAMP FOR UPDATE", + [lease.lock_key, lease.owner_id], + ) + return executor.fetchone() is not None + def _fetch_one_in( self, *, executor: Any, sql: str, params: Sequence[object] ) -> tuple[Any, ...] | None: @@ -544,66 +552,6 @@ def delete_virtual_environment( cursor.execute("ROLLBACK") raise - def upsert_virtual_environment_and_replace_node_ref_groups_if_locks_owned( - self, - *, - connection: Any, - schema: str, - record: VirtualEnvironmentRecord, - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], - leases: tuple[StateLockLease, ...], - checkpoint: VirtualEnvironmentCheckpointRecord | None = None, - checkpoint_refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...] = (), - checkpoint_function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...] = (), - checkpoint_seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...] = (), - ) -> bool: - validate_conditional_virtual_environment_publication( - record=record, - refs_by_node_type=refs_by_node_type, - checkpoint=checkpoint, - checkpoint_refs=checkpoint_refs, - checkpoint_function_refs=checkpoint_function_refs, - checkpoint_seed_refs=checkpoint_seed_refs, - ) - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - lease: StateLockLease - for lease in leases: - lock_table: str = self._qualified_name(schema=schema, table=LOCK_TABLE) - cursor.execute( - f"SELECT lock_key FROM {lock_table} " - "WHERE lock_key = %s AND owner_id = %s " - "AND expires_at > CURRENT_TIMESTAMP FOR UPDATE", - [lease.lock_key, lease.owner_id], - ) - if cursor.fetchone() is None: - cursor.execute("ROLLBACK") - return False - self._upsert_virtual_environment_record( - executor=cursor, schema=schema, record=record - ) - self._replace_virtual_environment_node_ref_groups( - executor=cursor, - schema=schema, - virtual_environment_name=record.virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - if checkpoint is not None: - self._insert_virtual_environment_checkpoint_rows( - cursor=cursor, - schema=schema, - checkpoint=checkpoint, - refs=checkpoint_refs, - function_refs=checkpoint_function_refs, - seed_refs=checkpoint_seed_refs, - ) - cursor.execute("COMMIT") - return True - except BaseException: - cursor.execute("ROLLBACK") - raise - def upsert_virtual_environment_node_ref( self, *, @@ -713,127 +661,6 @@ def replace_virtual_environment_source_freshness( cursor.execute("ROLLBACK") raise - def create_virtual_environment_checkpoint( - self, - *, - connection: Any, - schema: str, - checkpoint: VirtualEnvironmentCheckpointRecord, - refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...], - function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...] = (), - seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...] = (), - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - self._insert_virtual_environment_checkpoint_rows( - cursor=cursor, - schema=schema, - checkpoint=checkpoint, - refs=refs, - function_refs=function_refs, - seed_refs=seed_refs, - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def _insert_virtual_environment_checkpoint_rows( - self, - *, - cursor: Any, - schema: str, - checkpoint: VirtualEnvironmentCheckpointRecord, - refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...], - function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...], - seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...], - ) -> None: - checkpoint_table: str = self._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE - ) - model_ref_table: str = self._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE - ) - function_ref_table: str = self._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE - ) - seed_ref_table: str = self._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE - ) - cursor.execute( - f"INSERT INTO {checkpoint_table} " - "(checkpoint_id, virtual_environment_name, created_at) " - "VALUES (%s, %s, CURRENT_TIMESTAMP)", - [checkpoint.checkpoint_id, checkpoint.virtual_environment_name], - ) - ref: VirtualEnvironmentCheckpointModelRefRecord - for ref in refs: - cursor.execute( - f"INSERT INTO {model_ref_table} " - "(checkpoint_id, model_name, version_hash) VALUES (%s, %s, %s)", - [ref.checkpoint_id, ref.model_name, ref.version_hash], - ) - function_ref: VirtualEnvironmentCheckpointFunctionRefRecord - for function_ref in function_refs: - cursor.execute( - f"INSERT INTO {function_ref_table} " - "(checkpoint_id, function_name, version_hash) VALUES (%s, %s, %s)", - [function_ref.checkpoint_id, function_ref.function_name, function_ref.version_hash], - ) - seed_ref: VirtualEnvironmentCheckpointSeedRefRecord - for seed_ref in seed_refs: - cursor.execute( - f"INSERT INTO {seed_ref_table} " - "(checkpoint_id, seed_name, version_hash) VALUES (%s, %s, %s)", - [seed_ref.checkpoint_id, seed_ref.seed_name, seed_ref.version_hash], - ) - - def delete_virtual_environment_checkpoint( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - checkpoint_function_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, - ) - checkpoint_model_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, - ) - checkpoint_seed_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, - ) - cursor.execute( - f"DELETE FROM {checkpoint_seed_ref_table} WHERE checkpoint_id = %s", - [checkpoint_id], - ) - cursor.execute( - f"DELETE FROM {checkpoint_function_ref_table} WHERE checkpoint_id = %s", - [checkpoint_id], - ) - cursor.execute( - f"DELETE FROM {checkpoint_model_ref_table} WHERE checkpoint_id = %s", - [checkpoint_id], - ) - cursor.execute( - "DELETE FROM " - + self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE, - ) - + " " - "WHERE checkpoint_id = %s", - [checkpoint_id], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - def create_state_operation_event( self, *, connection: Any, schema: str, record: StateOperationEventRecord ) -> None: From 99c7ce6eb43bce05057a49ebc90386471dccedb2 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:40:15 +0100 Subject: [PATCH 17/57] refactor(state): share node result and event state SQL Move node result insert and read, the node result row mapping, the optional-equality and materialized storage helpers, and the state operation and reconcile event inserts into SqlStateBackend. SQL, parameters, ordering, and decoding were identical in both backends; the backends now only supply how a single autocommit statement runs. executor/node_results/main/_read.py keeps its own row mapping: it reads the direct-mode node results table, whose row layout (target columns, column positions, ISO timestamp fallback) differs from the virtual state node results query. (cherry picked from commit fe2de36514f8d779a1094c7088a0967d1222f3d3) --- .../state/classes/_sql_state_backend.py | 174 ++++++++++++++++++ src/sqlbuild/virtual/state/classes/duckdb.py | 158 ---------------- .../virtual/state/classes/postgres.py | 165 ----------------- 3 files changed, 174 insertions(+), 323 deletions(-) diff --git a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py index f8b732414c..584965fc58 100644 --- a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py +++ b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py @@ -9,6 +9,13 @@ from typing import Any, ClassVar from sqlbuild.compiler.compile.types import CompiledResourceType +from sqlbuild.executor.node_results.main.decode_json import decode_node_result_json +from sqlbuild.executor.node_results.main.encode_json import encode_node_result_json +from sqlbuild.executor.node_results.models import ( + NodeResultEnvelope, + NodeResultQuery, + NodeResultRecord, +) from sqlbuild.virtual.state._helpers.state_storage.datetime import ( from_naive_utc_wall_clock, ) @@ -20,11 +27,15 @@ FUNCTION_VERSION_TABLE, LOCK_TABLE, MODEL_VERSION_TABLE, + NODE_RESULTS_TABLE, PHYSICAL_RELATION_ANCESTRY_TABLE, PHYSICAL_RELATION_TABLE, PYTHON_NODE_VERSION_TABLE, + RECONCILE_EVENT_TABLE, SEED_VERSION_TABLE, SOURCE_FRESHNESS_OBSERVATION_TABLE, + STATE_BOOLEAN_TRUE, + STATE_OPERATION_EVENT_TABLE, STATE_OPERATION_TABLE, VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, @@ -40,10 +51,12 @@ PhysicalRelationAncestryRecord, PhysicalRelationRecord, PythonNodeVersionRecord, + ReconcileEventRecord, SeedVersionRecord, SourceFreshnessRecord, StateLockLease, StateLockRecord, + StateOperationEventRecord, StateOperationRecord, VirtualEnvironmentCheckpointFunctionRefRecord, VirtualEnvironmentCheckpointModelRefRecord, @@ -679,6 +692,167 @@ def _insert_virtual_environment_checkpoint_rows( params=[seed_ref.checkpoint_id, seed_ref.seed_name, seed_ref.version_hash], ) + def insert_node_result( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + record: NodeResultRecord, + ) -> None: + p: str = self._placeholder + self._execute( + connection=connection, + sql=f"INSERT INTO {self._qualified_name(schema=schema, table=NODE_RESULTS_TABLE)} " + "(virtual_environment_name, node_type, node_name, target_database, target_schema, " + "target_name, run_id, status, payload_json_b64, metadata_json_b64, error_message, " + "materialized, created_at) " + f"VALUES ({p}, {p}, {p}, {p}, {p}, {p}, {p}, {p}, {p}, {p}, {p}, {p}, {p})", + params=[ + virtual_environment_name, + record.node_type, + record.node_name, + record.target_database, + record.target_schema, + record.target_name, + record.run_id, + record.status, + encode_node_result_json( + value=record.payload, label="payload", node_name=record.node_name + ), + encode_node_result_json( + value=record.metadata, label="metadata", node_name=record.node_name + ), + record.error_message, + self._materialized_storage(record.materialized), + record.ts, + ], + ) + + def read_node_results( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + query: NodeResultQuery, + ) -> tuple[NodeResultEnvelope, ...]: + p: str = self._placeholder + if query.limit < 1: + return () + predicates: list[str] = [ + f"virtual_environment_name = {p}", + f"node_type = {p}", + f"node_name = {p}", + self._optional_equality_sql( + column="target_database", value=query.target_database, placeholder=f"{p}" + ), + self._optional_equality_sql( + column="target_schema", value=query.target_schema, placeholder=f"{p}" + ), + self._optional_equality_sql( + column="target_name", value=query.target_name, placeholder=f"{p}" + ), + ] + params: list[object] = [virtual_environment_name, query.node_type, query.node_name] + for value in (query.target_database, query.target_schema, query.target_name): + if value is not None: + params.append(value) + if query.statuses is not None: + placeholders: str = ", ".join(f"{p}" for _ in query.statuses) + predicates.append(f"status IN ({placeholders})") + params.extend(query.statuses) + if query.run_id is not None: + predicates.append(f"run_id = {p}") + params.append(query.run_id) + params.append(query.limit) + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT node_type, node_name, run_id, status, payload_json_b64, metadata_json_b64, " + "error_message, materialized, created_at " + f"FROM {self._qualified_name(schema=schema, table=NODE_RESULTS_TABLE)} " + f"WHERE {' AND '.join(predicates)} " + f"ORDER BY created_at DESC, run_id DESC LIMIT {p}", + params=params, + ) + return tuple(self._node_result_row_to_envelope(row) for row in rows) + + def create_state_operation_event( + self, *, connection: Any, schema: str, record: StateOperationEventRecord + ) -> None: + p: str = self._placeholder + self._execute( + connection=connection, + sql="INSERT INTO " + f"{self._qualified_name(schema=schema, table=STATE_OPERATION_EVENT_TABLE)} " + "(event_id, operation_id, action, status, message, created_at) " + f"VALUES ({p}, {p}, {p}, {p}, {p}, CURRENT_TIMESTAMP)", + params=[ + record.event_id, + record.operation_id, + record.action, + record.status.value, + record.message, + ], + ) + + def create_reconcile_event( + self, *, connection: Any, schema: str, record: ReconcileEventRecord + ) -> None: + p: str = self._placeholder + self._execute( + connection=connection, + sql=f"INSERT INTO {self._qualified_name(schema=schema, table=RECONCILE_EVENT_TABLE)} " + "(event_id, action, status, message, created_at) " + f"VALUES ({p}, {p}, {p}, {p}, CURRENT_TIMESTAMP)", + params=[record.event_id, record.action.value, record.status.value, record.message], + ) + + def _node_result_row_to_envelope(self, row: tuple[Any, ...]) -> NodeResultEnvelope: + node_name: str = str(row[1]) + metadata: object = decode_node_result_json( + value=str(row[5]), label="metadata", node_name=node_name + ) + normalized_metadata: dict[str, object] = ( + {str(key): value for key, value in metadata.items()} + if isinstance(metadata, dict) + else {} + ) + return NodeResultEnvelope( + node_type=str(row[0]), + node_name=node_name, + run_id=str(row[2]), + status=str(row[3]), + payload=decode_node_result_json( + value=str(row[4]), label="payload", node_name=node_name + ), + metadata=normalized_metadata, + error_message=str(row[6]) if row[6] is not None else None, + materialized=self._parse_materialized(row[7]), + ts=row[8], + ) + + def _optional_equality_sql(self, *, column: str, value: object | None, placeholder: str) -> str: + if value is None: + return f"{column} IS NULL" + return f"{column} = {placeholder}" + + def _materialized_storage(self, value: bool | None) -> str | None: + if value is None: + return None + return "true" if value else "false" + + def _parse_materialized(self, value: object) -> bool | None: + if value is None: + return None + return str(value).lower() == STATE_BOOLEAN_TRUE + + def _execute( + self, *, connection: Any, sql: str, params: Sequence[object] | None = None + ) -> None: + with self._statement_executor(connection=connection) as executor: + self._execute_in(executor=executor, sql=sql, params=params) + def _replace_row_preserving_created_at( self, *, diff --git a/src/sqlbuild/virtual/state/classes/duckdb.py b/src/sqlbuild/virtual/state/classes/duckdb.py index 53bd79d671..6e4b272b23 100644 --- a/src/sqlbuild/virtual/state/classes/duckdb.py +++ b/src/sqlbuild/virtual/state/classes/duckdb.py @@ -8,13 +8,6 @@ from typing import Any, ClassVar from sqlbuild.adapter.contract.classes.observed_connection import ObservedConnection -from sqlbuild.executor.node_results.main.decode_json import decode_node_result_json -from sqlbuild.executor.node_results.main.encode_json import encode_node_result_json -from sqlbuild.executor.node_results.models import ( - NodeResultEnvelope, - NodeResultQuery, - NodeResultRecord, -) from sqlbuild.microbatches.models import MicrobatchEvent, MicrobatchScope, MicrobatchWriteResult from sqlbuild.virtual.state._helpers.state_storage.datetime import ( to_naive_utc_wall_clock, @@ -36,14 +29,10 @@ DUCKDB_TIMESTAMP_TYPE_TOKEN, LOCK_TABLE, MICROBATCH_EVENT_TABLE, - NODE_RESULTS_TABLE, NON_UNIQUE_STATE_INDEXES, PYTHON_NODE_VERSION_TABLE, - RECONCILE_EVENT_TABLE, SOURCE_FRESHNESS_OBSERVATION_TABLE, - STATE_BOOLEAN_TRUE, STATE_MIGRATION_EVENTS_TABLE, - STATE_OPERATION_EVENT_TABLE, STATE_TABLE_COLUMNS, STATE_TABLE_INDEXES, STATE_TABLES, @@ -57,11 +46,9 @@ StateSchemaInvalidError, ) from sqlbuild.virtual.state.models import ( - ReconcileEventRecord, SourceFreshnessRecord, StateBackupRecord, StateLockLease, - StateOperationEventRecord, StateSchemaValidationResult, VirtualEnvironmentNodeRefRecord, ) @@ -298,86 +285,6 @@ def reset(self, *, connection: Any, schema: str) -> None: connection.execute("ROLLBACK") raise - def insert_node_result( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - record: NodeResultRecord, - ) -> None: - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=NODE_RESULTS_TABLE)} " - "(virtual_environment_name, node_type, node_name, target_database, target_schema, " - "target_name, run_id, status, payload_json_b64, metadata_json_b64, error_message, " - "materialized, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - [ - virtual_environment_name, - record.node_type, - record.node_name, - record.target_database, - record.target_schema, - record.target_name, - record.run_id, - record.status, - encode_node_result_json( - value=record.payload, label="payload", node_name=record.node_name - ), - encode_node_result_json( - value=record.metadata, label="metadata", node_name=record.node_name - ), - record.error_message, - self._materialized_storage(record.materialized), - record.ts, - ], - ) - - def read_node_results( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - query: NodeResultQuery, - ) -> tuple[NodeResultEnvelope, ...]: - if query.limit < 1: - return () - predicates: list[str] = [ - "virtual_environment_name = ?", - "node_type = ?", - "node_name = ?", - self._optional_equality_sql( - column="target_database", value=query.target_database, placeholder="?" - ), - self._optional_equality_sql( - column="target_schema", value=query.target_schema, placeholder="?" - ), - self._optional_equality_sql( - column="target_name", value=query.target_name, placeholder="?" - ), - ] - params: list[object] = [virtual_environment_name, query.node_type, query.node_name] - for value in (query.target_database, query.target_schema, query.target_name): - if value is not None: - params.append(value) - if query.statuses is not None: - placeholders: str = ", ".join("?" for _ in query.statuses) - predicates.append(f"status IN ({placeholders})") - params.extend(query.statuses) - if query.run_id is not None: - predicates.append("run_id = ?") - params.append(query.run_id) - params.append(query.limit) - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT node_type, node_name, run_id, status, payload_json_b64, metadata_json_b64, " - "error_message, materialized, created_at " - f"FROM {self._qualified_name(schema=schema, table=NODE_RESULTS_TABLE)} " - f"WHERE {' AND '.join(predicates)} " - "ORDER BY created_at DESC, run_id DESC LIMIT ?", - params, - ).fetchall() - return tuple(self._node_result_row_to_envelope(row) for row in rows) - def append_microbatch_event( self, *, connection: Any, schema: str, event: MicrobatchEvent ) -> None: @@ -573,32 +480,6 @@ def replace_virtual_environment_source_freshness( connection.execute("ROLLBACK") raise - def create_state_operation_event( - self, *, connection: Any, schema: str, record: StateOperationEventRecord - ) -> None: - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=STATE_OPERATION_EVENT_TABLE)} " - "(event_id, operation_id, action, status, message, created_at) " - "VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)", - [ - record.event_id, - record.operation_id, - record.action, - record.status.value, - record.message, - ], - ) - - def create_reconcile_event( - self, *, connection: Any, schema: str, record: ReconcileEventRecord - ) -> None: - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=RECONCILE_EVENT_TABLE)} " - "(event_id, action, status, message, created_at) " - "VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP)", - [record.event_id, record.action.value, record.status.value, record.message], - ) - def acquire_lock( self, *, @@ -776,45 +657,6 @@ def _state_column_sql_type(self, column_type: StateColumnType) -> str: return "TIMESTAMP" raise StateBackendConfigError(f"Unsupported state column type: {column_type}") - def _node_result_row_to_envelope(self, row: tuple[Any, ...]) -> NodeResultEnvelope: - node_name: str = str(row[1]) - metadata: object = decode_node_result_json( - value=str(row[5]), label="metadata", node_name=node_name - ) - normalized_metadata: dict[str, object] = ( - {str(key): value for key, value in metadata.items()} - if isinstance(metadata, dict) - else {} - ) - return NodeResultEnvelope( - node_type=str(row[0]), - node_name=node_name, - run_id=str(row[2]), - status=str(row[3]), - payload=decode_node_result_json( - value=str(row[4]), label="payload", node_name=node_name - ), - metadata=normalized_metadata, - error_message=str(row[6]) if row[6] is not None else None, - materialized=self._parse_materialized(row[7]), - ts=row[8], - ) - - def _optional_equality_sql(self, *, column: str, value: object | None, placeholder: str) -> str: - if value is None: - return f"{column} IS NULL" - return f"{column} = {placeholder}" - - def _materialized_storage(self, value: bool | None) -> str | None: - if value is None: - return None - return "true" if value else "false" - - def _parse_materialized(self, value: object) -> bool | None: - if value is None: - return None - return str(value).lower() == STATE_BOOLEAN_TRUE - def _validate_source_freshness_records( self, *, diff --git a/src/sqlbuild/virtual/state/classes/postgres.py b/src/sqlbuild/virtual/state/classes/postgres.py index e3e7bd9bbf..92dd2df62a 100644 --- a/src/sqlbuild/virtual/state/classes/postgres.py +++ b/src/sqlbuild/virtual/state/classes/postgres.py @@ -8,13 +8,6 @@ from typing import Any, ClassVar from sqlbuild.adapter.contract.classes.observed_connection import ObservedConnection -from sqlbuild.executor.node_results.main.decode_json import decode_node_result_json -from sqlbuild.executor.node_results.main.encode_json import encode_node_result_json -from sqlbuild.executor.node_results.models import ( - NodeResultEnvelope, - NodeResultQuery, - NodeResultRecord, -) from sqlbuild.microbatches.classes.event_codec import MicrobatchEventCodec from sqlbuild.microbatches.constants import ( MICROBATCH_COLUMNS, @@ -34,16 +27,12 @@ CURRENT_STATE_SCHEMA_VERSION, LOCK_TABLE, MICROBATCH_EVENT_TABLE, - NODE_RESULTS_TABLE, NON_UNIQUE_STATE_INDEXES, POSTGRES_INTEGER_TYPES, POSTGRES_TEXT_TYPES, PYTHON_NODE_VERSION_TABLE, - RECONCILE_EVENT_TABLE, SOURCE_FRESHNESS_OBSERVATION_TABLE, - STATE_BOOLEAN_TRUE, STATE_MIGRATION_EVENTS_TABLE, - STATE_OPERATION_EVENT_TABLE, STATE_TABLE_COLUMNS, STATE_TABLE_INDEXES, STATE_TABLES, @@ -57,11 +46,9 @@ StateSchemaInvalidError, ) from sqlbuild.virtual.state.models import ( - ReconcileEventRecord, SourceFreshnessRecord, StateBackupRecord, StateLockLease, - StateOperationEventRecord, StateSchemaValidationResult, VirtualEnvironmentNodeRefRecord, ) @@ -333,90 +320,6 @@ def reset(self, *, connection: Any, schema: str) -> None: cursor.execute("ROLLBACK") raise - def insert_node_result( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - record: NodeResultRecord, - ) -> None: - with connection.cursor() as cursor: - cursor.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=NODE_RESULTS_TABLE)} " - "(virtual_environment_name, node_type, node_name, target_database, target_schema, " - "target_name, run_id, status, payload_json_b64, metadata_json_b64, error_message, " - "materialized, created_at) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, " - "%s, %s)", - [ - virtual_environment_name, - record.node_type, - record.node_name, - record.target_database, - record.target_schema, - record.target_name, - record.run_id, - record.status, - encode_node_result_json( - value=record.payload, label="payload", node_name=record.node_name - ), - encode_node_result_json( - value=record.metadata, label="metadata", node_name=record.node_name - ), - record.error_message, - self._materialized_storage(record.materialized), - record.ts, - ], - ) - - def read_node_results( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - query: NodeResultQuery, - ) -> tuple[NodeResultEnvelope, ...]: - if query.limit < 1: - return () - predicates: list[str] = [ - "virtual_environment_name = %s", - "node_type = %s", - "node_name = %s", - self._optional_equality_sql( - column="target_database", value=query.target_database, placeholder="%s" - ), - self._optional_equality_sql( - column="target_schema", value=query.target_schema, placeholder="%s" - ), - self._optional_equality_sql( - column="target_name", value=query.target_name, placeholder="%s" - ), - ] - params: list[object] = [virtual_environment_name, query.node_type, query.node_name] - for value in (query.target_database, query.target_schema, query.target_name): - if value is not None: - params.append(value) - if query.statuses is not None: - placeholders: str = ", ".join("%s" for _ in query.statuses) - predicates.append(f"status IN ({placeholders})") - params.extend(query.statuses) - if query.run_id is not None: - predicates.append("run_id = %s") - params.append(query.run_id) - params.append(query.limit) - with connection.cursor() as cursor: - cursor.execute( - "SELECT node_type, node_name, run_id, status, payload_json_b64, " - "metadata_json_b64, error_message, materialized, created_at " - f"FROM {self._qualified_name(schema=schema, table=NODE_RESULTS_TABLE)} " - f"WHERE {' AND '.join(predicates)} " - "ORDER BY created_at DESC, run_id DESC LIMIT %s", - params, - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple(self._node_result_row_to_envelope(row) for row in rows) - def append_microbatch_event( self, *, connection: Any, schema: str, event: MicrobatchEvent ) -> None: @@ -661,35 +564,6 @@ def replace_virtual_environment_source_freshness( cursor.execute("ROLLBACK") raise - def create_state_operation_event( - self, *, connection: Any, schema: str, record: StateOperationEventRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=STATE_OPERATION_EVENT_TABLE)} " - "(event_id, operation_id, action, status, message, created_at) " - "VALUES (%s, %s, %s, %s, %s, CURRENT_TIMESTAMP)", - [ - record.event_id, - record.operation_id, - record.action, - record.status.value, - record.message, - ], - ) - - def create_reconcile_event( - self, *, connection: Any, schema: str, record: ReconcileEventRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=RECONCILE_EVENT_TABLE)} " - "(event_id, action, status, message, created_at) " - "VALUES (%s, %s, %s, %s, CURRENT_TIMESTAMP)", - [record.event_id, record.action.value, record.status.value, record.message], - ) - def acquire_lock( self, *, @@ -876,45 +750,6 @@ def _state_column_sql_type(self, column_type: StateColumnType) -> str: return "TIMESTAMP" raise StateBackendConfigError(f"Unsupported state column type: {column_type}") - def _node_result_row_to_envelope(self, row: tuple[Any, ...]) -> NodeResultEnvelope: - node_name: str = str(row[1]) - metadata: object = decode_node_result_json( - value=str(row[5]), label="metadata", node_name=node_name - ) - normalized_metadata: dict[str, object] = ( - {str(key): value for key, value in metadata.items()} - if isinstance(metadata, dict) - else {} - ) - return NodeResultEnvelope( - node_type=str(row[0]), - node_name=node_name, - run_id=str(row[2]), - status=str(row[3]), - payload=decode_node_result_json( - value=str(row[4]), label="payload", node_name=node_name - ), - metadata=normalized_metadata, - error_message=str(row[6]) if row[6] is not None else None, - materialized=self._parse_materialized(row[7]), - ts=row[8], - ) - - def _optional_equality_sql(self, *, column: str, value: object | None, placeholder: str) -> str: - if value is None: - return f"{column} IS NULL" - return f"{column} = {placeholder}" - - def _materialized_storage(self, value: bool | None) -> str | None: - if value is None: - return None - return "true" if value else "false" - - def _parse_materialized(self, value: object) -> bool | None: - if value is None: - return None - return str(value).lower() == STATE_BOOLEAN_TRUE - def _validate_source_freshness_records( self, *, From 9c44bbee7b4ce978bd291d5c9cd827af1dcc37b3 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:43:07 +0100 Subject: [PATCH 18/57] refactor(state): share relational state schema lifecycle Move initialize, create_backup, rollback, reset, table and index creation, migration event recording, backup listing and lookup, schema existence, column type mapping, and backup schema naming into SqlStateBackend. Each backend ran the same statements in the same order inside one BEGIN/COMMIT with ROLLBACK on failure. inspect_schema and type matching stay per backend because they read driver-specific catalog type names; delete_state_backup and delete_lock stay per backend because Postgres commits explicitly afterwards. (cherry picked from commit 22d0a942b3cce5aae844d98fe3eefd49b2cb0acb) --- .../state/classes/_sql_state_backend.py | 271 +++++++++++++++++- src/sqlbuild/virtual/state/classes/duckdb.py | 250 ---------------- .../virtual/state/classes/postgres.py | 266 ----------------- 3 files changed, 270 insertions(+), 517 deletions(-) diff --git a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py index 584965fc58..cd94c59cdf 100644 --- a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py +++ b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py @@ -19,15 +19,18 @@ from sqlbuild.virtual.state._helpers.state_storage.datetime import ( from_naive_utc_wall_clock, ) +from sqlbuild.virtual.state._helpers.state_storage.events import backup_id, event_id from sqlbuild.virtual.state._helpers.state_storage.validation import ( validate_conditional_virtual_environment_publication, ) from sqlbuild.virtual.state.classes.state_backend import StateBackend from sqlbuild.virtual.state.constants import ( + CURRENT_STATE_SCHEMA_VERSION, FUNCTION_VERSION_TABLE, LOCK_TABLE, MODEL_VERSION_TABLE, NODE_RESULTS_TABLE, + NON_UNIQUE_STATE_INDEXES, PHYSICAL_RELATION_ANCESTRY_TABLE, PHYSICAL_RELATION_TABLE, PYTHON_NODE_VERSION_TABLE, @@ -35,8 +38,13 @@ SEED_VERSION_TABLE, SOURCE_FRESHNESS_OBSERVATION_TABLE, STATE_BOOLEAN_TRUE, + STATE_MIGRATION_EVENTS_TABLE, STATE_OPERATION_EVENT_TABLE, STATE_OPERATION_TABLE, + STATE_TABLE_COLUMNS, + STATE_TABLE_INDEXES, + STATE_TABLES, + STATE_VERSION_TABLE, VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, @@ -44,7 +52,11 @@ VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, VIRTUAL_ENVIRONMENT_TABLE, ) -from sqlbuild.virtual.state.exceptions import StateBackendConfigError +from sqlbuild.virtual.state.exceptions import ( + StateBackendConfigError, + StateBackupNotFoundError, + StateSchemaInvalidError, +) from sqlbuild.virtual.state.models import ( FunctionVersionRecord, ModelVersionRecord, @@ -54,10 +66,12 @@ ReconcileEventRecord, SeedVersionRecord, SourceFreshnessRecord, + StateBackupRecord, StateLockLease, StateLockRecord, StateOperationEventRecord, StateOperationRecord, + StateSchemaValidationResult, VirtualEnvironmentCheckpointFunctionRefRecord, VirtualEnvironmentCheckpointModelRefRecord, VirtualEnvironmentCheckpointRecord, @@ -73,6 +87,9 @@ from sqlbuild.virtual.state.types import ( ModelVersionStatus, PhysicalArtifactType, + StateColumnType, + StateMigrationAction, + StateMigrationStatus, StateOperationStatus, StateOperationType, VirtualEnvironmentStatus, @@ -853,6 +870,258 @@ def _execute( with self._statement_executor(connection=connection) as executor: self._execute_in(executor=executor, sql=sql, params=params) + def initialize(self, *, connection: Any, schema: str, sqlbuild_version: str) -> None: + p: str = self._placeholder + with self._write_transaction(connection=connection) as executor: + self._execute_in( + executor=executor, + sql=f"CREATE SCHEMA IF NOT EXISTS {self._quote_identifier(schema)}", + ) + self._execute_in( + executor=executor, + sql="CREATE TABLE IF NOT EXISTS " + f"{self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)} (" + "schema_version INTEGER NOT NULL, " + "sqlbuild_version TEXT NOT NULL, " + "updated_at TIMESTAMP NOT NULL" + ")", + ) + self._execute_in( + executor=executor, + sql="CREATE TABLE IF NOT EXISTS " + f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} (" + "event_id TEXT NOT NULL, " + "action TEXT NOT NULL, " + "backup_id TEXT, " + "status TEXT NOT NULL, " + "message TEXT, " + "created_at TIMESTAMP NOT NULL" + ")", + ) + self._create_additional_state_tables(executor=executor, schema=schema) + self._execute_in( + executor=executor, + sql=f"DELETE FROM {self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)}", + ) + self._execute_in( + executor=executor, + sql=f"INSERT INTO {self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)} " + "(schema_version, sqlbuild_version, updated_at) " + f"VALUES ({p}, {p}, CURRENT_TIMESTAMP)", + params=[CURRENT_STATE_SCHEMA_VERSION, sqlbuild_version], + ) + self._record_event( + executor=executor, + schema=schema, + action=StateMigrationAction.INIT, + backup_id_value=None, + status=StateMigrationStatus.SUCCESS, + message=None, + ) + + def create_backup(self, *, connection: Any, schema: str) -> str: + validation: StateSchemaValidationResult = self.inspect_schema( + connection=connection, schema=schema + ) + if not validation.valid: + raise StateSchemaInvalidError("Cannot backup invalid state schema") + backup_id_value: str = backup_id() + backup_schema: str = self._backup_schema_name( + schema=schema, + backup_id_value=backup_id_value, + ) + with self._write_transaction(connection=connection) as executor: + self._execute_in( + executor=executor, sql=f"CREATE SCHEMA {self._quote_identifier(backup_schema)}" + ) + table_name: str + for table_name in STATE_TABLES: + self._execute_in( + executor=executor, + sql="CREATE TABLE " + f"{self._qualified_name(schema=backup_schema, table=table_name)} " + "AS " + f"SELECT * FROM {self._qualified_name(schema=schema, table=table_name)}", + ) + self._record_event( + executor=executor, + schema=schema, + action=StateMigrationAction.BACKUP, + backup_id_value=backup_id_value, + status=StateMigrationStatus.SUCCESS, + message=None, + ) + return backup_id_value + + def rollback(self, *, connection: Any, schema: str, backup_id: str | None = None) -> str: + backup_id_value: str = backup_id or self._latest_backup_id( + connection=connection, schema=schema + ) + backup_schema: str = self._backup_schema_name( + schema=schema, + backup_id_value=backup_id_value, + ) + if not self._schema_exists(connection=connection, schema=backup_schema): + raise StateBackupNotFoundError(f"State backup schema '{backup_schema}' does not exist") + with self._write_transaction(connection=connection) as executor: + table_name: str + for table_name in STATE_TABLES: + self._execute_in( + executor=executor, + sql="DROP TABLE IF EXISTS " + f"{self._qualified_name(schema=schema, table=table_name)}", + ) + self._execute_in( + executor=executor, + sql=f"CREATE SCHEMA IF NOT EXISTS {self._quote_identifier(schema)}", + ) + for table_name in STATE_TABLES: + self._execute_in( + executor=executor, + sql=f"CREATE TABLE {self._qualified_name(schema=schema, table=table_name)} AS " + f"SELECT * FROM {self._qualified_name(schema=backup_schema, table=table_name)}", + ) + self._create_state_indexes(executor=executor, schema=schema) + self._record_event( + executor=executor, + schema=schema, + action=StateMigrationAction.ROLLBACK, + backup_id_value=backup_id_value, + status=StateMigrationStatus.SUCCESS, + message=None, + ) + return backup_id_value + + def reset(self, *, connection: Any, schema: str) -> None: + with self._write_transaction(connection=connection) as executor: + for table_name in STATE_TABLES: + self._execute_in( + executor=executor, + sql="DROP TABLE IF EXISTS " + f"{self._qualified_name(schema=schema, table=table_name)}", + ) + + def _create_additional_state_tables(self, *, executor: Any, schema: str) -> None: + table_name: str + columns: dict[str, StateColumnType] + for table_name, columns in STATE_TABLE_COLUMNS.items(): + if table_name in {STATE_VERSION_TABLE, STATE_MIGRATION_EVENTS_TABLE}: + continue + column_sql: str = ", ".join( + f"{self._quote_identifier(column_name)} {self._state_column_sql_type(column_type)}" + for column_name, column_type in columns.items() + ) + self._execute_in( + executor=executor, + sql="CREATE TABLE IF NOT EXISTS " + f"{self._qualified_name(schema=schema, table=table_name)} " + f"({column_sql})", + ) + column_name: str + column_type: StateColumnType + for column_name, column_type in columns.items(): + self._execute_in( + executor=executor, + sql=f"ALTER TABLE {self._qualified_name(schema=schema, table=table_name)} " + f"ADD COLUMN IF NOT EXISTS {self._quote_identifier(column_name)} " + f"{self._state_column_sql_type(column_type)}", + ) + self._create_state_indexes(executor=executor, schema=schema) + + def _create_state_indexes(self, *, executor: Any, schema: str) -> None: + table_name: str + indexes: dict[str, tuple[str, ...]] + for table_name, indexes in STATE_TABLE_INDEXES.items(): + index_name: str + columns: tuple[str, ...] + for index_name, columns in indexes.items(): + column_sql: str = ", ".join(self._quote_identifier(column) for column in columns) + unique_sql: str = "" if index_name in NON_UNIQUE_STATE_INDEXES else "UNIQUE " + self._execute_in( + executor=executor, + sql=f"CREATE {unique_sql}INDEX IF NOT EXISTS " + f"{self._quote_identifier(index_name)} " + f"ON {self._qualified_name(schema=schema, table=table_name)} ({column_sql})", + ) + + def _record_event( + self, + *, + executor: Any, + schema: str, + action: StateMigrationAction, + backup_id_value: str | None, + status: StateMigrationStatus, + message: str | None, + ) -> None: + p: str = self._placeholder + self._execute_in( + executor=executor, + sql="INSERT INTO " + f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} " + "(event_id, action, backup_id, status, message, created_at) " + f"VALUES ({p}, {p}, {p}, {p}, {p}, CURRENT_TIMESTAMP)", + params=[event_id(), action.value, backup_id_value, status.value, message], + ) + + def list_state_backups(self, *, connection: Any, schema: str) -> tuple[StateBackupRecord, ...]: + p: str = self._placeholder + prefix: str = f"{schema}__backup_%" + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT s.schema_name, e.backup_id, MAX(e.created_at) " + "FROM information_schema.schemata s " + "LEFT JOIN " + f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} e " + f"ON s.schema_name = {p} || e.backup_id " + f"WHERE s.schema_name LIKE {p} " + "GROUP BY s.schema_name, e.backup_id ORDER BY s.schema_name DESC", + params=[f"{schema}__backup_", prefix], + ) + return tuple( + StateBackupRecord( + backup_id=row[1] or str(row[0]).removeprefix(f"{schema}__backup_"), + schema_name=row[0], + created_at=row[2], + ) + for row in rows + ) + + def _latest_backup_id(self, *, connection: Any, schema: str) -> str: + p: str = self._placeholder + prefix: str = f"{schema}__backup_%" + rows: list[tuple[str]] = self._fetch_all( + connection=connection, + sql=f"SELECT schema_name FROM information_schema.schemata WHERE schema_name LIKE {p} " + "ORDER BY schema_name DESC LIMIT 1", + params=[prefix], + ) + if not rows: + raise StateBackupNotFoundError("No state backup is available for rollback") + return rows[0][0].removeprefix(f"{schema}__backup_") + + def _schema_exists(self, *, connection: Any, schema: str) -> bool: + p: str = self._placeholder + rows: list[tuple[str]] = self._fetch_all( + connection=connection, + sql=f"SELECT schema_name FROM information_schema.schemata WHERE schema_name = {p}", + params=[schema], + ) + return bool(rows) + + def _state_column_sql_type(self, column_type: StateColumnType) -> str: + match column_type: + case StateColumnType.INTEGER: + return "INTEGER" + case StateColumnType.TEXT: + return "TEXT" + case StateColumnType.TIMESTAMP: + return "TIMESTAMP" + raise StateBackendConfigError(f"Unsupported state column type: {column_type}") + + def _backup_schema_name(self, *, schema: str, backup_id_value: str) -> str: + return f"{schema}__backup_{backup_id_value}" + def _replace_row_preserving_created_at( self, *, diff --git a/src/sqlbuild/virtual/state/classes/duckdb.py b/src/sqlbuild/virtual/state/classes/duckdb.py index 6e4b272b23..6a2df70a62 100644 --- a/src/sqlbuild/virtual/state/classes/duckdb.py +++ b/src/sqlbuild/virtual/state/classes/duckdb.py @@ -12,7 +12,6 @@ from sqlbuild.virtual.state._helpers.state_storage.datetime import ( to_naive_utc_wall_clock, ) -from sqlbuild.virtual.state._helpers.state_storage.events import backup_id, event_id from sqlbuild.virtual.state._helpers.state_storage.microbatch_events import ( append_duckdb_microbatch_event, append_duckdb_microbatch_events, @@ -23,39 +22,29 @@ from sqlbuild.virtual.state._helpers.state_storage.validation import build_validation_result from sqlbuild.virtual.state.classes._sql_state_backend import SqlStateBackend from sqlbuild.virtual.state.constants import ( - CURRENT_STATE_SCHEMA_VERSION, DUCKDB_DATETIME_TYPE_TOKEN, DUCKDB_INTEGER_TYPE_TOKEN, DUCKDB_TIMESTAMP_TYPE_TOKEN, LOCK_TABLE, MICROBATCH_EVENT_TABLE, - NON_UNIQUE_STATE_INDEXES, PYTHON_NODE_VERSION_TABLE, SOURCE_FRESHNESS_OBSERVATION_TABLE, - STATE_MIGRATION_EVENTS_TABLE, STATE_TABLE_COLUMNS, STATE_TABLE_INDEXES, - STATE_TABLES, - STATE_VERSION_TABLE, VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, VIRTUAL_ENVIRONMENT_TABLE, ) from sqlbuild.virtual.state.exceptions import ( StateBackendConfigError, - StateBackupNotFoundError, - StateSchemaInvalidError, ) from sqlbuild.virtual.state.models import ( SourceFreshnessRecord, - StateBackupRecord, StateLockLease, StateSchemaValidationResult, VirtualEnvironmentNodeRefRecord, ) from sqlbuild.virtual.state.types import ( StateColumnType, - StateMigrationAction, - StateMigrationStatus, ) @@ -125,51 +114,6 @@ def _fetch_one_in( ) -> tuple[Any, ...] | None: return executor.execute(sql, params).fetchone() - def initialize(self, *, connection: Any, schema: str, sqlbuild_version: str) -> None: - connection.execute("BEGIN") - try: - connection.execute(f"CREATE SCHEMA IF NOT EXISTS {self._quote_identifier(schema)}") - connection.execute( - "CREATE TABLE IF NOT EXISTS " - f"{self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)} (" - "schema_version INTEGER NOT NULL, " - "sqlbuild_version TEXT NOT NULL, " - "updated_at TIMESTAMP NOT NULL" - ")" - ) - connection.execute( - "CREATE TABLE IF NOT EXISTS " - f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} (" - "event_id TEXT NOT NULL, " - "action TEXT NOT NULL, " - "backup_id TEXT, " - "status TEXT NOT NULL, " - "message TEXT, " - "created_at TIMESTAMP NOT NULL" - ")" - ) - self._create_additional_state_tables(connection=connection, schema=schema) - connection.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)}" - ) - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)} " - "(schema_version, sqlbuild_version, updated_at) VALUES (?, ?, CURRENT_TIMESTAMP)", - [CURRENT_STATE_SCHEMA_VERSION, sqlbuild_version], - ) - self._record_event( - connection=connection, - schema=schema, - action=StateMigrationAction.INIT, - backup_id_value=None, - status=StateMigrationStatus.SUCCESS, - message=None, - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - def inspect_schema(self, *, connection: Any, schema: str) -> StateSchemaValidationResult: tables: set[str] = { row[0] @@ -200,91 +144,6 @@ def inspect_schema(self, *, connection: Any, schema: str) -> StateSchemaValidati existing_indexes_by_table=indexes_by_table, ) - def create_backup(self, *, connection: Any, schema: str) -> str: - validation: StateSchemaValidationResult = self.inspect_schema( - connection=connection, schema=schema - ) - if not validation.valid: - raise StateSchemaInvalidError("Cannot backup invalid state schema") - backup_id_value: str = backup_id() - backup_schema: str = self._backup_schema_name( - schema=schema, - backup_id_value=backup_id_value, - ) - connection.execute("BEGIN") - try: - connection.execute(f"CREATE SCHEMA {self._quote_identifier(backup_schema)}") - table_name: str - for table_name in STATE_TABLES: - connection.execute( - f"CREATE TABLE {self._qualified_name(schema=backup_schema, table=table_name)} " - "AS " - f"SELECT * FROM {self._qualified_name(schema=schema, table=table_name)}" - ) - self._record_event( - connection=connection, - schema=schema, - action=StateMigrationAction.BACKUP, - backup_id_value=backup_id_value, - status=StateMigrationStatus.SUCCESS, - message=None, - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - return backup_id_value - - def rollback(self, *, connection: Any, schema: str, backup_id: str | None = None) -> str: - backup_id_value: str = backup_id or self._latest_backup_id( - connection=connection, schema=schema - ) - backup_schema: str = self._backup_schema_name( - schema=schema, - backup_id_value=backup_id_value, - ) - if not self._schema_exists(connection=connection, schema=backup_schema): - raise StateBackupNotFoundError(f"State backup schema '{backup_schema}' does not exist") - connection.execute("BEGIN") - try: - table_name: str - for table_name in STATE_TABLES: - connection.execute( - f"DROP TABLE IF EXISTS {self._qualified_name(schema=schema, table=table_name)}" - ) - connection.execute(f"CREATE SCHEMA IF NOT EXISTS {self._quote_identifier(schema)}") - for table_name in STATE_TABLES: - connection.execute( - f"CREATE TABLE {self._qualified_name(schema=schema, table=table_name)} AS " - f"SELECT * FROM {self._qualified_name(schema=backup_schema, table=table_name)}" - ) - self._create_state_indexes(connection=connection, schema=schema) - self._record_event( - connection=connection, - schema=schema, - action=StateMigrationAction.ROLLBACK, - backup_id_value=backup_id_value, - status=StateMigrationStatus.SUCCESS, - message=None, - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - return backup_id_value - - def reset(self, *, connection: Any, schema: str) -> None: - connection.execute("BEGIN") - try: - for table_name in STATE_TABLES: - connection.execute( - f"DROP TABLE IF EXISTS {self._qualified_name(schema=schema, table=table_name)}" - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - def append_microbatch_event( self, *, connection: Any, schema: str, event: MicrobatchEvent ) -> None: @@ -547,116 +406,10 @@ def delete_lock(self, *, connection: Any, schema: str, lock_key: str) -> None: [lock_key], ) - def list_state_backups(self, *, connection: Any, schema: str) -> tuple[StateBackupRecord, ...]: - prefix: str = f"{schema}__backup_%" - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT s.schema_name, e.backup_id, MAX(e.created_at) " - "FROM information_schema.schemata s " - "LEFT JOIN " - f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} e " - "ON s.schema_name = ? || e.backup_id " - "WHERE s.schema_name LIKE ? " - "GROUP BY s.schema_name, e.backup_id ORDER BY s.schema_name DESC", - [f"{schema}__backup_", prefix], - ).fetchall() - return tuple( - StateBackupRecord( - backup_id=row[1] or str(row[0]).removeprefix(f"{schema}__backup_"), - schema_name=row[0], - created_at=row[2], - ) - for row in rows - ) - def delete_state_backup(self, *, connection: Any, schema: str, backup_id: str) -> None: backup_schema: str = self._backup_schema_name(schema=schema, backup_id_value=backup_id) connection.execute(f"DROP SCHEMA IF EXISTS {self._quote_identifier(backup_schema)} CASCADE") - def _latest_backup_id(self, *, connection: Any, schema: str) -> str: - prefix: str = f"{schema}__backup_%" - rows: list[tuple[str]] = connection.execute( - "SELECT schema_name FROM information_schema.schemata WHERE schema_name LIKE ? " - "ORDER BY schema_name DESC LIMIT 1", - [prefix], - ).fetchall() - if not rows: - raise StateBackupNotFoundError("No state backup is available for rollback") - return rows[0][0].removeprefix(f"{schema}__backup_") - - def _schema_exists(self, *, connection: Any, schema: str) -> bool: - rows: list[tuple[str]] = connection.execute( - "SELECT schema_name FROM information_schema.schemata WHERE schema_name = ?", - [schema], - ).fetchall() - return bool(rows) - - def _record_event( - self, - *, - connection: Any, - schema: str, - action: StateMigrationAction, - backup_id_value: str | None, - status: StateMigrationStatus, - message: str | None, - ) -> None: - connection.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} " - "(event_id, action, backup_id, status, message, created_at) " - "VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)", - [event_id(), action.value, backup_id_value, status.value, message], - ) - - def _create_additional_state_tables(self, *, connection: Any, schema: str) -> None: - table_name: str - columns: dict[str, StateColumnType] - for table_name, columns in STATE_TABLE_COLUMNS.items(): - if table_name in {STATE_VERSION_TABLE, STATE_MIGRATION_EVENTS_TABLE}: - continue - column_sql: str = ", ".join( - f"{self._quote_identifier(column_name)} {self._state_column_sql_type(column_type)}" - for column_name, column_type in columns.items() - ) - connection.execute( - "CREATE TABLE IF NOT EXISTS " - f"{self._qualified_name(schema=schema, table=table_name)} " - f"({column_sql})" - ) - column_name: str - column_type: StateColumnType - for column_name, column_type in columns.items(): - connection.execute( - f"ALTER TABLE {self._qualified_name(schema=schema, table=table_name)} " - f"ADD COLUMN IF NOT EXISTS {self._quote_identifier(column_name)} " - f"{self._state_column_sql_type(column_type)}" - ) - self._create_state_indexes(connection=connection, schema=schema) - - def _create_state_indexes(self, *, connection: Any, schema: str) -> None: - table_name: str - indexes: dict[str, tuple[str, ...]] - for table_name, indexes in STATE_TABLE_INDEXES.items(): - index_name: str - columns: tuple[str, ...] - for index_name, columns in indexes.items(): - column_sql: str = ", ".join(self._quote_identifier(column) for column in columns) - unique_sql: str = "" if index_name in NON_UNIQUE_STATE_INDEXES else "UNIQUE " - connection.execute( - f"CREATE {unique_sql}INDEX IF NOT EXISTS {self._quote_identifier(index_name)} " - f"ON {self._qualified_name(schema=schema, table=table_name)} ({column_sql})" - ) - - def _state_column_sql_type(self, column_type: StateColumnType) -> str: - match column_type: - case StateColumnType.INTEGER: - return "INTEGER" - case StateColumnType.TEXT: - return "TEXT" - case StateColumnType.TIMESTAMP: - return "TIMESTAMP" - raise StateBackendConfigError(f"Unsupported state column type: {column_type}") - def _validate_source_freshness_records( self, *, @@ -738,9 +491,6 @@ def _replace_virtual_environment_node_ref_groups( ) executor.execute(f"DROP TABLE IF EXISTS {temp_table_name}") - def _backup_schema_name(self, *, schema: str, backup_id_value: str) -> str: - return f"{schema}__backup_{backup_id_value}" - def _state_type_matches(self, *, actual_type: str, expected_type: StateColumnType) -> bool: actual: str = actual_type.lower() match expected_type: diff --git a/src/sqlbuild/virtual/state/classes/postgres.py b/src/sqlbuild/virtual/state/classes/postgres.py index 92dd2df62a..af6d689802 100644 --- a/src/sqlbuild/virtual/state/classes/postgres.py +++ b/src/sqlbuild/virtual/state/classes/postgres.py @@ -18,44 +18,33 @@ from sqlbuild.virtual.state._helpers.state_storage.datetime import ( to_naive_utc_wall_clock, ) -from sqlbuild.virtual.state._helpers.state_storage.events import backup_id, event_id from sqlbuild.virtual.state._helpers.state_storage.validation import ( build_validation_result, ) from sqlbuild.virtual.state.classes._sql_state_backend import SqlStateBackend from sqlbuild.virtual.state.constants import ( - CURRENT_STATE_SCHEMA_VERSION, LOCK_TABLE, MICROBATCH_EVENT_TABLE, - NON_UNIQUE_STATE_INDEXES, POSTGRES_INTEGER_TYPES, POSTGRES_TEXT_TYPES, PYTHON_NODE_VERSION_TABLE, SOURCE_FRESHNESS_OBSERVATION_TABLE, - STATE_MIGRATION_EVENTS_TABLE, STATE_TABLE_COLUMNS, STATE_TABLE_INDEXES, - STATE_TABLES, - STATE_VERSION_TABLE, VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, VIRTUAL_ENVIRONMENT_TABLE, ) from sqlbuild.virtual.state.exceptions import ( StateBackendConfigError, - StateBackupNotFoundError, - StateSchemaInvalidError, ) from sqlbuild.virtual.state.models import ( SourceFreshnessRecord, - StateBackupRecord, StateLockLease, StateSchemaValidationResult, VirtualEnvironmentNodeRefRecord, ) from sqlbuild.virtual.state.types import ( StateColumnType, - StateMigrationAction, - StateMigrationStatus, ) @@ -149,53 +138,6 @@ def _fetch_one_in( executor.execute(sql, params) return executor.fetchone() - def initialize(self, *, connection: Any, schema: str, sqlbuild_version: str) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - cursor.execute(f"CREATE SCHEMA IF NOT EXISTS {self._quote_identifier(schema)}") - cursor.execute( - "CREATE TABLE IF NOT EXISTS " - f"{self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)} (" - "schema_version INTEGER NOT NULL, " - "sqlbuild_version TEXT NOT NULL, " - "updated_at TIMESTAMP NOT NULL" - ")" - ) - cursor.execute( - "CREATE TABLE IF NOT EXISTS " - f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} (" - "event_id TEXT NOT NULL, " - "action TEXT NOT NULL, " - "backup_id TEXT, " - "status TEXT NOT NULL, " - "message TEXT, " - "created_at TIMESTAMP NOT NULL" - ")" - ) - self._create_additional_state_tables(cursor=cursor, schema=schema) - cursor.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)}" - ) - cursor.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)} " - "(schema_version, sqlbuild_version, updated_at) " - "VALUES (%s, %s, CURRENT_TIMESTAMP)", - [CURRENT_STATE_SCHEMA_VERSION, sqlbuild_version], - ) - self._record_event( - cursor=cursor, - schema=schema, - action=StateMigrationAction.INIT, - backup_id_value=None, - status=StateMigrationStatus.SUCCESS, - message=None, - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - def inspect_schema(self, *, connection: Any, schema: str) -> StateSchemaValidationResult: with connection.cursor() as cursor: cursor.execute( @@ -227,99 +169,6 @@ def inspect_schema(self, *, connection: Any, schema: str) -> StateSchemaValidati existing_indexes_by_table=indexes_by_table, ) - def create_backup(self, *, connection: Any, schema: str) -> str: - validation: StateSchemaValidationResult = self.inspect_schema( - connection=connection, schema=schema - ) - if not validation.valid: - raise StateSchemaInvalidError("Cannot backup invalid state schema") - backup_id_value: str = backup_id() - backup_schema: str = self._backup_schema_name( - schema=schema, - backup_id_value=backup_id_value, - ) - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - cursor.execute(f"CREATE SCHEMA {self._quote_identifier(backup_schema)}") - table_name: str - for table_name in STATE_TABLES: - cursor.execute( - "CREATE TABLE " - + self._qualified_name(schema=backup_schema, table=table_name) - + " AS " - f"SELECT * FROM {self._qualified_name(schema=schema, table=table_name)}" - ) - self._record_event( - cursor=cursor, - schema=schema, - action=StateMigrationAction.BACKUP, - backup_id_value=backup_id_value, - status=StateMigrationStatus.SUCCESS, - message=None, - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - return backup_id_value - - def rollback(self, *, connection: Any, schema: str, backup_id: str | None = None) -> str: - backup_id_value: str = backup_id or self._latest_backup_id( - connection=connection, schema=schema - ) - backup_schema: str = self._backup_schema_name( - schema=schema, - backup_id_value=backup_id_value, - ) - if not self._schema_exists(connection=connection, schema=backup_schema): - raise StateBackupNotFoundError(f"State backup schema '{backup_schema}' does not exist") - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - table_name: str - for table_name in STATE_TABLES: - cursor.execute( - "DROP TABLE IF EXISTS " - f"{self._qualified_name(schema=schema, table=table_name)}" - ) - cursor.execute(f"CREATE SCHEMA IF NOT EXISTS {self._quote_identifier(schema)}") - for table_name in STATE_TABLES: - cursor.execute( - f"CREATE TABLE {self._qualified_name(schema=schema, table=table_name)} AS " - "SELECT * FROM " - f"{self._qualified_name(schema=backup_schema, table=table_name)}" - ) - self._create_state_indexes(cursor=cursor, schema=schema) - self._record_event( - cursor=cursor, - schema=schema, - action=StateMigrationAction.ROLLBACK, - backup_id_value=backup_id_value, - status=StateMigrationStatus.SUCCESS, - message=None, - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - return backup_id_value - - def reset(self, *, connection: Any, schema: str) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - table_name: str - for table_name in STATE_TABLES: - cursor.execute( - "DROP TABLE IF EXISTS " - f"{self._qualified_name(schema=schema, table=table_name)}" - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - def append_microbatch_event( self, *, connection: Any, schema: str, event: MicrobatchEvent ) -> None: @@ -632,124 +481,12 @@ def delete_lock(self, *, connection: Any, schema: str, lock_key: str) -> None: ) connection.commit() - def list_state_backups(self, *, connection: Any, schema: str) -> tuple[StateBackupRecord, ...]: - prefix: str = f"{schema}__backup_%" - with connection.cursor() as cursor: - cursor.execute( - "SELECT s.schema_name, e.backup_id, MAX(e.created_at) " - "FROM information_schema.schemata s " - "LEFT JOIN " - f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} e " - "ON s.schema_name = %s || e.backup_id " - "WHERE s.schema_name LIKE %s " - "GROUP BY s.schema_name, e.backup_id ORDER BY s.schema_name DESC", - [f"{schema}__backup_", prefix], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - StateBackupRecord( - backup_id=row[1] or str(row[0]).removeprefix(f"{schema}__backup_"), - schema_name=row[0], - created_at=row[2], - ) - for row in rows - ) - def delete_state_backup(self, *, connection: Any, schema: str, backup_id: str) -> None: backup_schema: str = self._backup_schema_name(schema=schema, backup_id_value=backup_id) with connection.cursor() as cursor: cursor.execute(f"DROP SCHEMA IF EXISTS {self._quote_identifier(backup_schema)} CASCADE") connection.commit() - def _latest_backup_id(self, *, connection: Any, schema: str) -> str: - prefix: str = f"{schema}__backup_%" - with connection.cursor() as cursor: - cursor.execute( - "SELECT schema_name FROM information_schema.schemata WHERE schema_name LIKE %s " - "ORDER BY schema_name DESC LIMIT 1", - [prefix], - ) - row: tuple[str] | None = cursor.fetchone() - if row is None: - raise StateBackupNotFoundError("No state backup is available for rollback") - return row[0].removeprefix(f"{schema}__backup_") - - def _schema_exists(self, *, connection: Any, schema: str) -> bool: - with connection.cursor() as cursor: - cursor.execute( - "SELECT schema_name FROM information_schema.schemata WHERE schema_name = %s", - [schema], - ) - row: tuple[str] | None = cursor.fetchone() - return row is not None - - def _record_event( - self, - *, - cursor: Any, - schema: str, - action: StateMigrationAction, - backup_id_value: str | None, - status: StateMigrationStatus, - message: str | None, - ) -> None: - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} " - "(event_id, action, backup_id, status, message, created_at) " - "VALUES (%s, %s, %s, %s, %s, CURRENT_TIMESTAMP)", - [event_id(), action.value, backup_id_value, status.value, message], - ) - - def _create_additional_state_tables(self, *, cursor: Any, schema: str) -> None: - table_name: str - columns: dict[str, StateColumnType] - for table_name, columns in STATE_TABLE_COLUMNS.items(): - if table_name in {STATE_VERSION_TABLE, STATE_MIGRATION_EVENTS_TABLE}: - continue - column_sql: str = ", ".join( - f"{self._quote_identifier(column_name)} {self._state_column_sql_type(column_type)}" - for column_name, column_type in columns.items() - ) - cursor.execute( - "CREATE TABLE IF NOT EXISTS " - f"{self._qualified_name(schema=schema, table=table_name)} " - f"({column_sql})" - ) - column_name: str - column_type: StateColumnType - for column_name, column_type in columns.items(): - cursor.execute( - f"ALTER TABLE {self._qualified_name(schema=schema, table=table_name)} " - f"ADD COLUMN IF NOT EXISTS {self._quote_identifier(column_name)} " - f"{self._state_column_sql_type(column_type)}" - ) - self._create_state_indexes(cursor=cursor, schema=schema) - - def _create_state_indexes(self, *, cursor: Any, schema: str) -> None: - table_name: str - indexes: dict[str, tuple[str, ...]] - for table_name, indexes in STATE_TABLE_INDEXES.items(): - index_name: str - columns: tuple[str, ...] - for index_name, columns in indexes.items(): - column_sql: str = ", ".join(self._quote_identifier(column) for column in columns) - unique_sql: str = "" if index_name in NON_UNIQUE_STATE_INDEXES else "UNIQUE " - cursor.execute( - f"CREATE {unique_sql}INDEX IF NOT EXISTS {self._quote_identifier(index_name)} " - f"ON {self._qualified_name(schema=schema, table=table_name)} ({column_sql})" - ) - - def _state_column_sql_type(self, column_type: StateColumnType) -> str: - match column_type: - case StateColumnType.INTEGER: - return "INTEGER" - case StateColumnType.TEXT: - return "TEXT" - case StateColumnType.TIMESTAMP: - return "TIMESTAMP" - raise StateBackendConfigError(f"Unsupported state column type: {column_type}") - def _validate_source_freshness_records( self, *, @@ -814,9 +551,6 @@ def _replace_virtual_environment_node_ref_groups( ], ) - def _backup_schema_name(self, *, schema: str, backup_id_value: str) -> str: - return f"{schema}__backup_{backup_id_value}" - def _state_type_matches(self, *, actual_type: str, expected_type: StateColumnType) -> bool: actual: str = actual_type.lower() match expected_type: From 66364cc5a4070580985baeb04618bf68f8f556b1 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:58:36 +0100 Subject: [PATCH 19/57] test(state): characterize virtual environment deletion on DuckDB and Postgres state backends Extend the shared ref contract to delete one environment and check that its record and refs are removed while another environment's refs survive. (cherry picked from commit 350254fcc84e4dee6e6779568cfae82df1056734) --- .../sqlbuild/virtual/state/classes/helpers.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py b/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py index 1ea7658a61..34300313bd 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py @@ -301,6 +301,9 @@ class StateRefContractObservation(NamedTuple): grouped_seed_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] published_environment: VirtualEnvironmentRecord | None published_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] + deleted_environment: VirtualEnvironmentRecord | None + deleted_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] + surviving_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] def exercise_state_ref_contract( @@ -394,6 +397,22 @@ def exercise_state_ref_contract( connection=connection, schema=schema, virtual_environment_name="prod", node_type="model" ) ) + backend.delete_virtual_environment( + connection=connection, schema=schema, virtual_environment_name="prod" + ) + deleted_environment: VirtualEnvironmentRecord | None = backend.get_virtual_environment( + connection=connection, schema=schema, virtual_environment_name="prod" + ) + deleted_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="prod", node_type="model" + ) + ) + surviving_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="qa", node_type="model" + ) + ) return StateRefContractObservation( node_refs_after_first_replace=node_refs_after_first_replace, node_refs_after_second_replace=node_refs_after_second_replace, @@ -403,6 +422,9 @@ def exercise_state_ref_contract( grouped_seed_refs=grouped_seed_refs, published_environment=published_environment, published_model_refs=published_model_refs, + deleted_environment=deleted_environment, + deleted_model_refs=deleted_model_refs, + surviving_model_refs=surviving_model_refs, ) @@ -423,4 +445,7 @@ def exercise_state_ref_contract( grouped_seed_refs=(VirtualEnvironmentNodeRefRecord("qa", "seed", "countries", "seed-v2"),), published_environment=VirtualEnvironmentRecord("prod", VirtualEnvironmentStatus.ACTIVE), published_model_refs=(VirtualEnvironmentNodeRefRecord("prod", "model", "orders", "orders-v1"),), + deleted_environment=None, + deleted_model_refs=(), + surviving_model_refs=(VirtualEnvironmentNodeRefRecord("qa", "model", "orders", "orders-v3"),), ) From 4f5247b79090b8c81ba2fedbc7f2b6d5b968c69f Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:00:38 +0100 Subject: [PATCH 20/57] refactor(state): share environment deletion and Python node version count Move close, source freshness record validation, virtual environment deletion, and the unreferenced Python node version count into SqlStateBackend. Both backends ran the same statements in the same order and transaction; only placeholders and the executor differed. prune_unreferenced_python_node_versions stays per backend: the DELETE aliases the version table as "versions" on DuckDB and "AS versions" on Postgres, so the statement text is not identical. (cherry picked from commit 7c79874494c1bec9f097cc83071277e357b28122) --- .../state/classes/_sql_state_backend.py | 69 +++++++++++++++- src/sqlbuild/virtual/state/classes/duckdb.py | 67 ---------------- .../virtual/state/classes/postgres.py | 79 ------------------- 3 files changed, 68 insertions(+), 147 deletions(-) diff --git a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py index cd94c59cdf..0606be0831 100644 --- a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py +++ b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py @@ -6,7 +6,7 @@ from collections.abc import Sequence from contextlib import AbstractContextManager from datetime import datetime -from typing import Any, ClassVar +from typing import Any, ClassVar, cast from sqlbuild.compiler.compile.types import CompiledResourceType from sqlbuild.executor.node_results.main.decode_json import decode_node_result_json @@ -1122,6 +1122,73 @@ def _state_column_sql_type(self, column_type: StateColumnType) -> str: def _backup_schema_name(self, *, schema: str, backup_id_value: str) -> str: return f"{schema}__backup_{backup_id_value}" + def close(self, connection: Any) -> None: + connection.close() + + def delete_virtual_environment( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> None: + p: str = self._placeholder + table_names: tuple[str, ...] = ( + VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, + SOURCE_FRESHNESS_OBSERVATION_TABLE, + VIRTUAL_ENVIRONMENT_TABLE, + ) + with self._write_transaction(connection=connection) as executor: + table_name: str + for table_name in table_names: + self._execute_in( + executor=executor, + sql=( + f"DELETE FROM {self._qualified_name(schema=schema, table=table_name)} " + f"WHERE virtual_environment_name = {p}" + ), + params=[virtual_environment_name], + ) + + def count_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: + version_table: str = self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE) + ref_table: str = self._qualified_name( + schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE + ) + row: tuple[Any, ...] = cast( + tuple[Any, ...], + self._fetch_one( + connection=connection, + sql=( + "SELECT COUNT(*) " + f"FROM {version_table} versions " + "WHERE NOT EXISTS (" + "SELECT 1 " + f"FROM {ref_table} refs " + "WHERE refs.node_type = versions.node_type " + "AND refs.node_name = versions.node_name " + "AND refs.version_hash = versions.version_hash)" + ), + ), + ) + return int(row[0]) + + def _validate_source_freshness_records( + self, + *, + virtual_environment_name: str, + records: tuple[SourceFreshnessRecord, ...], + ) -> None: + seen_source_names: set[str] = set() + record: SourceFreshnessRecord + for record in records: + if record.virtual_environment_name != virtual_environment_name: + raise StateBackendConfigError( + "Source freshness record virtual_environment_name must match replacement " + "virtual_environment_name" + ) + if record.source_name in seen_source_names: + raise StateBackendConfigError( + f"Duplicate source freshness record for source '{record.source_name}'" + ) + seen_source_names.add(record.source_name) + def _replace_row_preserving_created_at( self, *, diff --git a/src/sqlbuild/virtual/state/classes/duckdb.py b/src/sqlbuild/virtual/state/classes/duckdb.py index 6a2df70a62..d4edb7945f 100644 --- a/src/sqlbuild/virtual/state/classes/duckdb.py +++ b/src/sqlbuild/virtual/state/classes/duckdb.py @@ -32,7 +32,6 @@ STATE_TABLE_COLUMNS, STATE_TABLE_INDEXES, VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, - VIRTUAL_ENVIRONMENT_TABLE, ) from sqlbuild.virtual.state.exceptions import ( StateBackendConfigError, @@ -61,9 +60,6 @@ def connect(self, config: dict[str, object]) -> Any: raise StateBackendConfigError("DuckDB state backend requires state.connection.database") return ObservedConnection(raw_connection=duckdb.connect(database), adapter="duckdb") - def close(self, connection: Any) -> None: - connection.close() - def _fetch_one( self, *, connection: Any, sql: str, params: Sequence[object] | None = None ) -> tuple[Any, ...] | None: @@ -188,34 +184,6 @@ def read_microbatch_model_history( scope=scope, ) - def delete_virtual_environment( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> None: - connection.execute("BEGIN") - try: - connection.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "WHERE virtual_environment_name = ?", - [virtual_environment_name], - ) - connection.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=SOURCE_FRESHNESS_OBSERVATION_TABLE)} " - "WHERE virtual_environment_name = ?", - [virtual_environment_name], - ) - connection.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "WHERE virtual_environment_name = ?", - [virtual_environment_name], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - def upsert_virtual_environment_node_ref( self, *, @@ -233,21 +201,6 @@ def upsert_virtual_environment_node_ref( [ref.virtual_environment_name, ref.node_type, ref.node_name, ref.version_hash], ) - def count_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: - row: tuple[Any, ...] = connection.execute( - "SELECT COUNT(*) " - f"FROM {self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} versions " - "WHERE NOT EXISTS (" - "SELECT 1 " - "FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "refs " - "WHERE refs.node_type = versions.node_type " - "AND refs.node_name = versions.node_name " - "AND refs.version_hash = versions.version_hash)" - ).fetchone() - return int(row[0]) - def prune_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: before_count: int = self.count_unreferenced_python_node_versions( connection=connection, schema=schema @@ -410,26 +363,6 @@ def delete_state_backup(self, *, connection: Any, schema: str, backup_id: str) - backup_schema: str = self._backup_schema_name(schema=schema, backup_id_value=backup_id) connection.execute(f"DROP SCHEMA IF EXISTS {self._quote_identifier(backup_schema)} CASCADE") - def _validate_source_freshness_records( - self, - *, - virtual_environment_name: str, - records: tuple[SourceFreshnessRecord, ...], - ) -> None: - seen_source_names: set[str] = set() - record: SourceFreshnessRecord - for record in records: - if record.virtual_environment_name != virtual_environment_name: - raise StateBackendConfigError( - "Source freshness record virtual_environment_name must match replacement " - "virtual_environment_name" - ) - if record.source_name in seen_source_names: - raise StateBackendConfigError( - f"Duplicate source freshness record for source '{record.source_name}'" - ) - seen_source_names.add(record.source_name) - def _replace_virtual_environment_node_ref_groups( self, *, diff --git a/src/sqlbuild/virtual/state/classes/postgres.py b/src/sqlbuild/virtual/state/classes/postgres.py index af6d689802..867e80760e 100644 --- a/src/sqlbuild/virtual/state/classes/postgres.py +++ b/src/sqlbuild/virtual/state/classes/postgres.py @@ -32,7 +32,6 @@ STATE_TABLE_COLUMNS, STATE_TABLE_INDEXES, VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, - VIRTUAL_ENVIRONMENT_TABLE, ) from sqlbuild.virtual.state.exceptions import ( StateBackendConfigError, @@ -77,9 +76,6 @@ def connect(self, config: dict[str, object]) -> Any: except Exception as error: raise StateBackendConfigError("Could not connect to Postgres state backend") from error - def close(self, connection: Any) -> None: - connection.close() - def _fetch_one( self, *, connection: Any, sql: str, params: Sequence[object] | None = None ) -> tuple[Any, ...] | None: @@ -267,43 +263,6 @@ def read_microbatch_model_history( rows: list[tuple[Any, ...]] = cursor.fetchall() return MicrobatchEventCodec.from_rows(tuple(row) for row in rows) - def delete_virtual_environment( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - cursor.execute( - "DELETE FROM " - + self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, - ) - + " " - "WHERE virtual_environment_name = %s", - [virtual_environment_name], - ) - cursor.execute( - "DELETE FROM " - + self._qualified_name( - schema=schema, - table=SOURCE_FRESHNESS_OBSERVATION_TABLE, - ) - + " " - "WHERE virtual_environment_name = %s", - [virtual_environment_name], - ) - cursor.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "WHERE virtual_environment_name = %s", - [virtual_environment_name], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - def upsert_virtual_environment_node_ref( self, *, @@ -323,24 +282,6 @@ def upsert_virtual_environment_node_ref( [ref.virtual_environment_name, ref.node_type, ref.node_name, ref.version_hash], ) - def count_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: - with connection.cursor() as cursor: - cursor.execute( - "SELECT COUNT(*) " - "FROM " - f"{self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} versions " - "WHERE NOT EXISTS (" - "SELECT 1 " - "FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "refs " - "WHERE refs.node_type = versions.node_type " - "AND refs.node_name = versions.node_name " - "AND refs.version_hash = versions.version_hash)" - ) - row: tuple[Any, ...] = cursor.fetchone() - return int(row[0]) - def prune_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: before_count: int = self.count_unreferenced_python_node_versions( connection=connection, schema=schema @@ -487,26 +428,6 @@ def delete_state_backup(self, *, connection: Any, schema: str, backup_id: str) - cursor.execute(f"DROP SCHEMA IF EXISTS {self._quote_identifier(backup_schema)} CASCADE") connection.commit() - def _validate_source_freshness_records( - self, - *, - virtual_environment_name: str, - records: tuple[SourceFreshnessRecord, ...], - ) -> None: - seen_source_names: set[str] = set() - record: SourceFreshnessRecord - for record in records: - if record.virtual_environment_name != virtual_environment_name: - raise StateBackendConfigError( - "Source freshness record virtual_environment_name must match replacement " - "virtual_environment_name" - ) - if record.source_name in seen_source_names: - raise StateBackendConfigError( - f"Duplicate source freshness record for source '{record.source_name}'" - ) - seen_source_names.add(record.source_name) - def _replace_virtual_environment_node_ref_groups( self, *, From 3ff0a9ee32e387c7f1da8d67ff4e4eb9a4c34a4a Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:26:04 +0100 Subject: [PATCH 21/57] refactor(scenario): share one scenario pipeline loop across modes Run the warehouse test, local test and capture scenario pipelines through one generic per-scenario loop and one target-connection context manager. Per-mode execution and failure-result construction stay in each entry point. Behaviour is unchanged: only the capture path observes the target connection with an operation lifecycle. (cherry picked from commit 53db84cab108c0d98c5e7918713d2bd3b4dac0be) --- .../executor/pipeline/_helpers/scenario.py | 372 +++++++++--------- 1 file changed, 191 insertions(+), 181 deletions(-) diff --git a/src/sqlbuild/executor/pipeline/_helpers/scenario.py b/src/sqlbuild/executor/pipeline/_helpers/scenario.py index cb0eb50b0f..7f8794923c 100644 --- a/src/sqlbuild/executor/pipeline/_helpers/scenario.py +++ b/src/sqlbuild/executor/pipeline/_helpers/scenario.py @@ -4,9 +4,10 @@ import logging import time -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from pathlib import Path -from typing import Any +from typing import Any, Protocol from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.cli.progress.classes.native_progress_projector import ( @@ -63,6 +64,18 @@ def _scenario_failure_help(exc: Exception) -> str | None: return _SCENARIO_INTERNAL_ERROR_HELP +class _ScenarioPipelineResult(Protocol): + @property + def status(self) -> ExecutionStatus: ... + + @property + def error_code(self) -> str | None: ... + + +class _ScenarioFailureResult[ResultT](Protocol): + def __call__(self, *, scenario_name: str, exc: Exception) -> ResultT: ... + + def run_scenario_test_pipeline( *, pipeline_result: CompilePipelineResult, @@ -80,12 +93,176 @@ def run_scenario_test_pipeline( ) -> tuple[ScenarioRunResult, ...]: """Execute selected scenarios from a compiled project.""" + with _scenario_target_connection( + adapter=adapter, + connection_config=connection_config, + connection_hooks=connection_hooks, + observe_connection=False, + ) as connection: + + def execute(scenario_plan: ScenarioExecutionPlan) -> ScenarioRunResult: + return execute_scenario_run( + scenario_plan=scenario_plan, + adapter=adapter, + connection=connection, + run_id=pipeline_result.project.run_id, + retain=retain, + ) + + def failed(*, scenario_name: str, exc: Exception) -> ScenarioRunResult: + return ScenarioRunResult( + scenario_name=scenario_name, + status=ExecutionStatus.FAILED, + retained=retain, + error_code=error_code(error=exc, fallback_code=SCENARIO_EXEC_INTERNAL), + error_help=_scenario_failure_help(exc), + error_message=error_message(exc), + ) + + return _run_scenarios( + pipeline_result=pipeline_result, + scenarios=scenarios, + adapter=adapter, + project_name=project_name, + execute=execute, + failed=failed, + on_scenario_start=on_scenario_start, + on_scenario_complete=on_scenario_complete, + ) + + +def run_scenario_local_test_pipeline( + *, + project_dir: Path, + pipeline_result: CompilePipelineResult, + scenarios: tuple[CompiledSqlScenario, ...], + adapter: BaseAdapter, + project_name: str, + strict: bool, + capture_adapter: str | None = None, + capture_dialect: str | None = None, + on_scenario_start: Callable[[CompiledSqlScenario], None] | None = None, + on_scenario_complete: Callable[ + [CompiledSqlScenario, ScenarioExecutionPlan | None, ScenarioRunResult], None + ] + | None = None, +) -> tuple[ScenarioRunResult, ...]: + """Load selected local scenario snapshots into run-scoped DuckDB databases.""" + + def execute(scenario_plan: ScenarioExecutionPlan) -> ScenarioRunResult: + return execute_local_scenario_load_only_run( + project_dir=project_dir, + scenario_plan=scenario_plan, + adapter=adapter, + strict=strict, + capture_adapter=capture_adapter, + capture_dialect=capture_dialect, + ) + + def failed(*, scenario_name: str, exc: Exception) -> ScenarioRunResult: + return ScenarioRunResult( + scenario_name=scenario_name, + status=ExecutionStatus.FAILED, + local_status=ScenarioLocalRunStatus.ERROR, + retained=False, + error_code=error_code(error=exc, fallback_code=SCENARIO_LOCAL_INTERNAL), + error_help=_scenario_failure_help(exc), + error_message=error_message(exc), + ) + + return _run_scenarios( + pipeline_result=pipeline_result, + scenarios=scenarios, + adapter=adapter, + project_name=project_name, + execute=execute, + failed=failed, + on_scenario_start=on_scenario_start, + on_scenario_complete=on_scenario_complete, + ) + + +def run_scenario_capture_pipeline( + *, + project_dir: Path, + pipeline_result: CompilePipelineResult, + scenarios: tuple[CompiledSqlScenario, ...], + connection_config: dict[str, object], + adapter: BaseAdapter, + project_name: str, + settings: ScenarioCaptureSettings, + connection_hooks: ConnectionHooks | None = None, + on_scenario_start: Callable[[CompiledSqlScenario], None] | None = None, + on_scenario_complete: Callable[ + [CompiledSqlScenario, ScenarioExecutionPlan | None, ScenarioSnapshotCaptureRunResult], None + ] + | None = None, +) -> tuple[ScenarioSnapshotCaptureRunResult, ...]: + """Capture selected scenario inputs into durable local snapshot files.""" + + with _scenario_target_connection( + adapter=adapter, + connection_config=connection_config, + connection_hooks=connection_hooks, + observe_connection=True, + ) as connection: + + def execute(scenario_plan: ScenarioExecutionPlan) -> ScenarioSnapshotCaptureRunResult: + return execute_scenario_snapshot_capture_run( + project_dir=project_dir, + scenario_plan=scenario_plan, + adapter=adapter, + connection=connection, + run_id=pipeline_result.project.run_id, + settings=settings, + local_type_overrides=scenario_local_type_overrides_for_dialect( + scenario_config=pipeline_result.project.scenario, + sql_analysis_dialect=adapter.sql_analysis_dialect(), + ), + ) + + def failed(*, scenario_name: str, exc: Exception) -> ScenarioSnapshotCaptureRunResult: + return ScenarioSnapshotCaptureRunResult( + scenario_name=scenario_name, + status=ExecutionStatus.FAILED, + retained=settings.retain, + error_code=error_code(error=exc, fallback_code=SCENARIO_EXEC_INTERNAL), + error_help=_scenario_failure_help(exc), + error_message=error_message(exc), + ) + + return _run_scenarios( + pipeline_result=pipeline_result, + scenarios=scenarios, + adapter=adapter, + project_name=project_name, + execute=execute, + failed=failed, + on_scenario_start=on_scenario_start, + on_scenario_complete=on_scenario_complete, + ) + + +@contextmanager +def _scenario_target_connection( + *, + adapter: BaseAdapter, + connection_config: dict[str, object], + connection_hooks: ConnectionHooks | None, + observe_connection: bool, +) -> Iterator[Any]: hooks: ConnectionHooks = connection_hooks if connection_hooks is not None else ConnectionHooks() if hooks.on_connection_start is not None: hooks.on_connection_start(1) start: float = time.monotonic() try: - connection: Any = adapter.connect(connection_config) + if observe_connection: + with OperationLifecycle( + operation_kind="scenario", operation_name="scenario_target_connection" + ): + connection: Any = adapter.connect(connection_config) + else: + connection = adapter.connect(connection_config) except Exception: if hooks.on_connection_error is not None: hooks.on_connection_error(1, elapsed_seconds=time.monotonic() - start) @@ -93,81 +270,26 @@ def run_scenario_test_pipeline( if hooks.on_connection_complete is not None: hooks.on_connection_complete(1, elapsed_seconds=time.monotonic() - start) try: - results: list[ScenarioRunResult] = [] - scenario: CompiledSqlScenario - for scenario in scenarios: - scenario_plan: ScenarioExecutionPlan | None = None - resource_id: str = f"sql_scenario:{scenario.name}" - projector: NativeProgressProjector | None = _prepare_scenario_presentation( - scenario_name=scenario.name, - has_completion_callback=on_scenario_complete is not None, - ) - with ResourceAttemptLifecycle( - resource_id=resource_id, - resource_kind="scenario", - resource_name=scenario.name, - run_id=pipeline_result.project.run_id, - ) as lifecycle: - if on_scenario_start is not None: - on_scenario_start(scenario) - try: - scenario_plan = build_scenario_plan( - scenario=scenario, - pipeline_result=pipeline_result, - adapter=adapter, - project_name=project_name, - ) - result: ScenarioRunResult = execute_scenario_run( - scenario_plan=scenario_plan, - adapter=adapter, - connection=connection, - run_id=pipeline_result.project.run_id, - retain=retain, - ) - except Exception as exc: - result = ScenarioRunResult( - scenario_name=scenario.name, - status=ExecutionStatus.FAILED, - retained=retain, - error_code=error_code(error=exc, fallback_code=SCENARIO_EXEC_INTERNAL), - error_help=_scenario_failure_help(exc), - error_message=error_message(exc), - ) - if result.status == ExecutionStatus.FAILED: - lifecycle.failed(error_code=result.error_code) - if projector is not None and on_scenario_complete is not None: - _ = projector.consume_resource_terminal( - resource_name=scenario.name, - resource_id=resource_id, - resource_attempt_id=lifecycle.resource_attempt_id, - ) - results.append(result) - if on_scenario_complete is not None: - on_scenario_complete(scenario, scenario_plan, result) - return tuple(results) + yield connection finally: adapter.close(connection) -def run_scenario_local_test_pipeline( +def _run_scenarios[ResultT: _ScenarioPipelineResult]( *, - project_dir: Path, pipeline_result: CompilePipelineResult, scenarios: tuple[CompiledSqlScenario, ...], adapter: BaseAdapter, project_name: str, - strict: bool, - capture_adapter: str | None = None, - capture_dialect: str | None = None, - on_scenario_start: Callable[[CompiledSqlScenario], None] | None = None, + execute: Callable[[ScenarioExecutionPlan], ResultT], + failed: _ScenarioFailureResult[ResultT], + on_scenario_start: Callable[[CompiledSqlScenario], None] | None, on_scenario_complete: Callable[ - [CompiledSqlScenario, ScenarioExecutionPlan | None, ScenarioRunResult], None + [CompiledSqlScenario, ScenarioExecutionPlan | None, ResultT], None ] - | None = None, -) -> tuple[ScenarioRunResult, ...]: - """Load selected local scenario snapshots into run-scoped DuckDB databases.""" - - results: list[ScenarioRunResult] = [] + | None, +) -> tuple[ResultT, ...]: + results: list[ResultT] = [] scenario: CompiledSqlScenario for scenario in scenarios: scenario_plan: ScenarioExecutionPlan | None = None @@ -191,24 +313,9 @@ def run_scenario_local_test_pipeline( adapter=adapter, project_name=project_name, ) - result: ScenarioRunResult = execute_local_scenario_load_only_run( - project_dir=project_dir, - scenario_plan=scenario_plan, - adapter=adapter, - strict=strict, - capture_adapter=capture_adapter, - capture_dialect=capture_dialect, - ) + result: ResultT = execute(scenario_plan) except Exception as exc: - result = ScenarioRunResult( - scenario_name=scenario.name, - status=ExecutionStatus.FAILED, - local_status=ScenarioLocalRunStatus.ERROR, - retained=False, - error_code=error_code(error=exc, fallback_code=SCENARIO_LOCAL_INTERNAL), - error_help=_scenario_failure_help(exc), - error_message=error_message(exc), - ) + result = failed(scenario_name=scenario.name, exc=exc) if result.status == ExecutionStatus.FAILED: lifecycle.failed(error_code=result.error_code) if projector is not None and on_scenario_complete is not None: @@ -223,103 +330,6 @@ def run_scenario_local_test_pipeline( return tuple(results) -def run_scenario_capture_pipeline( - *, - project_dir: Path, - pipeline_result: CompilePipelineResult, - scenarios: tuple[CompiledSqlScenario, ...], - connection_config: dict[str, object], - adapter: BaseAdapter, - project_name: str, - settings: ScenarioCaptureSettings, - connection_hooks: ConnectionHooks | None = None, - on_scenario_start: Callable[[CompiledSqlScenario], None] | None = None, - on_scenario_complete: Callable[ - [CompiledSqlScenario, ScenarioExecutionPlan | None, ScenarioSnapshotCaptureRunResult], None - ] - | None = None, -) -> tuple[ScenarioSnapshotCaptureRunResult, ...]: - """Capture selected scenario inputs into durable local snapshot files.""" - - hooks: ConnectionHooks = connection_hooks if connection_hooks is not None else ConnectionHooks() - if hooks.on_connection_start is not None: - hooks.on_connection_start(1) - start: float = time.monotonic() - try: - with OperationLifecycle( - operation_kind="scenario", operation_name="scenario_target_connection" - ): - connection: Any = adapter.connect(connection_config) - except Exception: - if hooks.on_connection_error is not None: - hooks.on_connection_error(1, elapsed_seconds=time.monotonic() - start) - raise - if hooks.on_connection_complete is not None: - hooks.on_connection_complete(1, elapsed_seconds=time.monotonic() - start) - try: - results: list[ScenarioSnapshotCaptureRunResult] = [] - scenario: CompiledSqlScenario - for scenario in scenarios: - scenario_plan: ScenarioExecutionPlan | None = None - resource_id: str = f"sql_scenario:{scenario.name}" - projector: NativeProgressProjector | None = _prepare_scenario_presentation( - scenario_name=scenario.name, - has_completion_callback=on_scenario_complete is not None, - ) - with ResourceAttemptLifecycle( - resource_id=resource_id, - resource_kind="scenario", - resource_name=scenario.name, - run_id=pipeline_result.project.run_id, - ) as lifecycle: - if on_scenario_start is not None: - on_scenario_start(scenario) - try: - scenario_plan = build_scenario_plan( - scenario=scenario, - pipeline_result=pipeline_result, - adapter=adapter, - project_name=project_name, - ) - result: ScenarioSnapshotCaptureRunResult = ( - execute_scenario_snapshot_capture_run( - project_dir=project_dir, - scenario_plan=scenario_plan, - adapter=adapter, - connection=connection, - run_id=pipeline_result.project.run_id, - settings=settings, - local_type_overrides=scenario_local_type_overrides_for_dialect( - scenario_config=pipeline_result.project.scenario, - sql_analysis_dialect=adapter.sql_analysis_dialect(), - ), - ) - ) - except Exception as exc: - result = ScenarioSnapshotCaptureRunResult( - scenario_name=scenario.name, - status=ExecutionStatus.FAILED, - retained=settings.retain, - error_code=error_code(error=exc, fallback_code=SCENARIO_EXEC_INTERNAL), - error_help=_scenario_failure_help(exc), - error_message=error_message(exc), - ) - if result.status == ExecutionStatus.FAILED: - lifecycle.failed(error_code=result.error_code) - if projector is not None and on_scenario_complete is not None: - _ = projector.consume_resource_terminal( - resource_name=scenario.name, - resource_id=resource_id, - resource_attempt_id=lifecycle.resource_attempt_id, - ) - results.append(result) - if on_scenario_complete is not None: - on_scenario_complete(scenario, scenario_plan, result) - return tuple(results) - finally: - adapter.close(connection) - - def _prepare_scenario_presentation( *, scenario_name: str, has_completion_callback: bool ) -> NativeProgressProjector | None: From 416ac204801b6be9d6844d37f9694ea54606af84 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:28:50 +0100 Subject: [PATCH 22/57] fix(scenario): observe the warehouse test target connection The capture pipeline wrapped adapter.connect in the scenario_target_connection operation lifecycle, but the warehouse scenario test pipeline connected to the same target without it, so observability consumers saw no connection operation for sqb test scenarios. Both warehouse pipelines now publish it. CLI text output is unchanged because the operation has no visible progress label. (cherry picked from commit e2e4e3e22a365075c7281988113a47b0e99bc715) --- .../executor/pipeline/_helpers/scenario.py | 14 +-- .../executor/pipeline/_helpers/_test_types.py | 10 ++ .../pipeline/_helpers/test_scenario.py | 100 +++++++++++++++++- 3 files changed, 111 insertions(+), 13 deletions(-) diff --git a/src/sqlbuild/executor/pipeline/_helpers/scenario.py b/src/sqlbuild/executor/pipeline/_helpers/scenario.py index 7f8794923c..12bc3f66d3 100644 --- a/src/sqlbuild/executor/pipeline/_helpers/scenario.py +++ b/src/sqlbuild/executor/pipeline/_helpers/scenario.py @@ -97,7 +97,6 @@ def run_scenario_test_pipeline( adapter=adapter, connection_config=connection_config, connection_hooks=connection_hooks, - observe_connection=False, ) as connection: def execute(scenario_plan: ScenarioExecutionPlan) -> ScenarioRunResult: @@ -204,7 +203,6 @@ def run_scenario_capture_pipeline( adapter=adapter, connection_config=connection_config, connection_hooks=connection_hooks, - observe_connection=True, ) as connection: def execute(scenario_plan: ScenarioExecutionPlan) -> ScenarioSnapshotCaptureRunResult: @@ -249,20 +247,16 @@ def _scenario_target_connection( adapter: BaseAdapter, connection_config: dict[str, object], connection_hooks: ConnectionHooks | None, - observe_connection: bool, ) -> Iterator[Any]: hooks: ConnectionHooks = connection_hooks if connection_hooks is not None else ConnectionHooks() if hooks.on_connection_start is not None: hooks.on_connection_start(1) start: float = time.monotonic() try: - if observe_connection: - with OperationLifecycle( - operation_kind="scenario", operation_name="scenario_target_connection" - ): - connection: Any = adapter.connect(connection_config) - else: - connection = adapter.connect(connection_config) + with OperationLifecycle( + operation_kind="scenario", operation_name="scenario_target_connection" + ): + connection: Any = adapter.connect(connection_config) except Exception: if hooks.on_connection_error is not None: hooks.on_connection_error(1, elapsed_seconds=time.monotonic() - start) diff --git a/tests/unit/src/sqlbuild/executor/pipeline/_helpers/_test_types.py b/tests/unit/src/sqlbuild/executor/pipeline/_helpers/_test_types.py index a56870ac9a..6222f977da 100644 --- a/tests/unit/src/sqlbuild/executor/pipeline/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/executor/pipeline/_helpers/_test_types.py @@ -91,6 +91,16 @@ class ScenarioFailureHelpTestCase: expected_help: str | None +@dataclass(frozen=True) +class ScenarioConnectionLifecycleTestCase: + """One warehouse scenario pipeline target-connection lifecycle case.""" + + description: str + expected_event_types: tuple[str, ...] + expected_operation_names: tuple[str, ...] + expected_connection_events: tuple[str, ...] + + @dataclass(frozen=True) class ScenarioTestPipelineTestCase: """One scenario test pipeline orchestration case.""" diff --git a/tests/unit/src/sqlbuild/executor/pipeline/_helpers/test_scenario.py b/tests/unit/src/sqlbuild/executor/pipeline/_helpers/test_scenario.py index 81c1caf395..795c0b8b82 100644 --- a/tests/unit/src/sqlbuild/executor/pipeline/_helpers/test_scenario.py +++ b/tests/unit/src/sqlbuild/executor/pipeline/_helpers/test_scenario.py @@ -17,11 +17,16 @@ SCENARIO_LOCAL_SNAPSHOT_MISSING, SCENARIO_LOCAL_SNAPSHOT_STALE, ) -from sqlbuild.executor.scenario.models import ScenarioRunResult, ScenarioSnapshotStateResult +from sqlbuild.executor.scenario.models import ( + ScenarioCaptureSettings, + ScenarioRunResult, + ScenarioSnapshotStateResult, +) from sqlbuild.executor.scenario.types import ScenarioLocalRunStatus, ScenarioSnapshotState from sqlbuild.executor.scheduling.types import ExecutionStatus from sqlbuild.observability import EventDispatcher, LifecycleEvent, dispatcher_scope from tests.unit.src.sqlbuild.executor.pipeline._helpers._test_types import ( + ScenarioConnectionLifecycleTestCase, ScenarioFailureHelpTestCase, ScenarioLocalPipelineTestCase, ScenarioPlanningLifecycleTestCase, @@ -119,19 +124,108 @@ def complete_scenario( assert tuple(adapter.events) == test_case.expected_connection_events assert results[-1].error_message is not None assert test_case.expected_error_fragment in results[-1].error_message + resource_events: list[LifecycleEvent] = lifecycle_events[2:] assert tuple(event.event_type for event in lifecycle_events) == ( + "operation_started", + "operation_completed", "resource_attempt_started", "resource_attempt_completed", "resource_attempt_started", "resource_attempt_failed", ) - assert tuple(event.resource_id for event in lifecycle_events) == ( + assert tuple(event.resource_id for event in resource_events) == ( "sql_scenario:passing_scenario", "sql_scenario:passing_scenario", "sql_scenario:planning_failure", "sql_scenario:planning_failure", ) - assert tuple(event.run_id for event in lifecycle_events) == ("scenario-test-run",) * 4 + assert tuple(event.run_id for event in resource_events) == ("scenario-test-run",) * 4 + + +@pytest.mark.parametrize( + "test_case", + ( + ScenarioConnectionLifecycleTestCase( + description="warehouse test pipeline observes the target connection", + expected_event_types=("operation_started", "operation_completed"), + expected_operation_names=("scenario_target_connection",) * 2, + expected_connection_events=("connect:pipeline.duckdb", "close"), + ), + ), + ids=lambda case: case.description, +) +def test_given_warehouse_test_pipeline_when_connecting_then_publishes_connection_operation( + test_case: ScenarioConnectionLifecycleTestCase, +) -> None: + pipeline_result: CompilePipelineResult = build_scenario_pipeline_result(scenario_names=()) + adapter: ScenarioPipelineTestAdapter = ScenarioPipelineTestAdapter() + lifecycle_events: list[LifecycleEvent] = [] + dispatcher: EventDispatcher = EventDispatcher() + dispatcher.subscribe_lifecycle(subscriber=lifecycle_events.append, accepts_opaque=False) + + with dispatcher_scope(dispatcher): + _ = scenario_pipeline.run_scenario_test_pipeline( + pipeline_result=pipeline_result, + scenarios=(), + connection_config={"database": "pipeline.duckdb"}, + adapter=adapter, + project_name="waffle_shop", + retain=False, + ) + + assert tuple(event.event_type for event in lifecycle_events) == test_case.expected_event_types + assert ( + tuple(event.payload.get("operation_name") for event in lifecycle_events) + == test_case.expected_operation_names + ) + assert tuple(adapter.events) == test_case.expected_connection_events + + +@pytest.mark.parametrize( + "test_case", + ( + ScenarioConnectionLifecycleTestCase( + description="capture pipeline observes the target connection", + expected_event_types=("operation_started", "operation_completed"), + expected_operation_names=("scenario_target_connection",) * 2, + expected_connection_events=("connect:pipeline.duckdb", "close"), + ), + ), + ids=lambda case: case.description, +) +def test_given_capture_pipeline_when_connecting_then_publishes_connection_operation( + tmp_path: Path, + test_case: ScenarioConnectionLifecycleTestCase, +) -> None: + pipeline_result: CompilePipelineResult = build_scenario_pipeline_result(scenario_names=()) + adapter: ScenarioPipelineTestAdapter = ScenarioPipelineTestAdapter() + lifecycle_events: list[LifecycleEvent] = [] + dispatcher: EventDispatcher = EventDispatcher() + dispatcher.subscribe_lifecycle(subscriber=lifecycle_events.append, accepts_opaque=False) + + with dispatcher_scope(dispatcher): + _ = scenario_pipeline.run_scenario_capture_pipeline( + project_dir=tmp_path, + pipeline_result=pipeline_result, + scenarios=(), + connection_config={"database": "pipeline.duckdb"}, + adapter=adapter, + project_name="waffle_shop", + settings=ScenarioCaptureSettings( + captured_at="2026-01-01T00:00:00Z", + capture_adapter="duckdb", + capture_dialect="duckdb", + sqlbuild_version="0.0.0", + retain=False, + ), + ) + + assert tuple(event.event_type for event in lifecycle_events) == test_case.expected_event_types + assert ( + tuple(event.payload.get("operation_name") for event in lifecycle_events) + == test_case.expected_operation_names + ) + assert tuple(adapter.events) == test_case.expected_connection_events @pytest.mark.parametrize( From b197a980e3216da57db01149605742f1235e215e Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:53:09 +0100 Subject: [PATCH 23/57] refactor(discovery): share per-file parse and module import helpers Route the SQL function, source, scenario, audit, SQL hook, enum and constant file discoverers through one fault-isolating parse loop, and the loader, Python node, provider, sink and materialization module loaders through one import helper. Per-kind rules are kept: SQL hooks skip underscore-prefixed files, sources scan top-level YAML only, and error types and messages are unchanged. Characterisation tests pin fault isolation and import failure messages per kind. (cherry picked from commit ca2cfb645961fbcc9dbcc4c64d79fbec581a5130) --- .../discovery/_helpers/filesystem/core.py | 497 +++++++++--------- .../discovery/_helpers/_test_types.py | 18 + .../compiler/discovery/_helpers/helpers.py | 10 + .../_helpers/test_filesystem_file_faults.py | 178 +++++++ 4 files changed, 448 insertions(+), 255 deletions(-) create mode 100644 tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_filesystem_file_faults.py diff --git a/src/sqlbuild/compiler/discovery/_helpers/filesystem/core.py b/src/sqlbuild/compiler/discovery/_helpers/filesystem/core.py index 7a954a3968..ebc6b82114 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/filesystem/core.py +++ b/src/sqlbuild/compiler/discovery/_helpers/filesystem/core.py @@ -5,7 +5,7 @@ import importlib.util import inspect import sys -from collections.abc import Callable +from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass, field from importlib.machinery import ModuleSpec from pathlib import Path @@ -57,7 +57,6 @@ SchemaParseError, ) from sqlbuild.compiler.discovery.models import ( - ConstantDeclaration, DiscoveredAdapterFile, DiscoveredAssetFunction, DiscoveredAuditFactory, @@ -87,7 +86,6 @@ DiscoveredSqlTestFile, DiscoveredTaskFunction, DiscoveryFileFault, - EnumDeclaration, ) from sqlbuild.compiler.resource_names.main._validate_resource_identity import ( validate_resource_identity, @@ -127,7 +125,7 @@ ) from sqlbuild.runtime.event_exporting.models import LifecycleEventSinkDefinition from sqlbuild.runtime.observability.models import LifecycleEvent -from sqlbuild.spec.contracts.models import SchemaModelEntry, SchemaSeedEntry, SourceEntry +from sqlbuild.spec.contracts.models import SchemaModelEntry, SchemaSeedEntry _PYTHON_NODE_KIND_FOLDERS: tuple[str, ...] = ("loaders", "tasks", "assets", "checks") _PYTHON_NODE_FACTORY_FOLDERS: tuple[str, ...] = (*_PYTHON_NODE_KIND_FOLDERS, "factories") @@ -465,39 +463,14 @@ def discover_enum_files( ) -> tuple[DiscoveredEnumFile, ...]: """Discover global and scoped enum declaration files.""" - discovered_files: list[DiscoveredEnumFile] = [] - facts: _DeclarationFileFacts - for facts in _discover_declaration_file_facts( + return _discover_declaration_kind_files( project_dir=project_dir, - declaration_kind=DeclarationKind.ENUM if isolate_declaration_kind else None, - ): - if facts.declaration_kind is not DeclarationKind.ENUM: - continue - try: - contents: str = facts.file_path.read_text(encoding="utf-8") - declarations: tuple[EnumDeclaration, ...] = parse_enum_declaration_file( - contents=contents, - file_path=facts.file_path, - relative_path=facts.relative_path, - ) - except (OSError, UnicodeError, ValueError, SyntaxError) as error: - if on_fault is None: - raise - on_fault(_discovery_fault(project_dir=project_dir, path=facts.file_path, error=error)) - continue - discovered_files.append( - DiscoveredEnumFile( - file_path=facts.file_path, - relative_path=facts.relative_path, - contents=contents, - declarations=declarations, - scope_kind=facts.scope_kind, - ownership_root=facts.ownership_root, - owning_path=facts.owning_path, - declaration_root=facts.declaration_root, - ) - ) - return tuple(discovered_files) + declaration_kind=DeclarationKind.ENUM, + isolate_declaration_kind=isolate_declaration_kind, + parse_declarations=parse_enum_declaration_file, + discovered_file_type=DiscoveredEnumFile, + on_fault=on_fault, + ) def discover_constant_files( @@ -508,39 +481,102 @@ def discover_constant_files( ) -> tuple[DiscoveredConstantFile, ...]: """Discover global and scoped constant declaration files.""" - discovered_files: list[DiscoveredConstantFile] = [] - facts: _DeclarationFileFacts - for facts in _discover_declaration_file_facts( + return _discover_declaration_kind_files( project_dir=project_dir, - declaration_kind=DeclarationKind.CONSTANT if isolate_declaration_kind else None, - ): - if facts.declaration_kind is not DeclarationKind.CONSTANT: - continue - try: - contents: str = facts.file_path.read_text(encoding="utf-8") - declarations: tuple[ConstantDeclaration, ...] = parse_constant_declaration_file( - contents=contents, - file_path=facts.file_path, - relative_path=facts.relative_path, + declaration_kind=DeclarationKind.CONSTANT, + isolate_declaration_kind=isolate_declaration_kind, + parse_declarations=parse_constant_declaration_file, + discovered_file_type=DiscoveredConstantFile, + on_fault=on_fault, + ) + + +def _discover_declaration_kind_files[DeclarationT, DiscoveredT]( + *, + project_dir: Path, + declaration_kind: DeclarationKind, + isolate_declaration_kind: bool, + parse_declarations: Callable[..., tuple[DeclarationT, ...]], + discovered_file_type: Callable[..., DiscoveredT], + on_fault: Callable[[DiscoveryFileFault], None] | None, +) -> tuple[DiscoveredT, ...]: + def parse(facts: _DeclarationFileFacts) -> DiscoveredT: + contents: str = facts.file_path.read_text(encoding="utf-8") + declarations: tuple[DeclarationT, ...] = parse_declarations( + contents=contents, + file_path=facts.file_path, + relative_path=facts.relative_path, + ) + return discovered_file_type( + file_path=facts.file_path, + relative_path=facts.relative_path, + contents=contents, + declarations=declarations, + scope_kind=facts.scope_kind, + ownership_root=facts.ownership_root, + owning_path=facts.owning_path, + declaration_root=facts.declaration_root, + ) + + return _parse_discovered_files( + project_dir=project_dir, + items=( + facts + for facts in _discover_declaration_file_facts( + project_dir=project_dir, + declaration_kind=declaration_kind if isolate_declaration_kind else None, ) + if facts.declaration_kind is declaration_kind + ), + item_path=lambda facts: facts.file_path, + parse=parse, + on_fault=on_fault, + ) + + +def _parse_discovered_files[ItemT, DiscoveredT]( + *, + project_dir: Path, + items: Iterable[ItemT], + item_path: Callable[[ItemT], Path], + parse: Callable[[ItemT], DiscoveredT], + on_fault: Callable[[DiscoveryFileFault], None] | None, +) -> tuple[DiscoveredT, ...]: + """Parse each discovered file, isolating per-file read and parse faults when requested.""" + + discovered: list[DiscoveredT] = [] + item: ItemT + for item in items: + try: + discovered.append(parse(item)) except (OSError, UnicodeError, ValueError, SyntaxError) as error: if on_fault is None: raise - on_fault(_discovery_fault(project_dir=project_dir, path=facts.file_path, error=error)) - continue - discovered_files.append( - DiscoveredConstantFile( - file_path=facts.file_path, - relative_path=facts.relative_path, - contents=contents, - declarations=declarations, - scope_kind=facts.scope_kind, - ownership_root=facts.ownership_root, - owning_path=facts.owning_path, - declaration_root=facts.declaration_root, - ) - ) - return tuple(discovered_files) + on_fault(_discovery_fault(project_dir=project_dir, path=item_path(item), error=error)) + return tuple(discovered) + + +def _parse_discovered_paths[DiscoveredT]( + *, + project_dir: Path, + file_paths: Iterable[Path], + parse: Callable[[Path], DiscoveredT], + on_fault: Callable[[DiscoveryFileFault], None] | None, +) -> tuple[DiscoveredT, ...]: + return _parse_discovered_files( + project_dir=project_dir, + items=file_paths, + item_path=lambda file_path: file_path, + parse=parse, + on_fault=on_fault, + ) + + +def _unscoped_files(*, root: Path, pattern: str, project_dir: Path) -> Iterator[Path]: + file_path: Path + for file_path in sorted(root.rglob(pattern)): + if not _is_in_scoped_declaration_tree(file_path=file_path, project_dir=project_dir): + yield file_path def discover_model_schema_files(*, project_dir: Path) -> tuple[DiscoveredModelSchemaFile, ...]: @@ -578,31 +614,25 @@ def discover_sql_function_files( if not function_root.is_dir(): return () - discovered_function_files: list[DiscoveredSqlFunctionFile] = [] - file_path: Path - for file_path in sorted(function_root.rglob("*.sql")): - if _is_in_scoped_declaration_tree(file_path=file_path, project_dir=project_dir): - continue - try: - contents: str = file_path.read_text(encoding="utf-8") - header_values: dict[str, object] - body_sql: str - header_values, body_sql = parse_function_sql(contents=contents, file_path=file_path) - except (OSError, UnicodeError, ValueError, SyntaxError) as error: - if on_fault is None: - raise - on_fault(_discovery_fault(project_dir=project_dir, path=file_path, error=error)) - continue - discovered_function_files.append( - DiscoveredSqlFunctionFile( - file_path=file_path, - relative_path=_project_relative_path(path=file_path, project_dir=project_dir), - contents=contents, - header_values=header_values, - body_sql=body_sql, - ) + def parse(file_path: Path) -> DiscoveredSqlFunctionFile: + contents: str = file_path.read_text(encoding="utf-8") + header_values: dict[str, object] + body_sql: str + header_values, body_sql = parse_function_sql(contents=contents, file_path=file_path) + return DiscoveredSqlFunctionFile( + file_path=file_path, + relative_path=_project_relative_path(path=file_path, project_dir=project_dir), + contents=contents, + header_values=header_values, + body_sql=body_sql, ) - return tuple(discovered_function_files) + + return _parse_discovered_paths( + project_dir=project_dir, + file_paths=_unscoped_files(root=function_root, pattern="*.sql", project_dir=project_dir), + parse=parse, + on_fault=on_fault, + ) def discover_python_function_files( @@ -691,28 +721,19 @@ def discover_source_files( yaml_paths: tuple[Path, ...] = tuple( sorted(path for path in sources_root.iterdir() if path.suffix in YAML_FILE_SUFFIXES) ) - discovered_source_files: list[DiscoveredSourceFile] = [] - file_path: Path - for file_path in yaml_paths: - try: - contents: str = file_path.read_text(encoding="utf-8") - source_entries: tuple[SourceEntry, ...] = parse_sources_yml( - contents=contents, file_path=file_path - ) - except (OSError, UnicodeError, ValueError, SyntaxError) as error: - if on_fault is None: - raise - on_fault(_discovery_fault(project_dir=project_dir, path=file_path, error=error)) - continue - discovered_source_files.append( - DiscoveredSourceFile( - file_path=file_path, - relative_path=_project_relative_path(path=file_path, project_dir=project_dir), - contents=contents, - source_entries=source_entries, - ) + + def parse(file_path: Path) -> DiscoveredSourceFile: + contents: str = file_path.read_text(encoding="utf-8") + return DiscoveredSourceFile( + file_path=file_path, + relative_path=_project_relative_path(path=file_path, project_dir=project_dir), + contents=contents, + source_entries=parse_sources_yml(contents=contents, file_path=file_path), ) - return tuple(discovered_source_files) + + return _parse_discovered_paths( + project_dir=project_dir, file_paths=yaml_paths, parse=parse, on_fault=on_fault + ) def discover_seed_files(*, project_dir: Path) -> tuple[DiscoveredSeedFile, ...]: @@ -766,25 +787,19 @@ def discover_scenario_files( if not scenarios_root.is_dir(): return () - discovered_scenario_files: list[DiscoveredSqlScenarioFile] = [] - file_path: Path - for file_path in sorted(scenarios_root.rglob("*.sql")): - if _is_in_scoped_declaration_tree(file_path=file_path, project_dir=project_dir): - continue - try: - contents: str = file_path.read_text(encoding="utf-8") - scenario: DiscoveredSqlScenarioFile = parse_sql_scenario_file( - contents=contents, - file_path=file_path, - relative_path=_project_relative_path(path=file_path, project_dir=project_dir), - ) - except (OSError, UnicodeError, ValueError, SyntaxError) as error: - if on_fault is None: - raise - on_fault(_discovery_fault(project_dir=project_dir, path=file_path, error=error)) - continue - discovered_scenario_files.append(scenario) - return tuple(discovered_scenario_files) + def parse(file_path: Path) -> DiscoveredSqlScenarioFile: + return parse_sql_scenario_file( + contents=file_path.read_text(encoding="utf-8"), + file_path=file_path, + relative_path=_project_relative_path(path=file_path, project_dir=project_dir), + ) + + return _parse_discovered_paths( + project_dir=project_dir, + file_paths=_unscoped_files(root=scenarios_root, pattern="*.sql", project_dir=project_dir), + parse=parse, + on_fault=on_fault, + ) def discover_audit_files( @@ -796,27 +811,21 @@ def discover_audit_files( if not audits_root.is_dir(): return () - discovered_audit_files: list[DiscoveredAuditFile] = [] - file_path: Path - for file_path in sorted(audits_root.rglob("*.sql")): - if _is_in_scoped_declaration_tree(file_path=file_path, project_dir=project_dir): - continue - try: - contents: str = file_path.read_text(encoding="utf-8") - discovered_audit_files.append( - DiscoveredAuditFile( - file_path=file_path, - relative_path=_project_relative_path(path=file_path, project_dir=project_dir), - contents=contents, - blocks=parse_sql_audit_file(contents=contents, file_path=file_path), - ) - ) - except (OSError, UnicodeError, ValueError, SyntaxError) as error: - if on_fault is None: - raise - on_fault(_discovery_fault(project_dir=project_dir, path=file_path, error=error)) - continue - return tuple(discovered_audit_files) + def parse(file_path: Path) -> DiscoveredAuditFile: + contents: str = file_path.read_text(encoding="utf-8") + return DiscoveredAuditFile( + file_path=file_path, + relative_path=_project_relative_path(path=file_path, project_dir=project_dir), + contents=contents, + blocks=parse_sql_audit_file(contents=contents, file_path=file_path), + ) + + return _parse_discovered_paths( + project_dir=project_dir, + file_paths=_unscoped_files(root=audits_root, pattern="*.sql", project_dir=project_dir), + parse=parse, + on_fault=on_fault, + ) def discover_macro_files( @@ -1003,29 +1012,26 @@ def discover_sql_hook_files( if not hooks_root.is_dir(): return () - discovered_hooks: list[DiscoveredSqlHookFile] = [] - file_path: Path - for file_path in sorted(hooks_root.rglob("*.sql")): - if file_path.name.startswith("_") or _is_in_scoped_declaration_tree( - file_path=file_path, project_dir=project_dir - ): - continue - try: - relative_path: Path = _project_relative_path(path=file_path, project_dir=project_dir) - contents: str = file_path.read_text(encoding="utf-8") - discovered_hooks.append( - parse_sql_hook_file( - contents=contents, - file_path=file_path, - relative_path=relative_path, - ) + def parse(file_path: Path) -> DiscoveredSqlHookFile: + relative_path: Path = _project_relative_path(path=file_path, project_dir=project_dir) + return parse_sql_hook_file( + contents=file_path.read_text(encoding="utf-8"), + file_path=file_path, + relative_path=relative_path, + ) + + return _parse_discovered_paths( + project_dir=project_dir, + file_paths=( + file_path + for file_path in _unscoped_files( + root=hooks_root, pattern="*.sql", project_dir=project_dir ) - except (OSError, UnicodeError, ValueError, SyntaxError) as error: - if on_fault is None: - raise - on_fault(_discovery_fault(project_dir=project_dir, path=file_path, error=error)) - continue - return tuple(discovered_hooks) + if not file_path.name.startswith("_") + ), + parse=parse, + on_fault=on_fault, + ) def discover_provider_classes(*, project_dir: Path) -> tuple[DiscoveredProvider, ...]: @@ -1804,82 +1810,56 @@ def _normalize_factory_result( def _load_loader_module(*, file_path: Path, project_dir: Path) -> ModuleType: - module_name: str = "sqlbuild_project_loader_" + "_".join( - _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts + return _exec_project_module( + module_name="sqlbuild_project_loader_" + + "_".join( + _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts + ), + file_path=file_path, + project_dir=project_dir, + file_label="source loader", + error_type=LoaderDiscoveryError, ) - spec: ModuleSpec | None = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None or spec.loader is None: - raise LoaderDiscoveryError(f"Could not load source loader file {file_path}") - module: ModuleType = importlib.util.module_from_spec(spec) - old_path: list[str] = list(sys.path) - sys.path.insert(0, str(project_dir)) - try: - sys.modules[module_name] = module - spec.loader.exec_module(module) - except Exception as error: - raise LoaderDiscoveryError( - f"Failed to import source loader file {file_path.relative_to(project_dir)}: {error}" - ) from error - finally: - sys.path = old_path - return module def _load_python_node_module(*, file_path: Path, project_dir: Path, node_folder: str) -> ModuleType: - module_name: str = "sqlbuild_project_python_node_" + "_".join( - _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts + return _exec_project_module( + module_name="sqlbuild_project_python_node_" + + "_".join( + _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts + ), + file_path=file_path, + project_dir=project_dir, + file_label="Python node", + error_type=PythonNodeDiscoveryError, ) - spec: ModuleSpec | None = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None or spec.loader is None: - raise PythonNodeDiscoveryError(f"Could not load Python node file {file_path}") - module: ModuleType = importlib.util.module_from_spec(spec) - old_path: list[str] = list(sys.path) - sys.path.insert(0, str(project_dir)) - try: - sys.modules[module_name] = module - spec.loader.exec_module(module) - except Exception as error: - raise PythonNodeDiscoveryError( - f"Failed to import Python node file {file_path.relative_to(project_dir)}: {error}" - ) from error - finally: - sys.path = old_path - return module def _load_provider_module(*, file_path: Path, project_dir: Path) -> ModuleType: - module_name: str = ".".join( - _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts - ) - _evict_stale_project_package_modules( - root_module=module_name.split(".", maxsplit=1)[0], + return _load_project_package_module( + file_path=file_path, project_dir=project_dir, + file_label="provider", + error_type=ProviderDiscoveryError, ) - existing_module: ModuleType | None = sys.modules.get(module_name) - if existing_module is not None: - existing_file: object = getattr(existing_module, "__file__", None) - if isinstance(existing_file, str) and Path(existing_file).resolve() == file_path.resolve(): - return existing_module - sys.modules.pop(module_name, None) - spec: ModuleSpec | None = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None or spec.loader is None: - raise ProviderDiscoveryError(f"Could not load provider file {file_path}") - module: ModuleType = importlib.util.module_from_spec(spec) - old_path: list[str] = list(sys.path) - sys.path.insert(0, str(project_dir)) - try: - sys.modules[module_name] = module - spec.loader.exec_module(module) - except Exception as error: - raise ProviderDiscoveryError( - f"Failed to import provider file {file_path.relative_to(project_dir)}: {error}" - ) from error - finally: - sys.path = old_path - return module def _load_sink_module(*, file_path: Path, project_dir: Path) -> ModuleType: + return _load_project_package_module( + file_path=file_path, + project_dir=project_dir, + file_label="sink", + error_type=EventExporterDiscoveryError, + ) + + +def _load_project_package_module( + *, + file_path: Path, + project_dir: Path, + file_label: str, + error_type: type[Exception], +) -> ModuleType: module_name: str = ".".join( _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts ) @@ -1893,9 +1873,26 @@ def _load_sink_module(*, file_path: Path, project_dir: Path) -> ModuleType: if isinstance(existing_file, str) and Path(existing_file).resolve() == file_path.resolve(): return existing_module sys.modules.pop(module_name, None) + return _exec_project_module( + module_name=module_name, + file_path=file_path, + project_dir=project_dir, + file_label=file_label, + error_type=error_type, + ) + + +def _exec_project_module( + *, + module_name: str, + file_path: Path, + project_dir: Path, + file_label: str, + error_type: type[Exception], +) -> ModuleType: spec: ModuleSpec | None = importlib.util.spec_from_file_location(module_name, file_path) if spec is None or spec.loader is None: - raise EventExporterDiscoveryError(f"Could not load sink file {file_path}") + raise error_type(f"Could not load {file_label} file {file_path}") module: ModuleType = importlib.util.module_from_spec(spec) old_path: list[str] = list(sys.path) sys.path.insert(0, str(project_dir)) @@ -1903,8 +1900,8 @@ def _load_sink_module(*, file_path: Path, project_dir: Path) -> ModuleType: sys.modules[module_name] = module spec.loader.exec_module(module) except Exception as error: - raise EventExporterDiscoveryError( - f"Failed to import sink file {file_path.relative_to(project_dir)}: {error}" + raise error_type( + f"Failed to import {file_label} file {file_path.relative_to(project_dir)}: {error}" ) from error finally: sys.path = old_path @@ -1936,25 +1933,15 @@ def _evict_stale_project_package_modules(*, root_module: str, project_dir: Path) def _load_materialization_module(*, file_path: Path, project_dir: Path) -> ModuleType: - module_name: str = ".".join( - _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts + return _exec_project_module( + module_name=".".join( + _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts + ), + file_path=file_path, + project_dir=project_dir, + file_label="materialization", + error_type=PythonNodeDiscoveryError, ) - spec: ModuleSpec | None = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None or spec.loader is None: - raise PythonNodeDiscoveryError(f"Could not load materialization file {file_path}") - module: ModuleType = importlib.util.module_from_spec(spec) - old_path: list[str] = list(sys.path) - sys.path.insert(0, str(project_dir)) - try: - sys.modules[module_name] = module - spec.loader.exec_module(module) - except Exception as error: - raise PythonNodeDiscoveryError( - f"Failed to import materialization file {file_path.relative_to(project_dir)}: {error}" - ) from error - finally: - sys.path = old_path - return module def discover_adapter_file(*, project_dir: Path) -> DiscoveredAdapterFile | None: diff --git a/tests/unit/src/sqlbuild/compiler/discovery/_helpers/_test_types.py b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/_test_types.py index db48cfebc5..2c10692b0a 100644 --- a/tests/unit/src/sqlbuild/compiler/discovery/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/_test_types.py @@ -827,3 +827,21 @@ class LoadTargetRetentionPoliciesErrorTestCase: description: str target_lines: tuple[str, ...] expected_error_fragment: str + + +@dataclass(frozen=True) +class DiscoverFileFaultsTestCase: + description: str + discovery_name: str + unreadable_files: tuple[str, ...] + expected_fault_paths: tuple[Path, ...] + expected_error_type: type[Exception] = UnicodeDecodeError + + +@dataclass(frozen=True) +class DiscoverModuleImportFailureTestCase: + description: str + discovery_name: str + broken_file: str + expected_error_type: type[Exception] + expected_error_fragment: str diff --git a/tests/unit/src/sqlbuild/compiler/discovery/_helpers/helpers.py b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/helpers.py index ac65e0bb96..55ce12db8c 100644 --- a/tests/unit/src/sqlbuild/compiler/discovery/_helpers/helpers.py +++ b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/helpers.py @@ -231,3 +231,13 @@ def discovered_test_case_values( for test_case in cases: case_values.append(tuple(value.value for _name, value in test_case.values)) return tuple(case_values) + + +def write_unreadable_files(*, project_dir: Path, relative_paths: tuple[str, ...]) -> None: + """Write files whose bytes are not valid UTF-8 under project_dir.""" + + relative_path: str + for relative_path in relative_paths: + file_path: Path = project_dir / relative_path + file_path.parent.mkdir(parents=True, exist_ok=True) + _ = file_path.write_bytes(b"\xff\xfe\xfa") diff --git a/tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_filesystem_file_faults.py b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_filesystem_file_faults.py new file mode 100644 index 0000000000..b82e41149b --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_filesystem_file_faults.py @@ -0,0 +1,178 @@ +"""Tests for per-file fault isolation and module import failures in filesystem discovery.""" + +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path + +import pytest + +from sqlbuild.compiler.discovery._helpers.filesystem.command_output_sinks import ( + discover_command_output_sink_declarations, +) +from sqlbuild.compiler.discovery._helpers.filesystem.core import ( + discover_audit_files, + discover_constant_files, + discover_enum_files, + discover_event_exporter_declarations, + discover_materialization_files, + discover_provider_classes, + discover_scenario_files, + discover_source_files, + discover_sql_function_files, + discover_sql_hook_files, +) +from sqlbuild.compiler.discovery.exceptions import ( + EventExporterDiscoveryError, + ProviderDiscoveryError, + PythonNodeDiscoveryError, +) +from sqlbuild.compiler.discovery.models import DiscoveryFileFault +from tests.unit.src.sqlbuild.compiler.discovery._helpers._test_types import ( + DiscoverFileFaultsTestCase, + DiscoverModuleImportFailureTestCase, +) +from tests.unit.src.sqlbuild.compiler.discovery._helpers.helpers import write_unreadable_files + +_FILE_DISCOVERIES: dict[str, Callable[..., tuple[object, ...]]] = { + "audits": discover_audit_files, + "constants": discover_constant_files, + "enums": discover_enum_files, + "scenarios": discover_scenario_files, + "sources": discover_source_files, + "sql_functions": discover_sql_function_files, + "sql_hooks": discover_sql_hook_files, +} +_MODULE_DISCOVERIES: dict[str, Callable[..., object]] = { + "command_output_sinks": discover_command_output_sink_declarations, + "event_exporters": discover_event_exporter_declarations, + "materializations": discover_materialization_files, + "providers": discover_provider_classes, +} + + +@pytest.mark.parametrize( + "test_case", + ( + DiscoverFileFaultsTestCase( + description="audits fault every nested SQL file", + discovery_name="audits", + unreadable_files=("audits/orders.sql", "audits/nested/_customers.sql"), + expected_fault_paths=(Path("audits/nested/_customers.sql"), Path("audits/orders.sql")), + ), + DiscoverFileFaultsTestCase( + description="constants fault declaration files", + discovery_name="constants", + unreadable_files=("constants/order_limits.sql",), + expected_fault_paths=(Path("constants/order_limits.sql"),), + ), + DiscoverFileFaultsTestCase( + description="enums fault declaration files", + discovery_name="enums", + unreadable_files=("enums/order_status.sql",), + expected_fault_paths=(Path("enums/order_status.sql"),), + ), + DiscoverFileFaultsTestCase( + description="scenarios fault nested SQL files", + discovery_name="scenarios", + unreadable_files=("tests/scenarios/orders/_refunds.sql",), + expected_fault_paths=(Path("tests/scenarios/orders/_refunds.sql"),), + ), + DiscoverFileFaultsTestCase( + description="sources fault only top-level YAML files", + discovery_name="sources", + unreadable_files=( + "sources/orders.yml", + "sources/customers.yaml", + "sources/notes.txt", + "sources/nested/products.yml", + ), + expected_fault_paths=(Path("sources/customers.yaml"), Path("sources/orders.yml")), + ), + DiscoverFileFaultsTestCase( + description="sql functions fault underscore-prefixed files", + discovery_name="sql_functions", + unreadable_files=("functions/sql/_order_total.sql", "functions/sql/tax.sql"), + expected_fault_paths=( + Path("functions/sql/_order_total.sql"), + Path("functions/sql/tax.sql"), + ), + ), + DiscoverFileFaultsTestCase( + description="sql hooks skip underscore-prefixed files", + discovery_name="sql_hooks", + unreadable_files=("hooks/sql/_shared.sql", "hooks/sql/grants/grant_access.sql"), + expected_fault_paths=(Path("hooks/sql/grants/grant_access.sql"),), + ), + ), + ids=lambda case: case.description, +) +def test_given_unreadable_files_when_discovering_then_faults_each_file_or_raises( + tmp_path: Path, + test_case: DiscoverFileFaultsTestCase, +) -> None: + write_unreadable_files(project_dir=tmp_path, relative_paths=test_case.unreadable_files) + discover: Callable[..., tuple[object, ...]] = _FILE_DISCOVERIES[test_case.discovery_name] + faults: list[DiscoveryFileFault] = [] + + discovered: tuple[object, ...] = discover(project_dir=tmp_path, on_fault=faults.append) + + assert discovered == () + assert tuple(fault.path for fault in faults) == test_case.expected_fault_paths + assert all("utf-8" in fault.message for fault in faults) + with pytest.raises(test_case.expected_error_type): + _ = discover(project_dir=tmp_path) + + +@pytest.mark.parametrize( + "test_case", + ( + DiscoverModuleImportFailureTestCase( + description="command output sink import failure identifies the sink file", + discovery_name="command_output_sinks", + broken_file="sinks/broken_orders.py", + expected_error_type=EventExporterDiscoveryError, + expected_error_fragment="Failed to import sink file sinks/broken_orders.py: ", + ), + DiscoverModuleImportFailureTestCase( + description="lifecycle sink import failure identifies the sink file", + discovery_name="event_exporters", + broken_file="sinks/broken_orders.py", + expected_error_type=EventExporterDiscoveryError, + expected_error_fragment="Failed to import sink file sinks/broken_orders.py: ", + ), + DiscoverModuleImportFailureTestCase( + description="materialization import failure identifies the materialization file", + discovery_name="materializations", + broken_file="materializations/broken_orders.py", + expected_error_type=PythonNodeDiscoveryError, + expected_error_fragment=( + "Failed to import materialization file materializations/broken_orders.py: " + ), + ), + DiscoverModuleImportFailureTestCase( + description="provider import failure identifies the provider file", + discovery_name="providers", + broken_file="providers/broken_orders.py", + expected_error_type=ProviderDiscoveryError, + expected_error_fragment="Failed to import provider file providers/broken_orders.py: ", + ), + ), + ids=lambda case: case.description, +) +def test_given_module_import_failure_when_discovering_then_raises_kind_specific_error( + tmp_path: Path, + test_case: DiscoverModuleImportFailureTestCase, +) -> None: + file_path: Path = tmp_path / test_case.broken_file + file_path.parent.mkdir(parents=True) + _ = file_path.write_text("raise RuntimeError('orders module failed')\n", encoding="utf-8") + + with pytest.raises(test_case.expected_error_type) as error_info: + _ = _MODULE_DISCOVERIES[test_case.discovery_name](project_dir=tmp_path) + + assert str(error_info.value) == f"{test_case.expected_error_fragment}orders module failed" + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) From 9bc506e20bf39225283464cc0a0b6aaaab16b0f2 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:58:20 +0100 Subject: [PATCH 24/57] refactor(python-nodes): build every Python node kind through one builder Loader, task, asset and check nodes now share one builder for the kind-independent fields, provider usages and identity, with a documented variant for kinds that carry tags, group, description and meta. Each kind still supplies exactly the same identity decorator_config fields as before; a characterisation test pins the per-kind field sets so version identities do not change. (cherry picked from commit b7504328ac306a13ee10688b85914709e0d0bb01) --- .../python_nodes/_helpers/inventory.py | 222 ++++++++++-------- .../python_nodes/_helpers/_test_types.py | 6 + .../python_nodes/_helpers/test_inventory.py | 125 +++++++++- 3 files changed, 249 insertions(+), 104 deletions(-) diff --git a/src/sqlbuild/compiler/python_nodes/_helpers/inventory.py b/src/sqlbuild/compiler/python_nodes/_helpers/inventory.py index f18448b7a1..fbe89b7bd4 100644 --- a/src/sqlbuild/compiler/python_nodes/_helpers/inventory.py +++ b/src/sqlbuild/compiler/python_nodes/_helpers/inventory.py @@ -4,7 +4,9 @@ import inspect from collections.abc import Callable -from typing import get_type_hints +from dataclasses import replace +from pathlib import Path +from typing import Protocol, get_type_hints from sqlbuild.compiler.discovery.models import ( DiscoveredAssetFunction, @@ -174,34 +176,45 @@ def build_python_node_path_index(*, nodes: tuple[DiscoveredPythonNode, ...]) -> return {node.name: node.relative_path.parent.as_posix() for node in nodes} +class _DiscoveredPythonFunction(Protocol): + @property + def file_path(self) -> Path: ... + + @property + def relative_path(self) -> Path: ... + + @property + def name(self) -> str: ... + + @property + def function(self) -> Callable[..., object]: ... + + @property + def depends_on(self) -> tuple[Callable[..., object] | SqlResourceRef, ...]: ... + + +class _DocumentedPythonFunction(_DiscoveredPythonFunction, Protocol): + @property + def tags(self) -> tuple[str, ...]: ... + + @property + def group(self) -> str | None: ... + + @property + def description(self) -> str | None: ... + + @property + def meta(self) -> dict[str, object] | None: ... + + def _build_loader_node( *, loader: DiscoveredLoaderFunction, provider_by_name: dict[str, DiscoveredProvider] ) -> DiscoveredPythonNode: - return DiscoveredPythonNode( - kind=PythonNodeKind.LOADER, - file_path=loader.file_path, - relative_path=loader.relative_path, - name=loader.name, - function=loader.function, - depends_on=loader.depends_on, - sql_deps=_sql_deps(loader.depends_on), - provider_usages=_provider_usages( - function=loader.function, provider_by_name=provider_by_name - ), - loader=DiscoveredPythonLoaderMetadata( - destination=loader.destination, - write_strategy=loader.write_strategy, - cursor_column=loader.cursor_column, - unique_key=loader.unique_key, - columns=loader.columns, - contract=loader.contract, - connection_mode=loader.connection_mode, - ), - identity=build_python_identity( - node_type=PythonNodeKind.LOADER.value, - node_name=loader.name, - function=loader.function, - project_dir=loader.file_path.parent, + return replace( + _build_python_node( + kind=PythonNodeKind.LOADER, + discovered=loader, + provider_by_name=provider_by_name, decorator_config={ "contract": loader.contract, "cursor_column": loader.cursor_column, @@ -212,112 +225,117 @@ def _build_loader_node( else None, }, ), + loader=DiscoveredPythonLoaderMetadata( + destination=loader.destination, + write_strategy=loader.write_strategy, + cursor_column=loader.cursor_column, + unique_key=loader.unique_key, + columns=loader.columns, + contract=loader.contract, + connection_mode=loader.connection_mode, + ), ) def _build_task_node( *, task: DiscoveredTaskFunction, provider_by_name: dict[str, DiscoveredProvider] ) -> DiscoveredPythonNode: - return DiscoveredPythonNode( - kind=PythonNodeKind.TASK, - file_path=task.file_path, - relative_path=task.relative_path, - name=task.name, - function=task.function, - depends_on=task.depends_on, - sql_deps=_sql_deps(task.depends_on), - tags=task.tags, - group=task.group, - description=task.description, - meta=task.meta, - provider_usages=_provider_usages(function=task.function, provider_by_name=provider_by_name), - task=DiscoveredPythonTaskMetadata(retry=task.retry), - identity=build_python_identity( - node_type=PythonNodeKind.TASK.value, - node_name=task.name, - function=task.function, - project_dir=task.file_path.parent, - decorator_config={ - "description": task.description, - "group": task.group, - "meta": task.meta, - "tags": task.tags, - }, + return replace( + _build_documented_python_node( + kind=PythonNodeKind.TASK, + discovered=task, + provider_by_name=provider_by_name, + decorator_config={}, ), + task=DiscoveredPythonTaskMetadata(retry=task.retry), ) def _build_asset_node( *, asset: DiscoveredAssetFunction, provider_by_name: dict[str, DiscoveredProvider] ) -> DiscoveredPythonNode: - return DiscoveredPythonNode( - kind=PythonNodeKind.ASSET, - file_path=asset.file_path, - relative_path=asset.relative_path, - name=asset.name, - function=asset.function, - depends_on=asset.depends_on, - sql_deps=_sql_deps(asset.depends_on), - tags=asset.tags, - group=asset.group, - description=asset.description, - meta=asset.meta, - provider_usages=_provider_usages( - function=asset.function, provider_by_name=provider_by_name + return replace( + _build_documented_python_node( + kind=PythonNodeKind.ASSET, + discovered=asset, + provider_by_name=provider_by_name, + decorator_config={"columns": asset.columns, "column_lineage": asset.column_lineage}, ), asset=DiscoveredPythonAssetMetadata( columns=asset.columns, column_lineage=asset.column_lineage, retry=asset.retry, ), - identity=build_python_identity( - node_type=PythonNodeKind.ASSET.value, - node_name=asset.name, - function=asset.function, - project_dir=asset.file_path.parent, + ) + + +def _build_check_node( + *, check: DiscoveredCheckFunction, provider_by_name: dict[str, DiscoveredProvider] +) -> DiscoveredPythonNode: + return replace( + _build_documented_python_node( + kind=PythonNodeKind.CHECK, + discovered=check, + provider_by_name=provider_by_name, + decorator_config={"severity": check.severity.value}, + ), + check=DiscoveredPythonCheckMetadata(severity=check.severity), + ) + + +def _build_documented_python_node( + *, + kind: PythonNodeKind, + discovered: _DocumentedPythonFunction, + provider_by_name: dict[str, DiscoveredProvider], + decorator_config: dict[str, object], +) -> DiscoveredPythonNode: + return replace( + _build_python_node( + kind=kind, + discovered=discovered, + provider_by_name=provider_by_name, decorator_config={ - "columns": asset.columns, - "column_lineage": asset.column_lineage, - "description": asset.description, - "group": asset.group, - "meta": asset.meta, - "tags": asset.tags, + **decorator_config, + "description": discovered.description, + "group": discovered.group, + "meta": discovered.meta, + "tags": discovered.tags, }, ), + tags=discovered.tags, + group=discovered.group, + description=discovered.description, + meta=discovered.meta, ) -def _build_check_node( - *, check: DiscoveredCheckFunction, provider_by_name: dict[str, DiscoveredProvider] +def _build_python_node( + *, + kind: PythonNodeKind, + discovered: _DiscoveredPythonFunction, + provider_by_name: dict[str, DiscoveredProvider], + decorator_config: dict[str, object], ) -> DiscoveredPythonNode: + """Build the kind-independent node view; decorator_config is exactly the identity input.""" + return DiscoveredPythonNode( - kind=PythonNodeKind.CHECK, - file_path=check.file_path, - relative_path=check.relative_path, - name=check.name, - function=check.function, - depends_on=check.depends_on, - sql_deps=_sql_deps(check.depends_on), - tags=check.tags, - group=check.group, - description=check.description, - meta=check.meta, + kind=kind, + file_path=discovered.file_path, + relative_path=discovered.relative_path, + name=discovered.name, + function=discovered.function, + depends_on=discovered.depends_on, + sql_deps=_sql_deps(discovered.depends_on), provider_usages=_provider_usages( - function=check.function, provider_by_name=provider_by_name + function=discovered.function, provider_by_name=provider_by_name ), - check=DiscoveredPythonCheckMetadata(severity=check.severity), identity=build_python_identity( - node_type=PythonNodeKind.CHECK.value, - node_name=check.name, - function=check.function, - project_dir=check.file_path.parent, - decorator_config={ - "description": check.description, - "group": check.group, - "meta": check.meta, - "severity": check.severity.value, - "tags": check.tags, - }, + node_type=kind.value, + node_name=discovered.name, + function=discovered.function, + project_dir=discovered.file_path.parent, + decorator_config=decorator_config, ), ) diff --git a/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/_test_types.py b/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/_test_types.py index 128296fa5b..968741bc70 100644 --- a/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/_test_types.py @@ -126,3 +126,9 @@ class PythonNodeIdentityChangeTestCase: expected_version_hash_changed: bool before_decorator_config: dict[str, object] | None = None after_decorator_config: dict[str, object] | None = None + + +@dataclass(frozen=True) +class PythonNodeIdentityConfigTestCase: + description: str + expected_decorator_config_keys_by_selector: dict[str, tuple[str, ...]] diff --git a/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/test_inventory.py b/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/test_inventory.py index 12056c1e06..b6eea07a48 100644 --- a/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/test_inventory.py +++ b/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/test_inventory.py @@ -2,8 +2,10 @@ from __future__ import annotations +import json from collections import defaultdict from pathlib import Path +from typing import cast import pytest @@ -15,14 +17,23 @@ DiscoveredProvider, DiscoveredTaskFunction, ) -from sqlbuild.compiler.python_nodes._helpers.inventory import build_python_node_graph -from sqlbuild.compiler.python_nodes.models import DiscoveredPythonNode, PythonNodeGraph +from sqlbuild.compiler.discovery.types import LoaderConnectionMode +from sqlbuild.compiler.python_nodes._helpers.inventory import ( + build_python_node_graph, + build_python_nodes, +) +from sqlbuild.compiler.python_nodes.models import ( + DiscoveredPythonNode, + PythonNodeGraph, + PythonNodeIdentity, +) from sqlbuild.compiler.python_nodes.types import PythonNodeKind from sqlbuild.python_nodes.models import RetryPolicy from sqlbuild.python_nodes.types import PythonCheckSeverity from sqlbuild.spec.contracts.models import LocalConfig, ProjectConfig, SourceColumnEntry from tests.unit.src.sqlbuild.compiler.python_nodes._helpers._test_types import ( PythonNodeGraphInventoryTestCase, + PythonNodeIdentityConfigTestCase, ) from tests.unit.src.sqlbuild.compiler.python_nodes._helpers.helpers import ( SlackProvider, @@ -217,3 +228,113 @@ def test_given_discovered_python_functions_when_building_graph_then_indexes_node for edge in graph.dependency_edges: edges_by_downstream[edge.downstream_name].append((edge.upstream_name, edge.downstream_name)) assert tuple(edges_by_downstream["notify_orders"]) == () + + +@pytest.mark.parametrize( + "test_case", + [ + PythonNodeIdentityConfigTestCase( + description="identity decorator config keeps each kind's authored field set", + expected_decorator_config_keys_by_selector={ + "loader:load_events": ( + "contract", + "cursor_column", + "destination", + "unique_key", + "write_strategy", + ), + "task:prepare_orders": ("description", "group", "meta", "tags"), + "asset:export_orders": ( + "column_lineage", + "columns", + "description", + "group", + "meta", + "tags", + ), + "check:check_orders_export": ( + "description", + "group", + "meta", + "severity", + "tags", + ), + }, + ) + ], + ids=lambda case: case.description, +) +def test_given_fully_configured_python_nodes_when_building_then_identity_config_is_per_kind( + test_case: PythonNodeIdentityConfigTestCase, +) -> None: + retry: RetryPolicy = RetryPolicy(max_attempts=2, retry_on=RuntimeError, jitter=False) + columns: tuple[SourceColumnEntry, ...] = (SourceColumnEntry(name="order_id", type="INTEGER"),) + + nodes: tuple[DiscoveredPythonNode, ...] = build_python_nodes( + loader_functions=( + DiscoveredLoaderFunction( + file_path=Path("/project/loaders/events.py"), + relative_path=Path("loaders/events.py"), + name="load_events", + function=load_events, + destination="raw.events", + cursor_column="updated_at", + unique_key=("event_id",), + columns=columns, + contract="enforced", + connection_mode=LoaderConnectionMode.EXTERNAL, + ), + ), + task_functions=( + DiscoveredTaskFunction( + file_path=Path("/project/tasks/orders.py"), + relative_path=Path("tasks/orders.py"), + name="prepare_orders", + function=prepare_orders, + tags=("orders",), + group="ingestion", + description="Prepare order export inputs.", + meta={"owner": "data-eng"}, + retry=retry, + ), + ), + asset_functions=( + DiscoveredAssetFunction( + file_path=Path("/project/assets/orders.py"), + relative_path=Path("assets/orders.py"), + name="export_orders", + function=export_orders, + tags=("orders",), + group="exports", + description="Export orders.", + meta={"owner": "data-eng"}, + columns=columns, + retry=retry, + ), + ), + check_functions=( + DiscoveredCheckFunction( + file_path=Path("/project/checks/orders.py"), + relative_path=Path("checks/orders.py"), + name="check_orders_export", + function=check_orders_export, + depends_on=(export_orders,), + severity=PythonCheckSeverity.WARN, + tags=("orders",), + group="exports", + description="Check exported orders.", + meta={"owner": "data-eng"}, + ), + ), + ) + + assert { + f"{node.kind.value}:{node.name}": tuple( + sorted( + json.loads(cast(PythonNodeIdentity, node.identity).definition_json)[ + "decorator_config" + ] + ) + ) + for node in nodes + } == test_case.expected_decorator_config_keys_by_selector From 92cb129fe77256a34a33e9100e069e6e9b63e4e6 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:03:28 +0100 Subject: [PATCH 25/57] refactor(sql-tests): classify direct-logic SQL test modes from one spec table The macro, UDF and table_fn CTE classifiers shared the same loop and checks and differed only in CTE names and foreign-CTE labels. They now run through one classifier driven by a per-mode spec. Every message is unchanged, including table_fn's generic direct-logic wording; a characterisation test pins the exact message per mode. (cherry picked from commit dca429a95b6b1a69da7a24766c695eb3262d27b1) --- .../compile/_helpers/sql_tests/core.py | 340 ++++++------------ .../compiler/compile/_helpers/_test_types.py | 8 + ...st_direct_logic_sql_test_classification.py | 263 ++++++++++++++ 3 files changed, 388 insertions(+), 223 deletions(-) create mode 100644 tests/unit/src/sqlbuild/compiler/compile/_helpers/test_direct_logic_sql_test_classification.py diff --git a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py index 6586f608d7..7a85a72f04 100644 --- a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py +++ b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py @@ -2,6 +2,7 @@ from __future__ import annotations +from dataclasses import dataclass from functools import lru_cache from typing import Any @@ -69,6 +70,51 @@ _SQL_IDENTIFIER_QUOTE_TOKENS: frozenset[str] = frozenset({'"', "`"}) +@dataclass(frozen=True) +class _DirectLogicModeSpec: + mode: SqlTestMode + actual_cte_name: str + expected_cte_name: str + foreign_ctes: tuple[tuple[frozenset[str], str], ...] + + +_MACRO_TEST_CTE_NAMES: frozenset[str] = frozenset( + {MACRO_ACTUAL_TEST_CTE_NAME, MACRO_EXPECTED_TEST_CTE_NAME} +) +_UDF_TEST_CTE_NAMES: frozenset[str] = frozenset( + {UDF_ACTUAL_TEST_CTE_NAME, UDF_EXPECTED_TEST_CTE_NAME} +) +_TABLE_FN_TEST_CTE_NAMES: frozenset[str] = frozenset( + {TABLE_FN_ACTUAL_TEST_CTE_NAME, TABLE_FN_EXPECTED_TEST_CTE_NAME} +) +_DIRECT_LOGIC_MODE_SPECS: dict[SqlTestMode, _DirectLogicModeSpec] = { + SqlTestMode.MACRO: _DirectLogicModeSpec( + mode=SqlTestMode.MACRO, + actual_cte_name=MACRO_ACTUAL_TEST_CTE_NAME, + expected_cte_name=MACRO_EXPECTED_TEST_CTE_NAME, + foreign_ctes=( + (_UDF_TEST_CTE_NAMES, "UDF-test"), + (_TABLE_FN_TEST_CTE_NAMES, "table_fn-test"), + ), + ), + SqlTestMode.UDF: _DirectLogicModeSpec( + mode=SqlTestMode.UDF, + actual_cte_name=UDF_ACTUAL_TEST_CTE_NAME, + expected_cte_name=UDF_EXPECTED_TEST_CTE_NAME, + foreign_ctes=( + (_MACRO_TEST_CTE_NAMES, "macro-test"), + (_TABLE_FN_TEST_CTE_NAMES, "table_fn-test"), + ), + ), + SqlTestMode.TABLE_FN: _DirectLogicModeSpec( + mode=SqlTestMode.TABLE_FN, + actual_cte_name=TABLE_FN_ACTUAL_TEST_CTE_NAME, + expected_cte_name=TABLE_FN_EXPECTED_TEST_CTE_NAME, + foreign_ctes=((_MACRO_TEST_CTE_NAMES | _UDF_TEST_CTE_NAMES, "another direct-logic"),), + ), +} + + def extract_sql_test_ctes( *, sql: str, file_label: str, mode: SqlTestMode = DEFAULT_SQL_TEST_MODE ) -> CompileSqlTestCtes: @@ -205,177 +251,86 @@ def _classify_sql_test_ctes( match mode: case SqlTestMode.MODEL: return _classify_model_sql_test_ctes(ctes=ctes, file_label=file_label) - case SqlTestMode.MACRO: - return _classify_macro_sql_test_ctes(ctes=ctes, file_label=file_label) - case SqlTestMode.UDF: - return _classify_udf_sql_test_ctes(ctes=ctes, file_label=file_label) - case SqlTestMode.TABLE_FN: - return _classify_table_fn_sql_test_ctes(ctes=ctes, file_label=file_label) + case SqlTestMode.MACRO | SqlTestMode.UDF | SqlTestMode.TABLE_FN: + return _classify_direct_logic_sql_test_ctes( + ctes=ctes, file_label=file_label, spec=_DIRECT_LOGIC_MODE_SPECS[mode] + ) case _: raise CompileInputError(f"SQL test '{file_label}' has unsupported mode '{mode}'") -def _classify_macro_sql_test_ctes( - *, ctes: tuple[CompileSqlTestCte, ...], file_label: str +def _classify_direct_logic_sql_test_ctes( + *, ctes: tuple[CompileSqlTestCte, ...], file_label: str, spec: _DirectLogicModeSpec ) -> CompileSqlTestCtes: + mode: str = spec.mode.value authored_ctes: list[CompileSqlTestCte] = [] - macro_actual_cte: CompileSqlTestCte | None = None - macro_expected_cte: CompileSqlTestCte | None = None + actual_cte: CompileSqlTestCte | None = None + expected_cte: CompileSqlTestCte | None = None cte: CompileSqlTestCte for cte in ctes: - if cte.name == MACRO_ACTUAL_TEST_CTE_NAME: - if macro_actual_cte is not None: + if cte.name == spec.actual_cte_name: + if actual_cte is not None: raise CompileInputError( - f"SQL test '{file_label}' mode 'macro' must define exactly one " - f"{MACRO_ACTUAL_TEST_CTE_NAME} CTE" + f"SQL test '{file_label}' mode '{mode}' must define exactly one " + f"{spec.actual_cte_name} CTE" ) - macro_actual_cte = cte + actual_cte = cte continue - if cte.name == MACRO_EXPECTED_TEST_CTE_NAME: - if macro_expected_cte is not None: + if cte.name == spec.expected_cte_name: + if expected_cte is not None: raise CompileInputError( - f"SQL test '{file_label}' mode 'macro' must define exactly one " - f"{MACRO_EXPECTED_TEST_CTE_NAME} CTE" + f"SQL test '{file_label}' mode '{mode}' must define exactly one " + f"{spec.expected_cte_name} CTE" ) _validate_expected_cte_query(cte=cte, file_label=file_label, label=cte.name) - macro_expected_cte = cte + expected_cte = cte continue if _is_model_mode_cte(cte.name): raise CompileInputError( - f"SQL test '{file_label}' is mode 'macro' but defines model-test CTE '{cte.name}'" - ) - if cte.name in {UDF_ACTUAL_TEST_CTE_NAME, UDF_EXPECTED_TEST_CTE_NAME}: - raise CompileInputError( - f"SQL test '{file_label}' is mode 'macro' but defines UDF-test CTE '{cte.name}'" - ) - if cte.name in {TABLE_FN_ACTUAL_TEST_CTE_NAME, TABLE_FN_EXPECTED_TEST_CTE_NAME}: - raise CompileInputError( - f"SQL test '{file_label}' is mode 'macro' but defines table_fn-test CTE " - f"'{cte.name}'" + f"SQL test '{file_label}' is mode '{mode}' but defines model-test CTE '{cte.name}'" ) - if cte.name in RESERVED_SQL_TEST_CTE_NAMES: - raise CompileInputError( - f"SQL test '{file_label}' uses reserved helper CTE name '{cte.name}'" - ) - authored_ctes.append(cte) - - direct_logic_payload: CompileDirectLogicSqlTestCtes = _validate_macro_test_ctes( - authored_ctes=tuple(authored_ctes), - macro_actual_cte=macro_actual_cte, - macro_expected_cte=macro_expected_cte, - file_label=file_label, - ) - return CompileSqlTestCtes(mode=SqlTestMode.MACRO, payload=direct_logic_payload) - - -def _classify_udf_sql_test_ctes( - *, ctes: tuple[CompileSqlTestCte, ...], file_label: str -) -> CompileSqlTestCtes: - authored_ctes: list[CompileSqlTestCte] = [] - udf_actual_cte: CompileSqlTestCte | None = None - udf_expected_cte: CompileSqlTestCte | None = None - - cte: CompileSqlTestCte - for cte in ctes: - if cte.name == UDF_ACTUAL_TEST_CTE_NAME: - if udf_actual_cte is not None: + foreign_names: frozenset[str] + foreign_label: str + for foreign_names, foreign_label in spec.foreign_ctes: + if cte.name in foreign_names: raise CompileInputError( - f"SQL test '{file_label}' mode 'udf' must define exactly one " - f"{UDF_ACTUAL_TEST_CTE_NAME} CTE" + f"SQL test '{file_label}' is mode '{mode}' but defines {foreign_label} CTE " + f"'{cte.name}'" ) - udf_actual_cte = cte - continue - if cte.name == UDF_EXPECTED_TEST_CTE_NAME: - if udf_expected_cte is not None: - raise CompileInputError( - f"SQL test '{file_label}' mode 'udf' must define exactly one " - f"{UDF_EXPECTED_TEST_CTE_NAME} CTE" - ) - _validate_expected_cte_query(cte=cte, file_label=file_label, label=cte.name) - udf_expected_cte = cte - continue - if _is_model_mode_cte(cte.name): - raise CompileInputError( - f"SQL test '{file_label}' is mode 'udf' but defines model-test CTE '{cte.name}'" - ) - if cte.name in {MACRO_ACTUAL_TEST_CTE_NAME, MACRO_EXPECTED_TEST_CTE_NAME}: - raise CompileInputError( - f"SQL test '{file_label}' is mode 'udf' but defines macro-test CTE '{cte.name}'" - ) - if cte.name in {TABLE_FN_ACTUAL_TEST_CTE_NAME, TABLE_FN_EXPECTED_TEST_CTE_NAME}: - raise CompileInputError( - f"SQL test '{file_label}' is mode 'udf' but defines table_fn-test CTE '{cte.name}'" - ) if cte.name in RESERVED_SQL_TEST_CTE_NAMES: raise CompileInputError( f"SQL test '{file_label}' uses reserved helper CTE name '{cte.name}'" ) authored_ctes.append(cte) - direct_logic_payload: CompileDirectLogicSqlTestCtes = _validate_udf_test_ctes( - authored_ctes=tuple(authored_ctes), - udf_actual_cte=udf_actual_cte, - udf_expected_cte=udf_expected_cte, - file_label=file_label, - ) - return CompileSqlTestCtes(mode=SqlTestMode.UDF, payload=direct_logic_payload) - - -def _classify_table_fn_sql_test_ctes( - *, ctes: tuple[CompileSqlTestCte, ...], file_label: str -) -> CompileSqlTestCtes: - authored_ctes: list[CompileSqlTestCte] = [] - table_fn_actual_cte: CompileSqlTestCte | None = None - table_fn_expected_cte: CompileSqlTestCte | None = None - - cte: CompileSqlTestCte - for cte in ctes: - if cte.name == TABLE_FN_ACTUAL_TEST_CTE_NAME: - if table_fn_actual_cte is not None: - raise CompileInputError( - f"SQL test '{file_label}' mode 'table_fn' must define exactly one " - f"{TABLE_FN_ACTUAL_TEST_CTE_NAME} CTE" - ) - table_fn_actual_cte = cte - continue - if cte.name == TABLE_FN_EXPECTED_TEST_CTE_NAME: - if table_fn_expected_cte is not None: - raise CompileInputError( - f"SQL test '{file_label}' mode 'table_fn' must define exactly one " - f"{TABLE_FN_EXPECTED_TEST_CTE_NAME} CTE" - ) - _validate_expected_cte_query(cte=cte, file_label=file_label, label=cte.name) - table_fn_expected_cte = cte - continue - if _is_model_mode_cte(cte.name): - raise CompileInputError( - f"SQL test '{file_label}' is mode 'table_fn' but defines model-test CTE " - f"'{cte.name}'" - ) - if cte.name in { - MACRO_ACTUAL_TEST_CTE_NAME, - MACRO_EXPECTED_TEST_CTE_NAME, - UDF_ACTUAL_TEST_CTE_NAME, - UDF_EXPECTED_TEST_CTE_NAME, - }: - raise CompileInputError( - f"SQL test '{file_label}' is mode 'table_fn' but defines another " - f"direct-logic CTE '{cte.name}'" - ) - if cte.name in RESERVED_SQL_TEST_CTE_NAMES: - raise CompileInputError( - f"SQL test '{file_label}' uses reserved helper CTE name '{cte.name}'" - ) - authored_ctes.append(cte) - - direct_logic_payload: CompileDirectLogicSqlTestCtes = _validate_table_fn_test_ctes( - authored_ctes=tuple(authored_ctes), - table_fn_actual_cte=table_fn_actual_cte, - table_fn_expected_cte=table_fn_expected_cte, - file_label=file_label, + if actual_cte is None or expected_cte is None: + raise CompileInputError( + f"SQL test '{file_label}' mode '{mode}' must define exactly one " + f"{spec.actual_cte_name} CTE and exactly one " + f"{spec.expected_cte_name} CTE" + ) + if spec.mode is SqlTestMode.MACRO: + _validate_macro_test_bodies( + helper_ctes=tuple(authored_ctes), expected_cte=expected_cte, file_label=file_label + ) + else: + _validate_call_free_direct_logic_bodies( + helper_ctes=tuple(authored_ctes), + expected_cte=expected_cte, + file_label=file_label, + mode=spec.mode, + actual_cte_name=spec.actual_cte_name, + ) + return CompileSqlTestCtes( + mode=spec.mode, + payload=CompileDirectLogicSqlTestCtes( + mode=spec.mode, + helper_ctes=tuple(authored_ctes), + actual_cte=actual_cte, + expected_cte=expected_cte, + ), ) - return CompileSqlTestCtes(mode=SqlTestMode.TABLE_FN, payload=direct_logic_payload) def _classify_model_sql_test_ctes( @@ -563,112 +518,51 @@ def _is_model_mode_cte(cte_name: str) -> bool: ) -def _validate_macro_test_ctes( +def _validate_macro_test_bodies( *, - authored_ctes: tuple[CompileSqlTestCte, ...], - macro_actual_cte: CompileSqlTestCte | None, - macro_expected_cte: CompileSqlTestCte | None, + helper_ctes: tuple[CompileSqlTestCte, ...], + expected_cte: CompileSqlTestCte, file_label: str, -) -> CompileDirectLogicSqlTestCtes: - if macro_actual_cte is None or macro_expected_cte is None: - raise CompileInputError( - f"SQL test '{file_label}' mode 'macro' must define exactly one " - f"{MACRO_ACTUAL_TEST_CTE_NAME} CTE and exactly one " - f"{MACRO_EXPECTED_TEST_CTE_NAME} CTE" - ) +) -> None: helper_cte: CompileSqlTestCte - for helper_cte in authored_ctes: + for helper_cte in helper_ctes: macro_names: tuple[str, ...] = find_macro_call_names(helper_cte.sql_body) if macro_names: raise CompileInputError( f"SQL test '{file_label}' mode 'macro' helper CTE '{helper_cte.name}' " "must not call macros; call macros only in __macro_actual__" ) - expected_macro_names: tuple[str, ...] = find_macro_call_names(macro_expected_cte.sql_body) + expected_macro_names: tuple[str, ...] = find_macro_call_names(expected_cte.sql_body) if expected_macro_names: raise CompileInputError( f"SQL test '{file_label}' mode 'macro' CTE {MACRO_EXPECTED_TEST_CTE_NAME} " "must not call macros" ) - return CompileDirectLogicSqlTestCtes( - mode=SqlTestMode.MACRO, - helper_ctes=authored_ctes, - actual_cte=macro_actual_cte, - expected_cte=macro_expected_cte, - ) -def _validate_udf_test_ctes( +def _validate_call_free_direct_logic_bodies( *, - authored_ctes: tuple[CompileSqlTestCte, ...], - udf_actual_cte: CompileSqlTestCte | None, - udf_expected_cte: CompileSqlTestCte | None, + helper_ctes: tuple[CompileSqlTestCte, ...], + expected_cte: CompileSqlTestCte, file_label: str, -) -> CompileDirectLogicSqlTestCtes: - if udf_actual_cte is None or udf_expected_cte is None: - raise CompileInputError( - f"SQL test '{file_label}' mode 'udf' must define exactly one " - f"{UDF_ACTUAL_TEST_CTE_NAME} CTE and exactly one " - f"{UDF_EXPECTED_TEST_CTE_NAME} CTE" - ) - helper_cte: CompileSqlTestCte - for helper_cte in authored_ctes: - _validate_no_direct_logic_calls( - sql=helper_cte.sql_body, - file_label=file_label, - mode=SqlTestMode.UDF, - cte_label=f"helper CTE '{helper_cte.name}'", - allowed_location=UDF_ACTUAL_TEST_CTE_NAME, - ) - _validate_no_direct_logic_calls( - sql=udf_expected_cte.sql_body, - file_label=file_label, - mode=SqlTestMode.UDF, - cte_label=f"CTE {UDF_EXPECTED_TEST_CTE_NAME}", - allowed_location=UDF_ACTUAL_TEST_CTE_NAME, - ) - return CompileDirectLogicSqlTestCtes( - mode=SqlTestMode.UDF, - helper_ctes=authored_ctes, - actual_cte=udf_actual_cte, - expected_cte=udf_expected_cte, - ) - - -def _validate_table_fn_test_ctes( - *, - authored_ctes: tuple[CompileSqlTestCte, ...], - table_fn_actual_cte: CompileSqlTestCte | None, - table_fn_expected_cte: CompileSqlTestCte | None, - file_label: str, -) -> CompileDirectLogicSqlTestCtes: - if table_fn_actual_cte is None or table_fn_expected_cte is None: - raise CompileInputError( - f"SQL test '{file_label}' mode 'table_fn' must define exactly one " - f"{TABLE_FN_ACTUAL_TEST_CTE_NAME} CTE and exactly one " - f"{TABLE_FN_EXPECTED_TEST_CTE_NAME} CTE" - ) + mode: SqlTestMode, + actual_cte_name: str, +) -> None: helper_cte: CompileSqlTestCte - for helper_cte in authored_ctes: + for helper_cte in helper_ctes: _validate_no_direct_logic_calls( sql=helper_cte.sql_body, file_label=file_label, - mode=SqlTestMode.TABLE_FN, + mode=mode, cte_label=f"helper CTE '{helper_cte.name}'", - allowed_location=TABLE_FN_ACTUAL_TEST_CTE_NAME, + allowed_location=actual_cte_name, ) _validate_no_direct_logic_calls( - sql=table_fn_expected_cte.sql_body, + sql=expected_cte.sql_body, file_label=file_label, - mode=SqlTestMode.TABLE_FN, - cte_label=f"CTE {TABLE_FN_EXPECTED_TEST_CTE_NAME}", - allowed_location=TABLE_FN_ACTUAL_TEST_CTE_NAME, - ) - return CompileDirectLogicSqlTestCtes( - mode=SqlTestMode.TABLE_FN, - helper_ctes=authored_ctes, - actual_cte=table_fn_actual_cte, - expected_cte=table_fn_expected_cte, + mode=mode, + cte_label=f"CTE {expected_cte.name}", + allowed_location=actual_cte_name, ) diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py index 1f5c1c2482..492e28a0c6 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py @@ -884,3 +884,11 @@ class CachedLineageRoundTripTestCase: transform_kind: ColumnTransformKind confidence: ColumnLineageConfidence expected_compact: bool + + +@dataclass(frozen=True) +class ClassifyDirectLogicSqlTestCtesErrorTestCase: + description: str + mode: SqlTestMode + ctes: tuple[tuple[str, str], ...] + expected_message: str diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_direct_logic_sql_test_classification.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_direct_logic_sql_test_classification.py new file mode 100644 index 0000000000..93032dd4c4 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_direct_logic_sql_test_classification.py @@ -0,0 +1,263 @@ +"""Characterisation tests for macro, UDF, and table_fn SQL-test CTE classification errors.""" + +from __future__ import annotations + +import pytest + +from sqlbuild.compiler.compile._helpers.sql_tests.core import classify_sql_test_ctes +from sqlbuild.compiler.compile.exceptions import CompileInputError +from sqlbuild.compiler.compile.models import CompileSqlTestCte +from sqlbuild.compiler.compile.types import SqlTestMode +from tests.unit.src.sqlbuild.compiler.compile._helpers._test_types import ( + ClassifyDirectLogicSqlTestCtesErrorTestCase, +) + +_FILE: str = "tests/unit/orders.sql" + + +@pytest.mark.parametrize( + "test_case", + [ + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro duplicate actual", + mode=SqlTestMode.MACRO, + ctes=(("__macro_actual__", "SELECT 1"), ("__macro_actual__", "SELECT 2")), + expected_message=( + f"SQL test '{_FILE}' mode 'macro' must define exactly one __macro_actual__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro duplicate expected", + mode=SqlTestMode.MACRO, + ctes=( + ("__macro_expected__", "SELECT 1 AS value"), + ("__macro_expected__", "SELECT 2 AS value"), + ), + expected_message=( + f"SQL test '{_FILE}' mode 'macro' must define exactly one __macro_expected__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro model-test cte", + mode=SqlTestMode.MACRO, + ctes=(("__ref__orders", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'macro' but defines model-test CTE '__ref__orders'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro udf-test cte", + mode=SqlTestMode.MACRO, + ctes=(("__udf_actual__", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'macro' but defines UDF-test CTE '__udf_actual__'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro table_fn-test cte", + mode=SqlTestMode.MACRO, + ctes=(("__table_fn_expected__", "SELECT 1 AS value"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'macro' but defines table_fn-test CTE " + "'__table_fn_expected__'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro reserved helper name", + mode=SqlTestMode.MACRO, + ctes=(("__missing__", "SELECT 1"),), + expected_message=f"SQL test '{_FILE}' uses reserved helper CTE name '__missing__'", + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro missing expected", + mode=SqlTestMode.MACRO, + ctes=(("__macro_actual__", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' mode 'macro' must define exactly one __macro_actual__ CTE " + "and exactly one __macro_expected__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro helper calls macro", + mode=SqlTestMode.MACRO, + ctes=( + ("helper", "SELECT @normalize_status('paid') AS status"), + ("__macro_actual__", "SELECT 1 AS status"), + ("__macro_expected__", "SELECT 1 AS status"), + ), + expected_message=( + f"SQL test '{_FILE}' mode 'macro' helper CTE 'helper' must not call macros; " + "call macros only in __macro_actual__" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro expected calls macro", + mode=SqlTestMode.MACRO, + ctes=( + ("__macro_actual__", "SELECT 1 AS status"), + ("__macro_expected__", "SELECT @normalize_status('paid') AS status"), + ), + expected_message=( + f"SQL test '{_FILE}' mode 'macro' CTE __macro_expected__ must not call macros" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf duplicate actual", + mode=SqlTestMode.UDF, + ctes=(("__udf_actual__", "SELECT 1"), ("__udf_actual__", "SELECT 2")), + expected_message=( + f"SQL test '{_FILE}' mode 'udf' must define exactly one __udf_actual__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf duplicate expected", + mode=SqlTestMode.UDF, + ctes=( + ("__udf_expected__", "SELECT 1 AS value"), + ("__udf_expected__", "SELECT 2 AS value"), + ), + expected_message=( + f"SQL test '{_FILE}' mode 'udf' must define exactly one __udf_expected__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf model-test cte", + mode=SqlTestMode.UDF, + ctes=(("__expected__orders", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'udf' but defines model-test CTE '__expected__orders'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf macro-test cte", + mode=SqlTestMode.UDF, + ctes=(("__macro_expected__", "SELECT 1 AS value"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'udf' but defines macro-test CTE '__macro_expected__'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf table_fn-test cte", + mode=SqlTestMode.UDF, + ctes=(("__table_fn_actual__", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'udf' but defines table_fn-test CTE " + "'__table_fn_actual__'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf reserved helper name", + mode=SqlTestMode.UDF, + ctes=(("__actual", "SELECT 1"),), + expected_message=f"SQL test '{_FILE}' uses reserved helper CTE name '__actual'", + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf missing actual", + mode=SqlTestMode.UDF, + ctes=(("__udf_expected__", "SELECT 1 AS value"),), + expected_message=( + f"SQL test '{_FILE}' mode 'udf' must define exactly one __udf_actual__ CTE " + "and exactly one __udf_expected__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf helper calls macro", + mode=SqlTestMode.UDF, + ctes=( + ("helper", "SELECT @normalize_status('paid') AS status"), + ("__udf_actual__", "SELECT 1 AS formatted"), + ("__udf_expected__", "SELECT 1 AS formatted"), + ), + expected_message=( + f"SQL test '{_FILE}' mode 'udf' helper CTE 'helper' must not call macros" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf expected calls udf", + mode=SqlTestMode.UDF, + ctes=( + ("__udf_actual__", "SELECT 1 AS formatted"), + ("__udf_expected__", 'SELECT __udf("format_cents")(1250) AS formatted'), + ), + expected_message=( + f"SQL test '{_FILE}' mode 'udf' CTE __udf_expected__ must not call udf; " + "call reusable logic only in __udf_actual__" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn duplicate actual", + mode=SqlTestMode.TABLE_FN, + ctes=(("__table_fn_actual__", "SELECT 1"), ("__table_fn_actual__", "SELECT 2")), + expected_message=( + f"SQL test '{_FILE}' mode 'table_fn' must define exactly one " + "__table_fn_actual__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn duplicate expected", + mode=SqlTestMode.TABLE_FN, + ctes=( + ("__table_fn_expected__", "SELECT 1 AS value"), + ("__table_fn_expected__", "SELECT 2 AS value"), + ), + expected_message=( + f"SQL test '{_FILE}' mode 'table_fn' must define exactly one " + "__table_fn_expected__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn model-test cte", + mode=SqlTestMode.TABLE_FN, + ctes=(("__source__raw_orders", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'table_fn' but defines model-test CTE " + "'__source__raw_orders'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn macro cte uses generic direct-logic wording", + mode=SqlTestMode.TABLE_FN, + ctes=(("__macro_actual__", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'table_fn' but defines another direct-logic CTE " + "'__macro_actual__'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn udf cte uses generic direct-logic wording", + mode=SqlTestMode.TABLE_FN, + ctes=(("__udf_expected__", "SELECT 1 AS value"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'table_fn' but defines another direct-logic CTE " + "'__udf_expected__'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn reserved helper name", + mode=SqlTestMode.TABLE_FN, + ctes=(("__unexpected__", "SELECT 1"),), + expected_message=f"SQL test '{_FILE}' uses reserved helper CTE name '__unexpected__'", + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn missing both", + mode=SqlTestMode.TABLE_FN, + ctes=(("helper", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' mode 'table_fn' must define exactly one " + "__table_fn_actual__ CTE and exactly one __table_fn_expected__ CTE" + ), + ), + ], + ids=lambda case: case.description, +) +def test_given_invalid_direct_logic_ctes_when_classifying_then_raises_mode_specific_message( + test_case: ClassifyDirectLogicSqlTestCtesErrorTestCase, +) -> None: + ctes: tuple[CompileSqlTestCte, ...] = tuple( + CompileSqlTestCte(name=name, sql_body=body) for name, body in test_case.ctes + ) + + with pytest.raises(CompileInputError) as error_info: + _ = classify_sql_test_ctes(ctes=ctes, file_label=_FILE, mode=test_case.mode) + + assert str(error_info.value) == test_case.expected_message From ad2cdef9d395af964fe4600856b7153a3d528250 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:10:05 +0100 Subject: [PATCH 26/57] refactor(cli): resolve audit, check, test and seed invocations through one helper The four invocation resolvers repeated the same discovery, adapter, connection, colour and progress-stream resolution. They now call a shared runtime helper; audit, check and test also build the shared connection and planning progress reporters, while seed keeps omitting them because its planning and execution phases build their own. An integration test pins the resolved context per command against a real DuckDB project. (cherry picked from commit bf3e65a12703e09e16414748391f72e9084b0868) --- .../cli/commands/_helpers/audit/invocation.py | 53 +------- .../cli/commands/_helpers/check/invocation.py | 53 +------- .../_helpers/runtime/command_invocation.py | 116 ++++++++++++++++++ .../cli/commands/_helpers/seed/invocation.py | 42 +------ .../cli/commands/_helpers/test/invocation.py | 53 +------- .../cli/commands/_helpers/runtime/__init__.py | 1 + .../commands/_helpers/runtime/_test_types.py | 20 +++ .../runtime/test_command_invocation.py | 89 ++++++++++++++ 8 files changed, 242 insertions(+), 185 deletions(-) create mode 100644 src/sqlbuild/cli/commands/_helpers/runtime/command_invocation.py create mode 100644 tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/__init__.py create mode 100644 tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/_test_types.py create mode 100644 tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/test_command_invocation.py diff --git a/src/sqlbuild/cli/commands/_helpers/audit/invocation.py b/src/sqlbuild/cli/commands/_helpers/audit/invocation.py index 85d92da554..a5ad81f7d5 100644 --- a/src/sqlbuild/cli/commands/_helpers/audit/invocation.py +++ b/src/sqlbuild/cli/commands/_helpers/audit/invocation.py @@ -2,58 +2,13 @@ from __future__ import annotations -import sys -from pathlib import Path -from typing import TextIO - -from sqlbuild.cli.commands._helpers.runtime.adapter_context import ( - resolve_adapter_connection_context, -) -from sqlbuild.cli.commands.models import ( - AdapterConnectionContext, - AuditCommandRequest, - AuditInvocation, -) -from sqlbuild.cli.progress.main._build_command_progress_reporters import ( - build_command_progress_reporters, +from sqlbuild.cli.commands._helpers.runtime.command_invocation import ( + resolve_reported_command_invocation, ) -from sqlbuild.cli.progress.models import CommandProgressReporters -from sqlbuild.compiler.discovery.main.discover import discover_project_inputs -from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs -from sqlbuild.presentation.main.supports_color import supports_color +from sqlbuild.cli.commands.models import AuditCommandRequest, AuditInvocation def resolve_audit_invocation(*, request: AuditCommandRequest) -> AuditInvocation: """Resolve discovery, adapter, connection, and reporters for audit.""" - effective_project_dir: Path = ( - request.project_dir if request.project_dir is not None else Path.cwd() - ) - discovered_inputs: DiscoveredProjectInputs = discover_project_inputs( - project_dir=effective_project_dir - ) - adapter_context: AdapterConnectionContext = resolve_adapter_connection_context( - discovered_inputs=discovered_inputs, - effective_project_dir=effective_project_dir, - selected_target=request.selected_target, - cli_vars=request.cli_vars, - ) - machine_output: bool = request.json_output - use_color: bool = not request.no_color and not machine_output and supports_color() - progress_stream: TextIO = sys.stderr if machine_output else sys.stdout - reporters: CommandProgressReporters = build_command_progress_reporters( - adapter_name=adapter_context.adapter_name, - stream=progress_stream, - use_color=use_color, - ) - return AuditInvocation( - effective_project_dir=effective_project_dir, - discovered_inputs=discovered_inputs, - adapter_name=adapter_context.adapter_name, - adapter=adapter_context.adapter, - connection_config=adapter_context.connection_config, - use_color=use_color, - progress_stream=progress_stream, - connection_progress=reporters.connection, - planning_progress=reporters.planning, - ) + return resolve_reported_command_invocation(request=request, invocation_type=AuditInvocation) diff --git a/src/sqlbuild/cli/commands/_helpers/check/invocation.py b/src/sqlbuild/cli/commands/_helpers/check/invocation.py index dc144d6c55..e0d0b8556d 100644 --- a/src/sqlbuild/cli/commands/_helpers/check/invocation.py +++ b/src/sqlbuild/cli/commands/_helpers/check/invocation.py @@ -2,58 +2,13 @@ from __future__ import annotations -import sys -from pathlib import Path -from typing import TextIO - -from sqlbuild.cli.commands._helpers.runtime.adapter_context import ( - resolve_adapter_connection_context, -) -from sqlbuild.cli.commands.models import ( - AdapterConnectionContext, - CheckCommandRequest, - CheckInvocation, -) -from sqlbuild.cli.progress.main._build_command_progress_reporters import ( - build_command_progress_reporters, +from sqlbuild.cli.commands._helpers.runtime.command_invocation import ( + resolve_reported_command_invocation, ) -from sqlbuild.cli.progress.models import CommandProgressReporters -from sqlbuild.compiler.discovery.main.discover import discover_project_inputs -from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs -from sqlbuild.presentation.main.supports_color import supports_color +from sqlbuild.cli.commands.models import CheckCommandRequest, CheckInvocation def resolve_check_invocation(*, request: CheckCommandRequest) -> CheckInvocation: """Resolve discovery, adapter, connection, and reporters for check.""" - effective_project_dir: Path = ( - request.project_dir if request.project_dir is not None else Path.cwd() - ) - discovered_inputs: DiscoveredProjectInputs = discover_project_inputs( - project_dir=effective_project_dir - ) - adapter_context: AdapterConnectionContext = resolve_adapter_connection_context( - discovered_inputs=discovered_inputs, - effective_project_dir=effective_project_dir, - selected_target=request.selected_target, - cli_vars=request.cli_vars, - ) - machine_output: bool = request.json_output - use_color: bool = not request.no_color and not machine_output and supports_color() - progress_stream: TextIO = sys.stderr if machine_output else sys.stdout - reporters: CommandProgressReporters = build_command_progress_reporters( - adapter_name=adapter_context.adapter_name, - stream=progress_stream, - use_color=use_color, - ) - return CheckInvocation( - effective_project_dir=effective_project_dir, - discovered_inputs=discovered_inputs, - adapter_name=adapter_context.adapter_name, - adapter=adapter_context.adapter, - connection_config=adapter_context.connection_config, - use_color=use_color, - progress_stream=progress_stream, - connection_progress=reporters.connection, - planning_progress=reporters.planning, - ) + return resolve_reported_command_invocation(request=request, invocation_type=CheckInvocation) diff --git a/src/sqlbuild/cli/commands/_helpers/runtime/command_invocation.py b/src/sqlbuild/cli/commands/_helpers/runtime/command_invocation.py new file mode 100644 index 0000000000..62e105d459 --- /dev/null +++ b/src/sqlbuild/cli/commands/_helpers/runtime/command_invocation.py @@ -0,0 +1,116 @@ +"""Shared invocation resolution for project commands that connect to one target.""" + +from __future__ import annotations + +import sys +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Protocol, TextIO + +from sqlbuild.cli.commands._helpers.runtime.adapter_context import ( + resolve_adapter_connection_context, +) +from sqlbuild.cli.commands.models import AdapterConnectionContext +from sqlbuild.cli.progress.main._build_command_progress_reporters import ( + build_command_progress_reporters, +) +from sqlbuild.cli.progress.models import CommandProgressReporters +from sqlbuild.compiler.discovery.main.discover import discover_project_inputs +from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs +from sqlbuild.presentation.main.supports_color import supports_color + + +class _CommandInvocationRequest(Protocol): + """CLI request fields needed to resolve a project command invocation.""" + + @property + def project_dir(self) -> Path | None: ... + + @property + def selected_target(self) -> str | None: ... + + @property + def cli_vars(self) -> dict[str, object] | None: ... + + @property + def json_output(self) -> bool: ... + + @property + def no_color(self) -> bool: ... + + +@dataclass(frozen=True) +class _CommandInvocationContext: + effective_project_dir: Path + discovered_inputs: DiscoveredProjectInputs + adapter_context: AdapterConnectionContext + use_color: bool + progress_stream: TextIO + + +def resolve_command_invocation[InvocationT]( + *, request: _CommandInvocationRequest, invocation_type: Callable[..., InvocationT] +) -> InvocationT: + """Resolve discovery, adapter, connection, and output context for one command.""" + + context: _CommandInvocationContext = _resolve_command_invocation_context(request=request) + return invocation_type( + effective_project_dir=context.effective_project_dir, + discovered_inputs=context.discovered_inputs, + adapter_name=context.adapter_context.adapter_name, + adapter=context.adapter_context.adapter, + connection_config=context.adapter_context.connection_config, + use_color=context.use_color, + progress_stream=context.progress_stream, + ) + + +def resolve_reported_command_invocation[InvocationT]( + *, request: _CommandInvocationRequest, invocation_type: Callable[..., InvocationT] +) -> InvocationT: + """Resolve command context plus shared connection and planning progress reporters.""" + + context: _CommandInvocationContext = _resolve_command_invocation_context(request=request) + reporters: CommandProgressReporters = build_command_progress_reporters( + adapter_name=context.adapter_context.adapter_name, + stream=context.progress_stream, + use_color=context.use_color, + ) + return invocation_type( + effective_project_dir=context.effective_project_dir, + discovered_inputs=context.discovered_inputs, + adapter_name=context.adapter_context.adapter_name, + adapter=context.adapter_context.adapter, + connection_config=context.adapter_context.connection_config, + use_color=context.use_color, + progress_stream=context.progress_stream, + connection_progress=reporters.connection, + planning_progress=reporters.planning, + ) + + +def _resolve_command_invocation_context( + *, request: _CommandInvocationRequest +) -> _CommandInvocationContext: + effective_project_dir: Path = ( + request.project_dir if request.project_dir is not None else Path.cwd() + ) + discovered_inputs: DiscoveredProjectInputs = discover_project_inputs( + project_dir=effective_project_dir + ) + adapter_context: AdapterConnectionContext = resolve_adapter_connection_context( + discovered_inputs=discovered_inputs, + effective_project_dir=effective_project_dir, + selected_target=request.selected_target, + cli_vars=request.cli_vars, + ) + machine_output: bool = request.json_output + use_color: bool = not request.no_color and not machine_output and supports_color() + return _CommandInvocationContext( + effective_project_dir=effective_project_dir, + discovered_inputs=discovered_inputs, + adapter_context=adapter_context, + use_color=use_color, + progress_stream=sys.stderr if machine_output else sys.stdout, + ) diff --git a/src/sqlbuild/cli/commands/_helpers/seed/invocation.py b/src/sqlbuild/cli/commands/_helpers/seed/invocation.py index 9e4a75332b..3fdb7028db 100644 --- a/src/sqlbuild/cli/commands/_helpers/seed/invocation.py +++ b/src/sqlbuild/cli/commands/_helpers/seed/invocation.py @@ -2,47 +2,13 @@ from __future__ import annotations -import sys -from pathlib import Path -from typing import TextIO - -from sqlbuild.cli.commands._helpers.runtime.adapter_context import ( - resolve_adapter_connection_context, -) -from sqlbuild.cli.commands.models import ( - AdapterConnectionContext, - SeedCommandRequest, - SeedInvocation, +from sqlbuild.cli.commands._helpers.runtime.command_invocation import ( + resolve_command_invocation, ) -from sqlbuild.compiler.discovery.main.discover import discover_project_inputs -from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs -from sqlbuild.presentation.main.supports_color import supports_color +from sqlbuild.cli.commands.models import SeedCommandRequest, SeedInvocation def resolve_seed_invocation(*, request: SeedCommandRequest) -> SeedInvocation: """Resolve discovery, adapter, connection, and output context for seed.""" - effective_project_dir: Path = ( - request.project_dir if request.project_dir is not None else Path.cwd() - ) - discovered_inputs: DiscoveredProjectInputs = discover_project_inputs( - project_dir=effective_project_dir - ) - adapter_context: AdapterConnectionContext = resolve_adapter_connection_context( - discovered_inputs=discovered_inputs, - effective_project_dir=effective_project_dir, - selected_target=request.selected_target, - cli_vars=request.cli_vars, - ) - machine_output: bool = request.json_output - use_color: bool = not request.no_color and not machine_output and supports_color() - progress_stream: TextIO = sys.stderr if machine_output else sys.stdout - return SeedInvocation( - effective_project_dir=effective_project_dir, - discovered_inputs=discovered_inputs, - adapter_name=adapter_context.adapter_name, - adapter=adapter_context.adapter, - connection_config=adapter_context.connection_config, - use_color=use_color, - progress_stream=progress_stream, - ) + return resolve_command_invocation(request=request, invocation_type=SeedInvocation) diff --git a/src/sqlbuild/cli/commands/_helpers/test/invocation.py b/src/sqlbuild/cli/commands/_helpers/test/invocation.py index f03cb38837..6572d84451 100644 --- a/src/sqlbuild/cli/commands/_helpers/test/invocation.py +++ b/src/sqlbuild/cli/commands/_helpers/test/invocation.py @@ -2,58 +2,13 @@ from __future__ import annotations -import sys -from pathlib import Path -from typing import TextIO - -from sqlbuild.cli.commands._helpers.runtime.adapter_context import ( - resolve_adapter_connection_context, -) -from sqlbuild.cli.commands.models import ( - AdapterConnectionContext, - TestCommandRequest, - TestInvocation, -) -from sqlbuild.cli.progress.main._build_command_progress_reporters import ( - build_command_progress_reporters, +from sqlbuild.cli.commands._helpers.runtime.command_invocation import ( + resolve_reported_command_invocation, ) -from sqlbuild.cli.progress.models import CommandProgressReporters -from sqlbuild.compiler.discovery.main.discover import discover_project_inputs -from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs -from sqlbuild.presentation.main.supports_color import supports_color +from sqlbuild.cli.commands.models import TestCommandRequest, TestInvocation def resolve_test_invocation(*, request: TestCommandRequest) -> TestInvocation: """Resolve discovery, adapter, connection, and reporters for test.""" - effective_project_dir: Path = ( - request.project_dir if request.project_dir is not None else Path.cwd() - ) - discovered_inputs: DiscoveredProjectInputs = discover_project_inputs( - project_dir=effective_project_dir - ) - adapter_context: AdapterConnectionContext = resolve_adapter_connection_context( - discovered_inputs=discovered_inputs, - effective_project_dir=effective_project_dir, - selected_target=request.selected_target, - cli_vars=request.cli_vars, - ) - machine_output: bool = request.json_output - use_color: bool = not request.no_color and not machine_output and supports_color() - progress_stream: TextIO = sys.stderr if machine_output else sys.stdout - reporters: CommandProgressReporters = build_command_progress_reporters( - adapter_name=adapter_context.adapter_name, - stream=progress_stream, - use_color=use_color, - ) - return TestInvocation( - effective_project_dir=effective_project_dir, - discovered_inputs=discovered_inputs, - adapter_name=adapter_context.adapter_name, - adapter=adapter_context.adapter, - connection_config=adapter_context.connection_config, - use_color=use_color, - progress_stream=progress_stream, - connection_progress=reporters.connection, - planning_progress=reporters.planning, - ) + return resolve_reported_command_invocation(request=request, invocation_type=TestInvocation) diff --git a/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/__init__.py b/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/__init__.py new file mode 100644 index 0000000000..e0de81c887 --- /dev/null +++ b/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/__init__.py @@ -0,0 +1 @@ +"""Integration tests for shared CLI command runtime helpers.""" diff --git a/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/_test_types.py b/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/_test_types.py new file mode 100644 index 0000000000..8d5e37e5b1 --- /dev/null +++ b/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/_test_types.py @@ -0,0 +1,20 @@ +from collections.abc import Callable +from dataclasses import dataclass + +from sqlbuild.cli.commands.models import ( + AuditInvocation, + CheckInvocation, + SeedInvocation, + TestInvocation, +) + + +@dataclass(frozen=True) +class CommandInvocationTestCase: + description: str + resolve: Callable[..., AuditInvocation | CheckInvocation | SeedInvocation | TestInvocation] + request_type: Callable[..., object] + json_output: bool + expected_invocation_type_name: str + expected_stream_name: str + expected_has_progress_reporters: bool diff --git a/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/test_command_invocation.py b/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/test_command_invocation.py new file mode 100644 index 0000000000..1f6b7b4de5 --- /dev/null +++ b/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/test_command_invocation.py @@ -0,0 +1,89 @@ +"""Integration tests for resolving CLI command invocations against a real DuckDB project.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +from sqlbuild.cli.commands import models as command_models +from sqlbuild.cli.commands._helpers.audit.invocation import resolve_audit_invocation +from sqlbuild.cli.commands._helpers.check.invocation import resolve_check_invocation +from sqlbuild.cli.commands._helpers.seed.invocation import resolve_seed_invocation +from sqlbuild.cli.commands._helpers.test.invocation import resolve_test_invocation +from tests.integration.src.sqlbuild.cli.commands._helpers.runtime._test_types import ( + CommandInvocationTestCase, +) + +_PROJECT_TOML: str = ( + 'name = "orders_shop"\nadapter = "duckdb"\n\n[connection]\ndatabase = "orders.duckdb"\n' +) + + +@pytest.mark.parametrize( + "test_case", + ( + CommandInvocationTestCase( + description="audit json output reports progress on stderr", + resolve=resolve_audit_invocation, + request_type=command_models.AuditCommandRequest, + json_output=True, + expected_invocation_type_name="AuditInvocation", + expected_stream_name="stderr", + expected_has_progress_reporters=True, + ), + CommandInvocationTestCase( + description="check text output reports progress on stdout", + resolve=resolve_check_invocation, + request_type=command_models.CheckCommandRequest, + json_output=False, + expected_invocation_type_name="CheckInvocation", + expected_stream_name="stdout", + expected_has_progress_reporters=True, + ), + CommandInvocationTestCase( + description="test json output reports progress on stderr", + resolve=resolve_test_invocation, + request_type=command_models.TestCommandRequest, + json_output=True, + expected_invocation_type_name="TestInvocation", + expected_stream_name="stderr", + expected_has_progress_reporters=True, + ), + CommandInvocationTestCase( + description="seed text output omits shared progress reporters", + resolve=resolve_seed_invocation, + request_type=command_models.SeedCommandRequest, + json_output=False, + expected_invocation_type_name="SeedInvocation", + expected_stream_name="stdout", + expected_has_progress_reporters=False, + ), + ), + ids=lambda case: case.description, +) +def test_given_duckdb_project_when_resolving_command_invocation_then_returns_command_context( + tmp_path: Path, + test_case: CommandInvocationTestCase, +) -> None: + _ = (tmp_path / "sqlbuild_project.toml").write_text(_PROJECT_TOML, encoding="utf-8") + request: object = test_case.request_type( + project_dir=tmp_path, no_color=True, json_output=test_case.json_output + ) + + invocation: ( + command_models.AuditInvocation + | command_models.CheckInvocation + | command_models.SeedInvocation + | command_models.TestInvocation + ) = test_case.resolve(request=request) + + assert type(invocation).__name__ == test_case.expected_invocation_type_name + assert invocation.effective_project_dir == tmp_path + assert invocation.adapter_name == "duckdb" + assert invocation.discovered_inputs.project_config.name == "orders_shop" + assert invocation.use_color is False + assert invocation.progress_stream is getattr(sys, test_case.expected_stream_name) + assert hasattr(invocation, "connection_progress") is test_case.expected_has_progress_reporters + assert hasattr(invocation, "planning_progress") is test_case.expected_has_progress_reporters From ea175a3f5730725b5a08d9e6f15cc8580d9a07ba Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:11:41 +0100 Subject: [PATCH 27/57] fix(scenario): fall back when a failed capture step has no message The capture path stringified a failed step's error_message directly, so a fixture or seed failure without a message reported "None" (or an empty message). It now uses the capture fallback message, matching the warehouse and local scenario runners. (cherry picked from commit df6a0bae12d02954408719444890c1340ef54f40) --- .../scenario/_helpers/capture/core.py | 5 +- .../executor/scenario/_helpers/_test_types.py | 7 +++ .../scenario/_helpers/test_capture.py | 61 +++++++++++++++++++ 3 files changed, 72 insertions(+), 1 deletion(-) diff --git a/src/sqlbuild/executor/scenario/_helpers/capture/core.py b/src/sqlbuild/executor/scenario/_helpers/capture/core.py index 73c8897f8f..4bd03a5db7 100644 --- a/src/sqlbuild/executor/scenario/_helpers/capture/core.py +++ b/src/sqlbuild/executor/scenario/_helpers/capture/core.py @@ -234,7 +234,10 @@ def _first_error(results: tuple[object, ...]) -> str | None: result: object for result in results: if getattr(result, "status", None) == ExecutionStatus.FAILED: - return str(getattr(result, "error_message", "scenario snapshot capture failed")) + error_message: object | None = getattr(result, "error_message", None) + if isinstance(error_message, str) and error_message: + return error_message + return "scenario snapshot capture failed" return None diff --git a/tests/unit/src/sqlbuild/executor/scenario/_helpers/_test_types.py b/tests/unit/src/sqlbuild/executor/scenario/_helpers/_test_types.py index 3da48e4ec9..507ea43d64 100644 --- a/tests/unit/src/sqlbuild/executor/scenario/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/executor/scenario/_helpers/_test_types.py @@ -168,3 +168,10 @@ class ScenarioCaptureRunIdentityTestCase: expected_run_id: str expected_seed_resource_ids: tuple[str, ...] expected_operation_name: str + + +@dataclass(frozen=True) +class ScenarioCaptureMissingFixtureMessageTestCase: + description: str + fixture_error_message: str | None + expected_error_message: str diff --git a/tests/unit/src/sqlbuild/executor/scenario/_helpers/test_capture.py b/tests/unit/src/sqlbuild/executor/scenario/_helpers/test_capture.py index 8ec24e5a4f..352c0adcb5 100644 --- a/tests/unit/src/sqlbuild/executor/scenario/_helpers/test_capture.py +++ b/tests/unit/src/sqlbuild/executor/scenario/_helpers/test_capture.py @@ -6,15 +6,19 @@ import pytest from sqlbuild.compiler.planner.models import ScenarioExecutionPlan, SeedPlanEntry +from sqlbuild.compiler.planner.types import ScenarioArtifactKind +from sqlbuild.executor.scenario._helpers.capture import core as capture_core from sqlbuild.executor.scenario._helpers.capture.core import execute_scenario_snapshot_capture_steps from sqlbuild.executor.scenario.models import ( ScenarioCaptureSettings, + ScenarioFixtureExecutionResult, ScenarioSnapshotCaptureRunResult, ) from sqlbuild.executor.scheduling.types import ExecutionStatus from sqlbuild.observability import EventDispatcher, LifecycleEvent, dispatcher_scope from tests.unit.src.sqlbuild.executor.scenario._helpers._test_types import ( ExecuteScenarioSnapshotCaptureStepsTestCase, + ScenarioCaptureMissingFixtureMessageTestCase, ScenarioCaptureRunIdentityTestCase, ) from tests.unit.src.sqlbuild.executor.scenario._helpers.helpers import ( @@ -205,3 +209,60 @@ def test_given_multiple_capture_seeds_when_executing_then_operation_and_resource test_case.expected_seed_resource_ids[1], test_case.expected_seed_resource_ids[1], ) + + +@pytest.mark.parametrize( + "test_case", + ( + ScenarioCaptureMissingFixtureMessageTestCase( + description="failed fixture without a message uses the capture fallback message", + fixture_error_message=None, + expected_error_message="scenario snapshot capture failed", + ), + ScenarioCaptureMissingFixtureMessageTestCase( + description="failed fixture with an empty message uses the capture fallback message", + fixture_error_message="", + expected_error_message="scenario snapshot capture failed", + ), + ScenarioCaptureMissingFixtureMessageTestCase( + description="failed fixture message is reported unchanged", + fixture_error_message="orders fixture failed", + expected_error_message="orders fixture failed", + ), + ), + ids=lambda case: case.description, +) +def test_given_failed_fixture_message_when_capturing_then_reports_message_or_fallback( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + test_case: ScenarioCaptureMissingFixtureMessageTestCase, +) -> None: + failed_fixture: ScenarioFixtureExecutionResult = ScenarioFixtureExecutionResult( + scenario_name=SCENARIO_PLAN.name, + kind=ScenarioArtifactKind.SOURCE, + logical_name="raw__orders", + target_relation="scenario_schema.raw__orders", + status=ExecutionStatus.FAILED, + error_message=test_case.fixture_error_message, + ) + monkeypatch.setattr( + capture_core, "execute_scenario_fixtures", lambda **_kwargs: (failed_fixture,) + ) + + result: ScenarioSnapshotCaptureRunResult = execute_scenario_snapshot_capture_steps( + project_dir=tmp_path, + scenario_plan=SCENARIO_PLAN, + adapter=ScenarioSnapshotCaptureStepsTestAdapter(), + connection=object(), + run_id="capture-test-run", + settings=ScenarioCaptureSettings( + captured_at="2026-05-09T00:00:00Z", + capture_adapter="duckdb", + capture_dialect="duckdb", + sqlbuild_version="0.1.0", + retain=False, + ), + ) + + assert result.status == ExecutionStatus.FAILED + assert result.error_message == test_case.expected_error_message From 924d5e9cfe497863ca85832e8fb1582879ab4880 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:15:40 +0100 Subject: [PATCH 28/57] refactor(scenario): project step failure details through one helper The warehouse, local and capture scenario runners each carried their own _first_error, _first_error_code and _first_error_help copies. They now share first_failure_details, keeping each runner's fallback message and fallback codes. Semantics are unchanged: the message comes from the first failed step, while code and help come from the first failed step that has them; a unit test pins that behaviour. (cherry picked from commit 535dd6eca2c5b720ad47f5209a4e2adce3470441) --- .../scenario/_helpers/capture/core.py | 58 +++--------- .../scenario/_helpers/execution/run.py | 62 +++---------- .../scenario/_helpers/lifecycle/failures.py | 41 +++++++++ .../scenario/_helpers/local/execution.py | 63 ++++--------- .../executor/scenario/_helpers/_test_types.py | 10 ++ .../scenario/_helpers/test_failures.py | 92 +++++++++++++++++++ 6 files changed, 186 insertions(+), 140 deletions(-) create mode 100644 src/sqlbuild/executor/scenario/_helpers/lifecycle/failures.py create mode 100644 tests/unit/src/sqlbuild/executor/scenario/_helpers/test_failures.py diff --git a/src/sqlbuild/executor/scenario/_helpers/capture/core.py b/src/sqlbuild/executor/scenario/_helpers/capture/core.py index 4bd03a5db7..b210b1bd9c 100644 --- a/src/sqlbuild/executor/scenario/_helpers/capture/core.py +++ b/src/sqlbuild/executor/scenario/_helpers/capture/core.py @@ -8,6 +8,7 @@ from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.compiler.planner.models import ScenarioExecutionPlan from sqlbuild.executor.build.models import SeedExecutionResult +from sqlbuild.executor.scenario._helpers.lifecycle.failures import first_failure_details from sqlbuild.executor.scenario._helpers.lifecycle.fixtures import ( execute_scenario_fixtures, execute_scenario_seed_entries, @@ -33,6 +34,8 @@ from sqlbuild.runtime.observability.classes.operation_lifecycle import OperationLifecycle from sqlbuild.runtime.observability.main.run_scope import run_scope +_CAPTURE_FAILED_MESSAGE: str = "scenario snapshot capture failed" + def execute_scenario_snapshot_capture_steps( *, @@ -97,8 +100,10 @@ def _execute_scenario_snapshot_capture_steps( adapter=adapter, connection=connection, ) - fixture_error: str | None = _first_error(fixture_results) - if fixture_error is not None: + fixture_failure: ScenarioFailureDetails = first_failure_details( + results=fixture_results, fallback_message=_CAPTURE_FAILED_MESSAGE + ) + if fixture_failure.error_message is not None: return _finish_capture_run( scenario_plan=scenario_plan, adapter=adapter, @@ -106,11 +111,7 @@ def _execute_scenario_snapshot_capture_steps( retain=settings.retain, prepare_cleanup_result=prepare_result, fixture_results=fixture_results, - failure=ScenarioFailureDetails( - error_code=_first_error_code(fixture_results), - error_help=_first_error_help(fixture_results), - error_message=fixture_error, - ), + failure=fixture_failure, ) seed_results: tuple[SeedExecutionResult, ...] = execute_scenario_seed_entries( @@ -120,8 +121,10 @@ def _execute_scenario_snapshot_capture_steps( connection=connection, run_id=run_id, ) - seed_error: str | None = _first_error(seed_results) - if seed_error is not None: + seed_failure: ScenarioFailureDetails = first_failure_details( + results=seed_results, fallback_message=_CAPTURE_FAILED_MESSAGE + ) + if seed_failure.error_message is not None: return _finish_capture_run( scenario_plan=scenario_plan, adapter=adapter, @@ -130,11 +133,7 @@ def _execute_scenario_snapshot_capture_steps( prepare_cleanup_result=prepare_result, fixture_results=fixture_results, seed_results=seed_results, - failure=ScenarioFailureDetails( - error_code=_first_error_code(seed_results), - error_help=_first_error_help(seed_results), - error_message=seed_error, - ), + failure=seed_failure, ) capture_plan: ScenarioSnapshotCapturePlan = build_scenario_snapshot_capture_plan( @@ -228,34 +227,3 @@ def _finish_capture_run( error_help=error_help, error_message=error_message, ) - - -def _first_error(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - error_message: object | None = getattr(result, "error_message", None) - if isinstance(error_message, str) and error_message: - return error_message - return "scenario snapshot capture failed" - return None - - -def _first_error_code(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - error_code: object | None = getattr(result, "error_code", None) - if isinstance(error_code, str) and error_code: - return error_code - return None - - -def _first_error_help(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - error_help: object | None = getattr(result, "error_help", None) - if isinstance(error_help, str) and error_help: - return error_help - return None diff --git a/src/sqlbuild/executor/scenario/_helpers/execution/run.py b/src/sqlbuild/executor/scenario/_helpers/execution/run.py index 5d7fab2665..666b58b377 100644 --- a/src/sqlbuild/executor/scenario/_helpers/execution/run.py +++ b/src/sqlbuild/executor/scenario/_helpers/execution/run.py @@ -13,6 +13,7 @@ execute_scenario_assertion_expectations, execute_scenario_expected_expectations, ) +from sqlbuild.executor.scenario._helpers.lifecycle.failures import first_failure_details from sqlbuild.executor.scenario._helpers.lifecycle.fixtures import ( execute_scenario_fixtures, execute_scenario_seed_entries, @@ -31,6 +32,8 @@ from sqlbuild.executor.scheduling.types import ExecutionStatus from sqlbuild.runtime.observability.classes.operation_lifecycle import OperationLifecycle +_STEP_FAILED_MESSAGE: str = "scenario step failed" + def execute_scenario_run_steps( *, @@ -91,10 +94,8 @@ def _execute_scenario_run_steps( retain=retain, prepare_cleanup_result=prepare_result, results=ScenarioStepResults(fixture_results=fixture_results), - failure=ScenarioFailureDetails( - error_code=_first_error_code(fixture_results), - error_help=_first_error_help(fixture_results), - error_message=_first_error(fixture_results), + failure=first_failure_details( + results=fixture_results, fallback_message=_STEP_FAILED_MESSAGE ), ) @@ -116,10 +117,8 @@ def _execute_scenario_run_steps( fixture_results=fixture_results, seed_results=seed_results, ), - failure=ScenarioFailureDetails( - error_code=_first_error_code(seed_results), - error_help=_first_error_help(seed_results), - error_message=_first_error(seed_results), + failure=first_failure_details( + results=seed_results, fallback_message=_STEP_FAILED_MESSAGE ), ) @@ -141,10 +140,8 @@ def _execute_scenario_run_steps( seed_results=seed_results, model_results=model_results, ), - failure=ScenarioFailureDetails( - error_code=_first_error_code(model_results), - error_help=_first_error_help(model_results), - error_message=_first_error(model_results), + failure=first_failure_details( + results=model_results, fallback_message=_STEP_FAILED_MESSAGE ), ) @@ -160,9 +157,6 @@ def _execute_scenario_run_steps( adapter=adapter, connection=connection, ) - failed_check_message: str | None = _first_error((*expected_results, *assertion_results)) - failed_check_code: str | None = _first_error_code((*expected_results, *assertion_results)) - failed_check_help: str | None = _first_error_help((*expected_results, *assertion_results)) return _finish_scenario( scenario_plan=scenario_plan, adapter=adapter, @@ -176,10 +170,9 @@ def _execute_scenario_run_steps( expected_results=expected_results, assertion_results=assertion_results, ), - failure=ScenarioFailureDetails( - error_code=failed_check_code, - error_help=failed_check_help, - error_message=failed_check_message, + failure=first_failure_details( + results=(*expected_results, *assertion_results), + fallback_message=_STEP_FAILED_MESSAGE, ), ) @@ -265,34 +258,3 @@ def _scenario_failure( def _has_failed(results: tuple[object, ...]) -> bool: return any(getattr(result, "status", None) == ExecutionStatus.FAILED for result in results) - - -def _first_error(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - error_message: object | None = getattr(result, "error_message", None) - if isinstance(error_message, str) and error_message: - return error_message - return "scenario step failed" - return None - - -def _first_error_code(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - error_code: object | None = getattr(result, "error_code", None) - if isinstance(error_code, str) and error_code: - return error_code - return None - - -def _first_error_help(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - error_help: object | None = getattr(result, "error_help", None) - if isinstance(error_help, str) and error_help: - return error_help - return None diff --git a/src/sqlbuild/executor/scenario/_helpers/lifecycle/failures.py b/src/sqlbuild/executor/scenario/_helpers/lifecycle/failures.py new file mode 100644 index 0000000000..575dc78a2f --- /dev/null +++ b/src/sqlbuild/executor/scenario/_helpers/lifecycle/failures.py @@ -0,0 +1,41 @@ +"""Failure-detail projection over scenario step results.""" + +from __future__ import annotations + +from sqlbuild.executor.scenario.models import ScenarioFailureDetails +from sqlbuild.executor.scheduling.types import ExecutionStatus + + +def first_failure_details( + *, + results: tuple[object, ...], + fallback_message: str, + fallback_code: str | None = None, +) -> ScenarioFailureDetails: + """Return the first failed step's message plus the first failed code and help present.""" + + failed: tuple[object, ...] = tuple( + result for result in results if getattr(result, "status", None) == ExecutionStatus.FAILED + ) + error_message: str | None = None + if failed: + first_message: str | None = _non_empty_text(getattr(failed[0], "error_message", None)) + error_message = first_message if first_message is not None else fallback_message + return ScenarioFailureDetails( + error_code=_first_text(results=failed, field_name="error_code") or fallback_code, + error_help=_first_text(results=failed, field_name="error_help"), + error_message=error_message, + ) + + +def _first_text(*, results: tuple[object, ...], field_name: str) -> str | None: + result: object + for result in results: + text: str | None = _non_empty_text(getattr(result, field_name, None)) + if text is not None: + return text + return None + + +def _non_empty_text(value: object) -> str | None: + return value if isinstance(value, str) and value else None diff --git a/src/sqlbuild/executor/scenario/_helpers/local/execution.py b/src/sqlbuild/executor/scenario/_helpers/local/execution.py index 21d774797d..52b73bfd98 100644 --- a/src/sqlbuild/executor/scenario/_helpers/local/execution.py +++ b/src/sqlbuild/executor/scenario/_helpers/local/execution.py @@ -30,6 +30,7 @@ execute_scenario_assertion_expectations, execute_scenario_expected_expectations, ) +from sqlbuild.executor.scenario._helpers.lifecycle.failures import first_failure_details from sqlbuild.executor.scenario._helpers.local.snapshots import ( load_scenario_snapshot_into_duckdb, local_snapshot_table_name, @@ -71,6 +72,8 @@ ) from sqlbuild.spec.contracts.models import SourceEntry +_LOCAL_STEP_FAILED_MESSAGE: str = "local scenario step failed" + def execute_local_scenario_load_only_run( *, @@ -302,10 +305,10 @@ def _execute_local_plan( duckdb_path=duckdb_path, loaded_relations=loaded_relations, function_results=function_results, - failure=ScenarioFailureDetails( - error_code=_first_error_code(function_results) or SCENARIO_LOCAL_FUNCTION_FAILED, - error_help=_first_error_help(function_results), - error_message=_first_error(function_results), + failure=first_failure_details( + results=function_results, + fallback_message=_LOCAL_STEP_FAILED_MESSAGE, + fallback_code=SCENARIO_LOCAL_FUNCTION_FAILED, ), ) @@ -327,10 +330,10 @@ def _execute_local_plan( loaded_relations=loaded_relations, function_results=function_results, model_results=local_model_results, - failure=ScenarioFailureDetails( - error_code=_first_error_code(local_model_results) or SCENARIO_LOCAL_MODEL_FAILED, - error_help=_first_error_help(local_model_results), - error_message=_first_error(local_model_results), + failure=first_failure_details( + results=local_model_results, + fallback_message=_LOCAL_STEP_FAILED_MESSAGE, + fallback_code=SCENARIO_LOCAL_MODEL_FAILED, ), ) @@ -362,10 +365,10 @@ def _execute_local_plan( model_results=model_results, expected_results=expected_results, assertion_results=assertion_results, - failure=ScenarioFailureDetails( - error_code=_first_error_code(check_results) or SCENARIO_LOCAL_MODEL_FAILED, - error_help=_first_error_help(check_results), - error_message=_first_error(check_results), + failure=first_failure_details( + results=check_results, + fallback_message=_LOCAL_STEP_FAILED_MESSAGE, + fallback_code=SCENARIO_LOCAL_MODEL_FAILED, ), ) if _has_failed(check_results): @@ -379,10 +382,9 @@ def _execute_local_plan( model_results=model_results, expected_results=expected_results, assertion_results=assertion_results, - failure=ScenarioFailureDetails( - error_code=_first_error_code(check_results), - error_help=_first_error_help(check_results), - error_message=_first_error(check_results), + failure=first_failure_details( + results=check_results, + fallback_message=_LOCAL_STEP_FAILED_MESSAGE, ), ) return _local_result( @@ -719,35 +721,6 @@ def _has_local_check_error(results: tuple[object, ...]) -> bool: ) -def _first_error(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - message: object | None = getattr(result, "error_message", None) - return message if isinstance(message, str) and message else "local scenario step failed" - return None - - -def _first_error_code(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - code: object | None = getattr(result, "error_code", None) - if isinstance(code, str) and code: - return code - return None - - -def _first_error_help(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - help_text: object | None = getattr(result, "error_help", None) - if isinstance(help_text, str) and help_text: - return help_text - return None - - def _remove_local_duckdb_files(duckdb_path: Path) -> None: duckdb_path.unlink(missing_ok=True) duckdb_path.with_name(f"{duckdb_path.name}.wal").unlink(missing_ok=True) diff --git a/tests/unit/src/sqlbuild/executor/scenario/_helpers/_test_types.py b/tests/unit/src/sqlbuild/executor/scenario/_helpers/_test_types.py index 507ea43d64..be60a39820 100644 --- a/tests/unit/src/sqlbuild/executor/scenario/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/executor/scenario/_helpers/_test_types.py @@ -175,3 +175,13 @@ class ScenarioCaptureMissingFixtureMessageTestCase: description: str fixture_error_message: str | None expected_error_message: str + + +@dataclass(frozen=True) +class FirstFailureDetailsTestCase: + description: str + step_outcomes: tuple[tuple[ExecutionStatus, str | None, str | None, str | None], ...] + fallback_code: str | None + expected_error_code: str | None + expected_error_help: str | None + expected_error_message: str | None diff --git a/tests/unit/src/sqlbuild/executor/scenario/_helpers/test_failures.py b/tests/unit/src/sqlbuild/executor/scenario/_helpers/test_failures.py new file mode 100644 index 0000000000..540d94ae32 --- /dev/null +++ b/tests/unit/src/sqlbuild/executor/scenario/_helpers/test_failures.py @@ -0,0 +1,92 @@ +"""Tests for projecting failure details from scenario step results.""" + +from __future__ import annotations + +import pytest + +from sqlbuild.compiler.planner.types import ScenarioArtifactKind +from sqlbuild.executor.scenario._helpers.lifecycle.failures import first_failure_details +from sqlbuild.executor.scenario.models import ( + ScenarioFailureDetails, + ScenarioFixtureExecutionResult, +) +from sqlbuild.executor.scheduling.types import ExecutionStatus +from tests.unit.src.sqlbuild.executor.scenario._helpers._test_types import ( + FirstFailureDetailsTestCase, +) + +_FAILED: ExecutionStatus = ExecutionStatus.FAILED +_SUCCESS: ExecutionStatus = ExecutionStatus.SUCCESS + + +@pytest.mark.parametrize( + "test_case", + ( + FirstFailureDetailsTestCase( + description="no failed step yields no message and the fallback code", + step_outcomes=((_SUCCESS, "S100", "help", "ignored"),), + fallback_code="S999", + expected_error_code="S999", + expected_error_help=None, + expected_error_message=None, + ), + FirstFailureDetailsTestCase( + description="first failed step supplies every detail", + step_outcomes=( + (_SUCCESS, "S100", "success help", "success message"), + (_FAILED, "S200", "orders help", "orders failed"), + (_FAILED, "S300", "customers help", "customers failed"), + ), + fallback_code=None, + expected_error_code="S200", + expected_error_help="orders help", + expected_error_message="orders failed", + ), + FirstFailureDetailsTestCase( + description="missing message uses the fallback message", + step_outcomes=((_FAILED, "S200", None, None),), + fallback_code=None, + expected_error_code="S200", + expected_error_help=None, + expected_error_message="step failed", + ), + FirstFailureDetailsTestCase( + description="code and help come from the first failed step that has them", + step_outcomes=( + (_FAILED, None, None, "orders failed"), + (_FAILED, "S300", "customers help", "customers failed"), + ), + fallback_code="S999", + expected_error_code="S300", + expected_error_help="customers help", + expected_error_message="orders failed", + ), + ), + ids=lambda case: case.description, +) +def test_given_step_results_when_projecting_failure_then_returns_expected_details( + test_case: FirstFailureDetailsTestCase, +) -> None: + results: tuple[ScenarioFixtureExecutionResult, ...] = tuple( + ScenarioFixtureExecutionResult( + scenario_name="orders_refund", + kind=ScenarioArtifactKind.SOURCE, + logical_name=f"fixture_{index}", + target_relation=f"scenario_schema.fixture_{index}", + status=status, + error_code=code, + error_help=help_text, + error_message=message, + ) + for index, (status, code, help_text, message) in enumerate(test_case.step_outcomes) + ) + + details: ScenarioFailureDetails = first_failure_details( + results=results, fallback_message="step failed", fallback_code=test_case.fallback_code + ) + + assert details == ScenarioFailureDetails( + error_code=test_case.expected_error_code, + error_help=test_case.expected_error_help, + error_message=test_case.expected_error_message, + ) From ae24a4ca3ad011277e7860bbd69c4e32310c9019 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:17:52 +0100 Subject: [PATCH 29/57] refactor(sql-tests): share the top-level CTE scanner with scenarios SQL tests and SQL scenarios each carried an identical cached scanner for the leading WITH clause. Scenarios now call the SQL-test scanner with their own context label and CTE model. Messages are unchanged. (cherry picked from commit 72df5a3de6ce0eef0b98a43c4eddeb9f85383b19) --- .../compile/_helpers/scenarios/core.py | 70 ++----------------- .../compile/_helpers/sql_tests/core.py | 38 ++++++---- 2 files changed, 30 insertions(+), 78 deletions(-) diff --git a/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py b/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py index 471e1b93f0..873419c1bc 100644 --- a/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py +++ b/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py @@ -2,18 +2,14 @@ from __future__ import annotations -from functools import lru_cache - from sqlbuild.compiler.compile._helpers.analysis.ctes import ( extract_top_level_ctes_with_sql_analysis, ) from sqlbuild.compiler.compile._helpers.sql_tests.core import ( - _consume_keyword, - _read_identifier, _require_prefixed_name, _skip_ignorable, _try_consume_keyword, - _validate_ceremonial_select, + extract_top_level_ctes_with_scanner, validate_independent_expected_and_assertion_ctes, ) from sqlbuild.compiler.compile.constants import ( @@ -24,8 +20,6 @@ REF_TEST_CTE_PREFIX, SEED_TEST_CTE_PREFIX, SOURCE_TEST_CTE_PREFIX, - SQL_ARGUMENT_SEPARATOR_TOKEN, - SQL_OPEN_PAREN_TOKEN, SQL_WITH_KEYWORD, ) from sqlbuild.compiler.compile.exceptions import CompileInputError @@ -33,7 +27,6 @@ CompileSqlScenarioCte, CompileSqlScenarioCtes, ) -from sqlbuild.compiler.sql_analysis.main._find_matching_paren import find_matching_paren _CONTEXT: str = "SQL scenario" @@ -42,9 +35,11 @@ def extract_sql_scenario_ctes(*, sql: str, file_label: str) -> CompileSqlScenari """Extract top-level SQL-native scenario fixture, expected, and assertion CTEs.""" try: - ctes: tuple[CompileSqlScenarioCte, ...] = _extract_sql_scenario_ctes_with_scanner( + ctes: tuple[CompileSqlScenarioCte, ...] = extract_top_level_ctes_with_scanner( sql=sql, file_label=file_label, + context_label=_CONTEXT, + cte_type=CompileSqlScenarioCte, ) except CompileInputError as scanner_error: cte_values: tuple[tuple[str, str], ...] | None = extract_top_level_ctes_with_sql_analysis( @@ -64,9 +59,11 @@ def extract_sql_scenario_expected_model_names(*, sql: str, file_label: str) -> t start: int = _skip_ignorable(sql=sql, start=0) if _try_consume_keyword(sql=sql, start=start, keyword=SQL_WITH_KEYWORD) is None: return () - ctes: tuple[CompileSqlScenarioCte, ...] = _extract_sql_scenario_ctes_with_scanner( + ctes: tuple[CompileSqlScenarioCte, ...] = extract_top_level_ctes_with_scanner( sql=sql, file_label=file_label, + context_label=_CONTEXT, + cte_type=CompileSqlScenarioCte, ) return tuple( _require_prefixed_name( @@ -80,59 +77,6 @@ def extract_sql_scenario_expected_model_names(*, sql: str, file_label: str) -> t ) -@lru_cache(maxsize=4096) -def _extract_sql_scenario_ctes_with_scanner( - *, sql: str, file_label: str -) -> tuple[CompileSqlScenarioCte, ...]: - index: int = _skip_ignorable(sql=sql, start=0) - index = _consume_keyword(sql=sql, start=index, keyword="WITH", file_label=file_label) - index = _skip_ignorable(sql=sql, start=index) - recursive_end: int | None = _try_consume_keyword(sql=sql, start=index, keyword="RECURSIVE") - if recursive_end is not None: - index = _skip_ignorable(sql=sql, start=recursive_end) - - ctes: list[CompileSqlScenarioCte] = [] - seen_cte_names: set[str] = set() - while True: - cte_name, index = _read_identifier(sql=sql, start=index, file_label=file_label) - if cte_name in seen_cte_names: - raise CompileInputError( - f"SQL scenario '{file_label}' defines duplicate CTE '{cte_name}'" - ) - seen_cte_names.add(cte_name) - - index = _skip_ignorable(sql=sql, start=index) - if index < len(sql) and sql[index] == SQL_OPEN_PAREN_TOKEN: - index = find_matching_paren(sql=sql, open_paren_index=index, context=_CONTEXT) + 1 - index = _skip_ignorable(sql=sql, start=index) - index = _consume_keyword(sql=sql, start=index, keyword="AS", file_label=file_label) - index = _skip_ignorable(sql=sql, start=index) - if index >= len(sql) or sql[index] != SQL_OPEN_PAREN_TOKEN: - raise CompileInputError( - f"SQL scenario '{file_label}' CTE '{cte_name}' must use AS (...)" - ) - cte_body_start: int = index + 1 - cte_body_end: int = find_matching_paren(sql=sql, open_paren_index=index, context=_CONTEXT) - ctes.append( - CompileSqlScenarioCte( - name=cte_name, - sql_body=sql[cte_body_start:cte_body_end].strip(), - ) - ) - index = _skip_ignorable(sql=sql, start=cte_body_end + 1) - if index < len(sql) and sql[index] == SQL_ARGUMENT_SEPARATOR_TOKEN: - index = _skip_ignorable(sql=sql, start=index + 1) - continue - break - - _validate_ceremonial_select( - sql=sql, - start=index, - file_label=file_label, - ) - return tuple(ctes) - - def _classify_sql_scenario_ctes( *, ctes: tuple[CompileSqlScenarioCte, ...], file_label: str ) -> CompileSqlScenarioCtes: diff --git a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py index 7a85a72f04..4619e251a8 100644 --- a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py +++ b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass from functools import lru_cache from typing import Any @@ -133,9 +134,11 @@ def extract_unclassified_sql_test_ctes( """Extract raw top-level CTEs before mode-specific classification.""" try: - ctes: tuple[CompileSqlTestCte, ...] = _extract_sql_test_ctes_with_scanner( + ctes: tuple[CompileSqlTestCte, ...] = extract_top_level_ctes_with_scanner( sql=sql, file_label=file_label, + context_label=_CONTEXT, + cte_type=CompileSqlTestCte, ) except CompileInputError as scanner_error: cte_values: tuple[tuple[str, str], ...] | None = extract_top_level_ctes_with_sql_analysis( @@ -167,9 +170,11 @@ def extract_sql_test_expected_model_names( start: int = _skip_ignorable(sql=sql, start=0) if _try_consume_keyword(sql=sql, start=start, keyword=SQL_WITH_KEYWORD) is None: return () - ctes: tuple[CompileSqlTestCte, ...] = _extract_sql_test_ctes_with_scanner( + ctes: tuple[CompileSqlTestCte, ...] = extract_top_level_ctes_with_scanner( sql=sql, file_label=file_label, + context_label=_CONTEXT, + cte_type=CompileSqlTestCte, ) return tuple( _require_prefixed_name( @@ -197,9 +202,11 @@ def extract_assertion_target_model_names(*, assertion_sql: tuple[str, ...]) -> t @lru_cache(maxsize=4096) -def _extract_sql_test_ctes_with_scanner( - *, sql: str, file_label: str -) -> tuple[CompileSqlTestCte, ...]: +def extract_top_level_ctes_with_scanner[CteT]( + *, sql: str, file_label: str, context_label: str, cte_type: Callable[..., CteT] +) -> tuple[CteT, ...]: + """Scan top-level `WITH` CTEs followed by the ceremonial `SELECT 1` of a test file.""" + index: int = _skip_ignorable(sql=sql, start=0) index = _consume_keyword(sql=sql, start=index, keyword="WITH", file_label=file_label) index = _skip_ignorable(sql=sql, start=index) @@ -207,30 +214,31 @@ def _extract_sql_test_ctes_with_scanner( if recursive_end is not None: index = _skip_ignorable(sql=sql, start=recursive_end) - ctes: list[CompileSqlTestCte] = [] + ctes: list[CteT] = [] seen_cte_names: set[str] = set() while True: cte_name, index = _read_identifier(sql=sql, start=index, file_label=file_label) if cte_name in seen_cte_names: - raise CompileInputError(f"SQL test '{file_label}' defines duplicate CTE '{cte_name}'") + raise CompileInputError( + f"{context_label} '{file_label}' defines duplicate CTE '{cte_name}'" + ) seen_cte_names.add(cte_name) index = _skip_ignorable(sql=sql, start=index) if index < len(sql) and sql[index] == SQL_OPEN_PAREN_TOKEN: - index = find_matching_paren(sql=sql, open_paren_index=index, context=_CONTEXT) + 1 + index = find_matching_paren(sql=sql, open_paren_index=index, context=context_label) + 1 index = _skip_ignorable(sql=sql, start=index) index = _consume_keyword(sql=sql, start=index, keyword="AS", file_label=file_label) index = _skip_ignorable(sql=sql, start=index) if index >= len(sql) or sql[index] != SQL_OPEN_PAREN_TOKEN: - raise CompileInputError(f"SQL test '{file_label}' CTE '{cte_name}' must use AS (...)") - cte_body_start: int = index + 1 - cte_body_end: int = find_matching_paren(sql=sql, open_paren_index=index, context=_CONTEXT) - ctes.append( - CompileSqlTestCte( - name=cte_name, - sql_body=sql[cte_body_start:cte_body_end].strip(), + raise CompileInputError( + f"{context_label} '{file_label}' CTE '{cte_name}' must use AS (...)" ) + cte_body_start: int = index + 1 + cte_body_end: int = find_matching_paren( + sql=sql, open_paren_index=index, context=context_label ) + ctes.append(cte_type(name=cte_name, sql_body=sql[cte_body_start:cte_body_end].strip())) index = _skip_ignorable(sql=sql, start=cte_body_end + 1) if index < len(sql) and sql[index] == SQL_ARGUMENT_SEPARATOR_TOKEN: index = _skip_ignorable(sql=sql, start=index + 1) From 11f10bf3fdfab2d8259e4077bc892558bf3215d2 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:23:24 +0100 Subject: [PATCH 30/57] fix(scenarios): name SQL scenarios in CTE scanner errors Scenario files reuse the SQL-test CTE scanner, so a missing WITH clause, an invalid CTE name, a missing AS keyword, a missing ceremonial SELECT 1 or an empty fixture or expected target reported "SQL test ''" wording, and the missing-WITH message asked for mock CTEs. The scanner and its helpers now take the caller's context label, so scenario errors say "SQL scenario ''" and ask for fixture plus expected or assertion CTEs. SQL-test messages are unchanged and pinned by tests. (cherry picked from commit c9927fc8d05bfe275fe0cf15466be283690b7b6d) --- .../compile/_helpers/scenarios/core.py | 12 +++ .../compile/_helpers/sql_tests/core.py | 67 +++++++++----- .../compiler/compile/_helpers/_test_types.py | 7 ++ .../compile/_helpers/test_scenarios.py | 91 ++++++++++++++++++- .../compiler/compile/_helpers/test_tests.py | 75 +++++++++++++++ 5 files changed, 229 insertions(+), 23 deletions(-) diff --git a/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py b/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py index 873419c1bc..1fc69eab28 100644 --- a/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py +++ b/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py @@ -29,6 +29,9 @@ ) _CONTEXT: str = "SQL scenario" +_WITH_REQUIREMENT: str = ( + "fixture CTEs and at least one __expected__ or __assert__ CTE" +) def extract_sql_scenario_ctes(*, sql: str, file_label: str) -> CompileSqlScenarioCtes: @@ -39,6 +42,7 @@ def extract_sql_scenario_ctes(*, sql: str, file_label: str) -> CompileSqlScenari sql=sql, file_label=file_label, context_label=_CONTEXT, + with_requirement=_WITH_REQUIREMENT, cte_type=CompileSqlScenarioCte, ) except CompileInputError as scanner_error: @@ -63,6 +67,7 @@ def extract_sql_scenario_expected_model_names(*, sql: str, file_label: str) -> t sql=sql, file_label=file_label, context_label=_CONTEXT, + with_requirement=_WITH_REQUIREMENT, cte_type=CompileSqlScenarioCte, ) return tuple( @@ -71,6 +76,7 @@ def extract_sql_scenario_expected_model_names(*, sql: str, file_label: str) -> t prefix=EXPECTED_TEST_CTE_PREFIX, label="__expected__", file_label=file_label, + context_label=_CONTEXT, ) for cte in ctes if cte.name.startswith(EXPECTED_TEST_CTE_PREFIX) @@ -106,6 +112,7 @@ def _classify_sql_scenario_ctes( prefix=SOURCE_TEST_CTE_PREFIX, label="__source__", file_label=file_label, + context_label=_CONTEXT, ) ) authored_ctes.append(cte) @@ -117,6 +124,7 @@ def _classify_sql_scenario_ctes( prefix=REF_TEST_CTE_PREFIX, label="__ref__", file_label=file_label, + context_label=_CONTEXT, ) ) authored_ctes.append(cte) @@ -128,6 +136,7 @@ def _classify_sql_scenario_ctes( prefix=SEED_TEST_CTE_PREFIX, label="__seed__", file_label=file_label, + context_label=_CONTEXT, ) ) authored_ctes.append(cte) @@ -139,6 +148,7 @@ def _classify_sql_scenario_ctes( prefix=DBT_REF_TEST_CTE_PREFIX, label="__dbt_ref__ or __dbt_ref____", file_label=file_label, + context_label=_CONTEXT, ) ) authored_ctes.append(cte) @@ -150,6 +160,7 @@ def _classify_sql_scenario_ctes( prefix=EXPECTED_TEST_CTE_PREFIX, label="__expected__", file_label=file_label, + context_label=_CONTEXT, ) ) expected_ctes.append(cte) @@ -161,6 +172,7 @@ def _classify_sql_scenario_ctes( prefix=ASSERT_SCENARIO_CTE_PREFIX, label="__assert__", file_label=file_label, + context_label=_CONTEXT, ) ) assertion_ctes.append(cte) diff --git a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py index 4619e251a8..f4b92bb6bc 100644 --- a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py +++ b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py @@ -67,6 +67,7 @@ from sqlbuild.compiler.sql_analysis.main.import_polyglot_sql import import_polyglot_sql _CONTEXT: str = "SQL test" +_SQL_TEST_WITH_REQUIREMENT: str = "mock CTEs and one __expected__ CTE" _DIRECT_DEPENDENCY_PATH_LENGTH: int = 2 _SQL_IDENTIFIER_QUOTE_TOKENS: frozenset[str] = frozenset({'"', "`"}) @@ -138,6 +139,7 @@ def extract_unclassified_sql_test_ctes( sql=sql, file_label=file_label, context_label=_CONTEXT, + with_requirement=_SQL_TEST_WITH_REQUIREMENT, cte_type=CompileSqlTestCte, ) except CompileInputError as scanner_error: @@ -174,6 +176,7 @@ def extract_sql_test_expected_model_names( sql=sql, file_label=file_label, context_label=_CONTEXT, + with_requirement=_SQL_TEST_WITH_REQUIREMENT, cte_type=CompileSqlTestCte, ) return tuple( @@ -203,13 +206,23 @@ def extract_assertion_target_model_names(*, assertion_sql: tuple[str, ...]) -> t @lru_cache(maxsize=4096) def extract_top_level_ctes_with_scanner[CteT]( - *, sql: str, file_label: str, context_label: str, cte_type: Callable[..., CteT] + *, + sql: str, + file_label: str, + context_label: str, + with_requirement: str, + cte_type: Callable[..., CteT], ) -> tuple[CteT, ...]: """Scan top-level `WITH` CTEs followed by the ceremonial `SELECT 1` of a test file.""" - index: int = _skip_ignorable(sql=sql, start=0) - index = _consume_keyword(sql=sql, start=index, keyword="WITH", file_label=file_label) - index = _skip_ignorable(sql=sql, start=index) + with_end: int | None = _try_consume_keyword( + sql=sql, start=_skip_ignorable(sql=sql, start=0), keyword=SQL_WITH_KEYWORD + ) + if with_end is None: + raise CompileInputError( + f"{context_label} '{file_label}' must declare {with_requirement} before `SELECT 1`" + ) + index: int = _skip_ignorable(sql=sql, start=with_end) recursive_end: int | None = _try_consume_keyword(sql=sql, start=index, keyword="RECURSIVE") if recursive_end is not None: index = _skip_ignorable(sql=sql, start=recursive_end) @@ -217,7 +230,9 @@ def extract_top_level_ctes_with_scanner[CteT]( ctes: list[CteT] = [] seen_cte_names: set[str] = set() while True: - cte_name, index = _read_identifier(sql=sql, start=index, file_label=file_label) + cte_name, index = _read_identifier( + sql=sql, start=index, file_label=file_label, context_label=context_label + ) if cte_name in seen_cte_names: raise CompileInputError( f"{context_label} '{file_label}' defines duplicate CTE '{cte_name}'" @@ -228,7 +243,13 @@ def extract_top_level_ctes_with_scanner[CteT]( if index < len(sql) and sql[index] == SQL_OPEN_PAREN_TOKEN: index = find_matching_paren(sql=sql, open_paren_index=index, context=context_label) + 1 index = _skip_ignorable(sql=sql, start=index) - index = _consume_keyword(sql=sql, start=index, keyword="AS", file_label=file_label) + index = _consume_keyword( + sql=sql, + start=index, + keyword="AS", + file_label=file_label, + context_label=context_label, + ) index = _skip_ignorable(sql=sql, start=index) if index >= len(sql) or sql[index] != SQL_OPEN_PAREN_TOKEN: raise CompileInputError( @@ -249,6 +270,7 @@ def extract_top_level_ctes_with_scanner[CteT]( sql=sql, start=index, file_label=file_label, + context_label=context_label, ) return tuple(ctes) @@ -885,11 +907,13 @@ def _contains_select_star(sql: str) -> bool: return False -def _require_prefixed_name(*, cte_name: str, prefix: str, label: str, file_label: str) -> str: +def _require_prefixed_name( + *, cte_name: str, prefix: str, label: str, file_label: str, context_label: str = _CONTEXT +) -> str: extracted_name: str = cte_name.removeprefix(prefix) if extracted_name: return extracted_name - raise CompileInputError(f"SQL test '{file_label}' must use {label} to identify a target") + raise CompileInputError(f"{context_label} '{file_label}' must use {label} to identify a target") def _validate_ceremonial_select( @@ -897,10 +921,14 @@ def _validate_ceremonial_select( sql: str, start: int, file_label: str, + context_label: str = _CONTEXT, ) -> None: if _is_ceremonial_select_statement(sql=sql, start=start): return - raise CompileInputError(_ceremonial_select_error(file_label)) + raise CompileInputError( + f"{context_label} '{file_label}' must end with a ceremonial top-level `SELECT 1` " + "after its CTEs" + ) def _is_ceremonial_select_statement(*, sql: str, start: int) -> bool: @@ -922,10 +950,6 @@ def _is_statement_end(*, sql: str, start: int) -> bool: return index == len(sql) -def _ceremonial_select_error(file_label: str) -> str: - return f"SQL test '{file_label}' must end with a ceremonial top-level `SELECT 1` after its CTEs" - - def validate_independent_expected_and_assertion_ctes( *, ctes: tuple[CompileSqlTestCte | CompileSqlScenarioCte, ...], @@ -1108,16 +1132,13 @@ def _dependency_path( return None -def _consume_keyword(*, sql: str, start: int, keyword: str, file_label: str) -> int: +def _consume_keyword( + *, sql: str, start: int, keyword: str, file_label: str, context_label: str = _CONTEXT +) -> int: keyword_end: int | None = _try_consume_keyword(sql=sql, start=start, keyword=keyword) if keyword_end is not None: return keyword_end - if keyword == SQL_WITH_KEYWORD: - raise CompileInputError( - f"SQL test '{file_label}' must declare mock CTEs and one __expected__ " - "CTE before `SELECT 1`" - ) - raise CompileInputError(f"SQL test '{file_label}' expected keyword {keyword}") + raise CompileInputError(f"{context_label} '{file_label}' expected keyword {keyword}") def _try_consume_keyword(*, sql: str, start: int, keyword: str) -> int | None: @@ -1131,9 +1152,11 @@ def _try_consume_keyword(*, sql: str, start: int, keyword: str) -> int | None: return keyword_end -def _read_identifier(*, sql: str, start: int, file_label: str) -> tuple[str, int]: +def _read_identifier( + *, sql: str, start: int, file_label: str, context_label: str = _CONTEXT +) -> tuple[str, int]: if start >= len(sql) or not is_identifier_start(sql[start]): - raise CompileInputError(f"SQL test '{file_label}' expected a CTE name") + raise CompileInputError(f"{context_label} '{file_label}' expected a CTE name") index: int = start + 1 while index < len(sql) and is_identifier_character(sql[index]): index += 1 diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py index 492e28a0c6..54aa5600f8 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py @@ -892,3 +892,10 @@ class ClassifyDirectLogicSqlTestCtesErrorTestCase: mode: SqlTestMode ctes: tuple[tuple[str, str], ...] expected_message: str + + +@dataclass(frozen=True) +class CteScannerMessageTestCase: + description: str + sql: str + expected_message: str diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_scenarios.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_scenarios.py index bb84908477..6031c6786d 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_scenarios.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_scenarios.py @@ -5,7 +5,11 @@ import pytest from sqlbuild.compiler.compile._helpers.attachment.sql_tests import build_scenario_inputs -from sqlbuild.compiler.compile._helpers.scenarios.core import extract_sql_scenario_ctes +from sqlbuild.compiler.compile._helpers.scenarios.core import ( + extract_sql_scenario_ctes, + extract_sql_scenario_expected_model_names, +) +from sqlbuild.compiler.compile.exceptions import CompileInputError from sqlbuild.compiler.compile.models import ( CompileSqlScenarioCtes, CompileSqlScenarioInput, @@ -20,6 +24,7 @@ from tests.unit.src.sqlbuild.compiler.compile._helpers._test_types import ( BuildScenarioInputsErrorTestCase, BuildScenarioInputsTestCase, + CteScannerMessageTestCase, ExtractSqlScenarioCtesErrorTestCase, ExtractSqlScenarioCtesTestCase, ) @@ -439,3 +444,87 @@ def test_given_invalid_scenario_source_refs_when_building_inputs_then_it_raises_ loaded_macros={}, declaration_expansion=DUCKDB_DECLARATION_EXPANSION_CONTEXT, ) + + +@pytest.mark.parametrize( + "test_case", + ( + CteScannerMessageTestCase( + description="scenario without a with clause uses scenario wording", + sql="SELECT * FROM orders", + expected_message=( + "SQL scenario 'tests/scenarios/orders.sql' must declare fixture CTEs and at least " + "one __expected__ or __assert__ CTE before `SELECT 1`" + ), + ), + CteScannerMessageTestCase( + description="scenario with an invalid cte name uses scenario wording", + sql="WITH 1orders AS (SELECT 1) SELECT 1", + expected_message="SQL scenario 'tests/scenarios/orders.sql' expected a CTE name", + ), + CteScannerMessageTestCase( + description="scenario without a ceremonial select uses scenario wording", + sql=( + "WITH __source__raw_orders AS (SELECT 1 AS order_id), " + "__expected__orders AS (SELECT 1 AS order_id), " + "result AS (SELECT 1) SELECT * FROM result" + ), + expected_message=( + "SQL scenario 'tests/scenarios/orders.sql' must end with a ceremonial top-level " + "`SELECT 1` after its CTEs" + ), + ), + CteScannerMessageTestCase( + description="scenario fixture without a target uses scenario wording", + sql=( + "WITH __source__ AS (SELECT 1 AS order_id), " + "__expected__orders AS (SELECT 1 AS order_id) SELECT 1" + ), + expected_message=( + "SQL scenario 'tests/scenarios/orders.sql' must use __source__ to " + "identify a target" + ), + ), + ), + ids=lambda case: case.description, +) +def test_given_malformed_scenario_ctes_when_extracting_then_messages_name_the_scenario( + test_case: CteScannerMessageTestCase, +) -> None: + with pytest.raises(CompileInputError) as error_info: + _ = extract_sql_scenario_ctes(sql=test_case.sql, file_label="tests/scenarios/orders.sql") + + assert str(error_info.value) == test_case.expected_message + + +@pytest.mark.parametrize( + "test_case", + ( + CteScannerMessageTestCase( + description="expected model scan without AS uses scenario wording", + sql="WITH __source__raw_orders (SELECT 1 AS order_id) SELECT 1", + expected_message="SQL scenario 'tests/scenarios/orders.sql' expected keyword AS", + ), + CteScannerMessageTestCase( + description="expected model scan without a target uses scenario wording", + sql=( + "WITH __source__raw_orders AS (SELECT 1 AS order_id), " + "__expected__ AS (SELECT 1 AS order_id) SELECT 1" + ), + expected_message=( + "SQL scenario 'tests/scenarios/orders.sql' must use __expected__ to " + "identify a target" + ), + ), + ), + ids=lambda case: case.description, +) +def test_given_malformed_scenario_ctes_when_scanning_expected_models_then_names_the_scenario( + test_case: CteScannerMessageTestCase, +) -> None: + with pytest.raises(CompileInputError) as error_info: + _ = extract_sql_scenario_expected_model_names( + sql=test_case.sql, file_label="tests/scenarios/orders.sql" + ) + + assert str(error_info.value) == test_case.expected_message diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_tests.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_tests.py index 6d9b6a2586..0476e95425 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_tests.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_tests.py @@ -5,13 +5,16 @@ from sqlbuild.compiler.compile._helpers.sql_tests.core import ( CompileSqlTestCtes, extract_sql_test_ctes, + extract_sql_test_expected_model_names, ) +from sqlbuild.compiler.compile.exceptions import CompileInputError from sqlbuild.compiler.compile.models import ( CompileDirectLogicSqlTestCtes, CompileModelSqlTestCtes, ) from sqlbuild.compiler.compile.types import SqlTestMode from tests.unit.src.sqlbuild.compiler.compile._helpers._test_types import ( + CteScannerMessageTestCase, ExtractSqlTestCtesErrorTestCase, ExtractSqlTestCtesTestCase, ) @@ -833,3 +836,75 @@ def test_given_invalid_sql_test_cte_variants_when_extracting_then_it_raises_clea file_label="tests/unit/orders.sql", mode=test_case.mode, ) + + +@pytest.mark.parametrize( + "test_case", + ( + CteScannerMessageTestCase( + description="sql test without a with clause keeps sql test wording", + sql="SELECT * FROM orders", + expected_message=( + "SQL test 'tests/unit/orders.sql' must declare mock CTEs and one " + "__expected__ CTE before `SELECT 1`" + ), + ), + CteScannerMessageTestCase( + description="sql test with an invalid cte name keeps sql test wording", + sql="WITH 1orders AS (SELECT 1) SELECT 1", + expected_message="SQL test 'tests/unit/orders.sql' expected a CTE name", + ), + CteScannerMessageTestCase( + description="sql test without a ceremonial select keeps sql test wording", + sql=( + "WITH __source__raw_orders AS (SELECT 1 AS order_id), " + "__expected__orders AS (SELECT 1 AS order_id), " + "result AS (SELECT 1) SELECT * FROM result" + ), + expected_message=( + "SQL test 'tests/unit/orders.sql' must end with a ceremonial top-level " + "`SELECT 1` after its CTEs" + ), + ), + ), + ids=lambda case: case.description, +) +def test_given_malformed_sql_test_ctes_when_extracting_then_messages_name_the_sql_test( + test_case: CteScannerMessageTestCase, +) -> None: + with pytest.raises(CompileInputError) as error_info: + _ = extract_sql_test_ctes(sql=test_case.sql, file_label="tests/unit/orders.sql") + + assert str(error_info.value) == test_case.expected_message + + +@pytest.mark.parametrize( + "test_case", + ( + CteScannerMessageTestCase( + description="expected model scan without AS keeps sql test wording", + sql="WITH __source__raw_orders (SELECT 1 AS order_id) SELECT 1", + expected_message="SQL test 'tests/unit/orders.sql' expected keyword AS", + ), + CteScannerMessageTestCase( + description="expected model scan without a target keeps sql test wording", + sql=( + "WITH __source__raw_orders AS (SELECT 1 AS order_id), " + "__expected__ AS (SELECT 1 AS order_id) SELECT 1" + ), + expected_message=( + "SQL test 'tests/unit/orders.sql' must use __expected__ to identify a target" + ), + ), + ), + ids=lambda case: case.description, +) +def test_given_malformed_sql_test_ctes_when_scanning_expected_models_then_names_the_sql_test( + test_case: CteScannerMessageTestCase, +) -> None: + with pytest.raises(CompileInputError) as error_info: + _ = extract_sql_test_expected_model_names( + sql=test_case.sql, file_label="tests/unit/orders.sql" + ) + + assert str(error_info.value) == test_case.expected_message From ec5c6e5f670ebd3f290b718ba6556420cee5245c Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:36:17 +0100 Subject: [PATCH 31/57] test(native): characterise every SQL scanner on quoting edge cases Pin the current results of the compiler matching-paren scanner, the rules table-function, Snowflake EXCLUDE/trailing-comma and numeric decision scanners, the SQL lint macro-site scanner and the SQL reference fast path on quotes containing parentheses, doubled quotes, backslash escapes, backticks, line/block/nested comments, dollar quotes, unterminated input and multibyte UTF-8. (cherry picked from commit 55fb3f520a9e37d736378aaac5387329b4bfd7ac) --- .../sqlbuild-rules-native/src/engine/tests.rs | 2 + .../src/engine/tests/test_sql_scanning.rs | 272 ++++++++++++++++++ .../src/engine/tests/test_types.rs | 17 ++ 3 files changed, 291 insertions(+) create mode 100644 crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs diff --git a/crates/sqlbuild-rules-native/src/engine/tests.rs b/crates/sqlbuild-rules-native/src/engine/tests.rs index 1907da7600..08d7e5b848 100644 --- a/crates/sqlbuild-rules-native/src/engine/tests.rs +++ b/crates/sqlbuild-rules-native/src/engine/tests.rs @@ -8,6 +8,8 @@ mod evaluation; mod helpers; #[path = "tests/test_scope_facts.rs"] mod scope_facts; +#[path = "tests/test_sql_scanning.rs"] +mod sql_scanning; #[path = "tests/test_sql_test_rules.rs"] mod sql_test_rules; #[path = "tests/test_types.rs"] diff --git a/crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs b/crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs new file mode 100644 index 0000000000..b2c85558b7 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs @@ -0,0 +1,272 @@ +use crate::compiler::_helpers::sql_references::extraction::extract; +use crate::compiler::_helpers::sql_tests::sql_scan::{self, Unclosed}; +use crate::engine::tests::test_types::SqlScannerTestCase; +use crate::rules::_helpers::evaluation::normalize_rules_sql; +use crate::rules::_helpers::numeric_decisions::compact_sql; +use crate::sql_lint::_helpers::preparation::prepare; + +#[test] +fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scanner_is_characterised() +-> Result<(), String> { + let test_cases = [ + SqlScannerTestCase { + description: "plain code", + fragment: "a", + expected_compiler_paren: Ok(4), + expected_table_function: "SELECT * FROM __table_fn(\"o\", a x, 1)", + expected_table_function_token: "SELECT a __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT a a FROM t", + expected_compact: "a=aandb", + expected_lint_site: Some("@m(a x)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "single quote containing open paren", + fragment: "'a(b'", + expected_compiler_paren: Ok(8), + expected_table_function: "SELECT * FROM __table_fn(\"o\", 'a(b' x, 1)", + expected_table_function_token: "SELECT 'a(b' __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT 'a(b' a FROM t", + expected_compact: "a='a(b'andb", + expected_lint_site: Some("@m('a(b' x)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "single quote containing close paren", + fragment: "'a)b'", + expected_compiler_paren: Ok(8), + expected_table_function: "SELECT * FROM __table_fn(\"o\", 'a)b' x, 1)", + expected_table_function_token: "SELECT 'a)b' __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT 'a)b' a FROM t", + expected_compact: "a='a)b'andb", + expected_lint_site: Some("@m('a)b' x)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "doubled single quote", + fragment: "'it''s )'", + expected_compiler_paren: Ok(12), + expected_table_function: "SELECT * FROM __table_fn(\"o\", 'it''s )' x, 1)", + expected_table_function_token: "SELECT 'it''s )' __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT 'it''s )' a FROM t", + expected_compact: "a='it''s )'andb", + expected_lint_site: Some("@m('it''s )' x)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "backslash escaped quote before close paren", + fragment: "'a\\')'", + expected_compiler_paren: Ok(5), + expected_table_function: "SELECT * FROM __table_fn(\"o\", 'a\\')' x)(1)", + expected_table_function_token: "SELECT 'a\\')' __table_fn(\"o\")(1)", + expected_snowflake_exclude: "SELECT * ' x) , a FROM t", + expected_snowflake_comma: "SELECT 'a\\')' a, FROM t", + expected_compact: "a='a\\')' and b", + expected_lint_site: Some("@m('a\\')' x)"), + expected_reference_fast_path: false, + }, + SqlScannerTestCase { + description: "backslash escaped quote before open paren", + fragment: "'a\\'('", + expected_compiler_paren: Err(Unclosed::Quote), + expected_table_function: "SELECT * FROM __table_fn(\"o\", 'a\\'(' x)(1)", + expected_table_function_token: "SELECT 'a\\'(' __table_fn(\"o\")(1)", + expected_snowflake_exclude: "SELECT * EXCLUDE ('a\\'(' x) , a FROM t", + expected_snowflake_comma: "SELECT 'a\\'(' a, FROM t", + expected_compact: "a='a\\'(' and b", + expected_lint_site: Some("@m('a\\'(' x)"), + expected_reference_fast_path: false, + }, + SqlScannerTestCase { + description: "double quote containing close paren", + fragment: "\"a)b\"", + expected_compiler_paren: Ok(8), + expected_table_function: "SELECT * FROM __table_fn(\"o\", \"a)b\" x, 1)", + expected_table_function_token: "SELECT \"a)b\" __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT \"a)b\" a FROM t", + expected_compact: "a=\"a)b\"andb", + expected_lint_site: Some("@m(\"a)b\" x)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "backtick containing close paren", + fragment: "`a)b`", + expected_compiler_paren: Ok(8), + expected_table_function: "SELECT * FROM __table_fn(\"o\", `a)b` x)(1)", + expected_table_function_token: "SELECT `a)b` __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * b` x) , a FROM t", + expected_snowflake_comma: "SELECT `a)b` a FROM t", + expected_compact: "a=`a)b`andb", + expected_lint_site: Some("@m(`a)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "doubled backtick", + fragment: "`a``)b`", + expected_compiler_paren: Ok(10), + expected_table_function: "SELECT * FROM __table_fn(\"o\", `a``)b` x)(1)", + expected_table_function_token: "SELECT `a``)b` __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * b` x) , a FROM t", + expected_snowflake_comma: "SELECT `a``)b` a FROM t", + expected_compact: "a=`a``)b`andb", + expected_lint_site: Some("@m(`a``)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "line comment", + fragment: "-- )\n", + expected_compiler_paren: Ok(8), + expected_table_function: "SELECT * FROM __table_fn(\"o\", -- )\n x, 1)", + expected_table_function_token: "SELECT -- )\n __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * \n , a FROM t", + expected_snowflake_comma: "SELECT -- )\n a FROM t", + expected_compact: "a=andb", + expected_lint_site: Some("@m(-- )\n x)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "block comment", + fragment: "/* ) */", + expected_compiler_paren: Ok(10), + expected_table_function: "SELECT * FROM __table_fn(\"o\", /* ) */ x, 1)", + expected_table_function_token: "SELECT /* ) */ __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT /* ) */ a FROM t", + expected_compact: "a=andb", + expected_lint_site: Some("@m(/* ) */ x)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "nested block comment", + fragment: "/* /* ) */ ) */", + expected_compiler_paren: Ok(12), + expected_table_function: "SELECT * FROM __table_fn(\"o\", /* /* ) */ ) */ x)(1)", + expected_table_function_token: "SELECT /* /* ) */ ) */ __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * */ x) , a FROM t", + expected_snowflake_comma: "SELECT /* /* ) */ ) */ a FROM t", + expected_compact: "a=)*/andb", + expected_lint_site: Some("@m(/* /* ) */ )"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "dollar quote", + fragment: "$$ ) $$", + expected_compiler_paren: Ok(4), + expected_table_function: "SELECT * FROM __table_fn(\"o\", $$ ) $$ x)(1)", + expected_table_function_token: "SELECT $$ ) $$ __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * $$ x) , a FROM t", + expected_snowflake_comma: "SELECT $$ ) $$ a FROM t", + expected_compact: "a=$$)$$andb", + expected_lint_site: Some("@m($$ )"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "unterminated quote", + fragment: "'abc", + expected_compiler_paren: Err(Unclosed::Quote), + expected_table_function: "SELECT * FROM __table_fn(\"o\", 'abc x)(1)", + expected_table_function_token: "SELECT 'abc __table_fn(\"o\")(1)", + expected_snowflake_exclude: "SELECT * EXCLUDE ('abc x) , a FROM t", + expected_snowflake_comma: "SELECT 'abc a, FROM t", + expected_compact: "a='abc and b", + expected_lint_site: Some(""), + expected_reference_fast_path: false, + }, + SqlScannerTestCase { + description: "unterminated block comment", + fragment: "/* abc", + expected_compiler_paren: Err(Unclosed::BlockComment), + expected_table_function: "SELECT * FROM __table_fn(\"o\", /* abc x)(1)", + expected_table_function_token: "SELECT /* abc __table_fn(\"o\")(1)", + expected_snowflake_exclude: "SELECT * EXCLUDE (/* abc x) , a FROM t", + expected_snowflake_comma: "SELECT /* abc a, FROM t", + expected_compact: "a=", + expected_lint_site: Some(""), + expected_reference_fast_path: false, + }, + SqlScannerTestCase { + description: "multibyte utf8", + fragment: "'é)ü' ñ", + expected_compiler_paren: Ok(13), + expected_table_function: "SELECT * FROM __table_fn(\"o\", 'é)ü' ñ x, 1)", + expected_table_function_token: "SELECT 'é)ü' ñ __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT 'é)ü' ñ a FROM t", + expected_compact: "a='é)ü'ñandb", + expected_lint_site: None, + expected_reference_fast_path: true, + }, + ]; + + for test_case in test_cases { + let fragment = test_case.fragment; + assert_eq!( + sql_scan::matching_paren(&format!("({fragment} x) y"), 0), + test_case.expected_compiler_paren, + "compiler matching paren: {}", + test_case.description + ); + assert_eq!( + normalize_rules_sql( + "duckdb", + &format!("SELECT * FROM __table_fn(\"o\", {fragment} x)(1)") + ), + test_case.expected_table_function, + "rules table-function parenthesis: {}", + test_case.description + ); + assert_eq!( + normalize_rules_sql("duckdb", &format!("SELECT {fragment} __table_fn(\"o\")(1)")), + test_case.expected_table_function_token, + "rules table-function token search: {}", + test_case.description + ); + assert_eq!( + normalize_rules_sql( + "snowflake", + &format!("SELECT * EXCLUDE ({fragment} x) , a FROM t") + ), + test_case.expected_snowflake_exclude, + "rules Snowflake EXCLUDE parenthesis: {}", + test_case.description + ); + assert_eq!( + normalize_rules_sql("snowflake", &format!("SELECT {fragment} a, FROM t")), + test_case.expected_snowflake_comma, + "rules Snowflake trailing comma: {}", + test_case.description + ); + assert_eq!( + compact_sql(&format!("A = {fragment} AND B")), + test_case.expected_compact, + "rules numeric decision compaction: {}", + test_case.description + ); + let lint_sql = format!("SELECT @m({fragment} x) AS y"); + let lint_site = prepare(&lint_sql, &lint_sql, &[])?.map(|prepared| { + prepared + .1 + .first() + .map(|site| site.5.clone()) + .unwrap_or_default() + }); + assert_eq!( + lint_site.as_deref(), + test_case.expected_lint_site, + "SQL lint macro site: {}", + test_case.description + ); + assert_eq!( + extract(&format!("SELECT {fragment} __ref('b')")).is_some(), + test_case.expected_reference_fast_path, + "SQL reference fast path: {}", + test_case.description + ); + } + Ok(()) +} diff --git a/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs b/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs index c3457bd571..573834b11c 100644 --- a/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs +++ b/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs @@ -1,5 +1,7 @@ use serde_json::Value; +use crate::compiler::_helpers::sql_tests::sql_scan::Unclosed; + pub(crate) struct NativeEvaluationTestCase { pub(crate) description: &'static str, pub(crate) config: Value, @@ -165,3 +167,18 @@ pub(crate) struct EmptyInputMinimumTestsTestCase { pub(crate) allowed_tests: Value, pub(crate) expected_messages: &'static [&'static str], } + +/// One SQL fragment scanned by every native quote- and comment-aware scanner. +pub(crate) struct SqlScannerTestCase { + pub(crate) description: &'static str, + pub(crate) fragment: &'static str, + pub(crate) expected_compiler_paren: Result, + pub(crate) expected_table_function: &'static str, + pub(crate) expected_table_function_token: &'static str, + pub(crate) expected_snowflake_exclude: &'static str, + pub(crate) expected_snowflake_comma: &'static str, + pub(crate) expected_compact: &'static str, + /// `None` when SQL lint declines non-ASCII input; `Some("")` when no macro site closes. + pub(crate) expected_lint_site: Option<&'static str>, + pub(crate) expected_reference_fast_path: bool, +} From 851de978d4cf8906d2df501d01f1824879e88efd Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:59:46 +0100 Subject: [PATCH 32/57] refactor(native): share one quote- and comment-aware SQL scanner Move the SQL-test planner's scanner to a crate-level sql_scan domain and route every former copy through it: the compiler SQL-test extraction, markers and CTE helpers, the SQL reference fast path, the rules table-function, Snowflake and generic-fallback normalisers, numeric decision compaction and SQL lint macro sites. A QuotePolicy captures the parts callers differed on (backtick identifiers, backslash escapes); each caller keeps its previous policy, so the characterisation tests are unchanged. (cherry picked from commit 8a6fde8d6737d46421c8eb526de9b3b01f83a8de) --- .../_helpers/sql_references/extraction.rs | 45 ++-- .../compiler/_helpers/sql_tests/cte_sql.rs | 5 +- .../compiler/_helpers/sql_tests/extraction.rs | 31 +-- .../compiler/_helpers/sql_tests/markers.rs | 7 +- .../src/compiler/_helpers/sql_tests/mod.rs | 1 - .../compiler/_helpers/sql_tests/sql_scan.rs | 83 ------- .../src/engine/tests/test_sql_scanning.rs | 10 +- .../src/engine/tests/test_types.rs | 2 +- crates/sqlbuild-rules-native/src/lib.rs | 1 + .../src/rules/_helpers/evaluation.rs | 206 ++++-------------- .../src/rules/_helpers/numeric_decisions.rs | 70 +++--- .../src/sql_lint/_helpers/preparation.rs | 50 +---- .../src/sql_scan/main/comment_end.rs | 21 ++ .../src/sql_scan/main/matching_paren.rs | 32 +++ .../src/sql_scan/main/mod.rs | 5 + .../src/sql_scan/main/non_code_end.rs | 20 ++ .../src/sql_scan/main/quote_end.rs | 23 ++ .../src/sql_scan/main/skip_whitespace.rs | 12 + .../sqlbuild-rules-native/src/sql_scan/mod.rs | 5 + .../src/sql_scan/models.rs | 61 ++++++ .../src/sql_scan/tests.rs | 4 + .../src/sql_scan/tests/test_sql_scan.rs | 149 +++++++++++++ .../src/sql_scan/tests/test_types.rs | 16 ++ 23 files changed, 463 insertions(+), 396 deletions(-) delete mode 100644 crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/sql_scan.rs create mode 100644 crates/sqlbuild-rules-native/src/sql_scan/main/comment_end.rs create mode 100644 crates/sqlbuild-rules-native/src/sql_scan/main/matching_paren.rs create mode 100644 crates/sqlbuild-rules-native/src/sql_scan/main/mod.rs create mode 100644 crates/sqlbuild-rules-native/src/sql_scan/main/non_code_end.rs create mode 100644 crates/sqlbuild-rules-native/src/sql_scan/main/quote_end.rs create mode 100644 crates/sqlbuild-rules-native/src/sql_scan/main/skip_whitespace.rs create mode 100644 crates/sqlbuild-rules-native/src/sql_scan/mod.rs create mode 100644 crates/sqlbuild-rules-native/src/sql_scan/models.rs create mode 100644 crates/sqlbuild-rules-native/src/sql_scan/tests.rs create mode 100644 crates/sqlbuild-rules-native/src/sql_scan/tests/test_sql_scan.rs create mode 100644 crates/sqlbuild-rules-native/src/sql_scan/tests/test_types.rs diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_references/extraction.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_references/extraction.rs index c7b0a8fd77..ac05233798 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_references/extraction.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_references/extraction.rs @@ -1,6 +1,9 @@ //! Conservative native fast path for logical SQL reference extraction. use crate::constants::{DBT_REFERENCE_KIND, TABLE_FUNCTION_REFERENCE_KIND}; +use crate::sql_scan::main::non_code_end::non_code_end; +use crate::sql_scan::main::quote_end::quote_end; +use crate::sql_scan::models::QuotePolicy; const PREFIXES: [(&str, &str); 6] = [ ("__dbt_ref(", "dbt_ref"), @@ -18,20 +21,13 @@ pub(crate) fn extract(sql: &str) -> Option> { let mut references: Vec = Vec::new(); let mut index = 0; while index < bytes.len() { - if bytes[index..].starts_with(b"--") { - index = sql[index + 2..] - .find('\n') - .map_or(bytes.len(), |offset| index + offset + 3); - continue; - } - if bytes[index..].starts_with(b"/*") { - let offset = sql[index + 2..].find("*/")?; - index += offset + 4; - continue; - } - if matches!(bytes[index], b'\'' | b'"' | b'`') { - index = quoted_end(bytes, index)?; - continue; + match non_code_end(bytes, index, QuotePolicy::COMPILER) { + Ok(Some(end)) => { + index = end; + continue; + } + Ok(None) => {} + Err(_) => return None, } let Some((prefix, kind)) = PREFIXES .iter() @@ -70,7 +66,9 @@ fn reference_name(sql: &str, start: usize) -> Option<(String, usize)> { let start = ascii_whitespace_end(bytes, start); let first = bytes.get(start).copied()?; if matches!(first, b'\'' | b'"') { - let end = quoted_end(bytes, start)?; + let Ok(end) = quote_end(bytes, start, QuotePolicy::COMPILER) else { + return None; + }; return Some((sql[start + 1..end - 1].to_owned(), end)); } if !is_identifier_start(first) { @@ -86,23 +84,6 @@ fn reference_name(sql: &str, start: usize) -> Option<(String, usize)> { Some((sql[start..end].to_owned(), end)) } -fn quoted_end(bytes: &[u8], start: usize) -> Option { - let quote = bytes[start]; - let mut index = start + 1; - while index < bytes.len() { - if bytes[index] != quote { - index += 1; - continue; - } - if matches!(quote, b'\'' | b'"') && bytes.get(index + 1).copied() == Some(quote) { - index += 2; - continue; - } - return Some(index + 1); - } - None -} - fn ascii_whitespace_end(bytes: &[u8], start: usize) -> usize { let mut end = start; while bytes diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/cte_sql.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/cte_sql.rs index 83e0e26683..03f0f43005 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/cte_sql.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/cte_sql.rs @@ -1,6 +1,7 @@ //! CTE definition and leading-WITH assembly shared by SQL-test planning and rendering. -use crate::compiler::_helpers::sql_tests::sql_scan::{comment_end, skip_whitespace}; +use crate::sql_scan::main::comment_end::comment_end; +use crate::sql_scan::main::skip_whitespace::skip_whitespace; /// Render one `name AS (body)` definition, closing a trailing line comment first. pub(crate) fn cte_definition_sql(name: &str, sql: &str) -> String { @@ -53,7 +54,7 @@ pub(crate) fn leading_with_prefix_end(sql: &str) -> Option { fn skip_leading_ignorable(sql: &str, mut index: usize) -> usize { loop { index = skip_whitespace(sql, index); - match comment_end(sql, index) { + match comment_end(sql.as_bytes(), index) { Ok(Some(end)) => index = end, Ok(None) | Err(_) => return index, } diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs index 625ab2d292..1bb1df6058 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs @@ -4,12 +4,16 @@ use std::collections::{HashMap, HashSet, VecDeque}; use serde::{Deserialize, Serialize}; -use crate::compiler::_helpers::sql_tests::sql_scan::{ - self, Unclosed, comment_end, quoted_end, skip_whitespace, -}; use crate::constants::{ DIRECT_DEPENDENCY_PATH_LENGTH, MACRO_TEST_MODE, TABLE_FUNCTION_TEST_MODE, UDF_TEST_MODE, }; +use crate::sql_scan::main::comment_end::comment_end; +use crate::sql_scan::main::matching_paren::matching_paren as scan_matching_paren; +use crate::sql_scan::main::non_code_end::non_code_end; +use crate::sql_scan::main::quote_end::quote_end; +use crate::sql_scan::main::skip_whitespace::skip_whitespace; +use crate::sql_scan::models::QuotePolicy; +use crate::sql_scan::models::Unclosed; const MODEL_PREFIXES: [&str; 8] = [ "__macro__", @@ -1060,7 +1064,9 @@ fn char_len(sql: &str, index: usize) -> usize { fn skip_ignorable(sql: &str, mut index: usize) -> Result { loop { index = skip_whitespace(sql, index); - match comment_end(sql, index).map_err(|error| scan_error_message(error, "SQL test"))? { + match comment_end(sql.as_bytes(), index) + .map_err(|error| scan_error_message(error, "SQL test"))? + { Some(end) => index = end, None => return Ok(index), } @@ -1068,26 +1074,21 @@ fn skip_ignorable(sql: &str, mut index: usize) -> Result { } fn skip_non_code(sql: &str, index: usize) -> Result { - if let Some(end) = - comment_end(sql, index).map_err(|error| scan_error_message(error, "SQL test"))? - { - return Ok(end); - } - if matches!(byte_at(sql, index), Some(b'\'') | Some(b'"') | Some(b'`')) { - return skip_quote(sql, index, "SQL test"); - } - Ok(index) + Ok(non_code_end(sql.as_bytes(), index, QuotePolicy::COMPILER) + .map_err(|error| scan_error_message(error, "SQL test"))? + .unwrap_or(index)) } fn skip_quote(sql: &str, start: usize, context: &str) -> Result { if byte_at(sql, start).is_none() { return Err(format!("{context} expected a quote")); } - quoted_end(sql, start).map_err(|error| scan_error_message(error, context)) + quote_end(sql.as_bytes(), start, QuotePolicy::COMPILER) + .map_err(|error| scan_error_message(error, context)) } fn matching_paren(sql: &str, open: usize, context: &str) -> Result { - sql_scan::matching_paren(sql, open).map_err(|error| match error { + scan_matching_paren(sql.as_bytes(), open, QuotePolicy::COMPILER).map_err(|error| match error { Unclosed::Parenthesis => scan_error_message(error, context), _ => scan_error_message(error, "SQL test"), }) diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/markers.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/markers.rs index da1babc312..ba1cda0325 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/markers.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/markers.rs @@ -5,7 +5,10 @@ use std::collections::HashSet; use regex::{Captures, Regex}; use crate::compiler::_helpers::sql_tests::planning::compile_error; -use crate::compiler::_helpers::sql_tests::sql_scan::{Unclosed, matching_paren, skip_whitespace}; +use crate::sql_scan::main::matching_paren::matching_paren; +use crate::sql_scan::main::skip_whitespace::skip_whitespace; +use crate::sql_scan::models::QuotePolicy; +use crate::sql_scan::models::Unclosed; pub(crate) fn replace_callable_markers( sql: &str, @@ -51,7 +54,7 @@ where } fn matching_paren_end(sql: &str, open: usize) -> Result { - matching_paren(sql, open) + matching_paren(sql.as_bytes(), open, QuotePolicy::COMPILER) .map(|close| close + 1) .map_err(|error| { compile_error(match error { diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/mod.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/mod.rs index 7969a2dd71..ca1796ea9f 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/mod.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/mod.rs @@ -5,4 +5,3 @@ pub(crate) mod helper_scope; pub(crate) mod markers; pub(crate) mod planning; pub(crate) mod rendering; -pub(crate) mod sql_scan; diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/sql_scan.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/sql_scan.rs deleted file mode 100644 index c1ba225599..0000000000 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/sql_scan.rs +++ /dev/null @@ -1,83 +0,0 @@ -//! Quote- and comment-aware SQL scanning shared by SQL-test extraction and planning. - -/// The SQL construct a scan reached the end of input inside. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(crate) enum Unclosed { - BlockComment, - Quote, - Parenthesis, -} - -/// Return the first index at or after `index` that is not whitespace. -pub(crate) fn skip_whitespace(sql: &str, mut index: usize) -> usize { - while let Some(character) = sql.get(index..).and_then(|rest| rest.chars().next()) { - if !character.is_whitespace() { - break; - } - index += character.len_utf8(); - } - index -} - -/// Return the end of a line or block comment starting at `index`, if one starts there. -pub(crate) fn comment_end(sql: &str, index: usize) -> Result, Unclosed> { - let bytes = sql.as_bytes(); - match (bytes.get(index), bytes.get(index + 1)) { - (Some(b'-'), Some(b'-')) => Ok(Some( - sql[index..] - .find('\n') - .map_or(sql.len(), |offset| index + offset + 1), - )), - (Some(b'/'), Some(b'*')) => sql[index + 2..] - .find("*/") - .map(|offset| Some(index + 2 + offset + 2)) - .ok_or(Unclosed::BlockComment), - _ => Ok(None), - } -} - -/// Return the index after the quoted text starting at `start`, treating doubled quotes as escapes. -pub(crate) fn quoted_end(sql: &str, start: usize) -> Result { - let bytes = sql.as_bytes(); - let quote = bytes[start]; - let mut index = start + 1; - while index < bytes.len() { - if bytes[index] != quote { - index += 1; - } else if bytes.get(index + 1) == Some("e) { - index += 2; - } else { - return Ok(index + 1); - } - } - Err(Unclosed::Quote) -} - -/// Return the index of the parenthesis closing the one at `open`, skipping quotes and comments. -pub(crate) fn matching_paren(sql: &str, open: usize) -> Result { - let bytes = sql.as_bytes(); - let mut depth = 0usize; - let mut index = open; - while index < bytes.len() { - if let Some(end) = comment_end(sql, index)? { - index = end; - continue; - } - match bytes[index] { - b'\'' | b'"' | b'`' => { - index = quoted_end(sql, index)?; - continue; - } - b'(' => depth += 1, - b')' => { - depth -= 1; - if depth == 0 { - return Ok(index); - } - } - _ => {} - } - index += 1; - } - Err(Unclosed::Parenthesis) -} diff --git a/crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs b/crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs index b2c85558b7..fd7e3e6086 100644 --- a/crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs +++ b/crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs @@ -1,9 +1,11 @@ use crate::compiler::_helpers::sql_references::extraction::extract; -use crate::compiler::_helpers::sql_tests::sql_scan::{self, Unclosed}; use crate::engine::tests::test_types::SqlScannerTestCase; use crate::rules::_helpers::evaluation::normalize_rules_sql; use crate::rules::_helpers::numeric_decisions::compact_sql; use crate::sql_lint::_helpers::preparation::prepare; +use crate::sql_scan::main::matching_paren::matching_paren; +use crate::sql_scan::models::QuotePolicy; +use crate::sql_scan::models::Unclosed; #[test] fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scanner_is_characterised() @@ -206,7 +208,11 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann for test_case in test_cases { let fragment = test_case.fragment; assert_eq!( - sql_scan::matching_paren(&format!("({fragment} x) y"), 0), + matching_paren( + format!("({fragment} x) y").as_bytes(), + 0, + QuotePolicy::COMPILER + ), test_case.expected_compiler_paren, "compiler matching paren: {}", test_case.description diff --git a/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs b/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs index 573834b11c..eb8a515cc1 100644 --- a/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs +++ b/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs @@ -1,6 +1,6 @@ use serde_json::Value; -use crate::compiler::_helpers::sql_tests::sql_scan::Unclosed; +use crate::sql_scan::models::Unclosed; pub(crate) struct NativeEvaluationTestCase { pub(crate) description: &'static str, diff --git a/crates/sqlbuild-rules-native/src/lib.rs b/crates/sqlbuild-rules-native/src/lib.rs index 773c00df7d..b37d85783b 100644 --- a/crates/sqlbuild-rules-native/src/lib.rs +++ b/crates/sqlbuild-rules-native/src/lib.rs @@ -12,3 +12,4 @@ mod scope_metadata; mod semantic_usage; mod semantic_validation; mod sql_lint; +mod sql_scan; diff --git a/crates/sqlbuild-rules-native/src/rules/_helpers/evaluation.rs b/crates/sqlbuild-rules-native/src/rules/_helpers/evaluation.rs index e997f1fdd4..ab0b7e9eda 100644 --- a/crates/sqlbuild-rules-native/src/rules/_helpers/evaluation.rs +++ b/crates/sqlbuild-rules-native/src/rules/_helpers/evaluation.rs @@ -11,6 +11,10 @@ use crate::rules::_helpers::{ use crate::rules::models::{ FaultCollector, ModelEvaluationRequest, ProjectEvaluationRequest, ResolvedThresholdOverride, }; +use crate::sql_scan::main::matching_paren::matching_paren; +use crate::sql_scan::main::non_code_end::non_code_end; +use crate::sql_scan::main::quote_end::quote_end; +use crate::sql_scan::models::QuotePolicy; use globset::{Glob, GlobSetBuilder}; use sqlparser::ast::{ BinaryOperator, Expr, GroupByExpr, JoinConstraint, JoinOperator, Query, Select, SelectItem, @@ -260,26 +264,20 @@ pub(super) fn parse_rule_statements( } fn normalize_generic_fallback(sql: &str) -> String { + let policy = QuotePolicy::RULES; let mut bytes = sql.as_bytes().to_vec(); - let mut quote: Option = None; let mut index = 0; while index < bytes.len() { let byte = bytes[index]; let next = bytes.get(index + 1).copied(); - if let Some(active_quote) = quote { - if byte == active_quote { - if next == Some(active_quote) { - index += 2; - continue; - } - quote = None; - } - index += 1; + if policy.is_quote(byte) { + let Ok(end) = quote_end(&bytes, index, policy) else { + break; + }; + index = end; continue; } - if matches!(byte, b'\'' | b'"') { - quote = Some(byte); - } else if byte == b'-' && next == Some(b'>') { + if byte == b'-' && next == Some(b'>') { bytes[index] = b','; bytes[index + 1] = b' '; index += 2; @@ -317,57 +315,22 @@ fn rules_dialect(name: &str) -> Box { } pub(crate) fn normalize_rules_sql(dialect: &str, sql: &str) -> String { - let sql = normalize_table_function_calls(sql); + let policy = QuotePolicy::RULES; + let sql = normalize_table_function_calls(sql, policy); if !dialect.eq_ignore_ascii_case("snowflake") { return sql; } let mut bytes = sql.as_bytes().to_vec(); let mut index = 0; - let mut quote: Option = None; - let mut line_comment = false; - let mut block_comment = false; while index < bytes.len() { let byte = bytes[index]; - let next = bytes.get(index + 1).copied(); - if line_comment { - line_comment = byte != b'\n'; - index += 1; - continue; - } - if block_comment { - if byte == b'*' && next == Some(b'/') { - block_comment = false; - index += 2; - } else { - index += 1; - } - continue; - } - if let Some(active_quote) = quote { - if byte == active_quote { - if next == Some(active_quote) { - index += 2; - continue; - } - quote = None; + match non_code_end(&bytes, index, policy) { + Ok(Some(end)) => { + index = end; + continue; } - index += 1; - continue; - } - if byte == b'-' && next == Some(b'-') { - line_comment = true; - index += 2; - continue; - } - if byte == b'/' && next == Some(b'*') { - block_comment = true; - index += 2; - continue; - } - if matches!(byte, b'\'' | b'"') { - quote = Some(byte); - index += 1; - continue; + Ok(None) => {} + Err(_) => break, } if byte == b',' && clause_follows(&bytes, index + 1) { bytes[index] = b' '; @@ -389,7 +352,7 @@ pub(crate) fn normalize_rules_sql(dialect: &str, sql: &str) -> String { if word.eq_ignore_ascii_case(b"EXCLUDE") && is_star_modifier_before(sql.as_bytes(), start) && bytes.get(following) == Some(&b'(') - && let Some(close) = matching_parenthesis(&bytes, following) + && let Ok(close) = matching_paren(&bytes, following, policy) { blank_preserving_newlines(&mut bytes[start..=close]); index = close + 1; @@ -398,7 +361,7 @@ pub(crate) fn normalize_rules_sql(dialect: &str, sql: &str) -> String { if word.eq_ignore_ascii_case(b"ARRAY") && is_cast_type_before(sql.as_bytes(), start) && bytes.get(following) == Some(&b'(') - && let Some(close) = matching_parenthesis(&bytes, following) + && let Ok(close) = matching_paren(&bytes, following, policy) { bytes[start..index].copy_from_slice(b"TEXT "); blank_preserving_newlines(&mut bytes[following..=close]); @@ -419,16 +382,17 @@ pub(crate) fn normalize_rules_sql(dialect: &str, sql: &str) -> String { String::from_utf8(bytes).unwrap_or(sql) } -fn normalize_table_function_calls(sql: &str) -> String { +fn normalize_table_function_calls(sql: &str, policy: QuotePolicy) -> String { const TABLE_FUNCTION: &[u8] = b"__table_fn("; let mut bytes = sql.as_bytes().to_vec(); let mut cursor = 0; while cursor + TABLE_FUNCTION.len() <= bytes.len() { - let Some(token_start) = find_token_outside_syntax(&bytes, TABLE_FUNCTION, cursor) else { + let Some(token_start) = find_token_outside_syntax(&bytes, TABLE_FUNCTION, cursor, policy) + else { break; }; let open = token_start + TABLE_FUNCTION.len() - 1; - let Some(close) = matching_parenthesis(&bytes, open) else { + let Ok(close) = matching_paren(&bytes, open, policy) else { break; }; let mut next = close + 1; @@ -452,53 +416,21 @@ fn blank_preserving_newlines(sql: &mut [u8]) { } } -fn find_token_outside_syntax(sql: &[u8], token: &[u8], start: usize) -> Option { +fn find_token_outside_syntax( + sql: &[u8], + token: &[u8], + start: usize, + policy: QuotePolicy, +) -> Option { let mut index = start; - let mut quote: Option = None; - let mut line_comment = false; - let mut block_comment = false; while index + token.len() <= sql.len() { - let byte = sql[index]; - let next = sql.get(index + 1).copied(); - if line_comment { - line_comment = byte != b'\n'; - index += 1; - continue; - } - if block_comment { - if byte == b'*' && next == Some(b'/') { - block_comment = false; - index += 2; - } else { - index += 1; - } - continue; - } - if let Some(active_quote) = quote { - if byte == active_quote { - if next == Some(active_quote) { - index += 2; - continue; - } - quote = None; + match non_code_end(sql, index, policy) { + Ok(Some(end)) => { + index = end; + continue; } - index += 1; - continue; - } - if byte == b'-' && next == Some(b'-') { - line_comment = true; - index += 2; - continue; - } - if byte == b'/' && next == Some(b'*') { - block_comment = true; - index += 2; - continue; - } - if matches!(byte, b'\'' | b'"') { - quote = Some(byte); - index += 1; - continue; + Ok(None) => {} + Err(_) => return None, } if sql[index..index + token.len()].eq_ignore_ascii_case(token) { return Some(index); @@ -548,70 +480,6 @@ fn is_cast_type_before(sql: &[u8], type_start: usize) -> bool { || !sql[cursor - BYTE_BEFORE_AS_OFFSET].is_ascii_alphanumeric()) } -fn matching_parenthesis(sql: &[u8], open: usize) -> Option { - let mut depth = 0; - let mut index = open; - let mut quote: Option = None; - let mut line_comment = false; - let mut block_comment = false; - while index < sql.len() { - let byte = sql[index]; - let next = sql.get(index + 1).copied(); - if line_comment { - line_comment = byte != b'\n'; - index += 1; - continue; - } - if block_comment { - if byte == b'*' && next == Some(b'/') { - block_comment = false; - index += 2; - } else { - index += 1; - } - continue; - } - if let Some(active_quote) = quote { - if byte == active_quote { - if next == Some(active_quote) { - index += 2; - continue; - } - quote = None; - } - index += 1; - continue; - } - if byte == b'-' && next == Some(b'-') { - line_comment = true; - index += 2; - continue; - } - if byte == b'/' && next == Some(b'*') { - block_comment = true; - index += 2; - continue; - } - if matches!(byte, b'\'' | b'"') { - quote = Some(byte); - index += 1; - continue; - } - match byte { - b'(' => depth += 1, - b')' => { - depth -= 1; - if depth == 0 { - return Some(index); - } - } - _ => {} - } - index += 1; - } - None -} - fn is_lambda_parameter_before(sql: &[u8], type_start: usize) -> bool { let mut cursor = type_start; while cursor > 0 && sql[cursor - 1].is_ascii_whitespace() { diff --git a/crates/sqlbuild-rules-native/src/rules/_helpers/numeric_decisions.rs b/crates/sqlbuild-rules-native/src/rules/_helpers/numeric_decisions.rs index fddcbcb86a..3ccda207f1 100644 --- a/crates/sqlbuild-rules-native/src/rules/_helpers/numeric_decisions.rs +++ b/crates/sqlbuild-rules-native/src/rules/_helpers/numeric_decisions.rs @@ -1,3 +1,7 @@ +use crate::sql_scan::main::comment_end::comment_end; +use crate::sql_scan::main::quote_end::quote_end; +use crate::sql_scan::models::QuotePolicy; + const CANONICAL_NUMERIC_DECISIONS: [&str; 3] = ["-1", "0", "1"]; pub(crate) struct AuthoredDecision<'a> { @@ -9,59 +13,35 @@ pub(crate) struct AuthoredDecision<'a> { } pub(crate) fn compact_sql(sql: &str) -> String { + let policy = QuotePolicy::RULES; let bytes = sql.as_bytes(); let mut compact = Vec::with_capacity(bytes.len()); let mut index = 0; - let mut quote: Option = None; - let mut line_comment = false; - let mut block_comment = false; while index < bytes.len() { let byte = bytes[index]; - let next = bytes.get(index + 1).copied(); - if line_comment { - line_comment = byte != b'\n'; - index += 1; - continue; - } - if block_comment { - if byte == b'*' && next == Some(b'/') { - block_comment = false; - index += 2; - } else { - index += 1; + let unclosed_end = match comment_end(bytes, index) { + Ok(Some(end)) => { + index = end; + continue; } - continue; - } - if let Some(active_quote) = quote { - compact.push(byte.to_ascii_lowercase()); - if byte == active_quote { - if next == Some(active_quote) { - compact.push(active_quote); - index += 2; - continue; + Ok(None) if policy.is_quote(byte) => { + quote_end(bytes, index, policy).unwrap_or(bytes.len()) + } + Ok(None) => { + if !byte.is_ascii_whitespace() { + compact.push(byte.to_ascii_lowercase()); } - quote = None; + index += 1; + continue; } - index += 1; - continue; - } - if byte == b'-' && next == Some(b'-') { - line_comment = true; - index += 2; - continue; - } - if byte == b'/' && next == Some(b'*') { - block_comment = true; - index += 2; - continue; - } - if matches!(byte, b'\'' | b'"') { - quote = Some(byte); - } - if !byte.is_ascii_whitespace() { - compact.push(byte.to_ascii_lowercase()); - } - index += 1; + Err(_) => break, + }; + compact.extend( + bytes[index..unclosed_end] + .iter() + .map(u8::to_ascii_lowercase), + ); + index = unclosed_end; } String::from_utf8(compact).unwrap_or_default() } diff --git a/crates/sqlbuild-rules-native/src/sql_lint/_helpers/preparation.rs b/crates/sqlbuild-rules-native/src/sql_lint/_helpers/preparation.rs index 8a9ddc97c2..872d157226 100644 --- a/crates/sqlbuild-rules-native/src/sql_lint/_helpers/preparation.rs +++ b/crates/sqlbuild-rules-native/src/sql_lint/_helpers/preparation.rs @@ -6,6 +6,8 @@ use std::sync::LazyLock; use regex::Regex; use crate::sql_lint::types::{InterpolationSite, PreparedSql}; +use crate::sql_scan::main::matching_paren::matching_paren as scan_matching_paren; +use crate::sql_scan::models::QuotePolicy; static SITES: LazyLock> = LazyLock::new(|| { Regex::new( @@ -135,49 +137,9 @@ fn scan_name(bytes: &[u8], mut index: usize, interpolation: bool) -> usize { index } -fn matching_paren(bytes: &[u8], mut index: usize) -> Option { - let mut depth = 0; - let mut quote = None; - while index < bytes.len() { - let byte = bytes[index]; - if let Some(quoted) = quote { - if byte == b'\\' && index + 1 < bytes.len() { - index += 2; - continue; - } - if byte == quoted { - if bytes.get(index + 1) == Some("ed) { - index += 2; - continue; - } - quote = None; - } - } else if bytes[index..].starts_with(b"--") { - index = bytes[index + 2..] - .iter() - .position(|byte| *byte == b'\n') - .map_or(bytes.len(), |offset| index + offset + 3); - continue; - } else if bytes[index..].starts_with(b"/*") { - index = bytes[index + 2..] - .windows(2) - .position(|pair| pair == b"*/") - .map_or(bytes.len(), |offset| index + offset + 4); - continue; - } else { - match byte { - b'\'' | b'"' => quote = Some(byte), - b'(' => depth += 1, - b')' => { - depth -= 1; - if depth == 0 { - return Some(index + 1); - } - } - _ => {} - } - } - index += 1; +fn matching_paren(bytes: &[u8], open: usize) -> Option { + match scan_matching_paren(bytes, open, QuotePolicy::SQL_LINT) { + Ok(close) => Some(close + 1), + Err(_) => None, } - None } diff --git a/crates/sqlbuild-rules-native/src/sql_scan/main/comment_end.rs b/crates/sqlbuild-rules-native/src/sql_scan/main/comment_end.rs new file mode 100644 index 0000000000..b90d31b65d --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/main/comment_end.rs @@ -0,0 +1,21 @@ +//! SQL line and block comment boundaries. + +use crate::sql_scan::models::Unclosed; + +/// Return the end of a line or block comment starting at `index`, if one starts there. +pub(crate) fn comment_end(sql: &[u8], index: usize) -> Result, Unclosed> { + match (sql.get(index), sql.get(index + 1)) { + (Some(b'-'), Some(b'-')) => Ok(Some( + sql[index + 2..] + .iter() + .position(|byte| *byte == b'\n') + .map_or(sql.len(), |offset| index + 2 + offset + 1), + )), + (Some(b'/'), Some(b'*')) => sql[index + 2..] + .windows(2) + .position(|pair| pair == b"*/") + .map(|offset| Some(index + 2 + offset + 2)) + .ok_or(Unclosed::BlockComment), + _ => Ok(None), + } +} diff --git a/crates/sqlbuild-rules-native/src/sql_scan/main/matching_paren.rs b/crates/sqlbuild-rules-native/src/sql_scan/main/matching_paren.rs new file mode 100644 index 0000000000..89ee9b98e3 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/main/matching_paren.rs @@ -0,0 +1,32 @@ +//! Quote- and comment-aware parenthesis matching. + +use crate::sql_scan::main::non_code_end::non_code_end; +use crate::sql_scan::models::{QuotePolicy, Unclosed}; + +/// Return the index of the parenthesis closing the one at `open`, skipping quotes and comments. +pub(crate) fn matching_paren( + sql: &[u8], + open: usize, + policy: QuotePolicy, +) -> Result { + let mut depth = 0isize; + let mut index = open; + while index < sql.len() { + if let Some(end) = non_code_end(sql, index, policy)? { + index = end; + continue; + } + match sql[index] { + b'(' => depth += 1, + b')' => { + depth -= 1; + if depth == 0 { + return Ok(index); + } + } + _ => {} + } + index += 1; + } + Err(Unclosed::Parenthesis) +} diff --git a/crates/sqlbuild-rules-native/src/sql_scan/main/mod.rs b/crates/sqlbuild-rules-native/src/sql_scan/main/mod.rs new file mode 100644 index 0000000000..4bf480ac79 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/main/mod.rs @@ -0,0 +1,5 @@ +pub(crate) mod comment_end; +pub(crate) mod matching_paren; +pub(crate) mod non_code_end; +pub(crate) mod quote_end; +pub(crate) mod skip_whitespace; diff --git a/crates/sqlbuild-rules-native/src/sql_scan/main/non_code_end.rs b/crates/sqlbuild-rules-native/src/sql_scan/main/non_code_end.rs new file mode 100644 index 0000000000..cb99e49124 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/main/non_code_end.rs @@ -0,0 +1,20 @@ +//! Boundaries of the comment or quoted text starting at one offset. + +use crate::sql_scan::main::comment_end::comment_end; +use crate::sql_scan::main::quote_end::quote_end; +use crate::sql_scan::models::{QuotePolicy, Unclosed}; + +/// Return the end of a comment or quoted text starting at `index`, if one starts there. +pub(crate) fn non_code_end( + sql: &[u8], + index: usize, + policy: QuotePolicy, +) -> Result, Unclosed> { + if let Some(end) = comment_end(sql, index)? { + return Ok(Some(end)); + } + match sql.get(index) { + Some(byte) if policy.is_quote(*byte) => quote_end(sql, index, policy).map(Some), + _ => Ok(None), + } +} diff --git a/crates/sqlbuild-rules-native/src/sql_scan/main/quote_end.rs b/crates/sqlbuild-rules-native/src/sql_scan/main/quote_end.rs new file mode 100644 index 0000000000..fdf256cea6 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/main/quote_end.rs @@ -0,0 +1,23 @@ +//! Quoted SQL text boundaries under a quoting policy. + +use crate::sql_scan::models::{QuotePolicy, Unclosed}; + +/// Return the index after the quoted text whose opening quote is at `start`. +pub(crate) fn quote_end(sql: &[u8], start: usize, policy: QuotePolicy) -> Result { + let quote = sql[start]; + let backslash_escapes = policy.backslash_escapes(quote); + let mut index = start + 1; + while index < sql.len() { + let byte = sql[index]; + if backslash_escapes && byte == b'\\' && index + 1 < sql.len() { + index += 2; + } else if byte != quote { + index += 1; + } else if sql.get(index + 1) == Some("e) { + index += 2; + } else { + return Ok(index + 1); + } + } + Err(Unclosed::Quote) +} diff --git a/crates/sqlbuild-rules-native/src/sql_scan/main/skip_whitespace.rs b/crates/sqlbuild-rules-native/src/sql_scan/main/skip_whitespace.rs new file mode 100644 index 0000000000..a60ea5e9c4 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/main/skip_whitespace.rs @@ -0,0 +1,12 @@ +//! Unicode whitespace skipping for SQL scanners. + +/// Return the first index at or after `index` that is not Unicode whitespace. +pub(crate) fn skip_whitespace(sql: &str, mut index: usize) -> usize { + while let Some(character) = sql.get(index..).and_then(|rest| rest.chars().next()) { + if !character.is_whitespace() { + break; + } + index += character.len_utf8(); + } + index +} diff --git a/crates/sqlbuild-rules-native/src/sql_scan/mod.rs b/crates/sqlbuild-rules-native/src/sql_scan/mod.rs new file mode 100644 index 0000000000..650595395f --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/mod.rs @@ -0,0 +1,5 @@ +pub(crate) mod main; +pub(crate) mod models; + +#[cfg(test)] +mod tests; diff --git a/crates/sqlbuild-rules-native/src/sql_scan/models.rs b/crates/sqlbuild-rules-native/src/sql_scan/models.rs new file mode 100644 index 0000000000..16ee59e865 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/models.rs @@ -0,0 +1,61 @@ +//! Lexical scanning outcomes and dialect-dependent quoting policy. + +/// The SQL construct a scan reached the end of input inside. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Unclosed { + BlockComment, + Quote, + Parenthesis, +} + +/// Dialect-dependent quoting rules; comments and doubled-quote escapes are fixed for every policy. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct QuotePolicy { + /// Backtick-delimited identifiers are quoted text. + pub(crate) backtick_identifiers: bool, + /// A backslash escapes the next byte inside single-quoted strings. + pub(crate) single_quote_backslash_escapes: bool, + /// A backslash escapes the next byte inside double-quoted text. + pub(crate) double_quote_backslash_escapes: bool, +} + +impl QuotePolicy { + /// Compiler policy, mirroring Python `sql_analysis`: backticks quote, backslashes are literal. + pub(crate) const COMPILER: Self = Self { + backtick_identifiers: true, + single_quote_backslash_escapes: false, + double_quote_backslash_escapes: false, + }; + + /// SQL lint policy, mirroring its site lexer: backslashes escape, backticks are code. + pub(crate) const SQL_LINT: Self = Self { + backtick_identifiers: false, + single_quote_backslash_escapes: true, + double_quote_backslash_escapes: true, + }; + + /// Rules policy: single and double quotes only, with backslashes treated as literal text. + pub(crate) const RULES: Self = Self { + backtick_identifiers: false, + single_quote_backslash_escapes: false, + double_quote_backslash_escapes: false, + }; + + /// Return whether `byte` opens quoted text under this policy. + pub(crate) fn is_quote(self, byte: u8) -> bool { + match byte { + b'\'' | b'"' => true, + b'`' => self.backtick_identifiers, + _ => false, + } + } + + /// Return whether a backslash escapes the next byte inside text opened by `quote`. + pub(crate) fn backslash_escapes(self, quote: u8) -> bool { + match quote { + b'\'' => self.single_quote_backslash_escapes, + b'"' => self.double_quote_backslash_escapes, + _ => false, + } + } +} diff --git a/crates/sqlbuild-rules-native/src/sql_scan/tests.rs b/crates/sqlbuild-rules-native/src/sql_scan/tests.rs new file mode 100644 index 0000000000..0b2156d36c --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/tests.rs @@ -0,0 +1,4 @@ +#[path = "tests/test_sql_scan.rs"] +mod sql_scan; +#[path = "tests/test_types.rs"] +mod test_types; diff --git a/crates/sqlbuild-rules-native/src/sql_scan/tests/test_sql_scan.rs b/crates/sqlbuild-rules-native/src/sql_scan/tests/test_sql_scan.rs new file mode 100644 index 0000000000..41f30993ba --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/tests/test_sql_scan.rs @@ -0,0 +1,149 @@ +use crate::sql_scan::main::matching_paren::matching_paren; +use crate::sql_scan::main::non_code_end::non_code_end; +use crate::sql_scan::models::QuotePolicy; +use crate::sql_scan::models::Unclosed; +use crate::sql_scan::tests::test_types::{MatchingParenPolicyTestCase, NonCodeEndTestCase}; + +const BACKSLASH_ONLY: QuotePolicy = QuotePolicy { + backtick_identifiers: false, + single_quote_backslash_escapes: true, + double_quote_backslash_escapes: false, +}; + +#[test] +fn given_quote_policies_when_matching_parentheses_then_policy_decides_quote_boundaries() { + let test_cases = [ + MatchingParenPolicyTestCase { + description: "compiler policy skips backtick identifiers", + sql: "(`a)b` x) y", + policy: QuotePolicy::COMPILER, + expected_close: Ok(8), + }, + MatchingParenPolicyTestCase { + description: "rules policy treats backticks as code", + sql: "(`a)b` x) y", + policy: QuotePolicy::RULES, + expected_close: Ok(3), + }, + MatchingParenPolicyTestCase { + description: "compiler policy ends a quote at a backslash-escaped quote", + sql: r"('a\')' x) y", + policy: QuotePolicy::COMPILER, + expected_close: Ok(5), + }, + MatchingParenPolicyTestCase { + description: "sql lint policy honours backslash escapes", + sql: r"('a\')' x) y", + policy: QuotePolicy::SQL_LINT, + expected_close: Ok(9), + }, + MatchingParenPolicyTestCase { + description: "single-quote backslash escapes do not apply to double quotes", + sql: r#"("a\")" x) y"#, + policy: BACKSLASH_ONLY, + expected_close: Ok(5), + }, + MatchingParenPolicyTestCase { + description: "doubled quotes are escapes under every policy", + sql: "('it''s )' x) y", + policy: QuotePolicy::RULES, + expected_close: Ok(12), + }, + MatchingParenPolicyTestCase { + description: "block comments do not nest", + sql: "(/* /* ) */ ) */ x) y", + policy: QuotePolicy::COMPILER, + expected_close: Ok(12), + }, + MatchingParenPolicyTestCase { + description: "dollar quotes are scanned as code", + sql: "($$ ) $$ x) y", + policy: QuotePolicy::COMPILER, + expected_close: Ok(4), + }, + MatchingParenPolicyTestCase { + description: "multibyte text keeps byte offsets on character boundaries", + sql: "('é)ü' ñ x) y", + policy: QuotePolicy::COMPILER, + expected_close: Ok(13), + }, + MatchingParenPolicyTestCase { + description: "unterminated quote reports the quote", + sql: "('abc x) y", + policy: QuotePolicy::COMPILER, + expected_close: Err(Unclosed::Quote), + }, + MatchingParenPolicyTestCase { + description: "unterminated block comment reports the comment", + sql: "(/* abc x) y", + policy: QuotePolicy::RULES, + expected_close: Err(Unclosed::BlockComment), + }, + MatchingParenPolicyTestCase { + description: "unclosed parenthesis reports the parenthesis", + sql: "(a (b)", + policy: QuotePolicy::SQL_LINT, + expected_close: Err(Unclosed::Parenthesis), + }, + ]; + + for test_case in test_cases { + assert_eq!( + matching_paren(test_case.sql.as_bytes(), 0, test_case.policy), + test_case.expected_close, + "{}", + test_case.description + ); + } +} + +#[test] +fn given_comment_or_quote_starts_when_finding_non_code_end_then_returns_end_offset() { + let test_cases = [ + NonCodeEndTestCase { + description: "line comment ends after its newline", + sql: "-- note\nSELECT", + policy: QuotePolicy::COMPILER, + expected_end: Ok(Some(8)), + }, + NonCodeEndTestCase { + description: "line comment without newline runs to the end", + sql: "-- note", + policy: QuotePolicy::COMPILER, + expected_end: Ok(Some(7)), + }, + NonCodeEndTestCase { + description: "block comment ends after its terminator", + sql: "/* a */ b", + policy: QuotePolicy::RULES, + expected_end: Ok(Some(7)), + }, + NonCodeEndTestCase { + description: "backtick is code when the policy disables it", + sql: "`a` b", + policy: QuotePolicy::SQL_LINT, + expected_end: Ok(None), + }, + NonCodeEndTestCase { + description: "trailing backslash leaves the quote unterminated", + sql: r"'a\", + policy: QuotePolicy::SQL_LINT, + expected_end: Err(Unclosed::Quote), + }, + NonCodeEndTestCase { + description: "plain code has no non-code end", + sql: "SELECT 1", + policy: QuotePolicy::COMPILER, + expected_end: Ok(None), + }, + ]; + + for test_case in test_cases { + assert_eq!( + non_code_end(test_case.sql.as_bytes(), 0, test_case.policy), + test_case.expected_end, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-rules-native/src/sql_scan/tests/test_types.rs b/crates/sqlbuild-rules-native/src/sql_scan/tests/test_types.rs new file mode 100644 index 0000000000..0951ec3e0e --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/tests/test_types.rs @@ -0,0 +1,16 @@ +use crate::sql_scan::models::QuotePolicy; +use crate::sql_scan::models::Unclosed; + +pub(crate) struct MatchingParenPolicyTestCase { + pub(crate) description: &'static str, + pub(crate) sql: &'static str, + pub(crate) policy: QuotePolicy, + pub(crate) expected_close: Result, +} + +pub(crate) struct NonCodeEndTestCase { + pub(crate) description: &'static str, + pub(crate) sql: &'static str, + pub(crate) policy: QuotePolicy, + pub(crate) expected_end: Result, Unclosed>, +} From 2669e801826cd18e8522fcbaf1234323d1bcf6e6 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:03:23 +0100 Subject: [PATCH 33/57] fix(rules): scan quotes the way the rules dialect tokenises them The rules normalisers and numeric-decision compaction treated only single and double quotes as quoted text and never honoured backslash escapes. In Snowflake or BigQuery a literal such as 'a\'(' or a backticked identifier containing ) could desynchronise the scanner, so table-function calls were left unnormalised and the model failed to parse for rules, or Snowflake EXCLUDE/trailing-comma rewrites were skipped. The rules quote policy now follows the sqlparser dialect used for the model: backticks where the dialect delimits identifiers with them, and backslash escapes where it supports them (double-quoted text only where it is a string). (cherry picked from commit 4dc640dc4b3f62b1081ba5b13b3bc1ea02e9205e) --- .../src/engine/tests/test_evaluation.rs | 51 +++++++++++++++++++ .../src/engine/tests/test_sql_scanning.rs | 48 +++++++++++++---- .../src/engine/tests/test_types.rs | 1 + .../src/rules/_helpers/evaluation.rs | 21 ++++++-- .../src/rules/_helpers/numeric_decisions.rs | 6 +-- .../src/sql_scan/models.rs | 7 --- .../src/sql_scan/tests/test_sql_scan.rs | 16 ++++-- 7 files changed, 122 insertions(+), 28 deletions(-) diff --git a/crates/sqlbuild-rules-native/src/engine/tests/test_evaluation.rs b/crates/sqlbuild-rules-native/src/engine/tests/test_evaluation.rs index e6af01313f..03d28d174c 100644 --- a/crates/sqlbuild-rules-native/src/engine/tests/test_evaluation.rs +++ b/crates/sqlbuild-rules-native/src/engine/tests/test_evaluation.rs @@ -428,6 +428,57 @@ fn given_snowflake_expression_when_evaluating_rules_then_uses_project_dialect() Ok(()) } +#[test] +fn given_dialect_quoting_around_table_functions_when_evaluating_rules_then_model_parses() +-> Result<(), String> { + let project_dir = TempDir::new().map_err(|error| error.to_string())?; + let test_cases = [ + test_types::DialectEvaluationTestCase { + description: "Snowflake backslash-escaped quote before a table function", + dialect: "snowflake", + query_sql: concat!( + "SELECT 'it\\'s (' AS label, item_id ", + "FROM __table_fn(\"inventory_items\")(42)" + ), + expected_code: "SQBRCONTRACT101", + }, + test_types::DialectEvaluationTestCase { + description: "BigQuery backslash-escaped quote inside table function arguments", + dialect: "bigquery", + query_sql: concat!( + "SELECT item_id ", + "FROM __table_fn(\"inventory_items\", 'a\\')')(42)" + ), + expected_code: "SQBRCONTRACT101", + }, + test_types::DialectEvaluationTestCase { + description: "BigQuery backtick identifier containing a parenthesis in arguments", + dialect: "bigquery", + query_sql: concat!( + "SELECT item_id ", + "FROM __table_fn(\"inventory_items\", `batch)size`)(42)" + ), + expected_code: "SQBRCONTRACT101", + }, + ]; + + for test_case in test_cases { + let config = json!({"select": ["SQBRCONTRACT101"], "cache": {"enabled": false}}); + let mut request: Value = serde_json::from_str(&helpers::request(&project_dir, &config)) + .map_err(|error| error.to_string())?; + request["dialect"] = json!(test_case.dialect); + request["models"][0]["query_sql"] = json!(test_case.query_sql); + let result: Value = serde_json::from_str(&evaluate_json(&request.to_string())?) + .map_err(|error| error.to_string())?; + assert_eq!( + result["faults"][0]["code"], test_case.expected_code, + "{}: {result}", + test_case.description + ); + } + Ok(()) +} + #[test] fn given_enforced_contract_outputs_when_evaluating_explicit_type_rule_then_returns_expected_faults() -> Result<(), String> { diff --git a/crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs b/crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs index fd7e3e6086..baea4d051b 100644 --- a/crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs +++ b/crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs @@ -1,6 +1,6 @@ use crate::compiler::_helpers::sql_references::extraction::extract; use crate::engine::tests::test_types::SqlScannerTestCase; -use crate::rules::_helpers::evaluation::normalize_rules_sql; +use crate::rules::_helpers::evaluation::{normalize_rules_sql, rules_quote_policy}; use crate::rules::_helpers::numeric_decisions::compact_sql; use crate::sql_lint::_helpers::preparation::prepare; use crate::sql_scan::main::matching_paren::matching_paren; @@ -20,6 +20,7 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann expected_snowflake_exclude: "SELECT * , a FROM t", expected_snowflake_comma: "SELECT a a FROM t", expected_compact: "a=aandb", + expected_snowflake_compact: "a=aandb", expected_lint_site: Some("@m(a x)"), expected_reference_fast_path: true, }, @@ -32,6 +33,7 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann expected_snowflake_exclude: "SELECT * , a FROM t", expected_snowflake_comma: "SELECT 'a(b' a FROM t", expected_compact: "a='a(b'andb", + expected_snowflake_compact: "a='a(b'andb", expected_lint_site: Some("@m('a(b' x)"), expected_reference_fast_path: true, }, @@ -44,6 +46,7 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann expected_snowflake_exclude: "SELECT * , a FROM t", expected_snowflake_comma: "SELECT 'a)b' a FROM t", expected_compact: "a='a)b'andb", + expected_snowflake_compact: "a='a)b'andb", expected_lint_site: Some("@m('a)b' x)"), expected_reference_fast_path: true, }, @@ -56,6 +59,7 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann expected_snowflake_exclude: "SELECT * , a FROM t", expected_snowflake_comma: "SELECT 'it''s )' a FROM t", expected_compact: "a='it''s )'andb", + expected_snowflake_compact: "a='it''s )'andb", expected_lint_site: Some("@m('it''s )' x)"), expected_reference_fast_path: true, }, @@ -65,9 +69,10 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann expected_compiler_paren: Ok(5), expected_table_function: "SELECT * FROM __table_fn(\"o\", 'a\\')' x)(1)", expected_table_function_token: "SELECT 'a\\')' __table_fn(\"o\")(1)", - expected_snowflake_exclude: "SELECT * ' x) , a FROM t", - expected_snowflake_comma: "SELECT 'a\\')' a, FROM t", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT 'a\\')' a FROM t", expected_compact: "a='a\\')' and b", + expected_snowflake_compact: "a='a\\')'andb", expected_lint_site: Some("@m('a\\')' x)"), expected_reference_fast_path: false, }, @@ -77,9 +82,10 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann expected_compiler_paren: Err(Unclosed::Quote), expected_table_function: "SELECT * FROM __table_fn(\"o\", 'a\\'(' x)(1)", expected_table_function_token: "SELECT 'a\\'(' __table_fn(\"o\")(1)", - expected_snowflake_exclude: "SELECT * EXCLUDE ('a\\'(' x) , a FROM t", - expected_snowflake_comma: "SELECT 'a\\'(' a, FROM t", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT 'a\\'(' a FROM t", expected_compact: "a='a\\'(' and b", + expected_snowflake_compact: "a='a\\'('andb", expected_lint_site: Some("@m('a\\'(' x)"), expected_reference_fast_path: false, }, @@ -92,6 +98,7 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann expected_snowflake_exclude: "SELECT * , a FROM t", expected_snowflake_comma: "SELECT \"a)b\" a FROM t", expected_compact: "a=\"a)b\"andb", + expected_snowflake_compact: "a=\"a)b\"andb", expected_lint_site: Some("@m(\"a)b\" x)"), expected_reference_fast_path: true, }, @@ -99,11 +106,12 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann description: "backtick containing close paren", fragment: "`a)b`", expected_compiler_paren: Ok(8), - expected_table_function: "SELECT * FROM __table_fn(\"o\", `a)b` x)(1)", + expected_table_function: "SELECT * FROM __table_fn(\"o\", `a)b` x, 1)", expected_table_function_token: "SELECT `a)b` __table_fn(\"o\", 1)", - expected_snowflake_exclude: "SELECT * b` x) , a FROM t", + expected_snowflake_exclude: "SELECT * , a FROM t", expected_snowflake_comma: "SELECT `a)b` a FROM t", expected_compact: "a=`a)b`andb", + expected_snowflake_compact: "a=`a)b`andb", expected_lint_site: Some("@m(`a)"), expected_reference_fast_path: true, }, @@ -111,11 +119,12 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann description: "doubled backtick", fragment: "`a``)b`", expected_compiler_paren: Ok(10), - expected_table_function: "SELECT * FROM __table_fn(\"o\", `a``)b` x)(1)", + expected_table_function: "SELECT * FROM __table_fn(\"o\", `a``)b` x, 1)", expected_table_function_token: "SELECT `a``)b` __table_fn(\"o\", 1)", - expected_snowflake_exclude: "SELECT * b` x) , a FROM t", + expected_snowflake_exclude: "SELECT * , a FROM t", expected_snowflake_comma: "SELECT `a``)b` a FROM t", expected_compact: "a=`a``)b`andb", + expected_snowflake_compact: "a=`a``)b`andb", expected_lint_site: Some("@m(`a``)"), expected_reference_fast_path: true, }, @@ -128,6 +137,7 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann expected_snowflake_exclude: "SELECT * \n , a FROM t", expected_snowflake_comma: "SELECT -- )\n a FROM t", expected_compact: "a=andb", + expected_snowflake_compact: "a=andb", expected_lint_site: Some("@m(-- )\n x)"), expected_reference_fast_path: true, }, @@ -140,6 +150,7 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann expected_snowflake_exclude: "SELECT * , a FROM t", expected_snowflake_comma: "SELECT /* ) */ a FROM t", expected_compact: "a=andb", + expected_snowflake_compact: "a=andb", expected_lint_site: Some("@m(/* ) */ x)"), expected_reference_fast_path: true, }, @@ -152,6 +163,7 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann expected_snowflake_exclude: "SELECT * */ x) , a FROM t", expected_snowflake_comma: "SELECT /* /* ) */ ) */ a FROM t", expected_compact: "a=)*/andb", + expected_snowflake_compact: "a=)*/andb", expected_lint_site: Some("@m(/* /* ) */ )"), expected_reference_fast_path: true, }, @@ -164,6 +176,7 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann expected_snowflake_exclude: "SELECT * $$ x) , a FROM t", expected_snowflake_comma: "SELECT $$ ) $$ a FROM t", expected_compact: "a=$$)$$andb", + expected_snowflake_compact: "a=$$)$$andb", expected_lint_site: Some("@m($$ )"), expected_reference_fast_path: true, }, @@ -176,6 +189,7 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann expected_snowflake_exclude: "SELECT * EXCLUDE ('abc x) , a FROM t", expected_snowflake_comma: "SELECT 'abc a, FROM t", expected_compact: "a='abc and b", + expected_snowflake_compact: "a='abc and b", expected_lint_site: Some(""), expected_reference_fast_path: false, }, @@ -188,6 +202,7 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann expected_snowflake_exclude: "SELECT * EXCLUDE (/* abc x) , a FROM t", expected_snowflake_comma: "SELECT /* abc a, FROM t", expected_compact: "a=", + expected_snowflake_compact: "a=", expected_lint_site: Some(""), expected_reference_fast_path: false, }, @@ -200,6 +215,7 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann expected_snowflake_exclude: "SELECT * , a FROM t", expected_snowflake_comma: "SELECT 'é)ü' ñ a FROM t", expected_compact: "a='é)ü'ñandb", + expected_snowflake_compact: "a='é)ü'ñandb", expected_lint_site: None, expected_reference_fast_path: true, }, @@ -248,11 +264,23 @@ fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scann test_case.description ); assert_eq!( - compact_sql(&format!("A = {fragment} AND B")), + compact_sql( + &format!("A = {fragment} AND B"), + rules_quote_policy("duckdb") + ), test_case.expected_compact, "rules numeric decision compaction: {}", test_case.description ); + assert_eq!( + compact_sql( + &format!("A = {fragment} AND B"), + rules_quote_policy("snowflake") + ), + test_case.expected_snowflake_compact, + "rules Snowflake numeric decision compaction: {}", + test_case.description + ); let lint_sql = format!("SELECT @m({fragment} x) AS y"); let lint_site = prepare(&lint_sql, &lint_sql, &[])?.map(|prepared| { prepared diff --git a/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs b/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs index eb8a515cc1..9c319a4f50 100644 --- a/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs +++ b/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs @@ -178,6 +178,7 @@ pub(crate) struct SqlScannerTestCase { pub(crate) expected_snowflake_exclude: &'static str, pub(crate) expected_snowflake_comma: &'static str, pub(crate) expected_compact: &'static str, + pub(crate) expected_snowflake_compact: &'static str, /// `None` when SQL lint declines non-ASCII input; `Some("")` when no macro site closes. pub(crate) expected_lint_site: Option<&'static str>, pub(crate) expected_reference_fast_path: bool, diff --git a/crates/sqlbuild-rules-native/src/rules/_helpers/evaluation.rs b/crates/sqlbuild-rules-native/src/rules/_helpers/evaluation.rs index ab0b7e9eda..6732ebc948 100644 --- a/crates/sqlbuild-rules-native/src/rules/_helpers/evaluation.rs +++ b/crates/sqlbuild-rules-native/src/rules/_helpers/evaluation.rs @@ -74,6 +74,7 @@ struct ParsedModel<'a> { model: &'a Model, classification: ModelClassification, authored_compact_sql: String, + quote_policy: QuotePolicy, } struct ModelClassification { @@ -149,12 +150,14 @@ fn evaluate_model_inner(request: ModelEvaluationRequest<'_>) -> Result String { - let policy = QuotePolicy::RULES; + let policy = rules_quote_policy("generic"); let mut bytes = sql.as_bytes().to_vec(); let mut index = 0; while index < bytes.len() { @@ -301,6 +304,17 @@ fn parse_with_dialect(sql: &str, dialect: &dyn Dialect) -> Result parser.parse_statements() } +pub(crate) fn rules_quote_policy(dialect_name: &str) -> QuotePolicy { + let dialect = rules_dialect(dialect_name); + let backslash_escapes = dialect.supports_string_literal_backslash_escape(); + QuotePolicy { + backtick_identifiers: dialect.is_delimited_identifier_start('`'), + single_quote_backslash_escapes: backslash_escapes, + double_quote_backslash_escapes: backslash_escapes + && !dialect.is_delimited_identifier_start('"'), + } +} + fn rules_dialect(name: &str) -> Box { match name.to_ascii_lowercase().as_str() { "bigquery" => Box::new(BigQueryDialect {}), @@ -315,7 +329,7 @@ fn rules_dialect(name: &str) -> Box { } pub(crate) fn normalize_rules_sql(dialect: &str, sql: &str) -> String { - let policy = QuotePolicy::RULES; + let policy = rules_quote_policy(dialect); let sql = normalize_table_function_calls(sql, policy); if !dialect.eq_ignore_ascii_case("snowflake") { return sql; @@ -1242,6 +1256,7 @@ fn evaluate_literal_rules( equality: comparison.equality, output_name: comparison.output_name.as_deref(), literal, + quote_policy: parsed.quote_policy, }) }); if magic { diff --git a/crates/sqlbuild-rules-native/src/rules/_helpers/numeric_decisions.rs b/crates/sqlbuild-rules-native/src/rules/_helpers/numeric_decisions.rs index 3ccda207f1..3d861d3fb3 100644 --- a/crates/sqlbuild-rules-native/src/rules/_helpers/numeric_decisions.rs +++ b/crates/sqlbuild-rules-native/src/rules/_helpers/numeric_decisions.rs @@ -10,10 +10,10 @@ pub(crate) struct AuthoredDecision<'a> { pub(crate) equality: bool, pub(crate) output_name: Option<&'a str>, pub(crate) literal: &'a str, + pub(crate) quote_policy: QuotePolicy, } -pub(crate) fn compact_sql(sql: &str) -> String { - let policy = QuotePolicy::RULES; +pub(crate) fn compact_sql(sql: &str, policy: QuotePolicy) -> String { let bytes = sql.as_bytes(); let mut compact = Vec::with_capacity(bytes.len()); let mut index = 0; @@ -50,7 +50,7 @@ pub(crate) fn is_authored_decision(decision: AuthoredDecision<'_>) -> bool { !CANONICAL_NUMERIC_DECISIONS.contains(&decision.literal) && decision .authored_compact_sql - .contains(&compact_sql(decision.comparison_sql)) + .contains(&compact_sql(decision.comparison_sql, decision.quote_policy)) && !decision.output_name.is_some_and(|name| { decision.equality && name.ends_with(&format!("_{}", decision.literal)) }) diff --git a/crates/sqlbuild-rules-native/src/sql_scan/models.rs b/crates/sqlbuild-rules-native/src/sql_scan/models.rs index 16ee59e865..74e0f5a908 100644 --- a/crates/sqlbuild-rules-native/src/sql_scan/models.rs +++ b/crates/sqlbuild-rules-native/src/sql_scan/models.rs @@ -34,13 +34,6 @@ impl QuotePolicy { double_quote_backslash_escapes: true, }; - /// Rules policy: single and double quotes only, with backslashes treated as literal text. - pub(crate) const RULES: Self = Self { - backtick_identifiers: false, - single_quote_backslash_escapes: false, - double_quote_backslash_escapes: false, - }; - /// Return whether `byte` opens quoted text under this policy. pub(crate) fn is_quote(self, byte: u8) -> bool { match byte { diff --git a/crates/sqlbuild-rules-native/src/sql_scan/tests/test_sql_scan.rs b/crates/sqlbuild-rules-native/src/sql_scan/tests/test_sql_scan.rs index 41f30993ba..99da540962 100644 --- a/crates/sqlbuild-rules-native/src/sql_scan/tests/test_sql_scan.rs +++ b/crates/sqlbuild-rules-native/src/sql_scan/tests/test_sql_scan.rs @@ -4,6 +4,12 @@ use crate::sql_scan::models::QuotePolicy; use crate::sql_scan::models::Unclosed; use crate::sql_scan::tests::test_types::{MatchingParenPolicyTestCase, NonCodeEndTestCase}; +const QUOTES_ONLY: QuotePolicy = QuotePolicy { + backtick_identifiers: false, + single_quote_backslash_escapes: false, + double_quote_backslash_escapes: false, +}; + const BACKSLASH_ONLY: QuotePolicy = QuotePolicy { backtick_identifiers: false, single_quote_backslash_escapes: true, @@ -20,9 +26,9 @@ fn given_quote_policies_when_matching_parentheses_then_policy_decides_quote_boun expected_close: Ok(8), }, MatchingParenPolicyTestCase { - description: "rules policy treats backticks as code", + description: "quotes-only policy treats backticks as code", sql: "(`a)b` x) y", - policy: QuotePolicy::RULES, + policy: QUOTES_ONLY, expected_close: Ok(3), }, MatchingParenPolicyTestCase { @@ -46,7 +52,7 @@ fn given_quote_policies_when_matching_parentheses_then_policy_decides_quote_boun MatchingParenPolicyTestCase { description: "doubled quotes are escapes under every policy", sql: "('it''s )' x) y", - policy: QuotePolicy::RULES, + policy: QUOTES_ONLY, expected_close: Ok(12), }, MatchingParenPolicyTestCase { @@ -76,7 +82,7 @@ fn given_quote_policies_when_matching_parentheses_then_policy_decides_quote_boun MatchingParenPolicyTestCase { description: "unterminated block comment reports the comment", sql: "(/* abc x) y", - policy: QuotePolicy::RULES, + policy: QUOTES_ONLY, expected_close: Err(Unclosed::BlockComment), }, MatchingParenPolicyTestCase { @@ -115,7 +121,7 @@ fn given_comment_or_quote_starts_when_finding_non_code_end_then_returns_end_offs NonCodeEndTestCase { description: "block comment ends after its terminator", sql: "/* a */ b", - policy: QuotePolicy::RULES, + policy: QUOTES_ONLY, expected_end: Ok(Some(7)), }, NonCodeEndTestCase { From e90c4cc7fef3518bbef943f0cea90000a340a9c1 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:05:12 +0100 Subject: [PATCH 34/57] test(sql-tests): characterise the expected-projection top-level scanners (cherry picked from commit e86b30ab70ec613d0a1fc42b66a1d03b37ed99da) --- .../compiler/compile/_helpers/_test_types.py | 11 +++ .../test_expected_projection_scanning.py | 85 +++++++++++++++++++ 2 files changed, 96 insertions(+) create mode 100644 tests/unit/src/sqlbuild/compiler/compile/_helpers/test_expected_projection_scanning.py diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py index 54aa5600f8..2981975de0 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py @@ -899,3 +899,14 @@ class CteScannerMessageTestCase: description: str sql: str expected_message: str + + +@dataclass(frozen=True) +class ExpectedProjectionScanTestCase: + description: str + sql: str + expected_branches: tuple[str, ...] + expected_select_list_end: int + expected_commas: tuple[str, ...] + expected_alias: str | None + expected_contains_select_star: bool diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_expected_projection_scanning.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_expected_projection_scanning.py new file mode 100644 index 0000000000..72ff6819e9 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_expected_projection_scanning.py @@ -0,0 +1,85 @@ +"""Characterisation tests for the top-level scanners behind SQL-test expected projections.""" + +from __future__ import annotations + +import pytest + +from sqlbuild.compiler.compile._helpers.sql_tests import core as sql_test_core +from tests.unit.src.sqlbuild.compiler.compile._helpers._test_types import ( + ExpectedProjectionScanTestCase, +) + + +@pytest.mark.parametrize( + "test_case", + ( + ExpectedProjectionScanTestCase( + description="plain select splits on top-level commas", + sql="SELECT a, b AS c FROM orders", + expected_branches=("SELECT a, b AS c FROM orders",), + expected_select_list_end=len("SELECT a, b AS c "), + expected_commas=("SELECT a", "b AS c FROM orders"), + expected_alias=None, + expected_contains_select_star=False, + ), + ExpectedProjectionScanTestCase( + description="quotes and comments hide keywords, commas and parentheses", + sql="SELECT 'x, FROM (' AS a, /* UNION , */ \"b)\" -- FROM\nFROM t", + expected_branches=("SELECT 'x, FROM (' AS a, /* UNION , */ \"b)\" -- FROM\nFROM t",), + expected_select_list_end=len("SELECT 'x, FROM (' AS a, /* UNION , */ \"b)\" -- FROM\n"), + expected_commas=("SELECT 'x, FROM (' AS a", '/* UNION , */ "b)" -- FROM\nFROM t'), + expected_alias=None, + expected_contains_select_star=False, + ), + ExpectedProjectionScanTestCase( + description="nested parentheses hide commas, FROM and AS", + sql="SELECT f(a, (SELECT 1 AS x FROM y)) AS total", + expected_branches=("SELECT f(a, (SELECT 1 AS x FROM y)) AS total",), + expected_select_list_end=len("SELECT f(a, (SELECT 1 AS x FROM y)) AS total"), + expected_commas=("SELECT f(a, (SELECT 1 AS x FROM y)) AS total",), + expected_alias="total", + expected_contains_select_star=False, + ), + ExpectedProjectionScanTestCase( + description="union all splits branches", + sql="SELECT 1 AS a UNION ALL SELECT 2 AS a", + expected_branches=("SELECT 1 AS a", "SELECT 2 AS a"), + expected_select_list_end=len("SELECT 1 AS a UNION ALL SELECT 2 AS a"), + expected_commas=("SELECT 1 AS a UNION ALL SELECT 2 AS a",), + expected_alias="a", + expected_contains_select_star=False, + ), + ExpectedProjectionScanTestCase( + description="backtick identifiers are quoted", + sql="SELECT `a, b` AS `c`, * FROM t", + expected_branches=("SELECT `a, b` AS `c`, * FROM t",), + expected_select_list_end=len("SELECT `a, b` AS `c`, * "), + expected_commas=("SELECT `a, b` AS `c`", "* FROM t"), + expected_alias=None, + expected_contains_select_star=False, + ), + ExpectedProjectionScanTestCase( + description="select star inside a subquery is detected", + sql="SELECT a FROM (SELECT * FROM t)", + expected_branches=("SELECT a FROM (SELECT * FROM t)",), + expected_select_list_end=len("SELECT a "), + expected_commas=("SELECT a FROM (SELECT * FROM t)",), + expected_alias=None, + expected_contains_select_star=True, + ), + ), + ids=lambda case: case.description, +) +def test_given_expected_cte_sql_when_scanning_top_level_then_results_are_characterised( + test_case: ExpectedProjectionScanTestCase, +) -> None: + assert sql_test_core._split_set_operation_branches(test_case.sql) == test_case.expected_branches + assert ( + sql_test_core._find_select_list_end(sql=test_case.sql, start=0) + == test_case.expected_select_list_end + ) + assert sql_test_core._split_top_level_commas(test_case.sql) == test_case.expected_commas + assert sql_test_core._extract_as_alias(test_case.sql) == test_case.expected_alias + assert sql_test_core._contains_select_star(test_case.sql) is ( + test_case.expected_contains_select_star + ) From 3db9ab5e9179d0551db9290a59c2a5bbb28120a5 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:06:57 +0100 Subject: [PATCH 35/57] refactor(sql-tests): scan expected projections through one code-position iterator The SQL-test set-operation splitter, select-list end finder, top-level comma splitter, AS-alias extractor and SELECT * detector each hand-rolled the same comment-, quote- and depth-aware loop over sql_analysis primitives. They now consume one sql_analysis iter_code_positions entrypoint. Results are unchanged and pinned by characterisation tests. (cherry picked from commit 9a0ab5da2c35adcbf73ce48020e52d03cd0bdf7b) --- .../compile/_helpers/sql_tests/core.py | 167 +++++------------- .../sql_analysis/_helpers/scanning.py | 26 +++ .../sql_analysis/main/_iter_code_positions.py | 11 ++ .../sql_analysis/_helpers/_test_types.py | 9 + .../sql_analysis/_helpers/test_scanning.py | 33 +++- 5 files changed, 125 insertions(+), 121 deletions(-) create mode 100644 src/sqlbuild/compiler/sql_analysis/main/_iter_code_positions.py diff --git a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py index f4b92bb6bc..9a8dba2e10 100644 --- a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py +++ b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py @@ -30,9 +30,7 @@ SOURCE_TEST_CTE_PREFIX, SQL_ARGUMENT_SEPARATOR_TOKEN, SQL_CEREMONIAL_SELECT_VALUE, - SQL_CLOSE_PAREN_TOKEN, SQL_OPEN_PAREN_TOKEN, - SQL_QUOTE_TOKENS, SQL_SINGLE_QUOTE_TOKEN, SQL_STATEMENT_TERMINATOR_TOKEN, SQL_WILDCARD_TOKEN, @@ -59,6 +57,7 @@ is_identifier_character, ) from sqlbuild.compiler.sql_analysis.main._is_identifier_start import is_identifier_start +from sqlbuild.compiler.sql_analysis.main._iter_code_positions import iter_code_positions from sqlbuild.compiler.sql_analysis.main._skip_block_comment import skip_block_comment from sqlbuild.compiler.sql_analysis.main._skip_line_comment import skip_line_comment from sqlbuild.compiler.sql_analysis.main._skip_quoted_text import ( @@ -708,38 +707,23 @@ def _extract_expected_branch_column_names( def _split_set_operation_branches(sql: str) -> tuple[str, ...]: branches: list[str] = [] branch_start: int = 0 - index: int = 0 - depth: int = 0 - while index < len(sql): - if sql.startswith("--", index): - index = skip_line_comment(sql=sql, start=index) - continue - if sql.startswith("/*", index): - index = skip_block_comment(sql=sql, start=index, context=_CONTEXT) - continue - if sql[index] in SQL_QUOTE_TOKENS: - index = skip_quoted_text(sql=sql, start=index, context=_CONTEXT) - continue - if sql[index] == SQL_OPEN_PAREN_TOKEN: - depth += 1 - index += 1 - continue - if sql[index] == SQL_CLOSE_PAREN_TOKEN: - depth -= 1 - index += 1 + resume: int = 0 + index: int + depth: int + for index, depth in iter_code_positions(sql=sql, context=_CONTEXT): + if index < resume or depth != 0: continue union_end: int | None = _try_consume_keyword(sql=sql, start=index, keyword="UNION") - if depth == 0 and union_end is not None: - branch_sql: str = sql[branch_start:index].strip() - if branch_sql: - branches.append(branch_sql) - index = _skip_ignorable(sql=sql, start=union_end) - all_end: int | None = _try_consume_keyword(sql=sql, start=index, keyword="ALL") - if all_end is not None: - index = _skip_ignorable(sql=sql, start=all_end) - branch_start = index + if union_end is None: continue - index += 1 + branch_sql: str = sql[branch_start:index].strip() + if branch_sql: + branches.append(branch_sql) + resume = _skip_ignorable(sql=sql, start=union_end) + all_end: int | None = _try_consume_keyword(sql=sql, start=resume, keyword="ALL") + if all_end is not None: + resume = _skip_ignorable(sql=sql, start=all_end) + branch_start = resume final_branch_sql: str = sql[branch_start:].strip() if final_branch_sql: @@ -769,57 +753,28 @@ def _extract_expected_select_column_names(*, branch_sql: str, file_label: str) - def _find_select_list_end(*, sql: str, start: int) -> int: - index: int = start - depth: int = 0 - while index < len(sql): - if sql.startswith("--", index): - index = skip_line_comment(sql=sql, start=index) - continue - if sql.startswith("/*", index): - index = skip_block_comment(sql=sql, start=index, context=_CONTEXT) - continue - if sql[index] in SQL_QUOTE_TOKENS: - index = skip_quoted_text(sql=sql, start=index, context=_CONTEXT) - continue - if sql[index] == SQL_OPEN_PAREN_TOKEN: - depth += 1 - index += 1 - continue - if sql[index] == SQL_CLOSE_PAREN_TOKEN: - depth -= 1 - index += 1 - continue - if depth == 0 and _try_consume_keyword(sql=sql, start=index, keyword="FROM") is not None: - return index - index += 1 + index: int + depth: int + for index, depth in iter_code_positions(sql=sql[start:], context=_CONTEXT): + if ( + depth == 0 + and _try_consume_keyword(sql=sql, start=start + index, keyword="FROM") is not None + ): + return start + index return len(sql) def _split_top_level_commas(raw_value: str) -> tuple[str, ...]: values: list[str] = [] value_start: int = 0 - index: int = 0 - depth: int = 0 - while index < len(raw_value): - if raw_value.startswith("--", index): - index = skip_line_comment(sql=raw_value, start=index) - continue - if raw_value.startswith("/*", index): - index = skip_block_comment(sql=raw_value, start=index, context=_CONTEXT) - continue - if raw_value[index] in SQL_QUOTE_TOKENS: - index = skip_quoted_text(sql=raw_value, start=index, context=_CONTEXT) - continue - if raw_value[index] == SQL_OPEN_PAREN_TOKEN: - depth += 1 - elif raw_value[index] == SQL_CLOSE_PAREN_TOKEN: - depth -= 1 - elif raw_value[index] == SQL_ARGUMENT_SEPARATOR_TOKEN and depth == 0: + index: int + depth: int + for index, depth in iter_code_positions(sql=raw_value, context=_CONTEXT): + if depth == 0 and raw_value[index] == SQL_ARGUMENT_SEPARATOR_TOKEN: item: str = raw_value[value_start:index].strip() if item: values.append(item) value_start = index + 1 - index += 1 final_item: str = raw_value[value_start:].strip() if final_item: @@ -840,41 +795,24 @@ def _extract_expected_projection_name(*, expression: str, file_label: str) -> st def _extract_as_alias(expression: str) -> str | None: - index: int = 0 - depth: int = 0 last_alias_name: str | None = None - while index < len(expression): - if expression.startswith("--", index): - index = skip_line_comment(sql=expression, start=index) - continue - if expression.startswith("/*", index): - index = skip_block_comment(sql=expression, start=index, context=_CONTEXT) - continue - if expression[index] in SQL_QUOTE_TOKENS: - index = skip_quoted_text(sql=expression, start=index, context=_CONTEXT) - continue - if expression[index] == SQL_OPEN_PAREN_TOKEN: - depth += 1 - index += 1 - continue - if expression[index] == SQL_CLOSE_PAREN_TOKEN: - depth -= 1 - index += 1 + index: int + depth: int + for index, depth in iter_code_positions(sql=expression, context=_CONTEXT): + if depth != 0: continue as_end: int | None = _try_consume_keyword(sql=expression, start=index, keyword="AS") - if depth == 0 and as_end is not None: - alias_index: int = _skip_ignorable(sql=expression, start=as_end) - if alias_index < len(expression) and is_identifier_start(expression[alias_index]): - alias_name, alias_end = _read_identifier( - sql=expression, - start=alias_index, - file_label="projection", - ) - if not expression[alias_end:].strip(): - last_alias_name = alias_name - index = as_end + if as_end is None: continue - index += 1 + alias_index: int = _skip_ignorable(sql=expression, start=as_end) + if alias_index < len(expression) and is_identifier_start(expression[alias_index]): + alias_name, alias_end = _read_identifier( + sql=expression, + start=alias_index, + file_label="projection", + ) + if not expression[alias_end:].strip(): + last_alias_name = alias_name return last_alias_name @@ -885,25 +823,14 @@ def _is_simple_identifier(value: str) -> bool: def _contains_select_star(sql: str) -> bool: - index: int = 0 - while index < len(sql): - if sql.startswith("--", index): - index = skip_line_comment(sql=sql, start=index) - continue - if sql.startswith("/*", index): - index = skip_block_comment(sql=sql, start=index, context=_CONTEXT) - continue - if sql[index] in SQL_QUOTE_TOKENS: - index = skip_quoted_text(sql=sql, start=index, context=_CONTEXT) - continue + index: int + for index, _depth in iter_code_positions(sql=sql, context=_CONTEXT): select_end: int | None = _try_consume_keyword(sql=sql, start=index, keyword="SELECT") - if select_end is not None: - value_index: int = _skip_ignorable(sql=sql, start=select_end) - if value_index < len(sql) and sql[value_index] == SQL_WILDCARD_TOKEN: - return True - index = select_end + if select_end is None: continue - index += 1 + value_index: int = _skip_ignorable(sql=sql, start=select_end) + if value_index < len(sql) and sql[value_index] == SQL_WILDCARD_TOKEN: + return True return False diff --git a/src/sqlbuild/compiler/sql_analysis/_helpers/scanning.py b/src/sqlbuild/compiler/sql_analysis/_helpers/scanning.py index f6253409c0..bfd1320426 100644 --- a/src/sqlbuild/compiler/sql_analysis/_helpers/scanning.py +++ b/src/sqlbuild/compiler/sql_analysis/_helpers/scanning.py @@ -3,6 +3,7 @@ from __future__ import annotations import re +from collections.abc import Iterator from sqlbuild.compiler.compile.exceptions import CompileInputError from sqlbuild.compiler.sql_analysis.constants import ( @@ -80,6 +81,31 @@ def find_matching_paren_impl(*, sql: str, open_paren_index: int, context: str = raise CompileInputError(f"{context} contains an unclosed parenthesis") +def iter_code_positions_impl(*, sql: str, context: str = "SQL") -> Iterator[tuple[int, int]]: + """Yield code offsets outside comments, quotes, and parentheses with their nesting depth.""" + + depth: int = 0 + index: int = 0 + while index < len(sql): + if sql.startswith("--", index): + index = skip_line_comment_impl(sql=sql, start=index) + continue + if sql.startswith("/*", index): + index = skip_block_comment_impl(sql=sql, start=index, context=context) + continue + character: str = sql[index] + if character in SQL_QUOTE_CHARACTERS: + index = skip_quoted_text_impl(sql=sql, start=index, context=context) + continue + if character == SQL_OPEN_PARENTHESIS: + depth += 1 + elif character == SQL_CLOSE_PARENTHESIS: + depth -= 1 + else: + yield index, depth + index += 1 + + def is_identifier_start_impl(character: str) -> bool: """Return whether a character can begin an SQL identifier.""" diff --git a/src/sqlbuild/compiler/sql_analysis/main/_iter_code_positions.py b/src/sqlbuild/compiler/sql_analysis/main/_iter_code_positions.py new file mode 100644 index 0000000000..b2a491333d --- /dev/null +++ b/src/sqlbuild/compiler/sql_analysis/main/_iter_code_positions.py @@ -0,0 +1,11 @@ +"""Top-level SQL code position scanning entrypoint.""" + +from collections.abc import Iterator + +from sqlbuild.compiler.sql_analysis._helpers.scanning import iter_code_positions_impl + + +def iter_code_positions(*, sql: str, context: str = "SQL") -> Iterator[tuple[int, int]]: + """Yield each code character offset outside comments and quotes with its parenthesis depth.""" + + return iter_code_positions_impl(sql=sql, context=context) diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/_test_types.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/_test_types.py index fcc46e88d0..4fa311faf2 100644 --- a/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/_test_types.py @@ -20,3 +20,12 @@ class SkipQuotedTextErrorTestCase: sql: str context: str expected_error: str + + +@dataclass(frozen=True) +class IterCodePositionsTestCase: + """Code-position scan case.""" + + description: str + sql: str + expected_positions: tuple[tuple[int, int], ...] diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/test_scanning.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/test_scanning.py index cdd6c6c355..3245e9b7aa 100644 --- a/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/test_scanning.py +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/test_scanning.py @@ -3,8 +3,12 @@ import pytest from sqlbuild.compiler.compile.exceptions import CompileInputError -from sqlbuild.compiler.sql_analysis._helpers.scanning import skip_quoted_text_impl +from sqlbuild.compiler.sql_analysis._helpers.scanning import ( + iter_code_positions_impl, + skip_quoted_text_impl, +) from tests.unit.src.sqlbuild.compiler.sql_analysis._helpers._test_types import ( + IterCodePositionsTestCase, SkipQuotedTextErrorTestCase, SkipQuotedTextSuccessTestCase, ) @@ -68,3 +72,30 @@ def test_given_unclosed_quote_when_skipping_then_raises_contextual_error( ) -> None: with pytest.raises(CompileInputError, match=test_case.expected_error): skip_quoted_text_impl(sql=test_case.sql, start=0, context=test_case.context) + + +@pytest.mark.parametrize( + "test_case", + ( + IterCodePositionsTestCase( + description="parentheses change depth and are not yielded", + sql="a(b)c", + expected_positions=((0, 0), (2, 1), (4, 0)), + ), + IterCodePositionsTestCase( + description="quotes and comments are skipped", + sql="a'(b'`c`--d\n/*e*/f", + expected_positions=((0, 0), (17, 0)), + ), + IterCodePositionsTestCase( + description="unbalanced close parenthesis goes below zero depth", + sql=")a", + expected_positions=((1, -1),), + ), + ), + ids=lambda case: case.description, +) +def test_given_sql_when_iterating_code_positions_then_yields_code_offsets_with_depth( + test_case: IterCodePositionsTestCase, +) -> None: + assert tuple(iter_code_positions_impl(sql=test_case.sql)) == test_case.expected_positions From a6294e5b17966ebfc54f5872e31b50db321831a4 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:10:42 +0100 Subject: [PATCH 36/57] fix(sql-tests): split UNION DISTINCT branches in expected CTEs The native and Python expected-CTE branch splitters only consumed an optional ALL after UNION, so an __expected__ CTE using UNION DISTINCT left "DISTINCT SELECT ..." as the next branch and sqb test failed with a false "must define each __expected__ set-operation branch as a SELECT query" error. Both splitters now accept ALL or DISTINCT. An end-to-end sqb test case covers a passing UNION DISTINCT expectation. (cherry picked from commit 08014ba6507b263c1cf0a197366234f10ed202a2) --- .../compiler/_helpers/sql_tests/extraction.rs | 6 ++- .../compile/_helpers/sql_tests/core.py | 8 ++-- .../cli/commands/main/test/helpers.py | 37 +++++++++++++++++++ .../cli/commands/main/test/test_test.py | 6 +++ .../test_expected_projection_scanning.py | 9 +++++ 5 files changed, 61 insertions(+), 5 deletions(-) diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs index 1bb1df6058..eef954157e 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs @@ -866,8 +866,10 @@ fn split_unions(sql: &str) -> Result, String> { values.push(value); } index = skip_ignorable(sql, end)?; - if let Some(all_end) = consume_keyword(sql, index, "ALL") { - index = skip_ignorable(sql, all_end)?; + if let Some(quantifier_end) = consume_keyword(sql, index, "ALL") + .or_else(|| consume_keyword(sql, index, "DISTINCT")) + { + index = skip_ignorable(sql, quantifier_end)?; } start = index; continue; diff --git a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py index 9a8dba2e10..2bbbfd2c15 100644 --- a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py +++ b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py @@ -720,9 +720,11 @@ def _split_set_operation_branches(sql: str) -> tuple[str, ...]: if branch_sql: branches.append(branch_sql) resume = _skip_ignorable(sql=sql, start=union_end) - all_end: int | None = _try_consume_keyword(sql=sql, start=resume, keyword="ALL") - if all_end is not None: - resume = _skip_ignorable(sql=sql, start=all_end) + quantifier_end: int | None = _try_consume_keyword( + sql=sql, start=resume, keyword="ALL" + ) or _try_consume_keyword(sql=sql, start=resume, keyword="DISTINCT") + if quantifier_end is not None: + resume = _skip_ignorable(sql=sql, start=quantifier_end) branch_start = resume final_branch_sql: str = sql[branch_start:].strip() diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py b/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py index e196fd3377..45639506f2 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py @@ -870,6 +870,43 @@ def build_star_mock_fixture_project_files() -> dict[str, str]: return files +def build_union_distinct_expected_project_files() -> dict[str, str]: + """Build a SQL test whose expected rows combine branches with UNION DISTINCT.""" + + return { + "sqlbuild_project.toml": ( + 'name = "union_distinct_expected"\n' + 'adapter = "duckdb"\n\n' + "[connection]\n" + 'database = "union_distinct_expected.duckdb"\n' + ), + "models/orders.sql": ( + 'MODEL ();\n\nSELECT DISTINCT order_id FROM __source("raw_orders")\n' + ), + "sources/raw_orders.yml": ( + "sources:\n" + " - name: raw_orders\n" + " schema: main\n" + " table: raw_orders\n" + " columns:\n" + " - name: order_id\n type: INTEGER\n" + ), + "tests/unit/test_orders.sql": ( + "TEST();\n\n" + "WITH\n" + "__source__raw_orders AS (\n" + " SELECT 1 AS order_id UNION ALL SELECT 1 AS order_id UNION ALL SELECT 2 AS order_id\n" + "),\n" + "__expected__orders AS (\n" + " SELECT 1 AS order_id\n" + " UNION DISTINCT\n" + " SELECT 2 AS order_id\n" + ")\n" + "SELECT 1\n" + ), + } + + def build_transformed_collection_project_files() -> dict[str, str]: """Build a valid aggregation that changes a scalar input into a collection output.""" diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/test/test_test.py b/tests/e2e/src/sqlbuild/cli/commands/main/test/test_test.py index 2ca60ccc35..3c4320c7fb 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/test/test_test.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/test/test_test.py @@ -62,6 +62,7 @@ build_star_partial_fixture_project_files, build_table_function_fixture_project_files, build_transformed_collection_project_files, + build_union_distinct_expected_project_files, build_unsatisfied_leaf_test_project_files, build_unspecified_nullability_fixture_project_files, build_untyped_null_fixture_project_files, @@ -316,6 +317,11 @@ def test_given_inspection_when_requesting_json_file_then_parser_rejects_conflict repo_files=build_star_mock_fixture_project_files(), expected_stdout_fragment="PASS=1", ), + FixtureCompatibilityE2ETestCase( + description="union distinct expected branches are validated per branch", + repo_files=build_union_distinct_expected_project_files(), + expected_stdout_fragment="PASS=1", + ), FixtureCompatibilityE2ETestCase( description="aggregation output does not inherit scalar input type", repo_files=build_transformed_collection_project_files(), diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_expected_projection_scanning.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_expected_projection_scanning.py index 72ff6819e9..af71f97607 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_expected_projection_scanning.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_expected_projection_scanning.py @@ -49,6 +49,15 @@ expected_alias="a", expected_contains_select_star=False, ), + ExpectedProjectionScanTestCase( + description="union distinct splits branches", + sql="SELECT 1 AS a UNION DISTINCT SELECT 2 AS a", + expected_branches=("SELECT 1 AS a", "SELECT 2 AS a"), + expected_select_list_end=len("SELECT 1 AS a UNION DISTINCT SELECT 2 AS a"), + expected_commas=("SELECT 1 AS a UNION DISTINCT SELECT 2 AS a",), + expected_alias="a", + expected_contains_select_star=False, + ), ExpectedProjectionScanTestCase( description="backtick identifiers are quoted", sql="SELECT `a, b` AS `c`, * FROM t", From 9b684af8fc8142ab9aa7514d8678864d698a3fbc Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:12:26 +0100 Subject: [PATCH 37/57] fix(sql-tests): name the direct-logic expected CTE in projection errors Expected-projection validation hard-coded __expected__ in its alias, empty-projection, set-operation branch and branch-mismatch errors, so macro, UDF and table_fn tests were told to fix an __expected__ CTE they do not have. The native extractor, the Python classifier and the Polyglot path now report the validated CTE (for example __macro_expected__); model-mode wording is unchanged. (cherry picked from commit 88c089023175eea7f3fa4e35e0f0c4344b637c45) --- .../compiler/_helpers/sql_tests/extraction.rs | 18 ++++++------ .../src/compiler/tests/helpers.rs | 22 +++++++++++++++ .../tests/test_sql_test_extraction.rs | 6 ++++ .../compile/_helpers/analysis/tests.py | 28 ++++++++++--------- .../compile/_helpers/sql_tests/core.py | 25 ++++++++++------- ...st_direct_logic_sql_test_classification.py | 23 +++++++++++++++ 6 files changed, 91 insertions(+), 31 deletions(-) diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs index eef954157e..ebcdd035c2 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs @@ -434,13 +434,13 @@ fn validate_expected( let branches = split_unions(&cte.1)?; let mut names: Vec> = Vec::new(); for branch in branches { - names.push(projection_names(branch, file)?); + names.push(projection_names(branch, file, label)?); } if let Some(first) = names.first() { for (index, branch) in names.iter().enumerate().skip(1) { if branch != first { return Err(format!( - "SQL test '{file}' must use the same __expected__ projection names and order in every set-operation branch; branch {} does not match branch 1", + "SQL test '{file}' must use the same {label} projection names and order in every set-operation branch; branch {} does not match branch 1", index + 1 )); } @@ -481,9 +481,11 @@ pub(crate) fn empty_fixture_marker_matches(sql: &str) -> Result { Ok(index == sql.len()) } -fn projection_names(branch: &str, file: &str) -> Result, String> { +fn projection_names(branch: &str, file: &str, label: &str) -> Result, String> { let start = skip_ignorable(branch, 0)?; - let select_end = consume_keyword(branch, start, "SELECT").ok_or_else(|| format!("SQL test '{file}' must define each __expected__ set-operation branch as a SELECT query"))?; + let select_end = consume_keyword(branch, start, "SELECT").ok_or_else(|| { + format!("SQL test '{file}' must define each {label} set-operation branch as a SELECT query") + })?; let mut end = branch.len(); for keyword in [ "FROM", "WHERE", "GROUP", "HAVING", "QUALIFY", "WINDOW", "ORDER", "LIMIT", "OFFSET", @@ -496,16 +498,16 @@ fn projection_names(branch: &str, file: &str) -> Result, String> { let expressions = split_top_level(&branch[select_end..end], b',')?; if expressions.is_empty() { return Err(format!( - "SQL test '{file}' must project at least one column in __expected__" + "SQL test '{file}' must project at least one column in {label}" )); } expressions .into_iter() - .map(|expression| projection_name(expression, file)) + .map(|expression| projection_name(expression, file, label)) .collect() } -fn projection_name(expression: &str, file: &str) -> Result { +fn projection_name(expression: &str, file: &str, label: &str) -> Result { if let Some(position) = find_last_top_level_keyword(expression, "AS")? { let alias_start = skip_ignorable(expression, position + 2)?; if let Some((alias, end)) = read_identifier(expression, alias_start) @@ -525,7 +527,7 @@ fn projection_name(expression: &str, file: &str) -> Result { return Ok(alias); } Err(format!( - "SQL test '{file}' must alias every non-trivial __expected__ projection" + "SQL test '{file}' must alias every non-trivial {label} projection" )) } diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs b/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs index ab82c272bd..422ac14668 100644 --- a/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs +++ b/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs @@ -200,6 +200,28 @@ pub(crate) fn empty_model_fixture_marker_preserves_direct_mode_validation() -> b true } +pub(crate) fn expected_projection_errors_name_the_expected_cte() -> bool { + let cases = [ + ( + r#"{"tests":[{"sql":"WITH __macro_actual__ AS (SELECT @status() AS status), __macro_expected__ AS (SELECT 1 + 1) SELECT 1","fileLabel":"tests/status.sql","mode":"macro"}]}"#, + "SQL test 'tests/status.sql' must alias every non-trivial __macro_expected__ projection", + ), + ( + r#"{"tests":[{"sql":"WITH __udf_actual__ AS (SELECT 1 AS value), __udf_expected__ AS (SELECT 1 AS value UNION ALL SELECT 2 AS other) SELECT 1","fileLabel":"tests/function.sql","mode":"udf"}]}"#, + "SQL test 'tests/function.sql' must use the same __udf_expected__ projection names and order in every set-operation branch; branch 2 does not match branch 1", + ), + ( + r#"{"tests":[{"sql":"WITH __source__raw_orders AS (SELECT 1 AS id), __expected__orders AS (SELECT 1 + 1) SELECT 1","fileLabel":"tests/orders.sql","mode":"model"}]}"#, + "SQL test 'tests/orders.sql' must alias every non-trivial __expected__ projection", + ), + ]; + for (request, expected_error) in cases { + let error = extract_batch_json(request).expect_err("expected projection is rejected"); + assert_eq!(error, expected_error); + } + true +} + pub(crate) fn concurrent_requests_initialize_shared_template_once() -> bool { let templates = Arc::new(Mutex::new(HashMap::new())); let starts = Arc::new(Barrier::new(4)); diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_extraction.rs b/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_extraction.rs index 6c3e4ea10d..671caec8be 100644 --- a/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_extraction.rs +++ b/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_extraction.rs @@ -1,6 +1,7 @@ use crate::compiler::tests::helpers::{ dependent_assertion_returns_authoritative_error, empty_model_fixture_marker_preserves_direct_mode_validation, + expected_projection_errors_name_the_expected_cte, mixed_expanded_tests_preserve_order_and_payloads, quoted_ctes_and_implicit_alias_preserve_payload, }; @@ -29,6 +30,11 @@ fn given_sql_test_cases_when_extracting_native_payloads_then_expected_behavior_h run: empty_model_fixture_marker_preserves_direct_mode_validation, expected_success: true, }, + SqlTestExtractionTestCase { + description: "expected projection errors name the expected CTE", + run: expected_projection_errors_name_the_expected_cte, + expected_success: true, + }, ]; for test_case in test_cases { diff --git a/src/sqlbuild/compiler/compile/_helpers/analysis/tests.py b/src/sqlbuild/compiler/compile/_helpers/analysis/tests.py index a0e46e240b..43294b01a3 100644 --- a/src/sqlbuild/compiler/compile/_helpers/analysis/tests.py +++ b/src/sqlbuild/compiler/compile/_helpers/analysis/tests.py @@ -25,7 +25,7 @@ def extract_expected_branch_column_names_with_sql_analysis( - *, sql: str, file_label: str + *, sql: str, file_label: str, label: str = "__expected__" ) -> tuple[tuple[str, ...], ...] | None: """Return expected SELECT branch names using required Polyglot analysis.""" @@ -40,32 +40,34 @@ def extract_expected_branch_column_names_with_sql_analysis( _extract_branch_names( expression=polyglot_module.parse_one(branch_sql, dialect="generic"), file_label=file_label, + label=label, )[0] for branch_sql in branches ) return None - return _extract_branch_names(expression=parsed_expression, file_label=file_label) + return _extract_branch_names(expression=parsed_expression, file_label=file_label, label=label) -def _extract_branch_names(*, expression: Any, file_label: str) -> tuple[tuple[str, ...], ...]: +def _extract_branch_names( + *, expression: Any, file_label: str, label: str +) -> tuple[tuple[str, ...], ...]: expression = _unwrap_expression(expression=expression) if expression.__class__.__name__ == POLYGLOT_UNION_EXPRESSION_NAME: left_expression: Any = expression.args["left"] right_expression: Any = expression.args["right"] return ( - *_extract_branch_names(expression=left_expression, file_label=file_label), - *_extract_branch_names(expression=right_expression, file_label=file_label), + *_extract_branch_names(expression=left_expression, file_label=file_label, label=label), + *_extract_branch_names(expression=right_expression, file_label=file_label, label=label), ) if expression.__class__.__name__ == POLYGLOT_SELECT_EXPRESSION_NAME: if _is_synthetic_values_set_branch(expression): raise CompileInputError( - f"SQL test '{file_label}' must define each __expected__ set-operation " + f"SQL test '{file_label}' must define each {label} set-operation " "branch as a SELECT query" ) - return (_extract_select_names(expression=expression, file_label=file_label),) + return (_extract_select_names(expression=expression, file_label=file_label, label=label),) raise CompileInputError( - f"SQL test '{file_label}' must define each __expected__ set-operation " - "branch as a SELECT query" + f"SQL test '{file_label}' must define each {label} set-operation branch as a SELECT query" ) @@ -88,13 +90,13 @@ def _unwrap_expression(*, expression: Any) -> Any: return expression -def _extract_select_names(*, expression: Any, file_label: str) -> tuple[str, ...]: +def _extract_select_names(*, expression: Any, file_label: str, label: str) -> tuple[str, ...]: names: list[str] = [] projection: Any for projection in expression.expressions: if projection.is_star: raise CompileInputError( - f"SQL test '{file_label}' must not use SELECT * in __expected__ CTEs" + f"SQL test '{file_label}' must not use SELECT * in {label} CTEs" ) alias_name: str = str(projection.alias or "") if alias_name: @@ -104,11 +106,11 @@ def _extract_select_names(*, expression: Any, file_label: str) -> tuple[str, ... names.append(str(projection.name)) continue raise CompileInputError( - f"SQL test '{file_label}' must alias every non-trivial __expected__ projection" + f"SQL test '{file_label}' must alias every non-trivial {label} projection" ) if not names: raise CompileInputError( - f"SQL test '{file_label}' must project at least one column in __expected__" + f"SQL test '{file_label}' must project at least one column in {label}" ) return tuple(names) diff --git a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py index 2bbbfd2c15..f1ca638c9d 100644 --- a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py +++ b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py @@ -677,29 +677,32 @@ def _validate_expected_cte_query( branch_column_names: tuple[tuple[str, ...], ...] = _extract_expected_branch_column_names( sql=cte.sql_body, file_label=file_label, + label=label, ) first_branch_column_names: tuple[str, ...] = branch_column_names[0] branch_index: int for branch_index, column_names in enumerate(branch_column_names[1:], start=2): if column_names != first_branch_column_names: raise CompileInputError( - f"SQL test '{file_label}' must use the same __expected__ " + f"SQL test '{file_label}' must use the same {label} " f"projection names and order in every set-operation branch; branch {branch_index} " "does not match branch 1" ) def _extract_expected_branch_column_names( - *, sql: str, file_label: str + *, sql: str, file_label: str, label: str ) -> tuple[tuple[str, ...], ...]: sql_analysis_column_names: tuple[tuple[str, ...], ...] | None = ( - extract_expected_branch_column_names_with_sql_analysis(sql=sql, file_label=file_label) + extract_expected_branch_column_names_with_sql_analysis( + sql=sql, file_label=file_label, label=label + ) ) if sql_analysis_column_names is not None: return sql_analysis_column_names branches: tuple[str, ...] = _split_set_operation_branches(sql) return tuple( - _extract_expected_select_column_names(branch_sql=branch, file_label=file_label) + _extract_expected_select_column_names(branch_sql=branch, file_label=file_label, label=label) for branch in branches ) @@ -733,12 +736,14 @@ def _split_set_operation_branches(sql: str) -> tuple[str, ...]: return tuple(branches) -def _extract_expected_select_column_names(*, branch_sql: str, file_label: str) -> tuple[str, ...]: +def _extract_expected_select_column_names( + *, branch_sql: str, file_label: str, label: str +) -> tuple[str, ...]: index: int = _skip_ignorable(sql=branch_sql, start=0) select_end: int | None = _try_consume_keyword(sql=branch_sql, start=index, keyword="SELECT") if select_end is None: raise CompileInputError( - f"SQL test '{file_label}' must define each __expected__ set-operation " + f"SQL test '{file_label}' must define each {label} set-operation " "branch as a SELECT query" ) select_list_end: int = _find_select_list_end(sql=branch_sql, start=select_end) @@ -746,10 +751,10 @@ def _extract_expected_select_column_names(*, branch_sql: str, file_label: str) - expressions: tuple[str, ...] = _split_top_level_commas(raw_select_list) if not expressions: raise CompileInputError( - f"SQL test '{file_label}' must project at least one column in __expected__" + f"SQL test '{file_label}' must project at least one column in {label}" ) return tuple( - _extract_expected_projection_name(expression=expression, file_label=file_label) + _extract_expected_projection_name(expression=expression, file_label=file_label, label=label) for expression in expressions ) @@ -784,7 +789,7 @@ def _split_top_level_commas(raw_value: str) -> tuple[str, ...]: return tuple(values) -def _extract_expected_projection_name(*, expression: str, file_label: str) -> str: +def _extract_expected_projection_name(*, expression: str, file_label: str, label: str) -> str: alias_name: str | None = _extract_as_alias(expression) if alias_name is not None: return alias_name @@ -792,7 +797,7 @@ def _extract_expected_projection_name(*, expression: str, file_label: str) -> st if _is_simple_identifier(stripped_expression): return stripped_expression raise CompileInputError( - f"SQL test '{file_label}' must alias every non-trivial __expected__ projection" + f"SQL test '{file_label}' must alias every non-trivial {label} projection" ) diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_direct_logic_sql_test_classification.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_direct_logic_sql_test_classification.py index 93032dd4c4..78559b4c50 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_direct_logic_sql_test_classification.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_direct_logic_sql_test_classification.py @@ -101,6 +101,29 @@ f"SQL test '{_FILE}' mode 'macro' CTE __macro_expected__ must not call macros" ), ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro expected unaliased projection names the expected cte", + mode=SqlTestMode.MACRO, + ctes=( + ("__macro_actual__", "SELECT 1 AS status"), + ("__macro_expected__", "SELECT 1 + 1"), + ), + expected_message=( + f"SQL test '{_FILE}' must alias every non-trivial __macro_expected__ projection" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn expected branch mismatch names the expected cte", + mode=SqlTestMode.TABLE_FN, + ctes=( + ("__table_fn_actual__", "SELECT 1 AS value"), + ("__table_fn_expected__", "SELECT 1 AS value UNION ALL SELECT 2 AS other"), + ), + expected_message=( + f"SQL test '{_FILE}' must use the same __table_fn_expected__ projection names " + "and order in every set-operation branch; branch 2 does not match branch 1" + ), + ), ClassifyDirectLogicSqlTestCtesErrorTestCase( description="udf duplicate actual", mode=SqlTestMode.UDF, From 948a5671ddb073e872405c4c3ed06a3044a7f505 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:33:05 +0100 Subject: [PATCH 38/57] fix(snapshots): reopen reappearing hard-deleted keys on DuckDB, MotherDuck, Postgres and SQL Server Historical check snapshots with invalidate_hard_deletes now open a new active version when a hard-deleted key reappears in a later observation with unchanged check columns, matching the documented behaviour and the BigQuery, Databricks and Snowflake adapters. Previously these four adapters left the key without an active version until a check column changed. Behaviour change: existing history rows are not rewritten. A key that is currently closed by a hard delete and is present in a later observation gets a new active version, starting at that observation, on the next build. - Initial build: a row also starts a version when the key was absent from the immediately preceding observation group. - Incremental build: the first observation after a hard-deleted latest version is added to the new changes. The latest version counts as hard-deleted only when its key is absent from the observation at its valid_to, so a version closed by a change in the same run is not reopened. Rows are deduplicated by unique key and observed_at, with no QUALIFY or whole-row UNION DISTINCT, so the SQL works on Postgres and SQL Server. (cherry picked from commit c9ef3edd189ad8484e29ae03bb5145d136bbe2a8) --- .../adapter/contract/classes/base_adapter.py | 100 ++++++++- .../contract/classes/duckdb_backed_adapter.py | 39 +++- .../postgres/classes/postgres_adapter.py | 21 +- .../sqlserver/classes/sqlserver_adapter.py | 58 +++-- .../cli/commands/main/build/_test_types.py | 13 ++ .../main/build/test_snapshot_build.py | 202 ++++++++++++++++++ .../commands/main/postgres/test_postgres.py | 2 + .../commands/main/sqlserver/test_sqlserver.py | 2 + .../snapshot_sql_rendering/_test_types.py | 9 + .../snapshot_sql_rendering/test_rendering.py | 145 +++++++++++++ 10 files changed, 550 insertions(+), 41 deletions(-) diff --git a/src/sqlbuild/adapter/contract/classes/base_adapter.py b/src/sqlbuild/adapter/contract/classes/base_adapter.py index 174f4b76a0..eeedd63d45 100644 --- a/src/sqlbuild/adapter/contract/classes/base_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/base_adapter.py @@ -2624,13 +2624,21 @@ def _historical_check_snapshot_select_sql( observed_at_column=observed_at_column, row_alias="__changes", ) + group_sequence_sql: str = _historical_observed_group_sequence_cte_sql( + origin=origin, observed_at_column=observed_at_column + ) + grouped_origin_sql: str = _historical_observed_group_sequence_from_sql( + origin=origin, observed_at_column=observed_at_column + ) return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" + f"WITH {group_sequence_sql}, __ordered AS (" + "SELECT __source.*, __observed_group_sequence.__prev_group_observed_at, " + f"LAG({observed_at_column}) OVER (" f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" + f") AS __prev_observed_at{previous_columns_sql} {grouped_origin_sql}" "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" + f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition}) " + "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at" "), __versions AS (" f"SELECT __changes.*, LEAD({observed_at_column}) OVER (" f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" @@ -2873,6 +2881,18 @@ def _historical_check_new_changes_cte_sql( observed_at_column=observed_at_column, row_alias="__target", ) + reappearing_new_changes_sql: str = _historical_reappearing_new_changes_ctes_sql( + changed_or_new_sql=( + "SELECT __delta_changes.* FROM __delta_changes " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{first_key} IS NULL " + f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " + f"AND ({latest_change_condition}))" + ), + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_to_column=valid_to_column, + ) return ( "__ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -2884,13 +2904,7 @@ def _historical_check_new_changes_cte_sql( f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - "), __hard_deletes AS (" + f"), {reappearing_new_changes_sql}, __hard_deletes AS (" f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " f"WHERE __target.{valid_to_column} IS NULL" @@ -2943,6 +2957,70 @@ def _historical_hard_deleted_at_sql( ) +def _historical_observed_group_sequence_cte_sql(*, origin: str, observed_at_column: str) -> str: + return ( + "__observed_group_sequence AS (" + "SELECT __observed_at, LAG(__observed_at) OVER (ORDER BY __observed_at) " + "AS __prev_group_observed_at " + f"FROM (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " + "AS __observed_group_values" + ")" + ) + + +def _historical_observed_group_sequence_from_sql(*, origin: str, observed_at_column: str) -> str: + return ( + f"FROM {origin} AS __source LEFT JOIN __observed_group_sequence " + f"ON __observed_group_sequence.__observed_at = __source.{observed_at_column}" + ) + + +def _historical_reappearing_new_changes_ctes_sql( + *, + changed_or_new_sql: str, + unique_key: tuple[str, ...], + observed_at_column: str, + valid_to_column: str, +) -> str: + """Render ``__new_changes`` plus the first observation after a hard-deleted latest version.""" + + ordered_key_sql: str = ", ".join(f"__ordered.{column}" for column in unique_key) + latest_condition: str = _snapshot_key_condition( + left_alias="__ordered", right_alias="__latest", unique_key=unique_key + ) + present_at_close_condition: str = _snapshot_key_condition( + left_alias="__present", right_alias="__latest", unique_key=unique_key + ) + reappearance_condition: str = _snapshot_key_condition( + left_alias="__ordered", right_alias="__reappearances", unique_key=unique_key + ) + already_changed_condition: str = _snapshot_key_condition( + left_alias="__changed_or_new", right_alias="__ordered", unique_key=unique_key + ) + return ( + f"__changed_or_new AS ({changed_or_new_sql}), __reappearances AS (" + f"SELECT {ordered_key_sql}, MIN(__ordered.{observed_at_column}) AS __reappeared_at " + f"FROM __ordered JOIN __latest ON {latest_condition} " + f"WHERE __latest.{valid_to_column} IS NOT NULL " + f"AND __ordered.{observed_at_column} > __latest.{valid_to_column} " + "AND NOT EXISTS (" + f"SELECT 1 FROM __ordered AS __present WHERE {present_at_close_condition} " + f"AND __present.{observed_at_column} = __latest.{valid_to_column}" + ") " + f"GROUP BY {ordered_key_sql}" + "), __new_changes AS (" + "SELECT * FROM __changed_or_new " + "UNION ALL " + "SELECT __ordered.* FROM __ordered " + f"JOIN __reappearances ON {reappearance_condition} " + f"AND __ordered.{observed_at_column} = __reappearances.__reappeared_at " + "WHERE NOT EXISTS (" + f"SELECT 1 FROM __changed_or_new WHERE {already_changed_condition} " + f"AND __changed_or_new.{observed_at_column} = __ordered.{observed_at_column}" + "))" + ) + + def _historical_snapshot_combined_close_sql( *, destination: str, diff --git a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py index 31f77c7009..9fd95d9708 100644 --- a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py @@ -16,6 +16,9 @@ from sqlbuild.adapter.contract.classes.base_adapter import ( BaseAdapter, _encode_typed_json, + _historical_observed_group_sequence_cte_sql, + _historical_observed_group_sequence_from_sql, + _historical_reappearing_new_changes_ctes_sql, _render_ansi_typed_scalar, _render_typed_value_list, ) @@ -2632,13 +2635,21 @@ def _historical_check_snapshot_select_sql( observed_at_column=observed_at_column, row_alias="__changes", ) + group_sequence_sql: str = _historical_observed_group_sequence_cte_sql( + origin=origin, observed_at_column=observed_at_column + ) + grouped_origin_sql: str = _historical_observed_group_sequence_from_sql( + origin=origin, observed_at_column=observed_at_column + ) return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" + f"WITH {group_sequence_sql}, __ordered AS (" + "SELECT __source.*, __observed_group_sequence.__prev_group_observed_at, " + f"LAG({observed_at_column}) OVER (" f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" + f") AS __prev_observed_at{previous_columns_sql} {grouped_origin_sql}" "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" + f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition}) " + "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at" "), __versions AS (" f"SELECT __changes.*, LEAD({observed_at_column}) OVER (" f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" @@ -2895,6 +2906,18 @@ def _historical_check_new_changes_cte_sql( observed_at_column=observed_at_column, row_alias="__target", ) + reappearing_new_changes_sql: str = _historical_reappearing_new_changes_ctes_sql( + changed_or_new_sql=( + "SELECT __delta_changes.* FROM __delta_changes " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{first_key} IS NULL " + f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " + f"AND ({latest_change_condition}))" + ), + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_to_column=valid_to_column, + ) return ( "__ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -2906,13 +2929,7 @@ def _historical_check_new_changes_cte_sql( f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - "), __hard_deletes AS (" + f"), {reappearing_new_changes_sql}, __hard_deletes AS (" f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " f"WHERE __target.{valid_to_column} IS NULL" diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py index 6b123118f8..f2fb3dbee9 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py @@ -17,6 +17,7 @@ _encode_typed_json, _historical_check_snapshot_select_sql, _historical_hard_deleted_at_sql, + _historical_reappearing_new_changes_ctes_sql, _historical_snapshot_combined_close_sql, _historical_timestamp_changes_select_sql, _historical_timestamp_snapshot_select_sql, @@ -2698,6 +2699,18 @@ def _pg_historical_check_new_changes_cte_sql( observed_at_column=observed_at_column, row_alias="__target", ) + reappearing_new_changes_sql: str = _historical_reappearing_new_changes_ctes_sql( + changed_or_new_sql=( + "SELECT __delta_changes.* FROM __delta_changes " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{first_key} IS NULL " + f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " + f"AND ({latest_change_condition}))" + ), + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_to_column=valid_to_column, + ) return ( "__ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -2705,13 +2718,7 @@ def _pg_historical_check_new_changes_cte_sql( f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" "), __delta_changes AS (" f"{changed_or_first_sql}" - f"), {latest_sql}, __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - "), __hard_deletes AS (" + f"), {latest_sql}, {reappearing_new_changes_sql}, __hard_deletes AS (" f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " f"WHERE __target.{valid_to_column} IS NULL" diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py index 4e60c013e1..7e2cee1f07 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py @@ -16,6 +16,9 @@ _build_schemas_filter, _encode_typed_json, _historical_hard_deleted_at_sql, + _historical_observed_group_sequence_cte_sql, + _historical_observed_group_sequence_from_sql, + _historical_reappearing_new_changes_ctes_sql, _historical_timestamp_changes_new_records_cte_sql, _quote_sql_string, _render_ansi_typed_scalar, @@ -293,6 +296,14 @@ def _historical_check_new_changes_cte_sql( for column in check_columns ) first_key: str = unique_key[0] + changed_or_new_sql: str = ( + "SELECT __delta_changes.* FROM __delta_changes " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{first_key} IS NULL OR (" + f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " + f"AND ({latest_change_condition}))" + ) + new_changes_sql: str = f"__new_changes AS ({changed_or_new_sql})" hard_deletes_sql: str = "" if invalidate_hard_deletes: hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( @@ -301,11 +312,18 @@ def _historical_check_new_changes_cte_sql( observed_at_column=observed_at_column, row_alias="__target", ) + new_changes_sql = _historical_reappearing_new_changes_ctes_sql( + changed_or_new_sql=changed_or_new_sql, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_to_column=valid_to_column, + ) hard_deletes_sql = ( - "), __hard_deletes AS (" + ", __hard_deletes AS (" f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " f"WHERE __target.{valid_to_column} IS NULL" + ")" ) changed_or_first_sql: str = ( "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " @@ -323,14 +341,7 @@ def _historical_check_new_changes_cte_sql( f"ORDER BY {valid_from_column} DESC) AS __rn FROM {destination}" "), __latest AS (" "SELECT * FROM __latest_ordered WHERE __rn = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - f"{hard_deletes_sql}" - ")" + f"), {new_changes_sql}{hard_deletes_sql}" ) def connect(self, config: dict[str, Any]) -> _SqlServerConnection: @@ -2203,12 +2214,35 @@ def render_create_initial_historical_check_snapshot_destination( "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " "ELSE __next_change_at END" ) - historical_sql: str = ( - f";WITH __ordered AS (SELECT *, LAG({observed_at_column}) OVER (" + ordered_sql: str = ( + f"SELECT *, LAG({observed_at_column}) OVER (" f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" + ) + changes_condition: str = f"__prev_observed_at IS NULL OR ({change_condition})" + leading_ctes_sql: str = "" + if invalidate_hard_deletes: + group_sequence_sql: str = _historical_observed_group_sequence_cte_sql( + origin=origin, observed_at_column=observed_at_column + ) + grouped_origin_sql: str = _historical_observed_group_sequence_from_sql( + origin=origin, observed_at_column=observed_at_column + ) + leading_ctes_sql = f"{group_sequence_sql}, " + ordered_sql = ( + "SELECT __source.*, __observed_group_sequence.__prev_group_observed_at, " + f"LAG({observed_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" + f") AS __prev_observed_at{previous_columns_sql} {grouped_origin_sql}" + ) + reappeared_condition: str = self._distinct_condition( + left="__prev_observed_at", right="__prev_group_observed_at" + ) + changes_condition = f"{changes_condition} OR {reappeared_condition}" + historical_sql: str = ( + f";WITH {leading_ctes_sql}__ordered AS ({ordered_sql}" "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" + f"SELECT * FROM __ordered WHERE {changes_condition}" "), __versions AS (" f"SELECT __changes.*, LEAD({observed_at_column}) OVER (" f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/build/_test_types.py b/tests/e2e/src/sqlbuild/cli/commands/main/build/_test_types.py index fc1ee9f4f3..5ae694aeed 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/build/_test_types.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/build/_test_types.py @@ -1055,6 +1055,19 @@ class RemoveColumnSemanticsBuildE2ETestCase: expected_warning_count: int +@dataclass(frozen=True) +class SnapshotReappearingKeyBuildE2ETestCase: + """Test case for hard-deleted keys that reappear in historical check snapshots.""" + + description: str + repo_files: dict[str, str] + initial_seed_sql: str + observation_sql_by_round: tuple[tuple[str, ...], ...] + history_query_template: str + expected_incremental_rows_by_round: tuple[tuple[tuple[object, ...], ...], ...] + expected_full_history_rows: tuple[tuple[object, ...], ...] + + @dataclass(frozen=True) class DynamicColumnContractE2ETestCase: description: str diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/build/test_snapshot_build.py b/tests/e2e/src/sqlbuild/cli/commands/main/build/test_snapshot_build.py index de079fee7d..b036f6cc79 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/build/test_snapshot_build.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/build/test_snapshot_build.py @@ -18,6 +18,7 @@ SnapshotHistoricalCheckBuildE2ETestCase, SnapshotHistoricalTimestampBuildE2ETestCase, SnapshotHookBuildE2ETestCase, + SnapshotReappearingKeyBuildE2ETestCase, SnapshotSelectorBuildE2ETestCase, SnapshotTimestampBuildE2ETestCase, SnapshotTimestampFailureBuildE2ETestCase, @@ -31,6 +32,30 @@ table_exists, ) +_REAPPEARING_KEY_FINAL_HISTORY: tuple[tuple[object, ...], ...] = ( + (1, "active", "2024-02-01", "2024-02-04"), + (1, "paused", "2024-02-04", None), + (2, "active", "2024-02-01", "2024-02-02"), + (2, "active", "2024-02-03", None), + (3, "active", "2024-02-01", "2024-02-04"), + (3, "active", "2024-02-05", None), +) +_REAPPEARING_KEY_OBSERVATIONS_SQL: str = ( + "SELECT 1 AS customer_id, 'active' AS status, TIMESTAMP '2024-02-01' AS observed_at " + "UNION ALL SELECT 2, 'active', TIMESTAMP '2024-02-01' " + "UNION ALL SELECT 3, 'active', TIMESTAMP '2024-02-01' " + "UNION ALL SELECT 1, 'active', TIMESTAMP '2024-02-02' " + "UNION ALL SELECT 3, 'active', TIMESTAMP '2024-02-02' " + "UNION ALL SELECT 1, 'active', TIMESTAMP '2024-02-03' " + "UNION ALL SELECT 2, 'active', TIMESTAMP '2024-02-03' " + "UNION ALL SELECT 3, 'active', TIMESTAMP '2024-02-03' " + "UNION ALL SELECT 1, 'paused', TIMESTAMP '2024-02-04' " + "UNION ALL SELECT 2, 'active', TIMESTAMP '2024-02-04' " + "UNION ALL SELECT 1, 'paused', TIMESTAMP '2024-02-05' " + "UNION ALL SELECT 2, 'active', TIMESTAMP '2024-02-05' " + "UNION ALL SELECT 3, 'active', TIMESTAMP '2024-02-05'" +) + @pytest.mark.parametrize( "test_case", @@ -1927,6 +1952,183 @@ def test_given_shallow_waffle_shop_snapshot_edges_when_sources_change_then_cli_t ) +@pytest.mark.parametrize( + "test_case", + [ + SnapshotReappearingKeyBuildE2ETestCase( + description=( + "historical check snapshot reopens hard-deleted keys that reappear unchanged" + ), + repo_files={ + "sqlbuild_project.toml": dedent( + """ + name = "snapshot_reappearing_keys" + adapter = "duckdb" + + [connection] + database = "snapshot_reappearing_keys.duckdb" + """ + ).strip() + + "\n", + "sources/raw.yml": dedent( + """ + sources: + - name: raw_membership_observations + schema: main + table: raw_membership_observations + - name: raw_membership_full_history + schema: main + table: raw_membership_full_history + """ + ).strip() + + "\n", + "models/membership_incremental_snapshot.sql": dedent( + """ + MODEL ( + materialized snapshot, + unique_key [customer_id], + snapshot_strategy check, + check_columns [status], + observed_at observed_at, + invalidate_hard_deletes true + ); + + SELECT customer_id, status, observed_at + FROM __source("raw_membership_observations") + """ + ).strip() + + "\n", + "models/membership_full_history_snapshot.sql": dedent( + """ + MODEL ( + materialized snapshot, + unique_key [customer_id], + snapshot_strategy check, + check_columns [status], + observed_at observed_at, + invalidate_hard_deletes true + ); + + SELECT customer_id, status, observed_at + FROM __source("raw_membership_full_history") + """ + ).strip() + + "\n", + }, + initial_seed_sql=( + "CREATE TABLE main.raw_membership_observations AS " + "SELECT 1 AS customer_id, 'active' AS status, TIMESTAMP '2024-02-01' AS observed_at " + "UNION ALL SELECT 2, 'active', TIMESTAMP '2024-02-01' " + "UNION ALL SELECT 3, 'active', TIMESTAMP '2024-02-01'; " + f"CREATE TABLE main.raw_membership_full_history AS {_REAPPEARING_KEY_OBSERVATIONS_SQL}" + ), + observation_sql_by_round=( + (), + ( + "INSERT INTO main.raw_membership_observations VALUES " + "(1, 'active', TIMESTAMP '2024-02-02'), (3, 'active', TIMESTAMP '2024-02-02')", + ), + ( + "INSERT INTO main.raw_membership_observations VALUES " + "(1, 'active', TIMESTAMP '2024-02-03'), (2, 'active', TIMESTAMP '2024-02-03'), " + "(3, 'active', TIMESTAMP '2024-02-03')", + ), + ( + "INSERT INTO main.raw_membership_observations VALUES " + "(1, 'paused', TIMESTAMP '2024-02-04'), (2, 'active', TIMESTAMP '2024-02-04'), " + "(1, 'paused', TIMESTAMP '2024-02-05'), (2, 'active', TIMESTAMP '2024-02-05'), " + "(3, 'active', TIMESTAMP '2024-02-05')", + ), + ), + history_query_template=( + "SELECT customer_id, status, strftime(valid_from, '%Y-%m-%d'), " + "strftime(valid_to, '%Y-%m-%d') FROM main.{table_name} " + "ORDER BY customer_id, valid_from" + ), + expected_incremental_rows_by_round=( + ( + (1, "active", "2024-02-01", None), + (2, "active", "2024-02-01", None), + (3, "active", "2024-02-01", None), + ), + ( + (1, "active", "2024-02-01", None), + (2, "active", "2024-02-01", "2024-02-02"), + (3, "active", "2024-02-01", None), + ), + ( + (1, "active", "2024-02-01", None), + (2, "active", "2024-02-01", "2024-02-02"), + (2, "active", "2024-02-03", None), + (3, "active", "2024-02-01", None), + ), + _REAPPEARING_KEY_FINAL_HISTORY, + ), + expected_full_history_rows=_REAPPEARING_KEY_FINAL_HISTORY, + ) + ], + ids=lambda case: case.description, +) +def test_given_hard_deleted_key_when_it_reappears_unchanged_then_snapshot_opens_new_version( + test_case: SnapshotReappearingKeyBuildE2ETestCase, + tmp_path: Path, +) -> None: + project_dir: Path = prepare_inline_project( + tmp_path=tmp_path, + project_name="snapshot_reappearing_keys", + repo_files=test_case.repo_files, + ) + db_path: Path = project_dir / "snapshot_reappearing_keys.duckdb" + + import duckdb + + connection: duckdb.DuckDBPyConnection = duckdb.connect(str(db_path)) + connection.execute(test_case.initial_seed_sql) + connection.close() + + observation_sql_round: tuple[str, ...] + expected_incremental_rows: tuple[tuple[object, ...], ...] + for observation_sql_round, expected_incremental_rows in zip( + test_case.observation_sql_by_round, + test_case.expected_incremental_rows_by_round, + strict=True, + ): + connection = duckdb.connect(str(db_path)) + observation_sql: str + for observation_sql in observation_sql_round: + connection.execute(observation_sql) + connection.close() + + result: object = run_sqb(command=("--no-color", "build"), project_dir=project_dir) + assert result.returncode == 0, result.stdout + result.stderr + + incremental_rows: tuple[tuple[object, ...], ...] = tuple( + tuple(row) + for row in query_duckdb( + db_path=db_path, + sql=test_case.history_query_template.format( + table_name="membership_incremental_snapshot" + ), + ) + ) + full_history_rows: tuple[tuple[object, ...], ...] = tuple( + tuple(row) + for row in query_duckdb( + db_path=db_path, + sql=test_case.history_query_template.format( + table_name="membership_full_history_snapshot" + ), + ) + ) + assert incremental_rows == expected_incremental_rows + assert full_history_rows == test_case.expected_full_history_rows + table_name: str + for table_name in ("membership_incremental_snapshot", "membership_full_history_snapshot"): + assert_snapshot_scd2_invariants( + db_path=db_path, table_name=table_name, key_columns=("customer_id",) + ) + + @pytest.mark.parametrize( "test_case", [ diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/postgres/test_postgres.py b/tests/e2e/src/sqlbuild/cli/commands/main/postgres/test_postgres.py index fde9573d92..0fdf5f26e5 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/postgres/test_postgres.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/postgres/test_postgres.py @@ -1116,6 +1116,7 @@ def test_given_intermediate_strategy_project_when_loading_twice_on_postgres_then ("1", "active", "2026-01-01", "2026-01-03"), ("1", "paused", "2026-01-03", None), ("2", "active", "2026-01-01", "2026-01-02"), + ("2", "active", "2026-01-03", None), ), expected_failure_fragments=( "current_customer_snapshot", @@ -1231,6 +1232,7 @@ def test_given_snapshot_project_when_building_on_postgres_then_scd2_history_is_v ("1", "active", "2026-01-01", "2026-01-03"), ("1", "paused", "2026-01-03", None), ("2", "active", "2026-01-01", "2026-01-02"), + ("2", "active", "2026-01-03", None), ), ) ], diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/sqlserver/test_sqlserver.py b/tests/e2e/src/sqlbuild/cli/commands/main/sqlserver/test_sqlserver.py index 346f8e23ff..7c0290f4ec 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/sqlserver/test_sqlserver.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/sqlserver/test_sqlserver.py @@ -1764,6 +1764,7 @@ def test_given_intermediate_strategy_project_when_loading_twice_on_sqlserver_the ("1", "active", "2026-01-01", "2026-01-03"), ("1", "paused", "2026-01-03", None), ("2", "active", "2026-01-01", "2026-01-02"), + ("2", "active", "2026-01-03", None), ), expected_failure_fragments=( "current_customer_snapshot", @@ -1885,6 +1886,7 @@ def test_given_snapshot_project_when_building_on_sqlserver_then_scd2_history_is_ ("1", "active", "2026-01-01", "2026-01-03"), ("1", "paused", "2026-01-03", None), ("2", "active", "2026-01-01", "2026-01-02"), + ("2", "active", "2026-01-03", None), ), ) ], diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py index cec73549b2..98f8b8aa86 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py @@ -13,3 +13,12 @@ class SnapshotSqlRenderingAdapterTestCase: expected_historical_timestamp_initial_hard_delete_fragments: tuple[str, ...] expected_historical_timestamp_apply_hard_delete_fragments: tuple[str, ...] expected_historical_check_apply_fragments: tuple[str, ...] + + +@dataclass(frozen=True) +class SnapshotReappearanceRenderingTestCase: + description: str + adapter: BaseAdapter + expected_initial_fragments: tuple[str, ...] + expected_apply_fragments: tuple[str, ...] + unexpected_apply_fragments: tuple[str, ...] diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py index 10d33e17db..53b9d86b0b 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py @@ -7,11 +7,30 @@ from sqlbuild.adapters.bigquery.classes.bigquery_adapter import BigQueryAdapter from sqlbuild.adapters.databricks.classes.databricks_adapter import DatabricksAdapter from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from sqlbuild.adapters.motherduck.classes.motherduck_adapter import MotherDuckAdapter +from sqlbuild.adapters.postgres.classes.postgres_adapter import PostgresAdapter from sqlbuild.adapters.snowflake.classes.snowflake_adapter import SnowflakeAdapter +from sqlbuild.adapters.sqlserver.classes.sqlserver_adapter import SqlServerAdapter from tests.integration.src.sqlbuild.adapters.snapshot_sql_rendering._test_types import ( + SnapshotReappearanceRenderingTestCase, SnapshotSqlRenderingAdapterTestCase, ) +_REAPPEARANCE_GROUP_SEQUENCE_FRAGMENT: str = ( + "__observed_group_sequence AS (SELECT __observed_at, " + "LAG(__observed_at) OVER (ORDER BY __observed_at) AS __prev_group_observed_at " + "FROM (SELECT DISTINCT observed_at AS __observed_at FROM source_table)" +) +_REAPPEARANCE_APPLY_FRAGMENTS: tuple[str, ...] = ( + "__reappearances AS (SELECT __ordered.customer_id, " + "MIN(__ordered.observed_at) AS __reappeared_at", + "WHERE __latest.valid_to IS NOT NULL AND __ordered.observed_at > __latest.valid_to", + "AND __present.observed_at = __latest.valid_to", + "__new_changes AS (SELECT * FROM __changed_or_new UNION ALL SELECT __ordered.* " + "FROM __ordered JOIN __reappearances", + "AND __changed_or_new.observed_at = __ordered.observed_at", +) + @pytest.mark.parametrize( "test_case", @@ -265,3 +284,129 @@ def test_given_builtin_adapter_when_rendering_snapshot_sql_then_covers_snapshot_ assert expected_fragment in historical_timestamp_apply_sql for expected_fragment in test_case.expected_historical_check_apply_fragments: assert expected_fragment in historical_check_apply_sql + + +@pytest.mark.parametrize( + "test_case", + [ + SnapshotReappearanceRenderingTestCase( + description="duckdb reopens reappearing historical check keys", + adapter=DuckDbAdapter(), + expected_initial_fragments=( + _REAPPEARANCE_GROUP_SEQUENCE_FRAGMENT, + "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at", + ), + expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, + unexpected_apply_fragments=("UNION DISTINCT",), + ), + SnapshotReappearanceRenderingTestCase( + description="motherduck reopens reappearing historical check keys", + adapter=MotherDuckAdapter(), + expected_initial_fragments=( + _REAPPEARANCE_GROUP_SEQUENCE_FRAGMENT, + "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at", + ), + expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, + unexpected_apply_fragments=("UNION DISTINCT",), + ), + SnapshotReappearanceRenderingTestCase( + description="postgres reopens reappearing historical check keys without qualify", + adapter=PostgresAdapter(), + expected_initial_fragments=( + _REAPPEARANCE_GROUP_SEQUENCE_FRAGMENT, + "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at", + ), + expected_apply_fragments=( + *_REAPPEARANCE_APPLY_FRAGMENTS, + ") AS __rn FROM target_table) AS __q WHERE __rn = 1)", + ), + unexpected_apply_fragments=("QUALIFY", "UNION DISTINCT"), + ), + SnapshotReappearanceRenderingTestCase( + description="sqlserver reopens reappearing historical check keys without qualify", + adapter=SqlServerAdapter(), + expected_initial_fragments=( + _REAPPEARANCE_GROUP_SEQUENCE_FRAGMENT, + "OR (__prev_observed_at <> __prev_group_observed_at " + "OR (__prev_observed_at IS NULL AND __prev_group_observed_at IS NOT NULL) " + "OR (__prev_observed_at IS NOT NULL AND __prev_group_observed_at IS NULL))", + ), + expected_apply_fragments=( + *_REAPPEARANCE_APPLY_FRAGMENTS, + "__latest AS (SELECT * FROM __latest_ordered WHERE __rn = 1)", + ), + unexpected_apply_fragments=("QUALIFY", "UNION DISTINCT", "IS DISTINCT FROM"), + ), + ], + ids=lambda case: case.description, +) +def test_given_hard_deletes_when_rendering_historical_check_snapshot_then_reopens_reappearing_keys( + test_case: SnapshotReappearanceRenderingTestCase, +) -> None: + initial_sql: str = "\n".join( + test_case.adapter.render_create_initial_historical_check_snapshot_destination( + table_type="table", + destination="target_table", + origin="source_table", + unique_key=("customer_id",), + check_columns=("plan",), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "plan", "observed_at"), + invalidate_hard_deletes=True, + ) + ) + apply_statements: tuple[str, ...] = ( + test_case.adapter.render_apply_historical_check_snapshot_changes( + destination="target_table", + origin="source_table", + unique_key=("customer_id",), + check_columns=("plan",), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "plan", "observed_at"), + invalidate_hard_deletes=True, + ) + ) + without_hard_deletes_sql: str = "\n".join( + ( + *test_case.adapter.render_create_initial_historical_check_snapshot_destination( + table_type="table", + destination="target_table", + origin="source_table", + unique_key=("customer_id",), + check_columns=("plan",), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "plan", "observed_at"), + invalidate_hard_deletes=False, + ), + *test_case.adapter.render_apply_historical_check_snapshot_changes( + destination="target_table", + origin="source_table", + unique_key=("customer_id",), + check_columns=("plan",), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "plan", "observed_at"), + invalidate_hard_deletes=False, + ), + ) + ) + + expected_fragment: str + for expected_fragment in test_case.expected_initial_fragments: + assert expected_fragment in initial_sql + statement: str + for statement in apply_statements: + for expected_fragment in test_case.expected_apply_fragments: + assert expected_fragment in statement + unexpected_fragment: str + for unexpected_fragment in test_case.unexpected_apply_fragments: + assert unexpected_fragment not in statement + assert "__prev_group_observed_at" not in without_hard_deletes_sql + assert "__reappearances" not in without_hard_deletes_sql From 4befbe0224f2f072ae2b54ce8e72d89bf884ac3b Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:54:32 +0100 Subject: [PATCH 39/57] test(snapshots): execute every adapter's snapshot SQL on DuckDB Render historical and current-state snapshot SQL for the check and timestamp strategies, with and without hard deletes, through the initial and incremental paths, and execute it on DuckDB for the DuckDB, MotherDuck, Postgres, BigQuery, Snowflake and Databricks adapters. Every adapter must produce the same expected history and keep the snapshot invariants: no overlapping windows and at most one active version per key. The only shims are BigQuery identifier quoting and a pinned execution clock for current-state histories. SQL Server renders T-SQL that DuckDB cannot run and keeps rendering assertions only. (cherry picked from commit dd8327e12dde193780fdfaec69ad649bfc204f3d) --- .../snapshot_sql_rendering/_test_types.py | 16 + .../snapshot_sql_rendering/helpers.py | 382 ++++++++++++++++++ .../test_duckdb_execution.py | 178 ++++++++ 3 files changed, 576 insertions(+) create mode 100644 tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py create mode 100644 tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py index 98f8b8aa86..9e5a812187 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py @@ -1,3 +1,4 @@ +from collections.abc import Callable from dataclasses import dataclass from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter @@ -22,3 +23,18 @@ class SnapshotReappearanceRenderingTestCase: expected_initial_fragments: tuple[str, ...] expected_apply_fragments: tuple[str, ...] unexpected_apply_fragments: tuple[str, ...] + + +@dataclass(frozen=True) +class SnapshotExecutionTestCase: + """One adapter executing one snapshot scenario through one build path on DuckDB.""" + + description: str + adapter_type: type[BaseAdapter] + normalize_sql: Callable[[str], str] + source_select_sql: Callable[[tuple[object, ...]], str] + render_initial: Callable[[BaseAdapter], tuple[str, ...]] + render_apply: Callable[[BaseAdapter], tuple[str, ...]] + history_sql: str + builds: tuple[tuple[tuple[object, ...], ...], ...] + expected_history: tuple[tuple[object, ...], ...] diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py new file mode 100644 index 0000000000..cc43637595 --- /dev/null +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py @@ -0,0 +1,382 @@ +"""Helpers that execute adapter-rendered snapshot SQL on DuckDB over shared scenarios. + +DuckDB runs the SQL rendered by the DuckDB, MotherDuck, Postgres, BigQuery, Snowflake and +Databricks adapters. The only shims are identifier quoting for BigQuery and a pinned execution +clock so current-state histories are deterministic; neither changes statement semantics. SQL +Server renders T-SQL that DuckDB cannot parse, so it is covered by rendering assertions only. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import duckdb + +from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter +from sqlbuild.adapter.contract.models import SnapshotChangeTarget +from sqlbuild.adapters.bigquery.classes.bigquery_adapter import BigQueryAdapter +from sqlbuild.adapters.databricks.classes.databricks_adapter import DatabricksAdapter +from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from sqlbuild.adapters.motherduck.classes.motherduck_adapter import MotherDuckAdapter +from sqlbuild.adapters.postgres.classes.postgres_adapter import PostgresAdapter +from sqlbuild.adapters.snowflake.classes.snowflake_adapter import SnowflakeAdapter +from sqlbuild.spec.contracts.types import TableType +from tests.integration.src.sqlbuild.adapters.snapshot_sql_rendering._test_types import ( + SnapshotExecutionTestCase, +) + + +@dataclass(frozen=True) +class SnapshotExecutionAdapter: + """Adapter whose rendered snapshot SQL is executed on DuckDB.""" + + name: str + adapter_type: type[BaseAdapter] + normalize_sql: Callable[[str], str] + + +@dataclass(frozen=True) +class SnapshotExecutionKind: + """Snapshot configuration family shared by execution scenarios.""" + + source_select_sql: Callable[[tuple[object, ...]], str] + render_initial: Callable[[BaseAdapter], tuple[str, ...]] + render_apply: Callable[[BaseAdapter], tuple[str, ...]] + history_sql: str + + +@dataclass(frozen=True) +class SnapshotExecutionScenario: + """Source rows per build and the history every adapter must produce.""" + + description: str + kind: SnapshotExecutionKind + builds: tuple[tuple[tuple[object, ...], ...], ...] + expected_history: tuple[tuple[object, ...], ...] + + +@dataclass(frozen=True) +class SnapshotExecutionRun: + """One adapter executing one scenario through incremental builds or one initial build.""" + + adapter: SnapshotExecutionAdapter + scenario: SnapshotExecutionScenario + path: str + builds: tuple[tuple[tuple[object, ...], ...], ...] + + +_TARGET: str = "snapshot_target" +_SOURCE: str = "snapshot_source" +FIRST_CLOCK_DAY: int = 20 +_KEY: tuple[str, ...] = ("customer_id",) + + +def _timestamp(day: object) -> str: + return f"TIMESTAMP '2026-01-{int(str(day)):02d}'" + + +def _check_history_row_sql(row: tuple[object, ...]) -> str: + return ( + f"SELECT {row[0]} AS customer_id, '{row[1]}' AS status, {_timestamp(row[2])} AS observed_at" + ) + + +def _timestamp_history_row_sql(row: tuple[object, ...]) -> str: + return ( + f"SELECT {row[0]} AS customer_id, '{row[1]}' AS plan, " + f"{_timestamp(row[2])} AS updated_at, {_timestamp(row[3])} AS observed_at" + ) + + +def _current_check_row_sql(row: tuple[object, ...]) -> str: + return f"SELECT {row[0]} AS customer_id, '{row[1]}' AS status" + + +def _current_timestamp_row_sql(row: tuple[object, ...]) -> str: + return f"SELECT {row[0]} AS customer_id, '{row[1]}' AS plan, {_timestamp(row[2])} AS updated_at" + + +def _historical_check_initial(adapter: BaseAdapter, hard_deletes: bool) -> tuple[str, ...]: + return adapter.render_create_initial_historical_check_snapshot_destination( + table_type=TableType.PERMANENT, + destination=_TARGET, + origin=_SOURCE, + unique_key=_KEY, + check_columns=("status",), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "status", "observed_at"), + invalidate_hard_deletes=hard_deletes, + ) + + +def _historical_check_apply(adapter: BaseAdapter, hard_deletes: bool) -> tuple[str, ...]: + return adapter.render_apply_historical_check_snapshot_changes( + destination=_TARGET, + origin=_SOURCE, + unique_key=_KEY, + check_columns=("status",), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "status", "observed_at"), + invalidate_hard_deletes=hard_deletes, + ) + + +def _historical_timestamp_initial(adapter: BaseAdapter, hard_deletes: bool) -> tuple[str, ...]: + return adapter.render_create_initial_historical_timestamp_snapshot_destination( + table_type=TableType.PERMANENT, + destination=_TARGET, + origin=_SOURCE, + unique_key=_KEY, + updated_at_column="updated_at", + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "plan", "updated_at", "observed_at"), + invalidate_hard_deletes=hard_deletes, + ) + + +def _historical_timestamp_apply(adapter: BaseAdapter, hard_deletes: bool) -> tuple[str, ...]: + return adapter.render_apply_historical_timestamp_snapshot_changes( + destination=_TARGET, + origin=_SOURCE, + unique_key=_KEY, + updated_at_column="updated_at", + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "plan", "updated_at", "observed_at"), + invalidate_hard_deletes=hard_deletes, + ) + + +def _current_initial( + adapter: BaseAdapter, strategy: str, updated_at_column: str | None +) -> tuple[str, ...]: + return adapter.render_create_initial_snapshot_destination( + table_type=TableType.PERMANENT, + destination=_TARGET, + origin=_SOURCE, + snapshot_strategy=strategy, + updated_at_column=updated_at_column, + observed_at_column=None, + valid_from_column="valid_from", + valid_to_column="valid_to", + initial_valid_from=None, + ) + + +def _current_timestamp_apply(adapter: BaseAdapter, hard_deletes: bool) -> tuple[str, ...]: + return adapter.render_apply_timestamp_snapshot_changes( + destination=_TARGET, + origin=_SOURCE, + unique_key=_KEY, + updated_at_column="updated_at", + observed_at_column=None, + valid_from_column="valid_from", + valid_to_column="valid_to", + initial_valid_from=None, + output_columns=("customer_id", "plan", "updated_at"), + invalidate_hard_deletes=hard_deletes, + ) + + +def _current_check_apply(adapter: BaseAdapter, hard_deletes: bool) -> tuple[str, ...]: + return adapter.render_apply_check_snapshot_changes( + target=SnapshotChangeTarget( + destination=_TARGET, + origin=_SOURCE, + unique_key=_KEY, + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "status"), + ), + check_columns=("status",), + updated_at_column=None, + observed_at_column=None, + initial_valid_from=None, + invalidate_hard_deletes=hard_deletes, + ) + + +def _history_sql(value_column: str) -> str: + return ( + f"SELECT customer_id, {value_column}, day(valid_from), day(valid_to) " + f"FROM {_TARGET} ORDER BY customer_id, valid_from" + ) + + +def _unquote_backticks(sql: str) -> str: + return sql.replace("`", '"') + + +def _unchanged(sql: str) -> str: + return sql + + +HISTORICAL_CHECK: SnapshotExecutionKind = SnapshotExecutionKind( + source_select_sql=_check_history_row_sql, + render_initial=lambda adapter: _historical_check_initial(adapter, False), + render_apply=lambda adapter: _historical_check_apply(adapter, False), + history_sql=_history_sql("status"), +) +HISTORICAL_CHECK_HARD_DELETES: SnapshotExecutionKind = SnapshotExecutionKind( + source_select_sql=_check_history_row_sql, + render_initial=lambda adapter: _historical_check_initial(adapter, True), + render_apply=lambda adapter: _historical_check_apply(adapter, True), + history_sql=_history_sql("status"), +) +HISTORICAL_TIMESTAMP: SnapshotExecutionKind = SnapshotExecutionKind( + source_select_sql=_timestamp_history_row_sql, + render_initial=lambda adapter: _historical_timestamp_initial(adapter, False), + render_apply=lambda adapter: _historical_timestamp_apply(adapter, False), + history_sql=_history_sql("plan"), +) +HISTORICAL_TIMESTAMP_HARD_DELETES: SnapshotExecutionKind = SnapshotExecutionKind( + source_select_sql=_timestamp_history_row_sql, + render_initial=lambda adapter: _historical_timestamp_initial(adapter, True), + render_apply=lambda adapter: _historical_timestamp_apply(adapter, True), + history_sql=_history_sql("plan"), +) +CURRENT_TIMESTAMP: SnapshotExecutionKind = SnapshotExecutionKind( + source_select_sql=_current_timestamp_row_sql, + render_initial=lambda adapter: _current_initial(adapter, "timestamp", "updated_at"), + render_apply=lambda adapter: _current_timestamp_apply(adapter, False), + history_sql=_history_sql("plan"), +) +CURRENT_TIMESTAMP_HARD_DELETES: SnapshotExecutionKind = SnapshotExecutionKind( + source_select_sql=_current_timestamp_row_sql, + render_initial=lambda adapter: _current_initial(adapter, "timestamp", "updated_at"), + render_apply=lambda adapter: _current_timestamp_apply(adapter, True), + history_sql=_history_sql("plan"), +) +CURRENT_CHECK_HARD_DELETES: SnapshotExecutionKind = SnapshotExecutionKind( + source_select_sql=_current_check_row_sql, + render_initial=lambda adapter: _current_initial(adapter, "check", None), + render_apply=lambda adapter: _current_check_apply(adapter, True), + history_sql=_history_sql("status"), +) + +ADAPTERS: tuple[SnapshotExecutionAdapter, ...] = ( + SnapshotExecutionAdapter(name="duckdb", adapter_type=DuckDbAdapter, normalize_sql=_unchanged), + SnapshotExecutionAdapter( + name="motherduck", adapter_type=MotherDuckAdapter, normalize_sql=_unchanged + ), + SnapshotExecutionAdapter( + name="postgres", adapter_type=PostgresAdapter, normalize_sql=_unchanged + ), + SnapshotExecutionAdapter( + name="bigquery", adapter_type=BigQueryAdapter, normalize_sql=_unquote_backticks + ), + SnapshotExecutionAdapter( + name="snowflake", adapter_type=SnowflakeAdapter, normalize_sql=_unchanged + ), + SnapshotExecutionAdapter( + name="databricks", adapter_type=DatabricksAdapter, normalize_sql=_unchanged + ), +) + +INVARIANT_VIOLATIONS_SQL: str = ( + "SELECT " + f"(SELECT COUNT(*) FROM (SELECT customer_id FROM {_TARGET} WHERE valid_to IS NULL " + "GROUP BY customer_id HAVING COUNT(*) > 1)), " + f"(SELECT COUNT(*) FROM {_TARGET} WHERE valid_to IS NOT NULL AND valid_from >= valid_to), " + f"(SELECT COUNT(*) FROM {_TARGET} AS left_row JOIN {_TARGET} AS right_row " + "ON left_row.customer_id = right_row.customer_id " + "AND left_row.valid_from < right_row.valid_from " + "AND right_row.valid_from < COALESCE(left_row.valid_to, TIMESTAMP '9999-12-31'))" +) + + +def build_execution_runs( + *, + incremental_scenarios: tuple[SnapshotExecutionScenario, ...], + full_history_scenarios: tuple[SnapshotExecutionScenario, ...], +) -> tuple[SnapshotExecutionRun, ...]: + """Run every scenario on every adapter, and historical scenarios as one initial build.""" + + runs: list[SnapshotExecutionRun] = [] + adapter: SnapshotExecutionAdapter + scenario: SnapshotExecutionScenario + for adapter in ADAPTERS: + for scenario in incremental_scenarios: + runs.append( + SnapshotExecutionRun( + adapter=adapter, + scenario=scenario, + path="incremental builds", + builds=scenario.builds, + ) + ) + for scenario in full_history_scenarios: + runs.append( + SnapshotExecutionRun( + adapter=adapter, + scenario=scenario, + path="initial full history", + builds=(scenario.builds[-1],), + ) + ) + return tuple(runs) + + +def run_snapshot_builds( + test_case: SnapshotExecutionTestCase, +) -> tuple[tuple[tuple[object, ...], ...], tuple[object, ...] | None]: + """Execute each build's rendered statements and return the history and invariant counts.""" + + connection: duckdb.DuckDBPyConnection = duckdb.connect() + _run_build( + connection=connection, + test_case=test_case, + rows=test_case.builds[0], + render=test_case.render_initial, + day=FIRST_CLOCK_DAY, + ) + build_index: int + rows: tuple[tuple[object, ...], ...] + for build_index, rows in enumerate(test_case.builds[1:], start=1): + _run_build( + connection=connection, + test_case=test_case, + rows=rows, + render=test_case.render_apply, + day=FIRST_CLOCK_DAY + build_index, + ) + history: tuple[tuple[object, ...], ...] = tuple( + tuple(row) for row in connection.execute(test_case.history_sql).fetchall() + ) + violations: tuple[object, ...] | None = connection.execute(INVARIANT_VIOLATIONS_SQL).fetchone() + connection.close() + return history, violations + + +def _pinned_clock_adapter(adapter_type: type[BaseAdapter], day: int) -> BaseAdapter: + clock_sql: str = _timestamp(day) + pinned_type: type[BaseAdapter] = type( + f"PinnedClock{adapter_type.__name__}", + (adapter_type,), + {"render_current_timestamp": lambda _self: clock_sql}, + ) + return pinned_type() + + +def _run_build( + *, + connection: duckdb.DuckDBPyConnection, + test_case: SnapshotExecutionTestCase, + rows: tuple[tuple[object, ...], ...], + render: Callable[[BaseAdapter], tuple[str, ...]], + day: int, +) -> None: + source_sql: str = " UNION ALL ".join(test_case.source_select_sql(row) for row in rows) + connection.execute(f"CREATE OR REPLACE TABLE {_SOURCE} AS {source_sql}") + adapter: BaseAdapter = _pinned_clock_adapter(test_case.adapter_type, day) + statement: str + for statement in render(adapter): + connection.execute(test_case.normalize_sql(statement)) diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py new file mode 100644 index 0000000000..4e79491358 --- /dev/null +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py @@ -0,0 +1,178 @@ +"""Cross-adapter snapshot SQL executed on DuckDB; see helpers for the adapter shims.""" + +from __future__ import annotations + +import pytest + +from tests.integration.src.sqlbuild.adapters.snapshot_sql_rendering._test_types import ( + SnapshotExecutionTestCase, +) +from tests.integration.src.sqlbuild.adapters.snapshot_sql_rendering.helpers import ( + CURRENT_CHECK_HARD_DELETES, + CURRENT_TIMESTAMP, + CURRENT_TIMESTAMP_HARD_DELETES, + FIRST_CLOCK_DAY, + HISTORICAL_CHECK, + HISTORICAL_CHECK_HARD_DELETES, + HISTORICAL_TIMESTAMP, + HISTORICAL_TIMESTAMP_HARD_DELETES, + SnapshotExecutionRun, + SnapshotExecutionScenario, + build_execution_runs, + run_snapshot_builds, +) + +_CHECK_DELETE_AND_REAPPEAR_ROWS: tuple[tuple[object, ...], ...] = ( + (1, "active", 1), + (2, "active", 1), + (1, "active", 2), + (1, "active", 3), + (2, "active", 3), +) +_TIMESTAMP_DELETE_AND_REAPPEAR_ROWS: tuple[tuple[object, ...], ...] = ( + (1, "basic", 1, 1), + (2, "pro", 1, 1), + (1, "basic", 1, 2), + (1, "basic", 1, 3), + (2, "pro", 1, 3), +) + +_HISTORICAL_SCENARIOS: tuple[SnapshotExecutionScenario, ...] = ( + SnapshotExecutionScenario( + description="check without hard deletes records changes and ignores missing keys", + kind=HISTORICAL_CHECK, + builds=( + ((1, "active", 1), (2, "active", 1)), + ((1, "active", 1), (2, "active", 1), (1, "paused", 2), (1, "paused", 3)), + ), + expected_history=((1, "active", 1, 2), (1, "paused", 2, None), (2, "active", 1, None)), + ), + SnapshotExecutionScenario( + description="check initial observation", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=(((1, "active", 1), (2, "active", 1)),), + expected_history=((1, "active", 1, None), (2, "active", 1, None)), + ), + SnapshotExecutionScenario( + description="check hard delete", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=(_CHECK_DELETE_AND_REAPPEAR_ROWS[:2], _CHECK_DELETE_AND_REAPPEAR_ROWS[:3]), + expected_history=((1, "active", 1, None), (2, "active", 1, 2)), + ), + SnapshotExecutionScenario( + description="check unchanged reappearance across builds", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=( + _CHECK_DELETE_AND_REAPPEAR_ROWS[:2], + _CHECK_DELETE_AND_REAPPEAR_ROWS[:3], + _CHECK_DELETE_AND_REAPPEAR_ROWS, + ), + expected_history=( + (1, "active", 1, None), + (2, "active", 1, 2), + (2, "active", 3, None), + ), + ), + SnapshotExecutionScenario( + description="check unchanged reappearance in the same build", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=(_CHECK_DELETE_AND_REAPPEAR_ROWS[:2], _CHECK_DELETE_AND_REAPPEAR_ROWS), + expected_history=( + (1, "active", 1, None), + (2, "active", 1, 2), + (2, "active", 3, None), + ), + ), + SnapshotExecutionScenario( + description="timestamp without hard deletes records updated_at changes", + kind=HISTORICAL_TIMESTAMP, + builds=( + ((1, "basic", 1, 1),), + ((1, "basic", 1, 1), (1, "pro", 2, 3), (1, "pro", 2, 4)), + ), + expected_history=((1, "basic", 1, 2), (1, "pro", 2, None)), + ), + SnapshotExecutionScenario( + description="timestamp initial observation", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=(_TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2],), + expected_history=((1, "basic", 1, None), (2, "pro", 1, None)), + ), + SnapshotExecutionScenario( + description="timestamp hard delete", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=( + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2], + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:3], + ), + expected_history=((1, "basic", 1, None), (2, "pro", 1, 2)), + ), + SnapshotExecutionScenario( + description="timestamp change followed by an unchanged repeat in one build", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=( + ((1, "basic", 1, 1),), + ((1, "basic", 1, 1), (1, "pro", 2, 3), (1, "pro", 2, 4)), + ), + expected_history=((1, "basic", 1, 2), (1, "pro", 2, None)), + ), +) + +_CURRENT_STATE_SCENARIOS: tuple[SnapshotExecutionScenario, ...] = ( + SnapshotExecutionScenario( + description="current timestamp without hard deletes records updated_at changes", + kind=CURRENT_TIMESTAMP, + builds=(((1, "basic", 1),), ((1, "pro", 5),)), + expected_history=((1, "basic", 1, 5), (1, "pro", 5, None)), + ), + SnapshotExecutionScenario( + description="current timestamp hard deletes keep updated_at for changes", + kind=CURRENT_TIMESTAMP_HARD_DELETES, + builds=(((1, "basic", 1),), ((1, "pro", 5),)), + expected_history=((1, "basic", 1, 5), (1, "pro", 5, None)), + ), + SnapshotExecutionScenario( + description="current check hard delete then unchanged reappearance", + kind=CURRENT_CHECK_HARD_DELETES, + builds=(((1, "active"), (2, "active")), ((1, "active"),), ((1, "active"), (2, "active"))), + expected_history=( + (1, "active", FIRST_CLOCK_DAY, None), + (2, "active", FIRST_CLOCK_DAY, FIRST_CLOCK_DAY + 1), + (2, "active", FIRST_CLOCK_DAY + 2, None), + ), + ), +) + +_RUNS: tuple[SnapshotExecutionRun, ...] = build_execution_runs( + incremental_scenarios=(*_HISTORICAL_SCENARIOS, *_CURRENT_STATE_SCENARIOS), + full_history_scenarios=_HISTORICAL_SCENARIOS, +) + + +@pytest.mark.parametrize( + "test_case", + [ + SnapshotExecutionTestCase( + description=f"{run.adapter.name}: {run.scenario.description} ({run.path})", + adapter_type=run.adapter.adapter_type, + normalize_sql=run.adapter.normalize_sql, + source_select_sql=run.scenario.kind.source_select_sql, + render_initial=run.scenario.kind.render_initial, + render_apply=run.scenario.kind.render_apply, + history_sql=run.scenario.kind.history_sql, + builds=run.builds, + expected_history=run.scenario.expected_history, + ) + for run in _RUNS + ], + ids=lambda case: case.description, +) +def test_given_snapshot_scenario_when_executing_adapter_sql_on_duckdb_then_history_matches( + test_case: SnapshotExecutionTestCase, +) -> None: + history: tuple[tuple[object, ...], ...] + violations: tuple[object, ...] | None + history, violations = run_snapshot_builds(test_case) + + assert history == test_case.expected_history + assert violations == (0, 0, 0) From 3546e3933e75a275d0a73988005ad664ae94bba1 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:55:48 +0100 Subject: [PATCH 40/57] fix(snapshots): stop duplicate versions after changes on BigQuery, Snowflake and Databricks Historical check snapshots with invalidate_hard_deletes on BigQuery, Snowflake and Databricks recorded a second, identical version when an incremental build contained a change followed by an unchanged observation of the new value: the reappearance branch treated the version closed by that change as hard-deleted. These adapters now use the shared reappearance rendering already used by DuckDB, MotherDuck, Postgres and SQL Server, which only counts a closed latest version as hard-deleted when its key is missing from the observation at valid_to. Existing duplicate rows are not rewritten. (cherry picked from commit 2b350ffbc6ef6c83ab0628c62d6d7a7c015a46ab) --- .../bigquery/classes/bigquery_adapter.py | 36 ++++++++----------- .../databricks/classes/databricks_adapter.py | 36 ++++++++----------- .../snowflake/classes/snowflake_adapter.py | 36 ++++++++----------- .../test_duckdb_execution.py | 9 +++++ 4 files changed, 51 insertions(+), 66 deletions(-) diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py index d37debc4bd..6c3532a8dd 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py @@ -13,6 +13,7 @@ _COLLECTION_SQL_VALUE_KINDS, BaseAdapter, _encode_typed_json, + _historical_reappearing_new_changes_ctes_sql, _render_ansi_typed_scalar, _render_typed_value_list, _typed_collection_items, @@ -3285,12 +3286,6 @@ def _historical_check_new_changes_cte_sql( latest_join_condition: str = cls._snapshot_key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) - latest_ordered_join_condition: str = cls._snapshot_key_condition( - left_alias="__ordered", right_alias="__latest", unique_key=unique_key - ) - reappearing_partition_sql: str = ", ".join( - f"__ordered.{column}" for column in unique_key - ) latest_change_condition: str = " OR ".join( f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" for column in check_columns @@ -3302,6 +3297,18 @@ def _historical_check_new_changes_cte_sql( observed_at_column=observed_at_column, valid_to_column=valid_to_column, ) + reappearing_new_changes_sql: str = _historical_reappearing_new_changes_ctes_sql( + changed_or_new_sql=( + "SELECT __delta_changes.* FROM __delta_changes " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{first_key} IS NULL " + f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " + f"AND ({latest_change_condition}))" + ), + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_to_column=valid_to_column, + ) return ( "__ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -3313,22 +3320,7 @@ def _historical_check_new_changes_cte_sql( f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - " UNION DISTINCT " - "SELECT __ordered.* FROM __ordered " - f"JOIN __latest ON {latest_ordered_join_condition} " - f"WHERE __latest.{valid_to_column} IS NOT NULL " - f"AND __ordered.{observed_at_column} > __latest.{valid_to_column} " - "QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {reappearing_partition_sql} " - f"ORDER BY __ordered.{observed_at_column}" - ") = 1" - "), __hard_deletes AS (" + f"), {reappearing_new_changes_sql}, __hard_deletes AS (" f"{hard_deletes_sql}" ")" ) diff --git a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py index 743883bbc2..cfc3ce8136 100644 --- a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py +++ b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py @@ -15,6 +15,7 @@ from sqlbuild.adapter.contract.classes.base_adapter import ( BaseAdapter, _encode_typed_json, + _historical_reappearing_new_changes_ctes_sql, _render_ansi_typed_scalar, _render_typed_value_list, ) @@ -3420,12 +3421,6 @@ def _historical_check_new_changes_cte_sql( latest_join_condition: str = cls._snapshot_key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) - latest_ordered_join_condition: str = cls._snapshot_key_condition( - left_alias="__ordered", right_alias="__latest", unique_key=unique_key - ) - reappearing_partition_sql: str = ", ".join( - f"__ordered.{column}" for column in unique_key - ) latest_change_condition: str = " OR ".join( f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" for column in check_columns @@ -3437,6 +3432,18 @@ def _historical_check_new_changes_cte_sql( observed_at_column=observed_at_column, valid_to_column=valid_to_column, ) + reappearing_new_changes_sql: str = _historical_reappearing_new_changes_ctes_sql( + changed_or_new_sql=( + "SELECT __delta_changes.* FROM __delta_changes " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{first_key} IS NULL " + f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " + f"AND ({latest_change_condition}))" + ), + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_to_column=valid_to_column, + ) return ( "__ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -3448,22 +3455,7 @@ def _historical_check_new_changes_cte_sql( f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - " UNION DISTINCT " - "SELECT __ordered.* FROM __ordered " - f"JOIN __latest ON {latest_ordered_join_condition} " - f"WHERE __latest.{valid_to_column} IS NOT NULL " - f"AND __ordered.{observed_at_column} > __latest.{valid_to_column} " - "QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {reappearing_partition_sql} " - f"ORDER BY __ordered.{observed_at_column}" - ") = 1" - "), __hard_deletes AS (" + f"), {reappearing_new_changes_sql}, __hard_deletes AS (" f"{hard_deletes_sql}" ")" ) diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index cbf3550519..c450bf2778 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -14,6 +14,7 @@ from sqlbuild.adapter.contract.classes.base_adapter import ( BaseAdapter, _encode_typed_json, + _historical_reappearing_new_changes_ctes_sql, _render_ansi_typed_scalar, _render_typed_value_list, ) @@ -3634,12 +3635,6 @@ def _historical_check_new_changes_cte_sql( latest_join_condition: str = cls._snapshot_key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) - latest_ordered_join_condition: str = cls._snapshot_key_condition( - left_alias="__ordered", right_alias="__latest", unique_key=unique_key - ) - reappearing_partition_sql: str = ", ".join( - f"__ordered.{column}" for column in unique_key - ) latest_change_condition: str = " OR ".join( f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" for column in check_columns @@ -3651,6 +3646,18 @@ def _historical_check_new_changes_cte_sql( observed_at_column=observed_at_column, valid_to_column=valid_to_column, ) + reappearing_new_changes_sql: str = _historical_reappearing_new_changes_ctes_sql( + changed_or_new_sql=( + "SELECT __delta_changes.* FROM __delta_changes " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{first_key} IS NULL " + f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " + f"AND ({latest_change_condition}))" + ), + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_to_column=valid_to_column, + ) return ( "__ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -3662,22 +3669,7 @@ def _historical_check_new_changes_cte_sql( f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - " UNION DISTINCT " - "SELECT __ordered.* FROM __ordered " - f"JOIN __latest ON {latest_ordered_join_condition} " - f"WHERE __latest.{valid_to_column} IS NOT NULL " - f"AND __ordered.{observed_at_column} > __latest.{valid_to_column} " - "QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {reappearing_partition_sql} " - f"ORDER BY __ordered.{observed_at_column}" - ") = 1" - "), __hard_deletes AS (" + f"), {reappearing_new_changes_sql}, __hard_deletes AS (" f"{hard_deletes_sql}" ")" ) diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py index 4e79491358..da96dfe6d5 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py @@ -116,6 +116,15 @@ ), expected_history=((1, "basic", 1, 2), (1, "pro", 2, None)), ), + SnapshotExecutionScenario( + description="check change followed by an unchanged repeat in one build", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=( + ((1, "active", 1),), + ((1, "active", 1), (1, "paused", 2), (1, "paused", 3)), + ), + expected_history=((1, "active", 1, 2), (1, "paused", 2, None)), + ), ) _CURRENT_STATE_SCENARIOS: tuple[SnapshotExecutionScenario, ...] = ( From 4f60e4ee70f2bc026cca0d7a07a959b40393876e Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:10:43 +0100 Subject: [PATCH 41/57] fix(snapshots): apply every hard delete within one incremental historical build Historical check and timestamp snapshots with invalidate_hard_deletes could only close the currently active version during an incremental build. Versions inserted by the same build ended at the next change only, so a key that was deleted and reappeared more than once in one build kept a version spanning the later deletions, and a check value that changed and changed back in one build lost the return to the earlier value. The final history was only correct after further builds. All adapters (DuckDB, MotherDuck, Postgres, SQL Server, BigQuery, Snowflake and Databricks) now render these paths from one shared, portable builder that uses window functions and joins only: no QUALIFY, no UNION DISTINCT and no correlated subqueries. Every new version now ends at the next version start or at the first later observation that is missing its key, on the initial and the incremental path alike, so one build produces the final history immediately. Initial full-history builds produce the same history as before. Existing history rows are not rewritten. (cherry picked from commit d6e86ba04b712815256006e10ded01f31295d29a) --- .../adapter/contract/classes/base_adapter.py | 282 ++++----------- .../contract/classes/duckdb_backed_adapter.py | 230 ++++--------- .../classes/historical_snapshot_sql.py | 278 +++++++++++++++ src/sqlbuild/adapter/contract/constants.py | 2 + .../bigquery/classes/bigquery_adapter.py | 320 ++++-------------- .../databricks/classes/databricks_adapter.py | 318 ++++------------- .../postgres/classes/postgres_adapter.py | 116 +++---- .../snowflake/classes/snowflake_adapter.py | 320 ++++-------------- .../sqlserver/classes/sqlserver_adapter.py | 151 ++++----- .../test_duckdb_execution.py | 40 +++ .../snapshot_sql_rendering/test_rendering.py | 150 ++++---- 11 files changed, 814 insertions(+), 1393 deletions(-) create mode 100644 src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py diff --git a/src/sqlbuild/adapter/contract/classes/base_adapter.py b/src/sqlbuild/adapter/contract/classes/base_adapter.py index eeedd63d45..e041af0d5f 100644 --- a/src/sqlbuild/adapter/contract/classes/base_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/base_adapter.py @@ -9,11 +9,16 @@ from pathlib import Path from typing import Any, ClassVar, cast +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + HistoricalSnapshotSql, + historical_insert_validity_sql, +) from sqlbuild.adapter.contract.classes.retention_adapter import RetentionAdapterMixin from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.strict_adapter import StrictAdapter from sqlbuild.adapter.contract.constants import ( INTEGER_TYPE_TOKEN, + SNAPSHOT_VERSION_START_COLUMN, TYPED_OBJECT_ENTRY_PART_COUNT, ) from sqlbuild.adapter.contract.exceptions import ( @@ -1088,7 +1093,7 @@ def render_apply_historical_timestamp_snapshot_changes( unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - change_time_column=updated_at_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: close_sql = ( @@ -1108,13 +1113,16 @@ def render_apply_historical_timestamp_snapshot_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{updated_at_column}, LEAD(__new_changes.{updated_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() insert_sql: str = ( f"WITH {new_changes_sql} " f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " + f"SELECT {output_select_sql}, {version_columns_sql} " f"FROM __new_changes" ) return (close_sql, insert_sql) @@ -1227,7 +1235,7 @@ def render_apply_historical_check_snapshot_changes( unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - change_time_column=observed_at_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: close_sql = ( @@ -1247,13 +1255,16 @@ def render_apply_historical_check_snapshot_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{observed_at_column}, LEAD(__new_changes.{observed_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() insert_sql: str = ( f"WITH {new_changes_sql} " f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{observed_at_column}, " - f"LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column}" - f") " + f"SELECT {output_select_sql}, {version_columns_sql} " f"FROM __new_changes" ) return (close_sql, insert_sql) @@ -2605,6 +2616,16 @@ def _historical_check_snapshot_select_sql( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + ).initial_select_sql(output_columns=output_columns) + partition_sql: str = ", ".join(unique_key) previous_columns_sql: str = ", ".join( f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " @@ -2617,42 +2638,6 @@ def _historical_check_snapshot_select_sql( f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns ) output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__changes", - ) - group_sequence_sql: str = _historical_observed_group_sequence_cte_sql( - origin=origin, observed_at_column=observed_at_column - ) - grouped_origin_sql: str = _historical_observed_group_sequence_from_sql( - origin=origin, observed_at_column=observed_at_column - ) - return ( - f"WITH {group_sequence_sql}, __ordered AS (" - "SELECT __source.*, __observed_group_sequence.__prev_group_observed_at, " - f"LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} {grouped_origin_sql}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition}) " - "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at" - "), __versions AS (" - f"SELECT __changes.*, LEAD({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __next_change_at, {hard_deleted_at_sql} AS __hard_deleted_at " - "FROM __changes" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) return ( "WITH __ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -2679,37 +2664,18 @@ def _historical_timestamp_snapshot_select_sql( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) if invalidate_hard_deletes: - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( + return HistoricalSnapshotSql( origin=origin, unique_key=unique_key, observed_at_column=observed_at_column, - row_alias="__changes", - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __versions AS (" - f"SELECT __changes.*, LEAD({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column}" - f") AS __next_change_at, {hard_deleted_at_sql} AS __hard_deleted_at " - "FROM __changes" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + ).initial_select_sql(output_columns=output_columns) + + partition_sql: str = ", ".join(unique_key) + output_select_sql: str = ", ".join(column for column in output_columns) return ( "WITH __ordered AS (" f"SELECT *, LAG({updated_at_column}) OVER (" @@ -2737,42 +2703,19 @@ def _historical_timestamp_new_changes_cte_sql( valid_to_column: str, invalidate_hard_deletes: bool, ) -> str: - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] if invalidate_hard_deletes: - latest_join_condition: str = _snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( + return HistoricalSnapshotSql( origin=origin, unique_key=unique_key, observed_at_column=observed_at_column, - row_alias="__target", - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - "), __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ")" - ) - latest_join_condition = _snapshot_key_condition( + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + ).new_changes_ctes_sql(destination=destination) + + partition_sql: str = ", ".join(unique_key) + first_key: str = unique_key[0] + latest_join_condition: str = _snapshot_key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) return ( @@ -2853,6 +2796,16 @@ def _historical_check_new_changes_cte_sql( valid_to_column: str, invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + ).new_changes_ctes_sql(destination=destination) + partition_sql: str = ", ".join(unique_key) previous_columns_sql: str = ", ".join( f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " @@ -2874,42 +2827,6 @@ def _historical_check_new_changes_cte_sql( changed_or_first_sql: str = ( f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({delta_change_condition})" ) - if invalidate_hard_deletes: - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - reappearing_new_changes_sql: str = _historical_reappearing_new_changes_ctes_sql( - changed_or_new_sql=( - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - ), - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - f"), {reappearing_new_changes_sql}, __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ")" - ) return ( "__ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -2938,89 +2855,6 @@ def _snapshot_key_condition( return " AND ".join(f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key) -def _historical_hard_deleted_at_sql( - *, origin: str, unique_key: tuple[str, ...], observed_at_column: str, row_alias: str -) -> str: - present_condition: str = _snapshot_key_condition( - left_alias="__present", right_alias=row_alias, unique_key=unique_key - ) - return ( - "(SELECT MIN(__observed_groups.__observed_at) " - f"FROM (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"WHERE __observed_groups.__observed_at > {row_alias}.{observed_at_column} " - "AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __present " - f"WHERE __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition}" - "))" - ) - - -def _historical_observed_group_sequence_cte_sql(*, origin: str, observed_at_column: str) -> str: - return ( - "__observed_group_sequence AS (" - "SELECT __observed_at, LAG(__observed_at) OVER (ORDER BY __observed_at) " - "AS __prev_group_observed_at " - f"FROM (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_group_values" - ")" - ) - - -def _historical_observed_group_sequence_from_sql(*, origin: str, observed_at_column: str) -> str: - return ( - f"FROM {origin} AS __source LEFT JOIN __observed_group_sequence " - f"ON __observed_group_sequence.__observed_at = __source.{observed_at_column}" - ) - - -def _historical_reappearing_new_changes_ctes_sql( - *, - changed_or_new_sql: str, - unique_key: tuple[str, ...], - observed_at_column: str, - valid_to_column: str, -) -> str: - """Render ``__new_changes`` plus the first observation after a hard-deleted latest version.""" - - ordered_key_sql: str = ", ".join(f"__ordered.{column}" for column in unique_key) - latest_condition: str = _snapshot_key_condition( - left_alias="__ordered", right_alias="__latest", unique_key=unique_key - ) - present_at_close_condition: str = _snapshot_key_condition( - left_alias="__present", right_alias="__latest", unique_key=unique_key - ) - reappearance_condition: str = _snapshot_key_condition( - left_alias="__ordered", right_alias="__reappearances", unique_key=unique_key - ) - already_changed_condition: str = _snapshot_key_condition( - left_alias="__changed_or_new", right_alias="__ordered", unique_key=unique_key - ) - return ( - f"__changed_or_new AS ({changed_or_new_sql}), __reappearances AS (" - f"SELECT {ordered_key_sql}, MIN(__ordered.{observed_at_column}) AS __reappeared_at " - f"FROM __ordered JOIN __latest ON {latest_condition} " - f"WHERE __latest.{valid_to_column} IS NOT NULL " - f"AND __ordered.{observed_at_column} > __latest.{valid_to_column} " - "AND NOT EXISTS (" - f"SELECT 1 FROM __ordered AS __present WHERE {present_at_close_condition} " - f"AND __present.{observed_at_column} = __latest.{valid_to_column}" - ") " - f"GROUP BY {ordered_key_sql}" - "), __new_changes AS (" - "SELECT * FROM __changed_or_new " - "UNION ALL " - "SELECT __ordered.* FROM __ordered " - f"JOIN __reappearances ON {reappearance_condition} " - f"AND __ordered.{observed_at_column} = __reappearances.__reappeared_at " - "WHERE NOT EXISTS (" - f"SELECT 1 FROM __changed_or_new WHERE {already_changed_condition} " - f"AND __changed_or_new.{observed_at_column} = __ordered.{observed_at_column}" - "))" - ) - - def _historical_snapshot_combined_close_sql( *, destination: str, diff --git a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py index 9fd95d9708..fd61347a66 100644 --- a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py @@ -16,12 +16,13 @@ from sqlbuild.adapter.contract.classes.base_adapter import ( BaseAdapter, _encode_typed_json, - _historical_observed_group_sequence_cte_sql, - _historical_observed_group_sequence_from_sql, - _historical_reappearing_new_changes_ctes_sql, _render_ansi_typed_scalar, _render_typed_value_list, ) +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + HistoricalSnapshotSql, + historical_insert_validity_sql, +) from sqlbuild.adapter.contract.classes.observed_connection import ObservedConnection from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin @@ -29,6 +30,7 @@ DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE, QUALIFIED_NAME_SEPARATOR, + SNAPSHOT_VERSION_START_COLUMN, ) from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapter.contract.models import ( @@ -571,7 +573,7 @@ def render_apply_historical_timestamp_snapshot_changes( unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - change_time_column=updated_at_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: close_sql = ( @@ -591,13 +593,16 @@ def render_apply_historical_timestamp_snapshot_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{updated_at_column}, LEAD(__new_changes.{updated_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() insert_sql: str = ( f"WITH {new_changes_sql} " f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " + f"SELECT {output_select_sql}, {version_columns_sql} " f"FROM __new_changes" ) return (close_sql, insert_sql) @@ -786,7 +791,7 @@ def render_apply_historical_check_snapshot_changes( unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - change_time_column=observed_at_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: close_sql = ( @@ -806,13 +811,16 @@ def render_apply_historical_check_snapshot_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{observed_at_column}, LEAD(__new_changes.{observed_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() insert_sql: str = ( f"WITH {new_changes_sql} " f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{observed_at_column}, " - f"LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column}" - f") " + f"SELECT {output_select_sql}, {version_columns_sql} " f"FROM __new_changes" ) return (close_sql, insert_sql) @@ -2616,6 +2624,16 @@ def _historical_check_snapshot_select_sql( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + ).initial_select_sql(output_columns=output_columns) + partition_sql: str = ", ".join(unique_key) previous_columns_sql: str = ", ".join( f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " @@ -2628,42 +2646,6 @@ def _historical_check_snapshot_select_sql( f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns ) output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_deleted_at_sql: str = cls._historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__changes", - ) - group_sequence_sql: str = _historical_observed_group_sequence_cte_sql( - origin=origin, observed_at_column=observed_at_column - ) - grouped_origin_sql: str = _historical_observed_group_sequence_from_sql( - origin=origin, observed_at_column=observed_at_column - ) - return ( - f"WITH {group_sequence_sql}, __ordered AS (" - "SELECT __source.*, __observed_group_sequence.__prev_group_observed_at, " - f"LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} {grouped_origin_sql}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition}) " - "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at" - "), __versions AS (" - f"SELECT __changes.*, LEAD({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __next_change_at, {hard_deleted_at_sql} AS __hard_deleted_at " - "FROM __changes" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) return ( "WITH __ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -2691,37 +2673,18 @@ def _historical_timestamp_snapshot_select_sql( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) if invalidate_hard_deletes: - hard_deleted_at_sql: str = cls._historical_hard_deleted_at_sql( + return HistoricalSnapshotSql( origin=origin, unique_key=unique_key, observed_at_column=observed_at_column, - row_alias="__changes", - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __versions AS (" - f"SELECT __changes.*, LEAD({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column}" - f") AS __next_change_at, {hard_deleted_at_sql} AS __hard_deleted_at " - "FROM __changes" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + ).initial_select_sql(output_columns=output_columns) + + partition_sql: str = ", ".join(unique_key) + output_select_sql: str = ", ".join(column for column in output_columns) return ( "WITH __ordered AS (" f"SELECT *, LAG({updated_at_column}) OVER (" @@ -2750,44 +2713,21 @@ def _historical_timestamp_new_changes_cte_sql( valid_to_column: str, invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + ).new_changes_ctes_sql(destination=destination) + partition_sql: str = ", ".join(unique_key) latest_join_condition: str = cls._snapshot_key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) first_key: str = unique_key[0] - if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - hard_deleted_at_sql: str = cls._historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - "), __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ")" - ) return ( "__ordered AS (" f"SELECT *, LAG({updated_at_column}) OVER (" @@ -2869,6 +2809,16 @@ def _historical_check_new_changes_cte_sql( valid_to_column: str, invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + ).new_changes_ctes_sql(destination=destination) + partition_sql: str = ", ".join(unique_key) previous_columns_sql: str = ", ".join( f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " @@ -2892,49 +2842,6 @@ def _historical_check_new_changes_cte_sql( "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " f"OR ({delta_change_condition})" ) - if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - hard_deleted_at_sql: str = cls._historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - reappearing_new_changes_sql: str = _historical_reappearing_new_changes_ctes_sql( - changed_or_new_sql=( - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - ), - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - f"), {reappearing_new_changes_sql}, __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ")" - ) return ( "__ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -2964,25 +2871,6 @@ def _snapshot_key_condition( f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key ) - @classmethod - def _historical_hard_deleted_at_sql( - cls, *, origin: str, unique_key: tuple[str, ...], observed_at_column: str, row_alias: str - ) -> str: - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias=row_alias, unique_key=unique_key - ) - return ( - "(SELECT MIN(__observed_groups.__observed_at) " - f"FROM (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"WHERE __observed_groups.__observed_at > {row_alias}.{observed_at_column} " - "AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __present " - f"WHERE __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition}" - "))" - ) - @classmethod def _historical_snapshot_combined_close_sql( cls, diff --git a/src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py new file mode 100644 index 0000000000..d2feaeb736 --- /dev/null +++ b/src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py @@ -0,0 +1,278 @@ +"""Portable SQL for historical snapshots that invalidate hard deletes.""" + +from __future__ import annotations + +from collections.abc import Callable + +from sqlbuild.adapter.contract.constants import SNAPSHOT_VERSION_START_COLUMN + + +def render_is_distinct_from(*, left: str, right: str) -> str: + """Render a null-safe inequality with the SQL standard predicate.""" + + return f"{left} IS DISTINCT FROM {right}" + + +def historical_insert_validity_sql() -> str: + """Render ``valid_from, valid_to`` expressions selected from ``__new_changes``.""" + + version_end_sql: str = _version_end_sql(alias="__new_changes") + return f"__new_changes.{SNAPSHOT_VERSION_START_COLUMN}, {version_end_sql}" + + +def _version_end_sql(*, alias: str | None) -> str: + prefix: str = f"{alias}." if alias is not None else "" + next_start: str = f"{prefix}__next_version_start" + next_absence: str = f"{prefix}__next_absence_at" + return ( + f"CASE WHEN {next_start} IS NULL THEN {next_absence} " + f"WHEN {next_absence} IS NULL THEN {next_start} " + f"WHEN {next_absence} < {next_start} THEN {next_absence} " + f"ELSE {next_start} END" + ) + + +class HistoricalSnapshotSql: + """Render hard-delete historical snapshot SQL without QUALIFY or correlated subqueries.""" + + def __init__( + self, + *, + origin: str, + unique_key: tuple[str, ...], + observed_at_column: str, + valid_from_column: str, + valid_to_column: str, + check_columns: tuple[str, ...] = (), + updated_at_column: str | None = None, + distinct_condition: Callable[..., str] = render_is_distinct_from, + ) -> None: + self.origin: str = origin + self.unique_key: tuple[str, ...] = unique_key + self.observed_at_column: str = observed_at_column + self.valid_from_column: str = valid_from_column + self.valid_to_column: str = valid_to_column + self.check_columns: tuple[str, ...] = check_columns + self.updated_at_column: str | None = updated_at_column + self.distinct_condition: Callable[..., str] = distinct_condition + + def initial_select_sql(self, *, output_columns: tuple[str, ...]) -> str: + """Render the complete initial-build query.""" + + return ( + f"WITH {self.initial_ctes_sql()} " + f"SELECT {self.initial_select_list_sql(output_columns=output_columns)} " + "FROM __versions" + ) + + def initial_ctes_sql(self) -> str: + """Render the initial-build CTE list ending in ``__versions``.""" + + partition_sql: str = self._key_sql(alias=None) + return ( + f"{self._tracked_source_ctes_sql()}, __changes AS (" + f"SELECT *, {self._initial_version_start_sql()} AS {SNAPSHOT_VERSION_START_COLUMN} " + "FROM __tracked " + f"WHERE __is_change = 1 OR {self._initial_reappearance_sql()}" + "), __versions AS (" + f"SELECT *, LEAD({SNAPSHOT_VERSION_START_COLUMN}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY {SNAPSHOT_VERSION_START_COLUMN}" + ") AS __next_version_start FROM __changes" + ")" + ) + + def initial_select_list_sql(self, *, output_columns: tuple[str, ...]) -> str: + """Render output columns and validity columns selected from ``__versions``.""" + + output_sql: str = ", ".join(output_columns) + return ( + f"{output_sql}, {SNAPSHOT_VERSION_START_COLUMN} AS {self.valid_from_column}, " + f"{_version_end_sql(alias=None)} AS {self.valid_to_column}" + ) + + def new_changes_ctes_sql(self, *, destination: str) -> str: + """Render the incremental CTE list defining ``__new_changes`` and ``__hard_deletes``.""" + + partition_sql: str = self._key_sql(alias=None) + first_key: str = self.unique_key[0] + latest_join_condition: str = self._key_condition(left="__tracked", right="__latest") + return ( + f"{self._tracked_source_ctes_sql()}, __latest AS (" + "SELECT * FROM (" + f"SELECT *, ROW_NUMBER() OVER (PARTITION BY {partition_sql} " + f"ORDER BY {self.valid_from_column} DESC) AS __latest_rn FROM {destination}" + ") AS __latest_ranked WHERE __latest_rn = 1" + "), __classified AS (" + "SELECT __tracked.*, " + f"CASE WHEN __latest.{first_key} IS NULL THEN 0 ELSE 1 END AS __has_latest, " + f"__latest.{self.valid_from_column} AS __latest_valid_from, " + f"{self._differs_from_latest_sql()} AS __differs_from_latest, " + f"{self._incremental_reappearance_sql()} AS __is_reappearance " + f"FROM __tracked LEFT JOIN __latest ON {latest_join_condition}" + "), __changed_or_new AS (" + f"SELECT __classified.*, {self._incremental_version_start_sql()} " + f"AS {SNAPSHOT_VERSION_START_COLUMN} FROM __classified " + f"WHERE {self._incremental_start_condition_sql()}" + f"), {self._new_change_starts_sql()}, __new_changes AS (" + f"SELECT *, LEAD({SNAPSHOT_VERSION_START_COLUMN}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY {SNAPSHOT_VERSION_START_COLUMN}" + ") AS __next_version_start FROM __new_change_starts" + f"), __hard_deletes AS ({self._open_version_hard_deletes_sql(destination=destination)})" + ) + + def _tracked_source_ctes_sql(self) -> str: + observed_at: str = self.observed_at_column + partition_sql: str = self._key_sql(alias=None) + window_sql: str = f"OVER (PARTITION BY {partition_sql} ORDER BY {observed_at})" + return ( + "__observed_group_sequence AS (" + "SELECT __observed_at, " + "LAG(__observed_at) OVER (ORDER BY __observed_at) AS __prev_group_observed_at, " + "LEAD(__observed_at) OVER (ORDER BY __observed_at) AS __next_group_observed_at " + f"FROM (SELECT DISTINCT {observed_at} AS __observed_at FROM {self.origin}) " + "AS __observed_group_values" + "), __ordered AS (" + "SELECT __source.*, __observed_group_sequence.__prev_group_observed_at, " + "__observed_group_sequence.__next_group_observed_at, " + f"LAG({observed_at}) {window_sql} AS __prev_observed_at, " + f"LEAD({observed_at}) {window_sql} AS __next_observed_at" + f"{self._previous_values_sql(window_sql=window_sql)} " + f"FROM {self.origin} AS __source LEFT JOIN __observed_group_sequence " + f"ON __observed_group_sequence.__observed_at = __source.{observed_at}" + "), __flagged AS (" + "SELECT *, " + "CASE WHEN __next_group_observed_at IS NOT NULL AND (__next_observed_at IS NULL " + "OR __next_observed_at <> __next_group_observed_at) " + "THEN __next_group_observed_at END AS __absent_after, " + f"CASE WHEN {self._change_condition_sql()} THEN 1 ELSE 0 END AS __is_change " + "FROM __ordered" + "), __tracked AS (" + "SELECT *, MIN(__absent_after) OVER (" + f"PARTITION BY {partition_sql} ORDER BY {observed_at} " + "ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING" + ") AS __next_absence_at FROM __flagged" + ")" + ) + + def _previous_values_sql(self, *, window_sql: str) -> str: + if self.updated_at_column is not None: + return f", LAG({self.updated_at_column}) {window_sql} AS __prev_updated_at" + return "".join( + f", LAG({column}) {window_sql} AS __prev_{column}" for column in self.check_columns + ) + + def _change_condition_sql(self) -> str: + if self.updated_at_column is not None: + return "__prev_updated_at IS NULL OR " + self.distinct_condition( + left=self.updated_at_column, right="__prev_updated_at" + ) + changed_sql: str = " OR ".join( + self.distinct_condition(left=column, right=f"__prev_{column}") + for column in self.check_columns + ) + return f"__prev_observed_at IS NULL OR ({changed_sql})" + + def _initial_reappearance_sql(self) -> str: + if self.updated_at_column is not None: + return "1 = 0" + return "(__prev_observed_at IS NOT NULL AND __prev_observed_at <> __prev_group_observed_at)" + + def _initial_version_start_sql(self) -> str: + return self.updated_at_column or self.observed_at_column + + def _incremental_reappearance_sql(self) -> str: + if self.updated_at_column is not None: + return "0" + return ( + "CASE WHEN __tracked.__prev_group_observed_at IS NOT NULL AND (" + f"(__tracked.__prev_observed_at IS NULL AND __latest.{self.unique_key[0]} IS NOT NULL) " + "OR __tracked.__prev_observed_at <> __tracked.__prev_group_observed_at" + ") THEN 1 ELSE 0 END" + ) + + def _incremental_version_start_sql(self) -> str: + return self.updated_at_column or self.observed_at_column + + def _differs_from_latest_sql(self) -> str: + if not self.check_columns: + return "0" + differs_sql: str = " OR ".join( + self.distinct_condition(left=f"__tracked.{column}", right=f"__latest.{column}") + for column in self.check_columns + ) + return f"CASE WHEN {differs_sql} THEN 1 ELSE 0 END" + + def _incremental_start_condition_sql(self) -> str: + return ( + "(__has_latest = 0 AND (__is_change = 1 OR __is_reappearance = 1)) " + "OR (__has_latest = 1 AND __is_reappearance = 1 " + f"AND {self.observed_at_column} > __latest_valid_from) " + "OR (__has_latest = 1 AND __is_reappearance = 0 AND __is_change = 1 " + f"AND {self._latest_change_condition_sql()})" + ) + + def _latest_change_condition_sql(self) -> str: + if self.updated_at_column is not None: + return f"{self.updated_at_column} > __latest_valid_from" + return ( + f"{self.observed_at_column} > __latest_valid_from " + "AND (__prev_observed_at >= __latest_valid_from OR __differs_from_latest = 1)" + ) + + def _new_change_starts_sql(self) -> str: + if self.updated_at_column is not None: + return "__new_change_starts AS (SELECT * FROM __changed_or_new)" + observed_at: str = self.observed_at_column + classified_key_sql: str = self._key_sql(alias="__classified") + return ( + "__first_new_starts AS (" + f"SELECT {self._key_sql(alias=None)}, MIN({SNAPSHOT_VERSION_START_COLUMN}) " + f"AS __first_version_start FROM __changed_or_new GROUP BY {self._key_sql(alias=None)}" + "), __reappearances AS (" + f"SELECT {classified_key_sql}, MIN(__classified.{observed_at}) AS __reappeared_at " + "FROM __classified " + f"JOIN __latest ON {self._key_condition(left='__classified', right='__latest')} " + "LEFT JOIN __observed_group_sequence AS __closing_group " + f"ON __closing_group.__observed_at = __latest.{self.valid_to_column} " + f"WHERE __latest.{self.valid_to_column} IS NOT NULL " + "AND __closing_group.__observed_at IS NULL " + f"AND __classified.{observed_at} > __latest.{self.valid_to_column} " + f"GROUP BY {classified_key_sql}" + "), __new_change_starts AS (" + "SELECT * FROM __changed_or_new " + "UNION ALL " + f"SELECT __classified.*, __classified.{observed_at} AS {SNAPSHOT_VERSION_START_COLUMN} " + "FROM __classified " + "JOIN __reappearances ON " + f"{self._key_condition(left='__classified', right='__reappearances')} " + f"AND __classified.{observed_at} = __reappearances.__reappeared_at " + "LEFT JOIN __first_new_starts ON " + f"{self._key_condition(left='__classified', right='__first_new_starts')} " + f"WHERE __first_new_starts.{self.unique_key[0]} IS NULL " + f"OR __classified.{observed_at} < __first_new_starts.__first_version_start" + ")" + ) + + def _open_version_hard_deletes_sql(self, *, destination: str) -> str: + target_key_sql: str = self._key_sql(alias="__target") + present_condition: str = self._key_condition(left="__present", right="__target") + observed_at: str = self.observed_at_column + return ( + f"SELECT {target_key_sql}, MIN(__observed_groups.__observed_at) AS __close_at " + f"FROM {destination} AS __target " + "JOIN __observed_group_sequence AS __observed_groups " + f"ON __observed_groups.__observed_at > __target.{observed_at} " + f"LEFT JOIN {self.origin} AS __present " + f"ON __present.{observed_at} = __observed_groups.__observed_at " + f"AND {present_condition} " + f"WHERE __target.{self.valid_to_column} IS NULL " + f"AND __present.{self.unique_key[0]} IS NULL " + f"GROUP BY {target_key_sql}" + ) + + def _key_sql(self, *, alias: str | None) -> str: + prefix: str = f"{alias}." if alias is not None else "" + return ", ".join(f"{prefix}{column}" for column in self.unique_key) + + def _key_condition(self, *, left: str, right: str) -> str: + return " AND ".join(f"{left}.{column} = {right}.{column}" for column in self.unique_key) diff --git a/src/sqlbuild/adapter/contract/constants.py b/src/sqlbuild/adapter/contract/constants.py index 1eb7331809..f4f7f1e53a 100644 --- a/src/sqlbuild/adapter/contract/constants.py +++ b/src/sqlbuild/adapter/contract/constants.py @@ -27,3 +27,5 @@ QUALIFIED_NAME_SEPARATOR: str = "." PYTHON_INIT_FILE_NAME: str = "__init__.py" PYTHON_IDENTIFIER_REPLACEMENT: str = "_" + +SNAPSHOT_VERSION_START_COLUMN: str = "__version_start" diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py index 6c3532a8dd..92e1eeef50 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py @@ -13,16 +13,23 @@ _COLLECTION_SQL_VALUE_KINDS, BaseAdapter, _encode_typed_json, - _historical_reappearing_new_changes_ctes_sql, _render_ansi_typed_scalar, _render_typed_value_list, _typed_collection_items, _typed_scalar_payload, ) +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + HistoricalSnapshotSql, + historical_insert_validity_sql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin -from sqlbuild.adapter.contract.constants import DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE +from sqlbuild.adapter.contract.constants import ( + DIFF_LEFT_SIDE, + DIFF_RIGHT_SIDE, + SNAPSHOT_VERSION_START_COLUMN, +) from sqlbuild.adapter.contract.exceptions import ( AdapterUserError, UnsupportedTypedSqlRenderingError, @@ -1154,7 +1161,7 @@ def render_apply_historical_timestamp_snapshot_changes( unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - change_time_column=updated_at_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: close_sql = self._historical_snapshot_close_sql( @@ -1171,16 +1178,17 @@ def render_apply_historical_timestamp_snapshot_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{updated_at_column}, LEAD(__new_changes.{updated_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() insert_sql: str = self._historical_snapshot_insert_sql( destination=destination, insert_column_sql=insert_column_sql, new_changes_sql=new_changes_sql, - select_sql=( - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") FROM __new_changes" - ), + select_sql=(f"SELECT {output_select_sql}, {version_columns_sql} FROM __new_changes"), ) return (close_sql, insert_sql) @@ -1359,7 +1367,7 @@ def render_apply_historical_check_snapshot_changes( unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - change_time_column=observed_at_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: close_sql = self._historical_snapshot_close_sql( @@ -1376,16 +1384,17 @@ def render_apply_historical_check_snapshot_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{observed_at_column}, LEAD(__new_changes.{observed_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() insert_sql: str = self._historical_snapshot_insert_sql( destination=destination, insert_column_sql=insert_column_sql, new_changes_sql=new_changes_sql, - select_sql=( - f"SELECT {output_select_sql}, __new_changes.{observed_at_column}, " - f"LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column}" - f") FROM __new_changes" - ), + select_sql=(f"SELECT {output_select_sql}, {version_columns_sql} FROM __new_changes"), ) return (close_sql, insert_sql) @@ -2955,6 +2964,16 @@ def _historical_check_snapshot_select_sql( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + ).initial_select_sql(output_columns=output_columns) + partition_sql: str = ", ".join(unique_key) previous_columns_sql: str = ", ".join( f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " @@ -2967,68 +2986,6 @@ def _historical_check_snapshot_select_sql( f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns ) output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_delete_join_condition: str = cls._snapshot_key_condition( - left_alias="__hard_delete_candidates", - right_alias="__changes", - unique_key=unique_key, - ) - hard_delete_key_sql: str = ", ".join(f"__changes.{column}" for column in unique_key) - hard_delete_group_sql: str = ", ".join( - [ - *(f"__changes.{column}" for column in unique_key), - f"__changes.{observed_at_column}", - ] - ) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__changes", unique_key=unique_key - ) - return ( - "WITH __ordered AS (" - "SELECT __source.*, " - f"(SELECT MAX(__groups.__observed_at) FROM (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - ") AS __groups " - f"WHERE __groups.__observed_at < __source.{observed_at_column}" - ") AS __prev_group_observed_at, " - f"LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin} AS __source" - "), __changes AS (" - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({change_condition}) " - "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at" - "), __observed_groups AS (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - "), __hard_delete_candidates AS (" - f"SELECT {hard_delete_key_sql}, __changes.{observed_at_column}, " - "MIN(__observed_groups.__observed_at) AS __hard_deleted_at " - "FROM __changes " - "JOIN __observed_groups " - f"ON __observed_groups.__observed_at > __changes.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __present.{unique_key[0]} IS NULL " - f"GROUP BY {hard_delete_group_sql}" - "), __versions AS (" - f"SELECT __changes.*, LEAD(__changes.{observed_at_column}) OVER (" - f"PARTITION BY {', '.join(f'__changes.{column}' for column in unique_key)} " - f"ORDER BY __changes.{observed_at_column}" - ") AS __next_change_at, __hard_delete_candidates.__hard_deleted_at " - "FROM __changes LEFT JOIN __hard_delete_candidates " - f"ON {hard_delete_join_condition} " - f"AND __hard_delete_candidates.{observed_at_column} = " - f"__changes.{observed_at_column}" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) return ( "WITH __ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -3056,63 +3013,18 @@ def _historical_timestamp_snapshot_select_sql( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) if invalidate_hard_deletes: - hard_delete_join_condition: str = cls._snapshot_key_condition( - left_alias="__hard_delete_candidates", - right_alias="__changes", + return HistoricalSnapshotSql( + origin=origin, unique_key=unique_key, - ) - hard_delete_key_sql: str = ", ".join(f"__changes.{column}" for column in unique_key) - hard_delete_group_sql: str = ", ".join( - [ - *(f"__changes.{column}" for column in unique_key), - f"__changes.{observed_at_column}", - ] - ) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__changes", unique_key=unique_key - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __observed_groups AS (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - "), __hard_delete_candidates AS (" - f"SELECT {hard_delete_key_sql}, __changes.{observed_at_column}, " - "MIN(__observed_groups.__observed_at) AS __hard_deleted_at " - "FROM __changes " - "JOIN __observed_groups " - f"ON __observed_groups.__observed_at > __changes.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __present.{unique_key[0]} IS NULL " - f"GROUP BY {hard_delete_group_sql}" - "), __versions AS (" - f"SELECT __changes.*, LEAD(__changes.{updated_at_column}) OVER (" - f"PARTITION BY {', '.join(f'__changes.{column}' for column in unique_key)} " - f"ORDER BY __changes.{updated_at_column}" - ") AS __next_change_at, __hard_delete_candidates.__hard_deleted_at " - "FROM __changes LEFT JOIN __hard_delete_candidates " - f"ON {hard_delete_join_condition} " - f"AND __hard_delete_candidates.{observed_at_column} = " - f"__changes.{observed_at_column}" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + ).initial_select_sql(output_columns=output_columns) + + partition_sql: str = ", ".join(unique_key) + output_select_sql: str = ", ".join(column for column in output_columns) return ( "WITH __ordered AS (" f"SELECT *, LAG({updated_at_column}) OVER (" @@ -3141,43 +3053,21 @@ def _historical_timestamp_new_changes_cte_sql( valid_to_column: str, invalidate_hard_deletes: bool, ) -> str: - partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - first_key: str = unique_key[0] if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - hard_deletes_sql: str = cls._historical_hard_deletes_select_sql( - destination=destination, + return HistoricalSnapshotSql( origin=origin, unique_key=unique_key, observed_at_column=observed_at_column, + valid_from_column=valid_from_column, valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - "), __hard_deletes AS (" - f"{hard_deletes_sql}" - ")" - ) + updated_at_column=updated_at_column, + ).new_changes_ctes_sql(destination=destination) + + partition_sql: str = ", ".join(unique_key) + latest_join_condition: str = cls._snapshot_key_condition( + left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key + ) + first_key: str = unique_key[0] return ( "__ordered AS (" f"SELECT *, LAG({updated_at_column}) OVER (" @@ -3259,6 +3149,16 @@ def _historical_check_new_changes_cte_sql( valid_to_column: str, invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + ).new_changes_ctes_sql(destination=destination) + partition_sql: str = ", ".join(unique_key) previous_columns_sql: str = ", ".join( f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " @@ -3282,48 +3182,6 @@ def _historical_check_new_changes_cte_sql( "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " f"OR ({delta_change_condition})" ) - if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - hard_deletes_sql: str = cls._historical_hard_deletes_select_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - reappearing_new_changes_sql: str = _historical_reappearing_new_changes_ctes_sql( - changed_or_new_sql=( - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - ), - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - f"), {reappearing_new_changes_sql}, __hard_deletes AS (" - f"{hard_deletes_sql}" - ")" - ) return ( "__ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -3353,54 +3211,6 @@ def _snapshot_key_condition( f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key ) - @classmethod - def _historical_hard_deleted_at_sql( - cls, *, origin: str, unique_key: tuple[str, ...], observed_at_column: str, row_alias: str - ) -> str: - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias=row_alias, unique_key=unique_key - ) - return ( - "(SELECT MIN(__observed_groups.__observed_at) " - f"FROM (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"WHERE __observed_groups.__observed_at > {row_alias}.{observed_at_column} " - "AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __present " - f"WHERE __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition}" - "))" - ) - - @classmethod - def _historical_hard_deletes_select_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - observed_at_column: str, - valid_to_column: str, - ) -> str: - target_key_sql: str = ", ".join(f"__target.{column}" for column in unique_key) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"SELECT {target_key_sql}, MIN(__observed_groups.__observed_at) AS __close_at " - f"FROM {destination} AS __target " - f"JOIN (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"ON __observed_groups.__observed_at > __target.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __present.{first_key} IS NULL " - f"GROUP BY {target_key_sql}" - ) - @classmethod def _historical_snapshot_combined_close_sql( cls, diff --git a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py index cfc3ce8136..b08f81cede 100644 --- a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py +++ b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py @@ -15,14 +15,21 @@ from sqlbuild.adapter.contract.classes.base_adapter import ( BaseAdapter, _encode_typed_json, - _historical_reappearing_new_changes_ctes_sql, _render_ansi_typed_scalar, _render_typed_value_list, ) +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + HistoricalSnapshotSql, + historical_insert_validity_sql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin -from sqlbuild.adapter.contract.constants import DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE +from sqlbuild.adapter.contract.constants import ( + DIFF_LEFT_SIDE, + DIFF_RIGHT_SIDE, + SNAPSHOT_VERSION_START_COLUMN, +) from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapter.contract.main.normalize_seed_csv_value import normalize_seed_csv_value from sqlbuild.adapter.contract.models import ( @@ -1007,7 +1014,7 @@ def render_apply_historical_timestamp_snapshot_changes( unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - change_time_column=updated_at_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: close_sql = ( @@ -1027,13 +1034,16 @@ def render_apply_historical_timestamp_snapshot_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{updated_at_column}, LEAD(__new_changes.{updated_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() insert_sql: str = ( f"WITH {new_changes_sql} " f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " + f"SELECT {output_select_sql}, {version_columns_sql} " f"FROM __new_changes" ) return (close_sql, insert_sql) @@ -1220,7 +1230,7 @@ def render_apply_historical_check_snapshot_changes( unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - change_time_column=observed_at_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: close_sql = ( @@ -1240,13 +1250,16 @@ def render_apply_historical_check_snapshot_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{observed_at_column}, LEAD(__new_changes.{observed_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() insert_sql: str = ( f"WITH {new_changes_sql} " f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{observed_at_column}, " - f"LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column}" - f") " + f"SELECT {output_select_sql}, {version_columns_sql} " f"FROM __new_changes" ) return (close_sql, insert_sql) @@ -3088,6 +3101,16 @@ def _historical_check_snapshot_select_sql( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + ).initial_select_sql(output_columns=output_columns) + partition_sql: str = ", ".join(unique_key) previous_columns_sql: str = ", ".join( f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " @@ -3100,70 +3123,6 @@ def _historical_check_snapshot_select_sql( f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns ) output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_delete_join_condition: str = cls._snapshot_key_condition( - left_alias="__hard_delete_candidates", - right_alias="__changes", - unique_key=unique_key, - ) - hard_delete_key_sql: str = ", ".join(f"__changes.{column}" for column in unique_key) - hard_delete_group_sql: str = ", ".join( - [ - *(f"__changes.{column}" for column in unique_key), - f"__changes.{observed_at_column}", - ] - ) - source_group_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__changes", unique_key=unique_key - ) - return ( - "WITH __observed_groups AS (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - "), __source_with_prev_group AS (" - "SELECT __source.*, MAX(__observed_groups.__observed_at) " - f"AS __prev_group_observed_at FROM {origin} AS __source " - "LEFT JOIN __observed_groups " - f"ON __observed_groups.__observed_at < __source.{observed_at_column} " - f"GROUP BY {source_group_sql}" - "), __ordered AS (" - "SELECT *, " - f"LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM __source_with_prev_group" - "), __changes AS (" - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({change_condition}) " - "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at" - "), __hard_delete_candidates AS (" - f"SELECT {hard_delete_key_sql}, __changes.{observed_at_column}, " - "MIN(__observed_groups.__observed_at) AS __hard_deleted_at " - "FROM __changes " - "JOIN __observed_groups " - f"ON __observed_groups.__observed_at > __changes.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __present.{unique_key[0]} IS NULL " - f"GROUP BY {hard_delete_group_sql}" - "), __versions AS (" - f"SELECT __changes.*, LEAD(__changes.{observed_at_column}) OVER (" - f"PARTITION BY {', '.join(f'__changes.{column}' for column in unique_key)} " - f"ORDER BY __changes.{observed_at_column}" - ") AS __next_change_at, __hard_delete_candidates.__hard_deleted_at " - "FROM __changes LEFT JOIN __hard_delete_candidates " - f"ON {hard_delete_join_condition} " - f"AND __hard_delete_candidates.{observed_at_column} = " - f"__changes.{observed_at_column}" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) return ( "WITH __ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -3191,63 +3150,18 @@ def _historical_timestamp_snapshot_select_sql( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) if invalidate_hard_deletes: - hard_delete_join_condition: str = cls._snapshot_key_condition( - left_alias="__hard_delete_candidates", - right_alias="__changes", + return HistoricalSnapshotSql( + origin=origin, unique_key=unique_key, - ) - hard_delete_key_sql: str = ", ".join(f"__changes.{column}" for column in unique_key) - hard_delete_group_sql: str = ", ".join( - [ - *(f"__changes.{column}" for column in unique_key), - f"__changes.{observed_at_column}", - ] - ) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__changes", unique_key=unique_key - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __observed_groups AS (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - "), __hard_delete_candidates AS (" - f"SELECT {hard_delete_key_sql}, __changes.{observed_at_column}, " - "MIN(__observed_groups.__observed_at) AS __hard_deleted_at " - "FROM __changes " - "JOIN __observed_groups " - f"ON __observed_groups.__observed_at > __changes.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __present.{unique_key[0]} IS NULL " - f"GROUP BY {hard_delete_group_sql}" - "), __versions AS (" - f"SELECT __changes.*, LEAD(__changes.{updated_at_column}) OVER (" - f"PARTITION BY {', '.join(f'__changes.{column}' for column in unique_key)} " - f"ORDER BY __changes.{updated_at_column}" - ") AS __next_change_at, __hard_delete_candidates.__hard_deleted_at " - "FROM __changes LEFT JOIN __hard_delete_candidates " - f"ON {hard_delete_join_condition} " - f"AND __hard_delete_candidates.{observed_at_column} = " - f"__changes.{observed_at_column}" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + ).initial_select_sql(output_columns=output_columns) + + partition_sql: str = ", ".join(unique_key) + output_select_sql: str = ", ".join(column for column in output_columns) return ( "WITH __ordered AS (" f"SELECT *, LAG({updated_at_column}) OVER (" @@ -3276,43 +3190,21 @@ def _historical_timestamp_new_changes_cte_sql( valid_to_column: str, invalidate_hard_deletes: bool, ) -> str: - partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - first_key: str = unique_key[0] if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - hard_deletes_sql: str = cls._historical_hard_deletes_select_sql( - destination=destination, + return HistoricalSnapshotSql( origin=origin, unique_key=unique_key, observed_at_column=observed_at_column, + valid_from_column=valid_from_column, valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - "), __hard_deletes AS (" - f"{hard_deletes_sql}" - ")" - ) + updated_at_column=updated_at_column, + ).new_changes_ctes_sql(destination=destination) + + partition_sql: str = ", ".join(unique_key) + latest_join_condition: str = cls._snapshot_key_condition( + left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key + ) + first_key: str = unique_key[0] return ( "__ordered AS (" f"SELECT *, LAG({updated_at_column}) OVER (" @@ -3394,6 +3286,16 @@ def _historical_check_new_changes_cte_sql( valid_to_column: str, invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + ).new_changes_ctes_sql(destination=destination) + partition_sql: str = ", ".join(unique_key) previous_columns_sql: str = ", ".join( f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " @@ -3417,48 +3319,6 @@ def _historical_check_new_changes_cte_sql( "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " f"OR ({delta_change_condition})" ) - if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - hard_deletes_sql: str = cls._historical_hard_deletes_select_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - reappearing_new_changes_sql: str = _historical_reappearing_new_changes_ctes_sql( - changed_or_new_sql=( - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - ), - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - f"), {reappearing_new_changes_sql}, __hard_deletes AS (" - f"{hard_deletes_sql}" - ")" - ) return ( "__ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -3488,54 +3348,6 @@ def _snapshot_key_condition( f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key ) - @classmethod - def _historical_hard_deleted_at_sql( - cls, *, origin: str, unique_key: tuple[str, ...], observed_at_column: str, row_alias: str - ) -> str: - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias=row_alias, unique_key=unique_key - ) - return ( - "(SELECT MIN(__observed_groups.__observed_at) " - f"FROM (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"WHERE __observed_groups.__observed_at > {row_alias}.{observed_at_column} " - "AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __present " - f"WHERE __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition}" - "))" - ) - - @classmethod - def _historical_hard_deletes_select_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - observed_at_column: str, - valid_to_column: str, - ) -> str: - target_key_sql: str = ", ".join(f"__target.{column}" for column in unique_key) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"SELECT {target_key_sql}, MIN(__observed_groups.__observed_at) AS __close_at " - f"FROM {destination} AS __target " - f"JOIN (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"ON __observed_groups.__observed_at > __target.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __present.{first_key} IS NULL " - f"GROUP BY {target_key_sql}" - ) - @classmethod def _historical_snapshot_combined_close_sql( cls, diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py index f2fb3dbee9..36c07208f6 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py @@ -16,8 +16,6 @@ _build_schemas_filter, _encode_typed_json, _historical_check_snapshot_select_sql, - _historical_hard_deleted_at_sql, - _historical_reappearing_new_changes_ctes_sql, _historical_snapshot_combined_close_sql, _historical_timestamp_changes_select_sql, _historical_timestamp_snapshot_select_sql, @@ -29,6 +27,10 @@ _snapshot_key_condition, _validate_rectangular_typed_array, ) +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + HistoricalSnapshotSql, + historical_insert_validity_sql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin @@ -36,6 +38,7 @@ DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE, QUALIFIED_NAME_SEPARATOR, + SNAPSHOT_VERSION_START_COLUMN, ) from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapter.contract.main.normalize_seed_csv_value import normalize_seed_csv_value @@ -2411,7 +2414,7 @@ def render_apply_historical_timestamp_snapshot_changes( unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - change_time_column=updated_at_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: close_sql = ( @@ -2431,13 +2434,16 @@ def render_apply_historical_timestamp_snapshot_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{updated_at_column}, LEAD(__new_changes.{updated_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() insert_sql: str = ( f"WITH {new_changes_sql} " f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " + f"SELECT {output_select_sql}, {version_columns_sql} " f"FROM __new_changes" ) return (close_sql, insert_sql) @@ -2524,7 +2530,7 @@ def render_apply_historical_check_snapshot_changes( unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - change_time_column=observed_at_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: close_sql = ( @@ -2544,13 +2550,16 @@ def render_apply_historical_check_snapshot_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{observed_at_column}, LEAD(__new_changes.{observed_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() insert_sql: str = ( f"WITH {new_changes_sql} " f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{observed_at_column}, " - f"LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column}" - f") " + f"SELECT {output_select_sql}, {version_columns_sql} " f"FROM __new_changes" ) return (close_sql, insert_sql) @@ -2567,6 +2576,16 @@ def _pg_historical_timestamp_new_changes_cte_sql( valid_to_column: str, invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + ).new_changes_ctes_sql(destination=destination) + partition_sql: str = ", ".join(unique_key) first_key: str = unique_key[0] latest_sql: str = ( @@ -2575,36 +2594,7 @@ def _pg_historical_timestamp_new_changes_cte_sql( f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" f") AS __rn FROM {destination}) AS __q WHERE __rn = 1)" ) - if invalidate_hard_deletes: - latest_join_condition: str = _snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - f"), {latest_sql}, __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - "), __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{col}' for col in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ")" - ) - latest_join_condition = _snapshot_key_condition( + latest_join_condition: str = _snapshot_key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) return ( @@ -2663,6 +2653,16 @@ def _pg_historical_check_new_changes_cte_sql( valid_to_column: str, invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + ).new_changes_ctes_sql(destination=destination) + partition_sql: str = ", ".join(unique_key) previous_columns_sql: str = ", ".join( f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " @@ -2692,38 +2692,6 @@ def _pg_historical_check_new_changes_cte_sql( f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" f") AS __rn FROM {destination}) AS __q WHERE __rn = 1)" ) - if invalidate_hard_deletes: - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - reappearing_new_changes_sql: str = _historical_reappearing_new_changes_ctes_sql( - changed_or_new_sql=( - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - ), - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - f"), {latest_sql}, {reappearing_new_changes_sql}, __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ")" - ) return ( "__ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index c450bf2778..63ce339ad1 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -14,14 +14,21 @@ from sqlbuild.adapter.contract.classes.base_adapter import ( BaseAdapter, _encode_typed_json, - _historical_reappearing_new_changes_ctes_sql, _render_ansi_typed_scalar, _render_typed_value_list, ) +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + HistoricalSnapshotSql, + historical_insert_validity_sql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin -from sqlbuild.adapter.contract.constants import DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE +from sqlbuild.adapter.contract.constants import ( + DIFF_LEFT_SIDE, + DIFF_RIGHT_SIDE, + SNAPSHOT_VERSION_START_COLUMN, +) from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapter.contract.main.normalize_seed_csv_value import normalize_seed_csv_value from sqlbuild.adapter.contract.models import ( @@ -1411,7 +1418,7 @@ def render_apply_historical_timestamp_snapshot_changes( unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - change_time_column=updated_at_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: close_sql = self._historical_snapshot_close_sql( @@ -1428,16 +1435,17 @@ def render_apply_historical_timestamp_snapshot_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{updated_at_column}, LEAD(__new_changes.{updated_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() insert_sql: str = self._historical_snapshot_insert_sql( destination=destination, insert_column_sql=insert_column_sql, new_changes_sql=new_changes_sql, - select_sql=( - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") FROM __new_changes" - ), + select_sql=(f"SELECT {output_select_sql}, {version_columns_sql} FROM __new_changes"), ) return (close_sql, insert_sql) @@ -1618,7 +1626,7 @@ def render_apply_historical_check_snapshot_changes( unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - change_time_column=observed_at_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: close_sql = self._historical_snapshot_close_sql( @@ -1635,16 +1643,17 @@ def render_apply_historical_check_snapshot_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{observed_at_column}, LEAD(__new_changes.{observed_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() insert_sql: str = self._historical_snapshot_insert_sql( destination=destination, insert_column_sql=insert_column_sql, new_changes_sql=new_changes_sql, - select_sql=( - f"SELECT {output_select_sql}, __new_changes.{observed_at_column}, " - f"LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column}" - f") FROM __new_changes" - ), + select_sql=(f"SELECT {output_select_sql}, {version_columns_sql} FROM __new_changes"), ) return (close_sql, insert_sql) @@ -3304,6 +3313,16 @@ def _historical_check_snapshot_select_sql( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + ).initial_select_sql(output_columns=output_columns) + partition_sql: str = ", ".join(unique_key) previous_columns_sql: str = ", ".join( f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " @@ -3316,68 +3335,6 @@ def _historical_check_snapshot_select_sql( f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns ) output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_delete_join_condition: str = cls._snapshot_key_condition( - left_alias="__hard_delete_candidates", - right_alias="__changes", - unique_key=unique_key, - ) - hard_delete_key_sql: str = ", ".join(f"__changes.{column}" for column in unique_key) - hard_delete_group_sql: str = ", ".join( - [ - *(f"__changes.{column}" for column in unique_key), - f"__changes.{observed_at_column}", - ] - ) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__changes", unique_key=unique_key - ) - return ( - "WITH __ordered AS (" - "SELECT __source.*, " - f"(SELECT MAX(__groups.__observed_at) FROM (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - ") AS __groups " - f"WHERE __groups.__observed_at < __source.{observed_at_column}" - ") AS __prev_group_observed_at, " - f"LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin} AS __source" - "), __changes AS (" - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({change_condition}) " - "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at" - "), __observed_groups AS (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - "), __hard_delete_candidates AS (" - f"SELECT {hard_delete_key_sql}, __changes.{observed_at_column}, " - "MIN(__observed_groups.__observed_at) AS __hard_deleted_at " - "FROM __changes " - "JOIN __observed_groups " - f"ON __observed_groups.__observed_at > __changes.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __present.{unique_key[0]} IS NULL " - f"GROUP BY {hard_delete_group_sql}" - "), __versions AS (" - f"SELECT __changes.*, LEAD(__changes.{observed_at_column}) OVER (" - f"PARTITION BY {', '.join(f'__changes.{column}' for column in unique_key)} " - f"ORDER BY __changes.{observed_at_column}" - ") AS __next_change_at, __hard_delete_candidates.__hard_deleted_at " - "FROM __changes LEFT JOIN __hard_delete_candidates " - f"ON {hard_delete_join_condition} " - f"AND __hard_delete_candidates.{observed_at_column} = " - f"__changes.{observed_at_column}" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) return ( "WITH __ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -3405,63 +3362,18 @@ def _historical_timestamp_snapshot_select_sql( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) if invalidate_hard_deletes: - hard_delete_join_condition: str = cls._snapshot_key_condition( - left_alias="__hard_delete_candidates", - right_alias="__changes", + return HistoricalSnapshotSql( + origin=origin, unique_key=unique_key, - ) - hard_delete_key_sql: str = ", ".join(f"__changes.{column}" for column in unique_key) - hard_delete_group_sql: str = ", ".join( - [ - *(f"__changes.{column}" for column in unique_key), - f"__changes.{observed_at_column}", - ] - ) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__changes", unique_key=unique_key - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __observed_groups AS (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - "), __hard_delete_candidates AS (" - f"SELECT {hard_delete_key_sql}, __changes.{observed_at_column}, " - "MIN(__observed_groups.__observed_at) AS __hard_deleted_at " - "FROM __changes " - "JOIN __observed_groups " - f"ON __observed_groups.__observed_at > __changes.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __present.{unique_key[0]} IS NULL " - f"GROUP BY {hard_delete_group_sql}" - "), __versions AS (" - f"SELECT __changes.*, LEAD(__changes.{updated_at_column}) OVER (" - f"PARTITION BY {', '.join(f'__changes.{column}' for column in unique_key)} " - f"ORDER BY __changes.{updated_at_column}" - ") AS __next_change_at, __hard_delete_candidates.__hard_deleted_at " - "FROM __changes LEFT JOIN __hard_delete_candidates " - f"ON {hard_delete_join_condition} " - f"AND __hard_delete_candidates.{observed_at_column} = " - f"__changes.{observed_at_column}" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + ).initial_select_sql(output_columns=output_columns) + + partition_sql: str = ", ".join(unique_key) + output_select_sql: str = ", ".join(column for column in output_columns) return ( "WITH __ordered AS (" f"SELECT *, LAG({updated_at_column}) OVER (" @@ -3490,43 +3402,21 @@ def _historical_timestamp_new_changes_cte_sql( valid_to_column: str, invalidate_hard_deletes: bool, ) -> str: - partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - first_key: str = unique_key[0] if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - hard_deletes_sql: str = cls._historical_hard_deletes_select_sql( - destination=destination, + return HistoricalSnapshotSql( origin=origin, unique_key=unique_key, observed_at_column=observed_at_column, + valid_from_column=valid_from_column, valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - "), __hard_deletes AS (" - f"{hard_deletes_sql}" - ")" - ) + updated_at_column=updated_at_column, + ).new_changes_ctes_sql(destination=destination) + + partition_sql: str = ", ".join(unique_key) + latest_join_condition: str = cls._snapshot_key_condition( + left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key + ) + first_key: str = unique_key[0] return ( "__ordered AS (" f"SELECT *, LAG({updated_at_column}) OVER (" @@ -3608,6 +3498,16 @@ def _historical_check_new_changes_cte_sql( valid_to_column: str, invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + ).new_changes_ctes_sql(destination=destination) + partition_sql: str = ", ".join(unique_key) previous_columns_sql: str = ", ".join( f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " @@ -3631,48 +3531,6 @@ def _historical_check_new_changes_cte_sql( "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " f"OR ({delta_change_condition})" ) - if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - hard_deletes_sql: str = cls._historical_hard_deletes_select_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - reappearing_new_changes_sql: str = _historical_reappearing_new_changes_ctes_sql( - changed_or_new_sql=( - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - ), - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - f"), {reappearing_new_changes_sql}, __hard_deletes AS (" - f"{hard_deletes_sql}" - ")" - ) return ( "__ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" @@ -3702,54 +3560,6 @@ def _snapshot_key_condition( f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key ) - @classmethod - def _historical_hard_deleted_at_sql( - cls, *, origin: str, unique_key: tuple[str, ...], observed_at_column: str, row_alias: str - ) -> str: - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias=row_alias, unique_key=unique_key - ) - return ( - "(SELECT MIN(__observed_groups.__observed_at) " - f"FROM (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"WHERE __observed_groups.__observed_at > {row_alias}.{observed_at_column} " - "AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __present " - f"WHERE __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition}" - "))" - ) - - @classmethod - def _historical_hard_deletes_select_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - observed_at_column: str, - valid_to_column: str, - ) -> str: - target_key_sql: str = ", ".join(f"__target.{column}" for column in unique_key) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"SELECT {target_key_sql}, MIN(__observed_groups.__observed_at) AS __close_at " - f"FROM {destination} AS __target " - f"JOIN (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"ON __observed_groups.__observed_at > __target.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __present.{first_key} IS NULL " - f"GROUP BY {target_key_sql}" - ) - @classmethod def _historical_snapshot_combined_close_sql( cls, diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py index 7e2cee1f07..17694f2d9b 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py @@ -15,10 +15,6 @@ _build_names_filter, _build_schemas_filter, _encode_typed_json, - _historical_hard_deleted_at_sql, - _historical_observed_group_sequence_cte_sql, - _historical_observed_group_sequence_from_sql, - _historical_reappearing_new_changes_ctes_sql, _historical_timestamp_changes_new_records_cte_sql, _quote_sql_string, _render_ansi_typed_scalar, @@ -27,10 +23,18 @@ _snapshot_key_condition, _typed_scalar_payload, ) +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + HistoricalSnapshotSql, + historical_insert_validity_sql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin -from sqlbuild.adapter.contract.constants import DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE +from sqlbuild.adapter.contract.constants import ( + DIFF_LEFT_SIDE, + DIFF_RIGHT_SIDE, + SNAPSHOT_VERSION_START_COLUMN, +) from sqlbuild.adapter.contract.exceptions import ( AdapterUserError, UnsupportedTypedSqlRenderingError, @@ -218,6 +222,16 @@ def _historical_timestamp_new_changes_cte_sql( valid_to_column: str, invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + distinct_condition=self._distinct_condition, + ).new_changes_ctes_sql(destination=destination) partition_sql: str = ", ".join(unique_key) first_key: str = unique_key[0] latest_join_condition: str = _snapshot_key_condition( @@ -226,20 +240,6 @@ def _historical_timestamp_new_changes_cte_sql( updated_changed: str = self._distinct_condition( left=f"{updated_at_column}", right="__prev_updated_at" ) - hard_deletes_sql: str = "" - if invalidate_hard_deletes: - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - hard_deletes_sql = ( - "), __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ) return ( "__ordered AS (" f"SELECT *, LAG({updated_at_column}) OVER (" @@ -257,7 +257,6 @@ def _historical_timestamp_new_changes_cte_sql( f"LEFT JOIN __latest ON {latest_join_condition} " f"WHERE __latest.{first_key} IS NULL " f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - f"{hard_deletes_sql}" ")" ) @@ -273,6 +272,16 @@ def _historical_check_new_changes_cte_sql( valid_to_column: str, invalidate_hard_deletes: bool, ) -> str: + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + distinct_condition=self._distinct_condition, + ).new_changes_ctes_sql(destination=destination) partition_sql: str = ", ".join(unique_key) previous_columns_sql: str = ", ".join( self._lag_expr( @@ -303,28 +312,6 @@ def _historical_check_new_changes_cte_sql( f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " f"AND ({latest_change_condition}))" ) - new_changes_sql: str = f"__new_changes AS ({changed_or_new_sql})" - hard_deletes_sql: str = "" - if invalidate_hard_deletes: - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - new_changes_sql = _historical_reappearing_new_changes_ctes_sql( - changed_or_new_sql=changed_or_new_sql, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - hard_deletes_sql = ( - ", __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ")" - ) changed_or_first_sql: str = ( "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " f"OR ({delta_change_condition})" @@ -341,7 +328,7 @@ def _historical_check_new_changes_cte_sql( f"ORDER BY {valid_from_column} DESC) AS __rn FROM {destination}" "), __latest AS (" "SELECT * FROM __latest_ordered WHERE __rn = 1" - f"), {new_changes_sql}{hard_deletes_sql}" + f"), __new_changes AS ({changed_or_new_sql})" ) def connect(self, config: dict[str, Any]) -> _SqlServerConnection: @@ -1889,7 +1876,7 @@ def render_apply_historical_check_snapshot_changes( unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - change_time_column=observed_at_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: close_sql = ( @@ -1910,13 +1897,16 @@ def render_apply_historical_check_snapshot_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{observed_at_column}, LEAD(__new_changes.{observed_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() insert_sql: str = ( f";WITH {new_changes_sql} " f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{observed_at_column}, " - f"LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column}" - f") " + f"SELECT {output_select_sql}, {version_columns_sql} " f"FROM __new_changes" ) return (close_sql, insert_sql) @@ -2004,7 +1994,7 @@ def render_apply_historical_timestamp_snapshot_changes( unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - change_time_column=updated_at_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: close_sql = ( @@ -2025,13 +2015,16 @@ def render_apply_historical_timestamp_snapshot_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{updated_at_column}, LEAD(__new_changes.{updated_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() insert_sql: str = ( f";WITH {new_changes_sql} " f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " + f"SELECT {output_select_sql}, {version_columns_sql} " f"FROM __new_changes" ) return (close_sql, insert_sql) @@ -2183,6 +2176,22 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: + if invalidate_hard_deletes: + historical: HistoricalSnapshotSql = HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + distinct_condition=self._distinct_condition, + ) + return ( + f"DROP TABLE IF EXISTS {destination}", + f";WITH {historical.initial_ctes_sql()} " + f"SELECT {historical.initial_select_list_sql(output_columns=output_columns)} " + f"INTO {destination} FROM __versions", + ) partition_sql: str = ", ".join(unique_key) previous_columns_sql: str = ", ".join( self._lag_expr( @@ -2200,13 +2209,6 @@ def render_create_initial_historical_check_snapshot_destination( ) output_select_sql: str = ", ".join(output_columns) hard_deleted_at_sql: str = "CAST(NULL AS DATETIME2)" - if invalidate_hard_deletes: - hard_deleted_at_sql = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__changes", - ) valid_to_expr: str = ( "CASE " "WHEN __next_change_at IS NULL THEN __hard_deleted_at " @@ -2214,35 +2216,12 @@ def render_create_initial_historical_check_snapshot_destination( "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " "ELSE __next_change_at END" ) - ordered_sql: str = ( - f"SELECT *, LAG({observed_at_column}) OVER (" + historical_sql: str = ( + f";WITH __ordered AS (SELECT *, LAG({observed_at_column}) OVER (" f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - ) - changes_condition: str = f"__prev_observed_at IS NULL OR ({change_condition})" - leading_ctes_sql: str = "" - if invalidate_hard_deletes: - group_sequence_sql: str = _historical_observed_group_sequence_cte_sql( - origin=origin, observed_at_column=observed_at_column - ) - grouped_origin_sql: str = _historical_observed_group_sequence_from_sql( - origin=origin, observed_at_column=observed_at_column - ) - leading_ctes_sql = f"{group_sequence_sql}, " - ordered_sql = ( - "SELECT __source.*, __observed_group_sequence.__prev_group_observed_at, " - f"LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} {grouped_origin_sql}" - ) - reappeared_condition: str = self._distinct_condition( - left="__prev_observed_at", right="__prev_group_observed_at" - ) - changes_condition = f"{changes_condition} OR {reappeared_condition}" - historical_sql: str = ( - f";WITH {leading_ctes_sql}__ordered AS ({ordered_sql}" "), __changes AS (" - f"SELECT * FROM __ordered WHERE {changes_condition}" + f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" "), __versions AS (" f"SELECT __changes.*, LEAD({observed_at_column}) OVER (" f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py index da96dfe6d5..f0cbbfd33d 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py @@ -125,6 +125,46 @@ ), expected_history=((1, "active", 1, 2), (1, "paused", 2, None)), ), + SnapshotExecutionScenario( + description="check multiple delete and reappear cycles in one build", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=( + _CHECK_DELETE_AND_REAPPEAR_ROWS[:2], + ( + *_CHECK_DELETE_AND_REAPPEAR_ROWS, + (1, "active", 4), + (1, "active", 5), + (2, "active", 5), + ), + ), + expected_history=( + (1, "active", 1, None), + (2, "active", 1, 2), + (2, "active", 3, 4), + (2, "active", 5, None), + ), + ), + SnapshotExecutionScenario( + description="check reappearance, change and change back in one build", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=( + _CHECK_DELETE_AND_REAPPEAR_ROWS[:2], + ( + *_CHECK_DELETE_AND_REAPPEAR_ROWS, + (1, "active", 4), + (2, "paused", 4), + (1, "active", 5), + (2, "active", 5), + ), + ), + expected_history=( + (1, "active", 1, None), + (2, "active", 1, 2), + (2, "active", 3, 4), + (2, "paused", 4, 5), + (2, "active", 5, None), + ), + ), ) _CURRENT_STATE_SCENARIOS: tuple[SnapshotExecutionScenario, ...] = ( diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py index 53b9d86b0b..0b512fbf3e 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py @@ -16,20 +16,27 @@ SnapshotSqlRenderingAdapterTestCase, ) -_REAPPEARANCE_GROUP_SEQUENCE_FRAGMENT: str = ( +_REAPPEARANCE_INITIAL_FRAGMENTS: tuple[str, ...] = ( "__observed_group_sequence AS (SELECT __observed_at, " - "LAG(__observed_at) OVER (ORDER BY __observed_at) AS __prev_group_observed_at " - "FROM (SELECT DISTINCT observed_at AS __observed_at FROM source_table)" + "LAG(__observed_at) OVER (ORDER BY __observed_at) AS __prev_group_observed_at, " + "LEAD(__observed_at) OVER (ORDER BY __observed_at) AS __next_group_observed_at " + "FROM (SELECT DISTINCT observed_at AS __observed_at FROM source_table)", + "(__prev_observed_at IS NOT NULL AND __prev_observed_at <> __prev_group_observed_at)", + "ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING) AS __next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ) _REAPPEARANCE_APPLY_FRAGMENTS: tuple[str, ...] = ( - "__reappearances AS (SELECT __ordered.customer_id, " - "MIN(__ordered.observed_at) AS __reappeared_at", - "WHERE __latest.valid_to IS NOT NULL AND __ordered.observed_at > __latest.valid_to", - "AND __present.observed_at = __latest.valid_to", - "__new_changes AS (SELECT * FROM __changed_or_new UNION ALL SELECT __ordered.* " - "FROM __ordered JOIN __reappearances", - "AND __changed_or_new.observed_at = __ordered.observed_at", + "AS __latest_rn FROM target_table) AS __latest_ranked WHERE __latest_rn = 1", + "__reappearances AS (SELECT __classified.customer_id, " + "MIN(__classified.observed_at) AS __reappeared_at", + "WHERE __latest.valid_to IS NOT NULL AND __closing_group.__observed_at IS NULL", + "OR __classified.observed_at < __first_new_starts.__first_version_start", + "__new_changes AS (SELECT *, LEAD(__version_start) OVER", + "customer_id, __version_start AS __close_at FROM __new_changes", + "__new_changes.__version_start, CASE WHEN __new_changes.__next_version_start IS NULL " + "THEN __new_changes.__next_absence_at", ) +_PORTABLE_APPLY_EXCLUSIONS: tuple[str, ...] = ("QUALIFY", "UNION DISTINCT", "NOT EXISTS") @pytest.mark.parametrize( @@ -50,14 +57,12 @@ "WHERE __target.effective_to IS NULL AND NOT EXISTS", ), expected_historical_check_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "NOT EXISTS", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "NOT EXISTS", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_apply_hard_delete_fragments=( "__hard_deletes AS (", @@ -87,14 +92,12 @@ "WHERE __target.effective_to IS NULL AND NOT EXISTS", ), expected_historical_check_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "__hard_delete_candidates AS (", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "__hard_delete_candidates AS (", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_apply_hard_delete_fragments=( "__hard_deletes AS (", @@ -124,14 +127,12 @@ "WHERE __target.effective_to IS NULL AND NOT EXISTS", ), expected_historical_check_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "__hard_delete_candidates AS (", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "__hard_delete_candidates AS (", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_apply_hard_delete_fragments=( "__hard_deletes AS (", @@ -161,14 +162,12 @@ "WHEN MATCHED THEN UPDATE SET effective_to = __source.updated_at", ), expected_historical_check_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "__hard_delete_candidates AS (", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "__hard_delete_candidates AS (", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_apply_hard_delete_fragments=( "__hard_deletes AS (", @@ -290,52 +289,53 @@ def test_given_builtin_adapter_when_rendering_snapshot_sql_then_covers_snapshot_ "test_case", [ SnapshotReappearanceRenderingTestCase( - description="duckdb reopens reappearing historical check keys", + description="duckdb renders portable hard-delete history SQL", adapter=DuckDbAdapter(), - expected_initial_fragments=( - _REAPPEARANCE_GROUP_SEQUENCE_FRAGMENT, - "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at", - ), + expected_initial_fragments=_REAPPEARANCE_INITIAL_FRAGMENTS, expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, - unexpected_apply_fragments=("UNION DISTINCT",), + unexpected_apply_fragments=_PORTABLE_APPLY_EXCLUSIONS, ), SnapshotReappearanceRenderingTestCase( - description="motherduck reopens reappearing historical check keys", + description="motherduck renders portable hard-delete history SQL", adapter=MotherDuckAdapter(), - expected_initial_fragments=( - _REAPPEARANCE_GROUP_SEQUENCE_FRAGMENT, - "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at", - ), + expected_initial_fragments=_REAPPEARANCE_INITIAL_FRAGMENTS, expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, - unexpected_apply_fragments=("UNION DISTINCT",), + unexpected_apply_fragments=_PORTABLE_APPLY_EXCLUSIONS, ), SnapshotReappearanceRenderingTestCase( - description="postgres reopens reappearing historical check keys without qualify", + description="postgres renders portable hard-delete history SQL", adapter=PostgresAdapter(), - expected_initial_fragments=( - _REAPPEARANCE_GROUP_SEQUENCE_FRAGMENT, - "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at", - ), - expected_apply_fragments=( - *_REAPPEARANCE_APPLY_FRAGMENTS, - ") AS __rn FROM target_table) AS __q WHERE __rn = 1)", - ), - unexpected_apply_fragments=("QUALIFY", "UNION DISTINCT"), + expected_initial_fragments=_REAPPEARANCE_INITIAL_FRAGMENTS, + expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, + unexpected_apply_fragments=_PORTABLE_APPLY_EXCLUSIONS, + ), + SnapshotReappearanceRenderingTestCase( + description="bigquery renders portable hard-delete history SQL", + adapter=BigQueryAdapter(), + expected_initial_fragments=_REAPPEARANCE_INITIAL_FRAGMENTS, + expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, + unexpected_apply_fragments=_PORTABLE_APPLY_EXCLUSIONS, + ), + SnapshotReappearanceRenderingTestCase( + description="snowflake renders portable hard-delete history SQL", + adapter=SnowflakeAdapter(), + expected_initial_fragments=_REAPPEARANCE_INITIAL_FRAGMENTS, + expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, + unexpected_apply_fragments=_PORTABLE_APPLY_EXCLUSIONS, + ), + SnapshotReappearanceRenderingTestCase( + description="databricks renders portable hard-delete history SQL", + adapter=DatabricksAdapter(), + expected_initial_fragments=_REAPPEARANCE_INITIAL_FRAGMENTS, + expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, + unexpected_apply_fragments=_PORTABLE_APPLY_EXCLUSIONS, ), SnapshotReappearanceRenderingTestCase( - description="sqlserver reopens reappearing historical check keys without qualify", + description="sqlserver renders portable hard-delete history SQL", adapter=SqlServerAdapter(), - expected_initial_fragments=( - _REAPPEARANCE_GROUP_SEQUENCE_FRAGMENT, - "OR (__prev_observed_at <> __prev_group_observed_at " - "OR (__prev_observed_at IS NULL AND __prev_group_observed_at IS NOT NULL) " - "OR (__prev_observed_at IS NOT NULL AND __prev_group_observed_at IS NULL))", - ), - expected_apply_fragments=( - *_REAPPEARANCE_APPLY_FRAGMENTS, - "__latest AS (SELECT * FROM __latest_ordered WHERE __rn = 1)", - ), - unexpected_apply_fragments=("QUALIFY", "UNION DISTINCT", "IS DISTINCT FROM"), + expected_initial_fragments=_REAPPEARANCE_INITIAL_FRAGMENTS, + expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, + unexpected_apply_fragments=(*_PORTABLE_APPLY_EXCLUSIONS, "IS DISTINCT FROM"), ), ], ids=lambda case: case.description, @@ -401,12 +401,12 @@ def test_given_hard_deletes_when_rendering_historical_check_snapshot_then_reopen expected_fragment: str for expected_fragment in test_case.expected_initial_fragments: assert expected_fragment in initial_sql - statement: str - for statement in apply_statements: - for expected_fragment in test_case.expected_apply_fragments: - assert expected_fragment in statement - unexpected_fragment: str - for unexpected_fragment in test_case.unexpected_apply_fragments: - assert unexpected_fragment not in statement + apply_sql: str = "\n".join(apply_statements) + for expected_fragment in test_case.expected_apply_fragments: + assert expected_fragment in apply_sql + unexpected_fragment: str + for unexpected_fragment in test_case.unexpected_apply_fragments: + assert unexpected_fragment not in apply_sql assert "__prev_group_observed_at" not in without_hard_deletes_sql assert "__reappearances" not in without_hard_deletes_sql + assert "__next_absence_at" not in without_hard_deletes_sql From 211f8e2f56da47128e275b89bc5bae29fe3389c0 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:21:46 +0100 Subject: [PATCH 42/57] fix(snapshots): reopen timestamp snapshot keys that reappear after a hard delete Historical timestamp snapshots with invalidate_hard_deletes left a hard-deleted key without an active version when it reappeared in a later observation with an unchanged updated_at, on every adapter. The key now gets a new active version whose valid_from is the observation time at which it reappeared, matching the observation-time boundary used when the deletion closed the previous version. A key that reappears with a changed updated_at still starts at that updated_at, and ordinary changes keep valid_from = updated_at. Applies to DuckDB, MotherDuck, Postgres, SQL Server, BigQuery, Snowflake and Databricks, on initial and incremental builds. Existing history rows are not rewritten; a key that is currently closed by a hard delete and present in a later observation is reopened on the next build. (cherry picked from commit 3e59b1ed5b6025eb762aa2dac599984eec55870c) --- .../classes/historical_snapshot_sql.py | 40 +++-- .../main/build/test_snapshot_build.py | 139 +++++++++++++++++- .../test_duckdb_execution.py | 78 +++++++++- 3 files changed, 246 insertions(+), 11 deletions(-) diff --git a/src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py index d2feaeb736..e2fe4878dc 100644 --- a/src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py +++ b/src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py @@ -107,7 +107,8 @@ def new_changes_ctes_sql(self, *, destination: str) -> str: f"CASE WHEN __latest.{first_key} IS NULL THEN 0 ELSE 1 END AS __has_latest, " f"__latest.{self.valid_from_column} AS __latest_valid_from, " f"{self._differs_from_latest_sql()} AS __differs_from_latest, " - f"{self._incremental_reappearance_sql()} AS __is_reappearance " + f"{self._incremental_reappearance_sql()} AS __is_reappearance, " + f"{self._unchanged_from_previous_sql()} AS __is_unchanged " f"FROM __tracked LEFT JOIN __latest ON {latest_join_condition}" "), __changed_or_new AS (" f"SELECT __classified.*, {self._incremental_version_start_sql()} " @@ -173,16 +174,14 @@ def _change_condition_sql(self) -> str: return f"__prev_observed_at IS NULL OR ({changed_sql})" def _initial_reappearance_sql(self) -> str: - if self.updated_at_column is not None: - return "1 = 0" return "(__prev_observed_at IS NOT NULL AND __prev_observed_at <> __prev_group_observed_at)" def _initial_version_start_sql(self) -> str: - return self.updated_at_column or self.observed_at_column + return self._version_start_sql( + unchanged_reappearance_sql=f"{self._initial_reappearance_sql()} AND __is_change = 0" + ) def _incremental_reappearance_sql(self) -> str: - if self.updated_at_column is not None: - return "0" return ( "CASE WHEN __tracked.__prev_group_observed_at IS NOT NULL AND (" f"(__tracked.__prev_observed_at IS NULL AND __latest.{self.unique_key[0]} IS NOT NULL) " @@ -191,7 +190,32 @@ def _incremental_reappearance_sql(self) -> str: ) def _incremental_version_start_sql(self) -> str: - return self.updated_at_column or self.observed_at_column + return self._version_start_sql( + unchanged_reappearance_sql="__is_reappearance = 1 AND __is_unchanged = 1" + ) + + def _version_start_sql(self, *, unchanged_reappearance_sql: str) -> str: + if self.updated_at_column is None: + return self.observed_at_column + return ( + f"CASE WHEN {unchanged_reappearance_sql} THEN {self.observed_at_column} " + f"ELSE {self.updated_at_column} END" + ) + + def _unchanged_from_previous_sql(self) -> str: + if self.updated_at_column is None: + return "0" + changed_from_latest_sql: str = self.distinct_condition( + left=f"__tracked.{self.updated_at_column}", + right=f"__latest.{self.updated_at_column}", + ) + return ( + "CASE WHEN __tracked.__prev_observed_at IS NOT NULL AND __tracked.__is_change = 0 " + "THEN 1 " + "WHEN __tracked.__prev_observed_at IS NULL AND __latest." + f"{self.unique_key[0]} IS NOT NULL AND NOT ({changed_from_latest_sql}) THEN 1 " + "ELSE 0 END" + ) def _differs_from_latest_sql(self) -> str: if not self.check_columns: @@ -220,8 +244,6 @@ def _latest_change_condition_sql(self) -> str: ) def _new_change_starts_sql(self) -> str: - if self.updated_at_column is not None: - return "__new_change_starts AS (SELECT * FROM __changed_or_new)" observed_at: str = self.observed_at_column classified_key_sql: str = self._key_sql(alias="__classified") return ( diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/build/test_snapshot_build.py b/tests/e2e/src/sqlbuild/cli/commands/main/build/test_snapshot_build.py index b036f6cc79..3268604f80 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/build/test_snapshot_build.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/build/test_snapshot_build.py @@ -32,6 +32,32 @@ table_exists, ) +_TIMESTAMP_REAPPEARING_KEY_FINAL_HISTORY: tuple[tuple[object, ...], ...] = ( + (1, "basic", "2024-01-15", "2024-02-03"), + (1, "pro", "2024-02-03", None), + (2, "basic", "2024-01-15", "2024-02-02"), + (2, "basic", "2024-02-03", None), + (3, "basic", "2024-01-15", "2024-02-04"), + (3, "basic", "2024-02-05", None), +) +_TIMESTAMP_REAPPEARING_KEY_ROUND_ROWS: tuple[str, ...] = ( + "(1, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-01'), " + "(2, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-01'), " + "(3, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-01')", + "(1, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-02'), " + "(3, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-02')", + "(1, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-03'), " + "(2, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-03'), " + "(3, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-03')", + "(1, 'pro', TIMESTAMP '2024-02-03', TIMESTAMP '2024-02-04'), " + "(2, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-04'), " + "(1, 'pro', TIMESTAMP '2024-02-03', TIMESTAMP '2024-02-05'), " + "(2, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-05'), " + "(3, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-05')", +) +_TIMESTAMP_REAPPEARING_KEY_TABLE_SQL: str = ( + "(customer_id INTEGER, status VARCHAR, updated_at TIMESTAMP, observed_at TIMESTAMP)" +) _REAPPEARING_KEY_FINAL_HISTORY: tuple[tuple[object, ...], ...] = ( (1, "active", "2024-02-01", "2024-02-04"), (1, "paused", "2024-02-04", None), @@ -2065,7 +2091,118 @@ def test_given_shallow_waffle_shop_snapshot_edges_when_sources_change_then_cli_t _REAPPEARING_KEY_FINAL_HISTORY, ), expected_full_history_rows=_REAPPEARING_KEY_FINAL_HISTORY, - ) + ), + SnapshotReappearingKeyBuildE2ETestCase( + description=( + "historical timestamp snapshot reopens hard-deleted keys at the observation time" + ), + repo_files={ + "sqlbuild_project.toml": dedent( + """ + name = "snapshot_reappearing_keys" + adapter = "duckdb" + + [connection] + database = "snapshot_reappearing_keys.duckdb" + """ + ).strip() + + "\n", + "sources/raw.yml": dedent( + """ + sources: + - name: raw_membership_observations + schema: main + table: raw_membership_observations + - name: raw_membership_full_history + schema: main + table: raw_membership_full_history + """ + ).strip() + + "\n", + "models/membership_incremental_snapshot.sql": dedent( + """ + MODEL ( + materialized snapshot, + unique_key [customer_id], + snapshot_strategy timestamp, + updated_at updated_at, + observed_at observed_at, + historical_input snapshot, + invalidate_hard_deletes true + ); + + SELECT customer_id, status, updated_at, observed_at + FROM __source("raw_membership_observations") + """ + ).strip() + + "\n", + "models/membership_full_history_snapshot.sql": dedent( + """ + MODEL ( + materialized snapshot, + unique_key [customer_id], + snapshot_strategy timestamp, + updated_at updated_at, + observed_at observed_at, + historical_input snapshot, + invalidate_hard_deletes true + ); + + SELECT customer_id, status, updated_at, observed_at + FROM __source("raw_membership_full_history") + """ + ).strip() + + "\n", + }, + initial_seed_sql=( + f"CREATE TABLE main.raw_membership_observations {_TIMESTAMP_REAPPEARING_KEY_TABLE_SQL}; " + "INSERT INTO main.raw_membership_observations VALUES " + f"{_TIMESTAMP_REAPPEARING_KEY_ROUND_ROWS[0]}; " + f"CREATE TABLE main.raw_membership_full_history {_TIMESTAMP_REAPPEARING_KEY_TABLE_SQL}; " + "INSERT INTO main.raw_membership_full_history VALUES " + f"{', '.join(_TIMESTAMP_REAPPEARING_KEY_ROUND_ROWS)}" + ), + observation_sql_by_round=( + (), + ( + "INSERT INTO main.raw_membership_observations VALUES " + f"{_TIMESTAMP_REAPPEARING_KEY_ROUND_ROWS[1]}", + ), + ( + "INSERT INTO main.raw_membership_observations VALUES " + f"{_TIMESTAMP_REAPPEARING_KEY_ROUND_ROWS[2]}", + ), + ( + "INSERT INTO main.raw_membership_observations VALUES " + f"{_TIMESTAMP_REAPPEARING_KEY_ROUND_ROWS[3]}", + ), + ), + history_query_template=( + "SELECT customer_id, status, strftime(valid_from, '%Y-%m-%d'), " + "strftime(valid_to, '%Y-%m-%d') FROM main.{table_name} " + "ORDER BY customer_id, valid_from" + ), + expected_incremental_rows_by_round=( + ( + (1, "basic", "2024-01-15", None), + (2, "basic", "2024-01-15", None), + (3, "basic", "2024-01-15", None), + ), + ( + (1, "basic", "2024-01-15", None), + (2, "basic", "2024-01-15", "2024-02-02"), + (3, "basic", "2024-01-15", None), + ), + ( + (1, "basic", "2024-01-15", None), + (2, "basic", "2024-01-15", "2024-02-02"), + (2, "basic", "2024-02-03", None), + (3, "basic", "2024-01-15", None), + ), + _TIMESTAMP_REAPPEARING_KEY_FINAL_HISTORY, + ), + expected_full_history_rows=_TIMESTAMP_REAPPEARING_KEY_FINAL_HISTORY, + ), ], ids=lambda case: case.description, ) diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py index f0cbbfd33d..545f2754f7 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py @@ -165,6 +165,78 @@ (2, "active", 5, None), ), ), + SnapshotExecutionScenario( + description="timestamp unchanged updated_at reappearance across builds", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=( + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2], + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:3], + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS, + ), + expected_history=((1, "basic", 1, None), (2, "pro", 1, 2), (2, "pro", 3, None)), + ), + SnapshotExecutionScenario( + description="timestamp unchanged updated_at reappearance in the same build", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=(_TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2], _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS), + expected_history=((1, "basic", 1, None), (2, "pro", 1, 2), (2, "pro", 3, None)), + ), + SnapshotExecutionScenario( + description="timestamp multiple delete and reappear cycles in one build", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=( + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2], + ( + *_TIMESTAMP_DELETE_AND_REAPPEAR_ROWS, + (1, "basic", 1, 4), + (1, "basic", 1, 5), + (2, "pro", 1, 5), + ), + ), + expected_history=( + (1, "basic", 1, None), + (2, "pro", 1, 2), + (2, "pro", 3, 4), + (2, "pro", 5, None), + ), + ), + SnapshotExecutionScenario( + description="timestamp changed reappearance starts at the new updated_at", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=( + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2], + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:3], + (*_TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:3], (1, "basic", 1, 4), (2, "team", 3, 4)), + ), + expected_history=((1, "basic", 1, None), (2, "pro", 1, 2), (2, "team", 3, None)), + ), +) + +_WINDOWED_SOURCE_SCENARIOS: tuple[SnapshotExecutionScenario, ...] = ( + SnapshotExecutionScenario( + description="check reappearance after the deletion observation left the source window", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=( + _CHECK_DELETE_AND_REAPPEAR_ROWS[:2], + _CHECK_DELETE_AND_REAPPEAR_ROWS[:3], + _CHECK_DELETE_AND_REAPPEAR_ROWS[3:], + ), + expected_history=( + (1, "active", 1, None), + (2, "active", 1, 2), + (2, "active", 3, None), + ), + ), + SnapshotExecutionScenario( + description="timestamp reappearance after the deletion observation left the source window", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=( + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2], + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:3], + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[3:], + ), + expected_history=((1, "basic", 1, None), (2, "pro", 1, 2), (2, "pro", 3, None)), + ), ) _CURRENT_STATE_SCENARIOS: tuple[SnapshotExecutionScenario, ...] = ( @@ -193,7 +265,11 @@ ) _RUNS: tuple[SnapshotExecutionRun, ...] = build_execution_runs( - incremental_scenarios=(*_HISTORICAL_SCENARIOS, *_CURRENT_STATE_SCENARIOS), + incremental_scenarios=( + *_HISTORICAL_SCENARIOS, + *_WINDOWED_SOURCE_SCENARIOS, + *_CURRENT_STATE_SCENARIOS, + ), full_history_scenarios=_HISTORICAL_SCENARIOS, ) From 13b17a51e0778a85146eac27ccd4f4d0af671aad Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:23:10 +0100 Subject: [PATCH 43/57] fix(snapshots): honour invalidate_hard_deletes on the first SQL Server historical timestamp build The first build of a SQL Server historical timestamp snapshot ignored invalidate_hard_deletes, so keys missing from later observations in the initial history stayed active and reappearing keys were not reopened until later incremental builds. The initial build now renders the same hard-delete history as the incremental path and the other adapters: deletions close a version at the observation where the key is missing, and a key that reappears with an unchanged updated_at opens a new version at that observation. Existing SQL Server snapshot tables are not rewritten. (cherry picked from commit 273dc4b486c89aae082e6305b4fb4c42b2bfdfec) --- .../sqlserver/classes/sqlserver_adapter.py | 17 ++++++++- .../snapshot_sql_rendering/test_rendering.py | 37 +++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py index 17694f2d9b..5dd7050a96 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py @@ -2269,7 +2269,22 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - del invalidate_hard_deletes + if invalidate_hard_deletes: + historical: HistoricalSnapshotSql = HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + distinct_condition=self._distinct_condition, + ) + return ( + f"DROP TABLE IF EXISTS {destination}", + f";WITH {historical.initial_ctes_sql()} " + f"SELECT {historical.initial_select_list_sql(output_columns=output_columns)} " + f"INTO {destination} FROM __versions", + ) partition_sql: str = ", ".join(unique_key) changed_condition: str = self._distinct_condition( left=updated_at_column, right="__prev_updated_at" diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py index 0b512fbf3e..734f74e101 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py @@ -182,6 +182,43 @@ "UNION ALL", ), ), + SnapshotSqlRenderingAdapterTestCase( + description="sqlserver renders complete snapshot SQL matrix", + adapter=SqlServerAdapter(), + expected_create_initial_fragments=( + "SELECT * INTO target_table FROM (SELECT *", + "updated_at AS valid_from", + "CAST(NULL AS DATETIME2) AS valid_to", + ), + expected_timestamp_hard_delete_fragments=( + "UPDATE __target SET effective_to = __source.updated_at", + "__target.customer_id = __source.customer_id", + "__target.region = __source.region", + "WHERE __target.effective_to IS NULL AND NOT EXISTS", + ), + expected_historical_check_initial_hard_delete_fragments=( + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", + "AS valid_to INTO target_table FROM __versions", + ), + expected_historical_timestamp_initial_hard_delete_fragments=( + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", + "AS valid_to INTO target_table FROM __versions", + ), + expected_historical_timestamp_apply_hard_delete_fragments=( + "__hard_deletes AS (", + "AS __close_at FROM target_table AS __target", + "UNION ALL", + ), + expected_historical_check_apply_fragments=( + "LAG(plan) OVER (PARTITION BY customer_id ORDER BY observed_at)", + "__hard_deletes AS (", + "INSERT INTO target_table (customer_id, plan, observed_at, valid_from, valid_to)", + "UPDATE __target SET valid_to = (SELECT MIN(__close_candidates.__close_at)", + "UNION ALL", + ), + ), ], ids=lambda case: case.description, ) From 66a5f4cc93ca5737b4a8aa2dcf0b156ea8bcb8c9 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:27:50 +0100 Subject: [PATCH 44/57] fix(snapshots): start reappearing current-state timestamp versions at execution time Current-state timestamp snapshots with invalidate_hard_deletes reopened a key that came back after a hard delete with valid_from = updated_at. Because updated_at usually predates the deletion, the new version overlapped the closed one. A reappearing key now starts at the execution time of the build that sees it again, the same boundary the hard delete used to close it. New keys and ordinary updated_at changes keep their existing valid_from. Applies to DuckDB, MotherDuck, Postgres, SQL Server, BigQuery, Snowflake and Databricks. Existing overlapping rows are not rewritten. (cherry picked from commit 7e77b82bcdc6c75206b37d157e44f116c84b2470) --- .../adapter/contract/classes/base_adapter.py | 18 ++++++++++++ .../contract/classes/duckdb_backed_adapter.py | 18 ++++++++++++ .../bigquery/classes/bigquery_adapter.py | 18 ++++++++++++ .../databricks/classes/databricks_adapter.py | 18 ++++++++++++ .../postgres/classes/postgres_adapter.py | 18 ++++++++++++ .../snowflake/classes/snowflake_adapter.py | 18 ++++++++++++ .../sqlserver/classes/sqlserver_adapter.py | 18 ++++++++++++ .../test_duckdb_execution.py | 28 +++++++++++++++++++ .../snapshot_sql_rendering/test_rendering.py | 15 ++++++++++ 9 files changed, 169 insertions(+) diff --git a/src/sqlbuild/adapter/contract/classes/base_adapter.py b/src/sqlbuild/adapter/contract/classes/base_adapter.py index e041af0d5f..d66a69cf7f 100644 --- a/src/sqlbuild/adapter/contract/classes/base_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/base_adapter.py @@ -912,12 +912,30 @@ def render_apply_timestamp_snapshot_changes( f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " f"ELSE __source.{updated_at_column} END" ) + history_join_sql: str = "" + if invalidate_hard_deletes: + key_sql: str = ", ".join(unique_key) + history_condition: str = " AND ".join( + f"__history.{column} = __source.{column}" for column in unique_key + ) + version_valid_from_expr = ( + f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " + f"AND __history.__closed_at <> __source.{updated_at_column} " + f"THEN {current_timestamp} " + f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " + f"ELSE __source.{updated_at_column} END" + ) + history_join_sql = ( + f"LEFT JOIN (SELECT {key_sql}, MAX({valid_to_column}) AS __closed_at " + f"FROM {destination} GROUP BY {key_sql}) AS __history ON {history_condition} " + ) insert_sql: str = ( f"INSERT INTO {destination} ({insert_column_sql}) " f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " f"FROM {origin} AS __source " f"LEFT JOIN {destination} AS __active " f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " + f"{history_join_sql}" f"WHERE __active.{first_key} IS NULL " f"OR __source.{updated_at_column} > __active.{updated_at_column}" ) diff --git a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py index fd61347a66..2576e1a99b 100644 --- a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py @@ -465,12 +465,30 @@ def render_apply_timestamp_snapshot_changes( f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " f"ELSE __source.{updated_at_column} END" ) + history_join_sql: str = "" + if invalidate_hard_deletes: + key_sql: str = ", ".join(unique_key) + history_condition: str = " AND ".join( + f"__history.{column} = __source.{column}" for column in unique_key + ) + version_valid_from_expr = ( + f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " + f"AND __history.__closed_at <> __source.{updated_at_column} " + f"THEN {current_timestamp} " + f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " + f"ELSE __source.{updated_at_column} END" + ) + history_join_sql = ( + f"LEFT JOIN (SELECT {key_sql}, MAX({valid_to_column}) AS __closed_at " + f"FROM {destination} GROUP BY {key_sql}) AS __history ON {history_condition} " + ) insert_sql: str = ( f"INSERT INTO {destination} ({insert_column_sql}) " f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " f"FROM {origin} AS __source " f"LEFT JOIN {destination} AS __active " f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " + f"{history_join_sql}" f"WHERE __active.{first_key} IS NULL " f"OR __source.{updated_at_column} > __active.{updated_at_column}" ) diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py index 92e1eeef50..5efdd42cde 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py @@ -1060,12 +1060,30 @@ def render_apply_timestamp_snapshot_changes( f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " f"ELSE __source.{updated_at_column} END" ) + history_join_sql: str = "" + if invalidate_hard_deletes: + key_sql: str = ", ".join(unique_key) + history_condition: str = " AND ".join( + f"__history.{column} = __source.{column}" for column in unique_key + ) + version_valid_from_expr = ( + f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " + f"AND __history.__closed_at <> __source.{updated_at_column} " + f"THEN {current_timestamp} " + f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " + f"ELSE __source.{updated_at_column} END" + ) + history_join_sql = ( + f"LEFT JOIN (SELECT {key_sql}, MAX({valid_to_column}) AS __closed_at " + f"FROM {destination} GROUP BY {key_sql}) AS __history ON {history_condition} " + ) insert_sql: str = ( f"INSERT INTO {destination} ({insert_column_sql}) " f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " f"FROM {origin} AS __source " f"LEFT JOIN {destination} AS __active " f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " + f"{history_join_sql}" f"WHERE __active.{first_key} IS NULL " f"OR __source.{updated_at_column} > __active.{updated_at_column}" ) diff --git a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py index b08f81cede..241e3737c9 100644 --- a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py +++ b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py @@ -910,12 +910,30 @@ def render_apply_timestamp_snapshot_changes( f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " f"ELSE __source.{updated_at_column} END" ) + history_join_sql: str = "" + if invalidate_hard_deletes: + key_sql: str = ", ".join(unique_key) + history_condition: str = " AND ".join( + f"__history.{column} = __source.{column}" for column in unique_key + ) + version_valid_from_expr = ( + f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " + f"AND __history.__closed_at <> __source.{updated_at_column} " + f"THEN {current_timestamp} " + f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " + f"ELSE __source.{updated_at_column} END" + ) + history_join_sql = ( + f"LEFT JOIN (SELECT {key_sql}, MAX({valid_to_column}) AS __closed_at " + f"FROM {destination} GROUP BY {key_sql}) AS __history ON {history_condition} " + ) insert_sql: str = ( f"INSERT INTO {destination} ({insert_column_sql}) " f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " f"FROM {origin} AS __source " f"LEFT JOIN {destination} AS __active " f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " + f"{history_join_sql}" f"WHERE __active.{first_key} IS NULL " f"OR __source.{updated_at_column} > __active.{updated_at_column}" ) diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py index 36c07208f6..a7358cf990 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py @@ -1174,12 +1174,30 @@ def render_apply_timestamp_snapshot_changes( f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " f"ELSE __source.{updated_at_column} END" ) + history_join_sql: str = "" + if invalidate_hard_deletes: + key_sql: str = ", ".join(unique_key) + history_condition: str = " AND ".join( + f"__history.{column} = __source.{column}" for column in unique_key + ) + version_valid_from_expr = ( + f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " + f"AND __history.__closed_at <> __source.{updated_at_column} " + f"THEN {current_timestamp} " + f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " + f"ELSE __source.{updated_at_column} END" + ) + history_join_sql = ( + f"LEFT JOIN (SELECT {key_sql}, MAX({valid_to_column}) AS __closed_at " + f"FROM {destination} GROUP BY {key_sql}) AS __history ON {history_condition} " + ) insert_sql: str = ( f"INSERT INTO {destination} ({insert_column_sql}) " f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " f"FROM {origin} AS __source " f"LEFT JOIN {destination} AS __active " f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " + f"{history_join_sql}" f"WHERE __active.{first_key} IS NULL " f"OR __source.{updated_at_column} > __active.{updated_at_column}" ) diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index 63ce339ad1..5a40211b9a 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -1313,12 +1313,30 @@ def render_apply_timestamp_snapshot_changes( f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " f"ELSE __source.{updated_at_column} END" ) + history_join_sql: str = "" + if invalidate_hard_deletes: + key_sql: str = ", ".join(unique_key) + history_condition: str = " AND ".join( + f"__history.{column} = __source.{column}" for column in unique_key + ) + version_valid_from_expr = ( + f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " + f"AND __history.__closed_at <> __source.{updated_at_column} " + f"THEN {current_timestamp} " + f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " + f"ELSE __source.{updated_at_column} END" + ) + history_join_sql = ( + f"LEFT JOIN (SELECT {key_sql}, MAX({valid_to_column}) AS __closed_at " + f"FROM {destination} GROUP BY {key_sql}) AS __history ON {history_condition} " + ) insert_sql: str = ( f"INSERT INTO {destination} ({insert_column_sql}) " f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " f"FROM {origin} AS __source " f"LEFT JOIN {destination} AS __active " f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " + f"{history_join_sql}" f"WHERE __active.{first_key} IS NULL " f"OR __source.{updated_at_column} > __active.{updated_at_column}" ) diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py index 5dd7050a96..9c3db579ac 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py @@ -2072,12 +2072,30 @@ def render_apply_timestamp_snapshot_changes( f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " f"ELSE __source.{updated_at_column} END" ) + history_join_sql: str = "" + if invalidate_hard_deletes: + key_sql: str = ", ".join(unique_key) + history_condition: str = " AND ".join( + f"__history.{column} = __source.{column}" for column in unique_key + ) + version_valid_from_expr = ( + f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " + f"AND __history.__closed_at <> __source.{updated_at_column} " + f"THEN {current_timestamp} " + f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " + f"ELSE __source.{updated_at_column} END" + ) + history_join_sql = ( + f"LEFT JOIN (SELECT {key_sql}, MAX({valid_to_column}) AS __closed_at " + f"FROM {destination} GROUP BY {key_sql}) AS __history ON {history_condition} " + ) insert_sql: str = ( f"INSERT INTO {destination} ({insert_column_sql}) " f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS DATETIME2) " f"FROM {origin} AS __source " f"LEFT JOIN {destination} AS __active " f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " + f"{history_join_sql}" f"WHERE __active.{first_key} IS NULL " f"OR __source.{updated_at_column} > __active.{updated_at_column}" ) diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py index 545f2754f7..8ec71fccc4 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py @@ -262,6 +262,34 @@ (2, "active", FIRST_CLOCK_DAY + 2, None), ), ), + SnapshotExecutionScenario( + description="current timestamp hard delete then unchanged reappearance", + kind=CURRENT_TIMESTAMP_HARD_DELETES, + builds=( + ((1, "basic", 1), (2, "pro", 1)), + ((1, "basic", 1),), + ((1, "basic", 1), (2, "pro", 1)), + ), + expected_history=( + (1, "basic", 1, None), + (2, "pro", 1, FIRST_CLOCK_DAY + 1), + (2, "pro", FIRST_CLOCK_DAY + 2, None), + ), + ), + SnapshotExecutionScenario( + description="current timestamp hard delete then reappearance with a newer updated_at", + kind=CURRENT_TIMESTAMP_HARD_DELETES, + builds=( + ((1, "basic", 1), (2, "pro", 1)), + ((1, "basic", 1),), + ((1, "basic", 1), (2, "team", 5)), + ), + expected_history=( + (1, "basic", 1, None), + (2, "pro", 1, FIRST_CLOCK_DAY + 1), + (2, "team", FIRST_CLOCK_DAY + 2, None), + ), + ), ) _RUNS: tuple[SnapshotExecutionRun, ...] = build_execution_runs( diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py index 734f74e101..1d859dabfa 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py @@ -55,6 +55,9 @@ "__target.customer_id = __source.customer_id", "__target.region = __source.region", "WHERE __target.effective_to IS NULL AND NOT EXISTS", + "MAX(effective_to) AS __closed_at FROM target_table " + "GROUP BY customer_id, region) AS __history", + "AND __history.__closed_at <> __source.updated_at THEN", ), expected_historical_check_initial_hard_delete_fragments=( "__next_absence_at", @@ -90,6 +93,9 @@ "__target.customer_id = __source.customer_id", "__target.region = __source.region", "WHERE __target.effective_to IS NULL AND NOT EXISTS", + "MAX(effective_to) AS __closed_at FROM target_table " + "GROUP BY customer_id, region) AS __history", + "AND __history.__closed_at <> __source.updated_at THEN", ), expected_historical_check_initial_hard_delete_fragments=( "__next_absence_at", @@ -125,6 +131,9 @@ "__target.customer_id = __source.customer_id", "__target.region = __source.region", "WHERE __target.effective_to IS NULL AND NOT EXISTS", + "MAX(effective_to) AS __closed_at FROM target_table " + "GROUP BY customer_id, region) AS __history", + "AND __history.__closed_at <> __source.updated_at THEN", ), expected_historical_check_initial_hard_delete_fragments=( "__next_absence_at", @@ -160,6 +169,9 @@ "__target.customer_id = __source.customer_id", "__target.region = __source.region", "WHEN MATCHED THEN UPDATE SET effective_to = __source.updated_at", + "MAX(effective_to) AS __closed_at FROM target_table " + "GROUP BY customer_id, region) AS __history", + "AND __history.__closed_at <> __source.updated_at THEN", ), expected_historical_check_initial_hard_delete_fragments=( "__next_absence_at", @@ -195,6 +207,9 @@ "__target.customer_id = __source.customer_id", "__target.region = __source.region", "WHERE __target.effective_to IS NULL AND NOT EXISTS", + "MAX(effective_to) AS __closed_at FROM target_table " + "GROUP BY customer_id, region) AS __history", + "AND __history.__closed_at <> __source.updated_at THEN", ), expected_historical_check_initial_hard_delete_fragments=( "__next_absence_at", From b7d6506cb5a60dad1f9a6638f9f894b39651358f Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:53:40 +0100 Subject: [PATCH 45/57] refactor(snapshots): share exact-copy historical check snapshot helpers Move private snapshot SQL helpers that were byte-for-byte copies across adapters into shared builders next to HistoricalSnapshotSql: - SnapshotSql.key_condition replaces seven copies of _snapshot_key_condition. - HistoricalCheckSnapshotSql holds the initial-build query and the QUALIFY-based incremental CTEs previously copied into the base, DuckDB, BigQuery, Snowflake and Databricks adapters. - HistoricalSnapshotStatementSql holds the UPDATE ... FROM close and INSERT ... WITH statements copied between BigQuery and Snowflake. Rendered SQL is unchanged for every adapter. (cherry picked from commit d650c86abfb75561640d32f2fbddd2084c7c1850) --- .../adapter/contract/classes/base_adapter.py | 148 ++--------- .../contract/classes/duckdb_backed_adapter.py | 154 ++---------- .../classes/historical_check_snapshot_sql.py | 128 ++++++++++ .../historical_snapshot_statement_sql.py | 78 ++++++ .../adapter/contract/classes/snapshot_sql.py | 15 ++ .../bigquery/classes/bigquery_adapter.py | 230 ++---------------- .../databricks/classes/databricks_adapter.py | 154 ++---------- .../postgres/classes/postgres_adapter.py | 28 ++- .../snowflake/classes/snowflake_adapter.py | 230 ++---------------- .../sqlserver/classes/sqlserver_adapter.py | 24 +- 10 files changed, 347 insertions(+), 842 deletions(-) create mode 100644 src/sqlbuild/adapter/contract/classes/historical_check_snapshot_sql.py create mode 100644 src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py create mode 100644 src/sqlbuild/adapter/contract/classes/snapshot_sql.py diff --git a/src/sqlbuild/adapter/contract/classes/base_adapter.py b/src/sqlbuild/adapter/contract/classes/base_adapter.py index d66a69cf7f..25da4219dd 100644 --- a/src/sqlbuild/adapter/contract/classes/base_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/base_adapter.py @@ -9,11 +9,15 @@ from pathlib import Path from typing import Any, ClassVar, cast +from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( + HistoricalCheckSnapshotSql, +) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( HistoricalSnapshotSql, historical_insert_validity_sql, ) from sqlbuild.adapter.contract.classes.retention_adapter import RetentionAdapterMixin +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.strict_adapter import StrictAdapter from sqlbuild.adapter.contract.constants import ( @@ -891,7 +895,7 @@ def render_apply_timestamp_snapshot_changes( source_alias="__source", current_timestamp=current_timestamp, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__source", unique_key=unique_key ) close_sql: str = ( @@ -904,7 +908,7 @@ def render_apply_timestamp_snapshot_changes( ) insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = _snapshot_key_condition( + active_join_condition: str = SnapshotSql.key_condition( left_alias="__active", right_alias="__source", unique_key=unique_key ) first_key: str = unique_key[0] @@ -978,7 +982,7 @@ def render_apply_check_snapshot_changes( source_alias="__source", current_timestamp=current_timestamp, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__source", unique_key=unique_key ) change_condition: str = " OR ".join( @@ -994,7 +998,7 @@ def render_apply_check_snapshot_changes( ) insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = _snapshot_key_condition( + active_join_condition: str = SnapshotSql.key_condition( left_alias="__active", right_alias="__source", unique_key=unique_key ) active_change_condition: str = " OR ".join( @@ -1101,7 +1105,7 @@ def render_apply_historical_timestamp_snapshot_changes( valid_from_column=valid_from_column, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) if invalidate_hard_deletes: @@ -1163,7 +1167,7 @@ def render_apply_historical_timestamp_changes( updated_at_column=updated_at_column, valid_to_column=valid_to_column, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) close_sql: str = ( @@ -1208,7 +1212,7 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = _historical_check_snapshot_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql.initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -1233,7 +1237,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = _historical_check_new_changes_cte_sql( + new_changes_sql: str = HistoricalCheckSnapshotSql.new_changes_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1243,7 +1247,7 @@ def render_apply_historical_check_snapshot_changes( valid_to_column=valid_to_column, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) if invalidate_hard_deletes: @@ -2608,7 +2612,7 @@ def _snapshot_hard_delete_close_sql( valid_to_column: str, current_timestamp: str, ) -> str: - missing_key_condition: str = _snapshot_key_condition( + missing_key_condition: str = SnapshotSql.key_condition( left_alias="__source", right_alias="__target", unique_key=unique_key ) first_key: str = unique_key[0] @@ -2623,54 +2627,6 @@ def _snapshot_hard_delete_close_sql( ) -def _historical_check_snapshot_select_sql( - *, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, -) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - check_columns=check_columns, - ).initial_select_sql(output_columns=output_columns) - - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - f"LEAD({observed_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {observed_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - def _historical_timestamp_snapshot_select_sql( *, origin: str, @@ -2733,7 +2689,7 @@ def _historical_timestamp_new_changes_cte_sql( partition_sql: str = ", ".join(unique_key) first_key: str = unique_key[0] - latest_join_condition: str = _snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) return ( @@ -2784,7 +2740,7 @@ def _historical_timestamp_changes_new_records_cte_sql( updated_at_column: str, valid_to_column: str, ) -> str: - latest_join_condition: str = _snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__source", right_alias="__latest", unique_key=unique_key ) partition_sql: str = ", ".join(unique_key) @@ -2803,76 +2759,6 @@ def _historical_timestamp_changes_new_records_cte_sql( ) -def _historical_check_new_changes_cte_sql( - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, -) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - check_columns=check_columns, - ).new_changes_ctes_sql(destination=destination) - - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - delta_change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - latest_join_condition: str = _snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" for column in check_columns - ) - first_key: str = unique_key[0] - changed_or_first_sql: str = ( - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({delta_change_condition})" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition})" - ")" - ")" - ) - - -def _snapshot_key_condition( - *, left_alias: str, right_alias: str, unique_key: tuple[str, ...] -) -> str: - return " AND ".join(f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key) - - def _historical_snapshot_combined_close_sql( *, destination: str, @@ -2882,7 +2768,7 @@ def _historical_snapshot_combined_close_sql( valid_to_column: str, change_time_column: str, ) -> str: - close_candidate_condition: str = _snapshot_key_condition( + close_candidate_condition: str = SnapshotSql.key_condition( left_alias="__close_candidates", right_alias="__target", unique_key=unique_key ) candidate_key_sql: str = ", ".join(unique_key) diff --git a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py index 2576e1a99b..ee240327e5 100644 --- a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py @@ -19,11 +19,15 @@ _render_ansi_typed_scalar, _render_typed_value_list, ) +from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( + HistoricalCheckSnapshotSql, +) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( HistoricalSnapshotSql, historical_insert_validity_sql, ) from sqlbuild.adapter.contract.classes.observed_connection import ObservedConnection +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin from sqlbuild.adapter.contract.constants import ( @@ -444,7 +448,7 @@ def render_apply_timestamp_snapshot_changes( source_alias="__source", current_timestamp=current_timestamp, ) - key_condition: str = self._snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__source", unique_key=unique_key ) close_sql: str = ( @@ -457,7 +461,7 @@ def render_apply_timestamp_snapshot_changes( ) insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( + active_join_condition: str = SnapshotSql.key_condition( left_alias="__active", right_alias="__source", unique_key=unique_key ) first_key: str = unique_key[0] @@ -581,7 +585,7 @@ def render_apply_historical_timestamp_snapshot_changes( valid_to_column=valid_to_column, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = self._snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) if invalidate_hard_deletes: @@ -643,7 +647,7 @@ def render_apply_historical_timestamp_changes( updated_at_column=updated_at_column, valid_to_column=valid_to_column, ) - key_condition: str = self._snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) close_sql: str = ( @@ -699,7 +703,7 @@ def render_apply_check_snapshot_changes( source_alias="__source", current_timestamp=current_timestamp, ) - key_condition: str = self._snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__source", unique_key=unique_key ) change_condition: str = " OR ".join( @@ -715,7 +719,7 @@ def render_apply_check_snapshot_changes( ) insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( + active_join_condition: str = SnapshotSql.key_condition( left_alias="__active", right_alias="__source", unique_key=unique_key ) active_change_condition: str = " OR ".join( @@ -764,7 +768,7 @@ def render_create_initial_historical_check_snapshot_destination( ) -> tuple[str, ...]: return self.render_create_table_as( destination=destination, - sql=self._historical_check_snapshot_select_sql( + sql=HistoricalCheckSnapshotSql.initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -789,7 +793,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_check_new_changes_cte_sql( + new_changes_sql: str = HistoricalCheckSnapshotSql.new_changes_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -799,7 +803,7 @@ def render_apply_historical_check_snapshot_changes( valid_to_column=valid_to_column, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = self._snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) if invalidate_hard_deletes: @@ -2615,7 +2619,7 @@ def _snapshot_hard_delete_close_sql( valid_to_column: str, current_timestamp: str, ) -> str: - missing_key_condition: str = cls._snapshot_key_condition( + missing_key_condition: str = SnapshotSql.key_condition( left_alias="__source", right_alias="__target", unique_key=unique_key ) first_key: str = unique_key[0] @@ -2629,55 +2633,6 @@ def _snapshot_hard_delete_close_sql( f")" ) - @classmethod - def _historical_check_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - check_columns=check_columns, - ).initial_select_sql(output_columns=output_columns) - - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - f"LEAD({observed_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {observed_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - @classmethod def _historical_timestamp_snapshot_select_sql( cls, @@ -2742,7 +2697,7 @@ def _historical_timestamp_new_changes_cte_sql( ).new_changes_ctes_sql(destination=destination) partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = cls._snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) first_key: str = unique_key[0] @@ -2796,7 +2751,7 @@ def _historical_timestamp_changes_new_records_cte_sql( updated_at_column: str, valid_to_column: str, ) -> str: - latest_join_condition: str = cls._snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__source", right_alias="__latest", unique_key=unique_key ) partition_sql: str = ", ".join(unique_key) @@ -2814,81 +2769,6 @@ def _historical_timestamp_changes_new_records_cte_sql( ")" ) - @classmethod - def _historical_check_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - check_columns=check_columns, - ).new_changes_ctes_sql(destination=destination) - - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - delta_change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - first_key: str = unique_key[0] - changed_or_first_sql: str = ( - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({delta_change_condition})" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition})" - ")" - ")" - ) - - @staticmethod - def _snapshot_key_condition( - *, left_alias: str, right_alias: str, unique_key: tuple[str, ...] - ) -> str: - return " AND ".join( - f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key - ) - @classmethod def _historical_snapshot_combined_close_sql( cls, @@ -2900,7 +2780,7 @@ def _historical_snapshot_combined_close_sql( valid_to_column: str, change_time_column: str, ) -> str: - close_candidate_condition: str = cls._snapshot_key_condition( + close_candidate_condition: str = SnapshotSql.key_condition( left_alias="__close_candidates", right_alias="__target", unique_key=unique_key ) candidate_key_sql: str = ", ".join(unique_key) diff --git a/src/sqlbuild/adapter/contract/classes/historical_check_snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/historical_check_snapshot_sql.py new file mode 100644 index 0000000000..220053ada3 --- /dev/null +++ b/src/sqlbuild/adapter/contract/classes/historical_check_snapshot_sql.py @@ -0,0 +1,128 @@ +"""Shared SQL for historical check snapshots.""" + +from __future__ import annotations + +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import HistoricalSnapshotSql +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql + + +class HistoricalCheckSnapshotSql: + """Render SQL for historical check snapshots.""" + + @staticmethod + def initial_select_sql( + *, + origin: str, + unique_key: tuple[str, ...], + check_columns: tuple[str, ...], + observed_at_column: str, + valid_from_column: str, + valid_to_column: str, + output_columns: tuple[str, ...], + invalidate_hard_deletes: bool, + ) -> str: + """Render the initial-build query for a historical check snapshot.""" + + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + ).initial_select_sql(output_columns=output_columns) + + partition_sql: str = ", ".join(unique_key) + previous_columns_sql: str = ", ".join( + f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " + f"AS __prev_{column}" + for column in check_columns + ) + if previous_columns_sql: + previous_columns_sql = f", {previous_columns_sql}" + change_condition: str = " OR ".join( + f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns + ) + output_select_sql: str = ", ".join(column for column in output_columns) + return ( + "WITH __ordered AS (" + f"SELECT *, LAG({observed_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" + f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" + "), __changes AS (" + f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" + ") " + f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " + f"LEAD({observed_at_column}) OVER (PARTITION BY {partition_sql} " + f"ORDER BY {observed_at_column}) AS {valid_to_column} " + "FROM __changes" + ) + + @staticmethod + def new_changes_ctes_sql( + *, + destination: str, + origin: str, + unique_key: tuple[str, ...], + check_columns: tuple[str, ...], + observed_at_column: str, + valid_from_column: str, + valid_to_column: str, + invalidate_hard_deletes: bool, + ) -> str: + """Render incremental historical check CTEs using ``QUALIFY`` for the latest versions.""" + + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + ).new_changes_ctes_sql(destination=destination) + + partition_sql: str = ", ".join(unique_key) + previous_columns_sql: str = ", ".join( + f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " + f"AS __prev_{column}" + for column in check_columns + ) + if previous_columns_sql: + previous_columns_sql = f", {previous_columns_sql}" + delta_change_condition: str = " OR ".join( + f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns + ) + latest_join_condition: str = SnapshotSql.key_condition( + left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key + ) + latest_change_condition: str = " OR ".join( + f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" + for column in check_columns + ) + first_key: str = unique_key[0] + changed_or_first_sql: str = ( + "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " + f"OR ({delta_change_condition})" + ) + return ( + "__ordered AS (" + f"SELECT *, LAG({observed_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" + f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" + "), __delta_changes AS (" + f"{changed_or_first_sql}" + "), __latest AS (" + f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" + f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" + ") = 1" + "), __new_changes AS (" + "SELECT __delta_changes.* FROM __delta_changes " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{first_key} IS NULL OR (" + f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " + f"AND ({latest_change_condition})" + ")" + ")" + ) diff --git a/src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py b/src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py new file mode 100644 index 0000000000..af27517117 --- /dev/null +++ b/src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py @@ -0,0 +1,78 @@ +"""Shared close and insert statements for historical snapshots.""" + +from __future__ import annotations + +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql + + +class HistoricalSnapshotStatementSql: + """Render close and insert statements shared by historical snapshots.""" + + @staticmethod + def grouped_combined_close_sql( + *, + destination: str, + new_changes_sql: str, + unique_key: tuple[str, ...], + valid_from_column: str, + valid_to_column: str, + change_time_column: str, + ) -> str: + """Render an ``UPDATE ... FROM`` close over new versions and hard deletes.""" + + candidate_key_sql: str = ", ".join(unique_key) + return HistoricalSnapshotStatementSql.grouped_close_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + close_candidates_sql=( + f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at " + "FROM __new_changes " + "UNION ALL " + f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " + "WHERE __close_at IS NOT NULL" + ), + ) + + @staticmethod + def grouped_close_sql( + *, + destination: str, + new_changes_sql: str, + unique_key: tuple[str, ...], + valid_from_column: str, + valid_to_column: str, + close_candidates_sql: str, + ) -> str: + """Render an ``UPDATE ... FROM`` close that uses the earliest candidate per key.""" + + close_candidate_condition: str = SnapshotSql.key_condition( + left_alias="__close_candidates", right_alias="__target", unique_key=unique_key + ) + candidate_key_sql: str = ", ".join(unique_key) + close_candidates_query: str = ( + f"WITH {new_changes_sql}, __close_candidates AS ({close_candidates_sql}) " + f"SELECT {candidate_key_sql}, MIN(__close_at) AS __close_at " + "FROM __close_candidates GROUP BY " + f"{candidate_key_sql}" + ) + return ( + f"UPDATE {destination} AS __target " + f"SET {valid_to_column} = __close_candidates.__close_at " + f"FROM ({close_candidates_query}) AS __close_candidates " + f"WHERE __target.{valid_to_column} IS NULL " + f"AND __target.{valid_from_column} < __close_candidates.__close_at " + f"AND {close_candidate_condition}" + ) + + @staticmethod + def insert_with_cte_sql( + *, destination: str, insert_column_sql: str, new_changes_sql: str, select_sql: str + ) -> str: + """Render an insert whose CTEs follow the ``INSERT INTO`` column list.""" + + return ( + f"INSERT INTO {destination} ({insert_column_sql}) WITH {new_changes_sql} {select_sql}" + ) diff --git a/src/sqlbuild/adapter/contract/classes/snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/snapshot_sql.py new file mode 100644 index 0000000000..aea802d795 --- /dev/null +++ b/src/sqlbuild/adapter/contract/classes/snapshot_sql.py @@ -0,0 +1,15 @@ +"""Shared SQL fragments for snapshot materializations.""" + +from __future__ import annotations + + +class SnapshotSql: + """Render SQL fragments shared by snapshot materializations.""" + + @staticmethod + def key_condition(*, left_alias: str, right_alias: str, unique_key: tuple[str, ...]) -> str: + """Render the equality join between two aliases over the snapshot unique key.""" + + return " AND ".join( + f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key + ) diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py index 5efdd42cde..7ee868be40 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py @@ -18,11 +18,18 @@ _typed_collection_items, _typed_scalar_payload, ) +from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( + HistoricalCheckSnapshotSql, +) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( HistoricalSnapshotSql, historical_insert_validity_sql, ) +from sqlbuild.adapter.contract.classes.historical_snapshot_statement_sql import ( + HistoricalSnapshotStatementSql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin from sqlbuild.adapter.contract.constants import ( @@ -1039,7 +1046,7 @@ def render_apply_timestamp_snapshot_changes( source_alias="__source", current_timestamp=current_timestamp, ) - key_condition: str = self._snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__source", unique_key=unique_key ) close_sql: str = ( @@ -1052,7 +1059,7 @@ def render_apply_timestamp_snapshot_changes( ) insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( + active_join_condition: str = SnapshotSql.key_condition( left_alias="__active", right_alias="__source", unique_key=unique_key ) first_key: str = unique_key[0] @@ -1173,7 +1180,7 @@ def render_apply_historical_timestamp_snapshot_changes( invalidate_hard_deletes=invalidate_hard_deletes, ) if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( + close_sql: str = HistoricalSnapshotStatementSql.grouped_combined_close_sql( destination=destination, new_changes_sql=new_changes_sql, unique_key=unique_key, @@ -1182,7 +1189,7 @@ def render_apply_historical_timestamp_snapshot_changes( change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: - close_sql = self._historical_snapshot_close_sql( + close_sql = HistoricalSnapshotStatementSql.grouped_close_sql( destination=destination, new_changes_sql=new_changes_sql, unique_key=unique_key, @@ -1202,7 +1209,7 @@ def render_apply_historical_timestamp_snapshot_changes( ) if invalidate_hard_deletes: version_columns_sql = historical_insert_validity_sql() - insert_sql: str = self._historical_snapshot_insert_sql( + insert_sql: str = HistoricalSnapshotStatementSql.insert_with_cte_sql( destination=destination, insert_column_sql=insert_column_sql, new_changes_sql=new_changes_sql, @@ -1228,7 +1235,7 @@ def render_apply_historical_timestamp_changes( updated_at_column=updated_at_column, valid_to_column=valid_to_column, ) - close_sql: str = self._historical_snapshot_close_sql( + close_sql: str = HistoricalSnapshotStatementSql.grouped_close_sql( destination=destination, new_changes_sql=new_changes_sql, unique_key=unique_key, @@ -1242,7 +1249,7 @@ def render_apply_historical_timestamp_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = self._historical_snapshot_insert_sql( + insert_sql: str = HistoricalSnapshotStatementSql.insert_with_cte_sql( destination=destination, insert_column_sql=insert_column_sql, new_changes_sql=new_changes_sql, @@ -1280,7 +1287,7 @@ def render_apply_check_snapshot_changes( source_alias="__source", current_timestamp=current_timestamp, ) - key_condition: str = self._snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__source", unique_key=unique_key ) change_condition: str = " OR ".join( @@ -1296,7 +1303,7 @@ def render_apply_check_snapshot_changes( ) insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( + active_join_condition: str = SnapshotSql.key_condition( left_alias="__active", right_alias="__source", unique_key=unique_key ) active_change_condition: str = " OR ".join( @@ -1343,7 +1350,7 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = self._historical_check_snapshot_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql.initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -1368,7 +1375,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_check_new_changes_cte_sql( + new_changes_sql: str = HistoricalCheckSnapshotSql.new_changes_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1379,7 +1386,7 @@ def render_apply_historical_check_snapshot_changes( invalidate_hard_deletes=invalidate_hard_deletes, ) if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( + close_sql: str = HistoricalSnapshotStatementSql.grouped_combined_close_sql( destination=destination, new_changes_sql=new_changes_sql, unique_key=unique_key, @@ -1388,7 +1395,7 @@ def render_apply_historical_check_snapshot_changes( change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: - close_sql = self._historical_snapshot_close_sql( + close_sql = HistoricalSnapshotStatementSql.grouped_close_sql( destination=destination, new_changes_sql=new_changes_sql, unique_key=unique_key, @@ -1408,7 +1415,7 @@ def render_apply_historical_check_snapshot_changes( ) if invalidate_hard_deletes: version_columns_sql = historical_insert_validity_sql() - insert_sql: str = self._historical_snapshot_insert_sql( + insert_sql: str = HistoricalSnapshotStatementSql.insert_with_cte_sql( destination=destination, insert_column_sql=insert_column_sql, new_changes_sql=new_changes_sql, @@ -2955,7 +2962,7 @@ def _snapshot_hard_delete_close_sql( valid_to_column: str, current_timestamp: str, ) -> str: - missing_key_condition: str = cls._snapshot_key_condition( + missing_key_condition: str = SnapshotSql.key_condition( left_alias="__source", right_alias="__target", unique_key=unique_key ) first_key: str = unique_key[0] @@ -2969,55 +2976,6 @@ def _snapshot_hard_delete_close_sql( f")" ) - @classmethod - def _historical_check_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - check_columns=check_columns, - ).initial_select_sql(output_columns=output_columns) - - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - f"LEAD({observed_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {observed_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - @classmethod def _historical_timestamp_snapshot_select_sql( cls, @@ -3082,7 +3040,7 @@ def _historical_timestamp_new_changes_cte_sql( ).new_changes_ctes_sql(destination=destination) partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = cls._snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) first_key: str = unique_key[0] @@ -3136,7 +3094,7 @@ def _historical_timestamp_changes_new_records_cte_sql( updated_at_column: str, valid_to_column: str, ) -> str: - latest_join_condition: str = cls._snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__source", right_alias="__latest", unique_key=unique_key ) partition_sql: str = ", ".join(unique_key) @@ -3153,143 +3111,3 @@ def _historical_timestamp_changes_new_records_cte_sql( f"OR __source.{updated_at_column} > __latest.{updated_at_column}" ")" ) - - @classmethod - def _historical_check_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - check_columns=check_columns, - ).new_changes_ctes_sql(destination=destination) - - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - delta_change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - first_key: str = unique_key[0] - changed_or_first_sql: str = ( - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({delta_change_condition})" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition})" - ")" - ")" - ) - - @staticmethod - def _snapshot_key_condition( - *, left_alias: str, right_alias: str, unique_key: tuple[str, ...] - ) -> str: - return " AND ".join( - f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key - ) - - @classmethod - def _historical_snapshot_combined_close_sql( - cls, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - change_time_column: str, - ) -> str: - candidate_key_sql: str = ", ".join(unique_key) - return cls._historical_snapshot_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at " - "FROM __new_changes " - "UNION ALL " - f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " - "WHERE __close_at IS NOT NULL" - ), - ) - - @classmethod - def _historical_snapshot_close_sql( - cls, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - close_candidates_sql: str, - ) -> str: - close_candidate_condition: str = cls._snapshot_key_condition( - left_alias="__close_candidates", right_alias="__target", unique_key=unique_key - ) - candidate_key_sql: str = ", ".join(unique_key) - close_candidates_query: str = ( - f"WITH {new_changes_sql}, __close_candidates AS ({close_candidates_sql}) " - f"SELECT {candidate_key_sql}, MIN(__close_at) AS __close_at " - "FROM __close_candidates GROUP BY " - f"{candidate_key_sql}" - ) - return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = __close_candidates.__close_at " - f"FROM ({close_candidates_query}) AS __close_candidates " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < __close_candidates.__close_at " - f"AND {close_candidate_condition}" - ) - - @staticmethod - def _historical_snapshot_insert_sql( - *, destination: str, insert_column_sql: str, new_changes_sql: str, select_sql: str - ) -> str: - return ( - f"INSERT INTO {destination} ({insert_column_sql}) WITH {new_changes_sql} {select_sql}" - ) diff --git a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py index 241e3737c9..88884152fa 100644 --- a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py +++ b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py @@ -18,11 +18,15 @@ _render_ansi_typed_scalar, _render_typed_value_list, ) +from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( + HistoricalCheckSnapshotSql, +) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( HistoricalSnapshotSql, historical_insert_validity_sql, ) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin from sqlbuild.adapter.contract.constants import ( @@ -889,7 +893,7 @@ def render_apply_timestamp_snapshot_changes( source_alias="__source", current_timestamp=current_timestamp, ) - key_condition: str = self._snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__source", unique_key=unique_key ) close_sql: str = ( @@ -902,7 +906,7 @@ def render_apply_timestamp_snapshot_changes( ) insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( + active_join_condition: str = SnapshotSql.key_condition( left_alias="__active", right_alias="__source", unique_key=unique_key ) first_key: str = unique_key[0] @@ -1022,7 +1026,7 @@ def render_apply_historical_timestamp_snapshot_changes( valid_to_column=valid_to_column, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = self._snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) if invalidate_hard_deletes: @@ -1084,7 +1088,7 @@ def render_apply_historical_timestamp_changes( updated_at_column=updated_at_column, valid_to_column=valid_to_column, ) - key_condition: str = self._snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) close_sql: str = ( @@ -1140,7 +1144,7 @@ def render_apply_check_snapshot_changes( source_alias="__source", current_timestamp=current_timestamp, ) - key_condition: str = self._snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__source", unique_key=unique_key ) change_condition: str = " OR ".join( @@ -1156,7 +1160,7 @@ def render_apply_check_snapshot_changes( ) insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( + active_join_condition: str = SnapshotSql.key_condition( left_alias="__active", right_alias="__source", unique_key=unique_key ) active_change_condition: str = " OR ".join( @@ -1203,7 +1207,7 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = self._historical_check_snapshot_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql.initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -1228,7 +1232,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_check_new_changes_cte_sql( + new_changes_sql: str = HistoricalCheckSnapshotSql.new_changes_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1238,7 +1242,7 @@ def render_apply_historical_check_snapshot_changes( valid_to_column=valid_to_column, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = self._snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) if invalidate_hard_deletes: @@ -3092,7 +3096,7 @@ def _snapshot_hard_delete_close_sql( valid_to_column: str, current_timestamp: str, ) -> str: - missing_key_condition: str = cls._snapshot_key_condition( + missing_key_condition: str = SnapshotSql.key_condition( left_alias="__source", right_alias="__target", unique_key=unique_key ) first_key: str = unique_key[0] @@ -3106,55 +3110,6 @@ def _snapshot_hard_delete_close_sql( f")" ) - @classmethod - def _historical_check_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - check_columns=check_columns, - ).initial_select_sql(output_columns=output_columns) - - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - f"LEAD({observed_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {observed_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - @classmethod def _historical_timestamp_snapshot_select_sql( cls, @@ -3219,7 +3174,7 @@ def _historical_timestamp_new_changes_cte_sql( ).new_changes_ctes_sql(destination=destination) partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = cls._snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) first_key: str = unique_key[0] @@ -3273,7 +3228,7 @@ def _historical_timestamp_changes_new_records_cte_sql( updated_at_column: str, valid_to_column: str, ) -> str: - latest_join_condition: str = cls._snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__source", right_alias="__latest", unique_key=unique_key ) partition_sql: str = ", ".join(unique_key) @@ -3291,81 +3246,6 @@ def _historical_timestamp_changes_new_records_cte_sql( ")" ) - @classmethod - def _historical_check_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - check_columns=check_columns, - ).new_changes_ctes_sql(destination=destination) - - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - delta_change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - first_key: str = unique_key[0] - changed_or_first_sql: str = ( - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({delta_change_condition})" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition})" - ")" - ")" - ) - - @staticmethod - def _snapshot_key_condition( - *, left_alias: str, right_alias: str, unique_key: tuple[str, ...] - ) -> str: - return " AND ".join( - f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key - ) - @classmethod def _historical_snapshot_combined_close_sql( cls, @@ -3377,7 +3257,7 @@ def _historical_snapshot_combined_close_sql( valid_to_column: str, change_time_column: str, ) -> str: - close_candidate_condition: str = cls._snapshot_key_condition( + close_candidate_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__close_candidates", unique_key=unique_key ) candidate_key_sql: str = ", ".join(unique_key) diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py index a7358cf990..822794c5f1 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py @@ -15,7 +15,6 @@ _build_names_filter, _build_schemas_filter, _encode_typed_json, - _historical_check_snapshot_select_sql, _historical_snapshot_combined_close_sql, _historical_timestamp_changes_select_sql, _historical_timestamp_snapshot_select_sql, @@ -24,14 +23,17 @@ _render_typed_value_list, _snapshot_hard_delete_close_sql, _snapshot_initial_valid_from_expr, - _snapshot_key_condition, _validate_rectangular_typed_array, ) +from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( + HistoricalCheckSnapshotSql, +) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( HistoricalSnapshotSql, historical_insert_validity_sql, ) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin from sqlbuild.adapter.contract.constants import ( @@ -1153,7 +1155,7 @@ def render_apply_timestamp_snapshot_changes( source_alias="__source", current_timestamp=current_timestamp, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__source", unique_key=unique_key ) close_sql: str = ( @@ -1166,7 +1168,7 @@ def render_apply_timestamp_snapshot_changes( ) insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = _snapshot_key_condition( + active_join_condition: str = SnapshotSql.key_condition( left_alias="__active", right_alias="__source", unique_key=unique_key ) first_key: str = unique_key[0] @@ -1240,7 +1242,7 @@ def render_apply_check_snapshot_changes( source_alias="__source", current_timestamp=current_timestamp, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__source", unique_key=unique_key ) change_condition: str = " OR ".join( @@ -1256,7 +1258,7 @@ def render_apply_check_snapshot_changes( ) insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = _snapshot_key_condition( + active_join_condition: str = SnapshotSql.key_condition( left_alias="__active", right_alias="__source", unique_key=unique_key ) active_change_condition: str = " OR ".join( @@ -1351,7 +1353,7 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = _historical_check_snapshot_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql.initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -2422,7 +2424,7 @@ def render_apply_historical_timestamp_snapshot_changes( valid_to_column=valid_to_column, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) if invalidate_hard_deletes: @@ -2484,7 +2486,7 @@ def render_apply_historical_timestamp_changes( updated_at_column=updated_at_column, valid_to_column=valid_to_column, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) close_sql: str = ( @@ -2538,7 +2540,7 @@ def render_apply_historical_check_snapshot_changes( valid_to_column=valid_to_column, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) if invalidate_hard_deletes: @@ -2612,7 +2614,7 @@ def _pg_historical_timestamp_new_changes_cte_sql( f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" f") AS __rn FROM {destination}) AS __q WHERE __rn = 1)" ) - latest_join_condition: str = _snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) return ( @@ -2641,7 +2643,7 @@ def _pg_historical_timestamp_changes_new_records_cte_sql( valid_to_column: str, ) -> str: del valid_to_column - latest_join_condition: str = _snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__source", right_alias="__latest", unique_key=unique_key ) partition_sql: str = ", ".join(unique_key) @@ -2692,7 +2694,7 @@ def _pg_historical_check_new_changes_cte_sql( delta_change_condition: str = " OR ".join( f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns ) - latest_join_condition: str = _snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) latest_change_condition: str = " OR ".join( diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index 5a40211b9a..80bbad19ff 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -17,11 +17,18 @@ _render_ansi_typed_scalar, _render_typed_value_list, ) +from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( + HistoricalCheckSnapshotSql, +) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( HistoricalSnapshotSql, historical_insert_validity_sql, ) +from sqlbuild.adapter.contract.classes.historical_snapshot_statement_sql import ( + HistoricalSnapshotStatementSql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin from sqlbuild.adapter.contract.constants import ( @@ -1292,7 +1299,7 @@ def render_apply_timestamp_snapshot_changes( source_alias="__source", current_timestamp=current_timestamp, ) - key_condition: str = self._snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__source", unique_key=unique_key ) close_sql: str = ( @@ -1305,7 +1312,7 @@ def render_apply_timestamp_snapshot_changes( ) insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( + active_join_condition: str = SnapshotSql.key_condition( left_alias="__active", right_alias="__source", unique_key=unique_key ) first_key: str = unique_key[0] @@ -1430,7 +1437,7 @@ def render_apply_historical_timestamp_snapshot_changes( invalidate_hard_deletes=invalidate_hard_deletes, ) if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( + close_sql: str = HistoricalSnapshotStatementSql.grouped_combined_close_sql( destination=destination, new_changes_sql=new_changes_sql, unique_key=unique_key, @@ -1439,7 +1446,7 @@ def render_apply_historical_timestamp_snapshot_changes( change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: - close_sql = self._historical_snapshot_close_sql( + close_sql = HistoricalSnapshotStatementSql.grouped_close_sql( destination=destination, new_changes_sql=new_changes_sql, unique_key=unique_key, @@ -1459,7 +1466,7 @@ def render_apply_historical_timestamp_snapshot_changes( ) if invalidate_hard_deletes: version_columns_sql = historical_insert_validity_sql() - insert_sql: str = self._historical_snapshot_insert_sql( + insert_sql: str = HistoricalSnapshotStatementSql.insert_with_cte_sql( destination=destination, insert_column_sql=insert_column_sql, new_changes_sql=new_changes_sql, @@ -1485,7 +1492,7 @@ def render_apply_historical_timestamp_changes( updated_at_column=updated_at_column, valid_to_column=valid_to_column, ) - close_sql: str = self._historical_snapshot_close_sql( + close_sql: str = HistoricalSnapshotStatementSql.grouped_close_sql( destination=destination, new_changes_sql=new_changes_sql, unique_key=unique_key, @@ -1499,7 +1506,7 @@ def render_apply_historical_timestamp_changes( insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = self._historical_snapshot_insert_sql( + insert_sql: str = HistoricalSnapshotStatementSql.insert_with_cte_sql( destination=destination, insert_column_sql=insert_column_sql, new_changes_sql=new_changes_sql, @@ -1537,7 +1544,7 @@ def render_apply_check_snapshot_changes( source_alias="__source", current_timestamp=current_timestamp, ) - key_condition: str = self._snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__source", unique_key=unique_key ) change_condition: str = " OR ".join( @@ -1553,7 +1560,7 @@ def render_apply_check_snapshot_changes( ) insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( + active_join_condition: str = SnapshotSql.key_condition( left_alias="__active", right_alias="__source", unique_key=unique_key ) active_change_condition: str = " OR ".join( @@ -1600,7 +1607,7 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = self._historical_check_snapshot_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql.initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -1627,7 +1634,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_check_new_changes_cte_sql( + new_changes_sql: str = HistoricalCheckSnapshotSql.new_changes_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1638,7 +1645,7 @@ def render_apply_historical_check_snapshot_changes( invalidate_hard_deletes=invalidate_hard_deletes, ) if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( + close_sql: str = HistoricalSnapshotStatementSql.grouped_combined_close_sql( destination=destination, new_changes_sql=new_changes_sql, unique_key=unique_key, @@ -1647,7 +1654,7 @@ def render_apply_historical_check_snapshot_changes( change_time_column=SNAPSHOT_VERSION_START_COLUMN, ) else: - close_sql = self._historical_snapshot_close_sql( + close_sql = HistoricalSnapshotStatementSql.grouped_close_sql( destination=destination, new_changes_sql=new_changes_sql, unique_key=unique_key, @@ -1667,7 +1674,7 @@ def render_apply_historical_check_snapshot_changes( ) if invalidate_hard_deletes: version_columns_sql = historical_insert_validity_sql() - insert_sql: str = self._historical_snapshot_insert_sql( + insert_sql: str = HistoricalSnapshotStatementSql.insert_with_cte_sql( destination=destination, insert_column_sql=insert_column_sql, new_changes_sql=new_changes_sql, @@ -3304,7 +3311,7 @@ def _snapshot_hard_delete_close_sql( valid_to_column: str, current_timestamp: str, ) -> str: - missing_key_condition: str = cls._snapshot_key_condition( + missing_key_condition: str = SnapshotSql.key_condition( left_alias="__source", right_alias="__target", unique_key=unique_key ) first_key: str = unique_key[0] @@ -3318,55 +3325,6 @@ def _snapshot_hard_delete_close_sql( f")" ) - @classmethod - def _historical_check_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - check_columns=check_columns, - ).initial_select_sql(output_columns=output_columns) - - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - f"LEAD({observed_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {observed_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - @classmethod def _historical_timestamp_snapshot_select_sql( cls, @@ -3431,7 +3389,7 @@ def _historical_timestamp_new_changes_cte_sql( ).new_changes_ctes_sql(destination=destination) partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = cls._snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) first_key: str = unique_key[0] @@ -3485,7 +3443,7 @@ def _historical_timestamp_changes_new_records_cte_sql( updated_at_column: str, valid_to_column: str, ) -> str: - latest_join_condition: str = cls._snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__source", right_alias="__latest", unique_key=unique_key ) partition_sql: str = ", ".join(unique_key) @@ -3502,143 +3460,3 @@ def _historical_timestamp_changes_new_records_cte_sql( f"OR __source.{updated_at_column} > __latest.{updated_at_column}" ")" ) - - @classmethod - def _historical_check_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - check_columns=check_columns, - ).new_changes_ctes_sql(destination=destination) - - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - delta_change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - first_key: str = unique_key[0] - changed_or_first_sql: str = ( - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({delta_change_condition})" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition})" - ")" - ")" - ) - - @staticmethod - def _snapshot_key_condition( - *, left_alias: str, right_alias: str, unique_key: tuple[str, ...] - ) -> str: - return " AND ".join( - f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key - ) - - @classmethod - def _historical_snapshot_combined_close_sql( - cls, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - change_time_column: str, - ) -> str: - candidate_key_sql: str = ", ".join(unique_key) - return cls._historical_snapshot_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at " - "FROM __new_changes " - "UNION ALL " - f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " - "WHERE __close_at IS NOT NULL" - ), - ) - - @classmethod - def _historical_snapshot_close_sql( - cls, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - close_candidates_sql: str, - ) -> str: - close_candidate_condition: str = cls._snapshot_key_condition( - left_alias="__close_candidates", right_alias="__target", unique_key=unique_key - ) - candidate_key_sql: str = ", ".join(unique_key) - close_candidates_query: str = ( - f"WITH {new_changes_sql}, __close_candidates AS ({close_candidates_sql}) " - f"SELECT {candidate_key_sql}, MIN(__close_at) AS __close_at " - "FROM __close_candidates GROUP BY " - f"{candidate_key_sql}" - ) - return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = __close_candidates.__close_at " - f"FROM ({close_candidates_query}) AS __close_candidates " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < __close_candidates.__close_at " - f"AND {close_candidate_condition}" - ) - - @staticmethod - def _historical_snapshot_insert_sql( - *, destination: str, insert_column_sql: str, new_changes_sql: str, select_sql: str - ) -> str: - return ( - f"INSERT INTO {destination} ({insert_column_sql}) WITH {new_changes_sql} {select_sql}" - ) diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py index 9c3db579ac..cf3cb0fefd 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py @@ -20,7 +20,6 @@ _render_ansi_typed_scalar, _render_typed_value_list, _snapshot_initial_valid_from_expr, - _snapshot_key_condition, _typed_scalar_payload, ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( @@ -28,6 +27,7 @@ historical_insert_validity_sql, ) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin from sqlbuild.adapter.contract.constants import ( @@ -160,7 +160,7 @@ def _historical_snapshot_combined_close_sql( valid_to_column: str, change_time_column: str, ) -> str: - close_candidate_condition: str = _snapshot_key_condition( + close_candidate_condition: str = SnapshotSql.key_condition( left_alias="__close_candidates", right_alias="__target", unique_key=unique_key ) candidate_key_sql: str = ", ".join(unique_key) @@ -196,7 +196,7 @@ def _snapshot_hard_delete_close_sql( valid_to_column: str, current_timestamp: str, ) -> str: - missing_key_condition: str = _snapshot_key_condition( + missing_key_condition: str = SnapshotSql.key_condition( left_alias="__source", right_alias="__target", unique_key=unique_key ) first_key: str = unique_key[0] @@ -234,7 +234,7 @@ def _historical_timestamp_new_changes_cte_sql( ).new_changes_ctes_sql(destination=destination) partition_sql: str = ", ".join(unique_key) first_key: str = unique_key[0] - latest_join_condition: str = _snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) updated_changed: str = self._distinct_condition( @@ -297,7 +297,7 @@ def _historical_check_new_changes_cte_sql( self._distinct_condition(left=column, right=f"__prev_{column}") for column in check_columns ) - latest_join_condition: str = _snapshot_key_condition( + latest_join_condition: str = SnapshotSql.key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) latest_change_condition: str = " OR ".join( @@ -1793,7 +1793,7 @@ def render_apply_check_snapshot_changes( source_alias="__source", current_timestamp=current_timestamp, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__source", unique_key=unique_key ) change_condition: str = " OR ".join( @@ -1809,7 +1809,7 @@ def render_apply_check_snapshot_changes( ) insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = _snapshot_key_condition( + active_join_condition: str = SnapshotSql.key_condition( left_alias="__active", right_alias="__source", unique_key=unique_key ) active_change_condition: str = " OR ".join( @@ -1866,7 +1866,7 @@ def render_apply_historical_check_snapshot_changes( valid_to_column=valid_to_column, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) if invalidate_hard_deletes: @@ -1929,7 +1929,7 @@ def render_apply_historical_timestamp_changes( updated_at_column=updated_at_column, valid_to_column=valid_to_column, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) close_sql: str = ( @@ -1984,7 +1984,7 @@ def render_apply_historical_timestamp_snapshot_changes( valid_from_column=valid_from_column, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__new_changes", unique_key=unique_key ) if invalidate_hard_deletes: @@ -2052,7 +2052,7 @@ def render_apply_timestamp_snapshot_changes( source_alias="__source", current_timestamp=current_timestamp, ) - key_condition: str = _snapshot_key_condition( + key_condition: str = SnapshotSql.key_condition( left_alias="__target", right_alias="__source", unique_key=unique_key ) close_sql: str = ( @@ -2064,7 +2064,7 @@ def render_apply_timestamp_snapshot_changes( ) insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = _snapshot_key_condition( + active_join_condition: str = SnapshotSql.key_condition( left_alias="__active", right_alias="__source", unique_key=unique_key ) first_key: str = unique_key[0] From fa167da74c1e457a79c3176994ce3d05f04ff398 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:55:39 +0100 Subject: [PATCH 46/57] refactor(snapshots): share exact-copy historical timestamp snapshot helpers Move the historical timestamp and historical changes helpers that were copied into the base, DuckDB, BigQuery, Snowflake and Databricks adapters into HistoricalTimestampSnapshotSql: the initial-build queries and the QUALIFY-based incremental CTEs. Postgres and SQL Server now call the same shared builder where they previously imported the base adapter's module-level copy. Rendered SQL is unchanged for every adapter. (cherry picked from commit 5a817113318b1ed80116697bcb407f62fa6f8ec0) --- .../adapter/contract/classes/base_adapter.py | 144 +---------------- .../contract/classes/duckdb_backed_adapter.py | 148 +---------------- .../historical_timestamp_snapshot_sql.py | 150 ++++++++++++++++++ .../bigquery/classes/bigquery_adapter.py | 148 +---------------- .../databricks/classes/databricks_adapter.py | 148 +---------------- .../postgres/classes/postgres_adapter.py | 9 +- .../snowflake/classes/snowflake_adapter.py | 148 +---------------- .../sqlserver/classes/sqlserver_adapter.py | 6 +- 8 files changed, 194 insertions(+), 707 deletions(-) create mode 100644 src/sqlbuild/adapter/contract/classes/historical_timestamp_snapshot_sql.py diff --git a/src/sqlbuild/adapter/contract/classes/base_adapter.py b/src/sqlbuild/adapter/contract/classes/base_adapter.py index 25da4219dd..bbfc635571 100644 --- a/src/sqlbuild/adapter/contract/classes/base_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/base_adapter.py @@ -13,9 +13,11 @@ HistoricalCheckSnapshotSql, ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( - HistoricalSnapshotSql, historical_insert_validity_sql, ) +from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( + HistoricalTimestampSnapshotSql, +) from sqlbuild.adapter.contract.classes.retention_adapter import RetentionAdapterMixin from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder @@ -1048,7 +1050,7 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = _historical_timestamp_snapshot_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1072,7 +1074,7 @@ def render_create_initial_historical_timestamp_changes_destination( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - historical_sql: str = _historical_timestamp_changes_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.changes_initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1095,7 +1097,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = _historical_timestamp_new_changes_cte_sql( + new_changes_sql: str = HistoricalTimestampSnapshotSql.new_changes_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1160,7 +1162,7 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = _historical_timestamp_changes_new_records_cte_sql( + new_changes_sql: str = HistoricalTimestampSnapshotSql.changes_new_records_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -2627,138 +2629,6 @@ def _snapshot_hard_delete_close_sql( ) -def _historical_timestamp_snapshot_select_sql( - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, -) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - updated_at_column=updated_at_column, - ).initial_select_sql(output_columns=output_columns) - - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - -def _historical_timestamp_new_changes_cte_sql( - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, -) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - updated_at_column=updated_at_column, - ).new_changes_ctes_sql(destination=destination) - - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - ")" - ) - - -def _historical_timestamp_changes_select_sql( - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], -) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - f"FROM {origin}" - ) - - -def _historical_timestamp_changes_new_records_cte_sql( - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_to_column: str, -) -> str: - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__source", right_alias="__latest", unique_key=unique_key - ) - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - return ( - "__latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column} DESC" - ") = 1" - "), __new_changes AS (" - f"SELECT __source.* FROM {origin} AS __source " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __latest.{updated_at_column}" - ")" - ) - - def _historical_snapshot_combined_close_sql( *, destination: str, diff --git a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py index ee240327e5..8f3c39291c 100644 --- a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py @@ -23,9 +23,11 @@ HistoricalCheckSnapshotSql, ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( - HistoricalSnapshotSql, historical_insert_validity_sql, ) +from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( + HistoricalTimestampSnapshotSql, +) from sqlbuild.adapter.contract.classes.observed_connection import ObservedConnection from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder @@ -526,7 +528,7 @@ def render_create_initial_historical_timestamp_snapshot_destination( ) -> tuple[str, ...]: return self.render_create_table_as( destination=destination, - sql=self._historical_timestamp_snapshot_select_sql( + sql=HistoricalTimestampSnapshotSql.initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -552,7 +554,7 @@ def render_create_initial_historical_timestamp_changes_destination( ) -> tuple[str, ...]: return self.render_create_table_as( destination=destination, - sql=self._historical_timestamp_changes_select_sql( + sql=HistoricalTimestampSnapshotSql.changes_initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -575,7 +577,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_new_changes_cte_sql( + new_changes_sql: str = HistoricalTimestampSnapshotSql.new_changes_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -640,7 +642,7 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_changes_new_records_cte_sql( + new_changes_sql: str = HistoricalTimestampSnapshotSql.changes_new_records_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -2633,142 +2635,6 @@ def _snapshot_hard_delete_close_sql( f")" ) - @classmethod - def _historical_timestamp_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - updated_at_column=updated_at_column, - ).initial_select_sql(output_columns=output_columns) - - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - @classmethod - def _historical_timestamp_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - updated_at_column=updated_at_column, - ).new_changes_ctes_sql(destination=destination) - - partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - ")" - ) - - @classmethod - def _historical_timestamp_changes_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - f"FROM {origin}" - ) - - @classmethod - def _historical_timestamp_changes_new_records_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_to_column: str, - ) -> str: - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__source", right_alias="__latest", unique_key=unique_key - ) - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - return ( - "__latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column} DESC" - ") = 1" - "), __new_changes AS (" - f"SELECT __source.* FROM {origin} AS __source " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __latest.{updated_at_column}" - ")" - ) - @classmethod def _historical_snapshot_combined_close_sql( cls, diff --git a/src/sqlbuild/adapter/contract/classes/historical_timestamp_snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/historical_timestamp_snapshot_sql.py new file mode 100644 index 0000000000..a2eeb8d532 --- /dev/null +++ b/src/sqlbuild/adapter/contract/classes/historical_timestamp_snapshot_sql.py @@ -0,0 +1,150 @@ +"""Shared SQL for historical timestamp snapshots.""" + +from __future__ import annotations + +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import HistoricalSnapshotSql +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql + + +class HistoricalTimestampSnapshotSql: + """Render SQL for historical timestamp snapshots.""" + + @staticmethod + def initial_select_sql( + *, + origin: str, + unique_key: tuple[str, ...], + updated_at_column: str, + observed_at_column: str, + valid_from_column: str, + valid_to_column: str, + output_columns: tuple[str, ...], + invalidate_hard_deletes: bool, + ) -> str: + """Render the initial-build query for a historical timestamp snapshot.""" + + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + ).initial_select_sql(output_columns=output_columns) + + partition_sql: str = ", ".join(unique_key) + output_select_sql: str = ", ".join(column for column in output_columns) + return ( + "WITH __ordered AS (" + f"SELECT *, LAG({updated_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" + f") AS __prev_updated_at FROM {origin}" + "), __changes AS (" + f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " + f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" + ") " + f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " + f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " + f"ORDER BY {updated_at_column}) AS {valid_to_column} " + "FROM __changes" + ) + + @staticmethod + def new_changes_ctes_sql( + *, + destination: str, + origin: str, + unique_key: tuple[str, ...], + updated_at_column: str, + observed_at_column: str, + valid_from_column: str, + valid_to_column: str, + invalidate_hard_deletes: bool, + ) -> str: + """Render incremental historical timestamp CTEs using QUALIFY for the latest versions.""" + + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + ).new_changes_ctes_sql(destination=destination) + + partition_sql: str = ", ".join(unique_key) + first_key: str = unique_key[0] + latest_join_condition: str = SnapshotSql.key_condition( + left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key + ) + return ( + "__ordered AS (" + f"SELECT *, LAG({updated_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" + f") AS __prev_updated_at FROM {origin}" + "), __delta_changes AS (" + f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " + f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" + "), __latest AS (" + f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" + f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" + ") = 1" + "), __new_changes AS (" + "SELECT __delta_changes.* FROM __delta_changes " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{first_key} IS NULL " + f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" + ")" + ) + + @staticmethod + def changes_initial_select_sql( + *, + origin: str, + unique_key: tuple[str, ...], + updated_at_column: str, + valid_from_column: str, + valid_to_column: str, + output_columns: tuple[str, ...], + ) -> str: + """Render the initial-build query for a historical changes snapshot.""" + + partition_sql: str = ", ".join(unique_key) + output_select_sql: str = ", ".join(column for column in output_columns) + return ( + f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " + f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " + f"ORDER BY {updated_at_column}) AS {valid_to_column} " + f"FROM {origin}" + ) + + @staticmethod + def changes_new_records_ctes_sql( + *, + destination: str, + origin: str, + unique_key: tuple[str, ...], + updated_at_column: str, + valid_to_column: str, + ) -> str: + """Render incremental historical changes CTEs using QUALIFY for the latest versions.""" + + latest_join_condition: str = SnapshotSql.key_condition( + left_alias="__source", right_alias="__latest", unique_key=unique_key + ) + partition_sql: str = ", ".join(unique_key) + first_key: str = unique_key[0] + return ( + "__latest AS (" + f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" + f"PARTITION BY {partition_sql} ORDER BY {updated_at_column} DESC" + ") = 1" + "), __new_changes AS (" + f"SELECT __source.* FROM {origin} AS __source " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{first_key} IS NULL " + f"OR __source.{updated_at_column} > __latest.{updated_at_column}" + ")" + ) diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py index 7ee868be40..c748e754dc 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py @@ -22,12 +22,14 @@ HistoricalCheckSnapshotSql, ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( - HistoricalSnapshotSql, historical_insert_validity_sql, ) from sqlbuild.adapter.contract.classes.historical_snapshot_statement_sql import ( HistoricalSnapshotStatementSql, ) +from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( + HistoricalTimestampSnapshotSql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder @@ -1122,7 +1124,7 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = self._historical_timestamp_snapshot_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1146,7 +1148,7 @@ def render_create_initial_historical_timestamp_changes_destination( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - historical_sql: str = self._historical_timestamp_changes_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.changes_initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1169,7 +1171,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_new_changes_cte_sql( + new_changes_sql: str = HistoricalTimestampSnapshotSql.new_changes_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1228,7 +1230,7 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_changes_new_records_cte_sql( + new_changes_sql: str = HistoricalTimestampSnapshotSql.changes_new_records_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -2975,139 +2977,3 @@ def _snapshot_hard_delete_close_sql( f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" f")" ) - - @classmethod - def _historical_timestamp_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - updated_at_column=updated_at_column, - ).initial_select_sql(output_columns=output_columns) - - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - @classmethod - def _historical_timestamp_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - updated_at_column=updated_at_column, - ).new_changes_ctes_sql(destination=destination) - - partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - ")" - ) - - @classmethod - def _historical_timestamp_changes_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - f"FROM {origin}" - ) - - @classmethod - def _historical_timestamp_changes_new_records_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_to_column: str, - ) -> str: - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__source", right_alias="__latest", unique_key=unique_key - ) - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - return ( - "__latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column} DESC" - ") = 1" - "), __new_changes AS (" - f"SELECT __source.* FROM {origin} AS __source " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __latest.{updated_at_column}" - ")" - ) diff --git a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py index 88884152fa..86f7d5bb72 100644 --- a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py +++ b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py @@ -22,9 +22,11 @@ HistoricalCheckSnapshotSql, ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( - HistoricalSnapshotSql, historical_insert_validity_sql, ) +from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( + HistoricalTimestampSnapshotSql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder @@ -969,7 +971,7 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = self._historical_timestamp_snapshot_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -993,7 +995,7 @@ def render_create_initial_historical_timestamp_changes_destination( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - historical_sql: str = self._historical_timestamp_changes_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.changes_initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1016,7 +1018,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_new_changes_cte_sql( + new_changes_sql: str = HistoricalTimestampSnapshotSql.new_changes_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1081,7 +1083,7 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_changes_new_records_cte_sql( + new_changes_sql: str = HistoricalTimestampSnapshotSql.changes_new_records_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -3110,142 +3112,6 @@ def _snapshot_hard_delete_close_sql( f")" ) - @classmethod - def _historical_timestamp_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - updated_at_column=updated_at_column, - ).initial_select_sql(output_columns=output_columns) - - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - @classmethod - def _historical_timestamp_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - updated_at_column=updated_at_column, - ).new_changes_ctes_sql(destination=destination) - - partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - ")" - ) - - @classmethod - def _historical_timestamp_changes_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - f"FROM {origin}" - ) - - @classmethod - def _historical_timestamp_changes_new_records_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_to_column: str, - ) -> str: - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__source", right_alias="__latest", unique_key=unique_key - ) - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - return ( - "__latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column} DESC" - ") = 1" - "), __new_changes AS (" - f"SELECT __source.* FROM {origin} AS __source " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __latest.{updated_at_column}" - ")" - ) - @classmethod def _historical_snapshot_combined_close_sql( cls, diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py index 822794c5f1..65f336f680 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py @@ -16,8 +16,6 @@ _build_schemas_filter, _encode_typed_json, _historical_snapshot_combined_close_sql, - _historical_timestamp_changes_select_sql, - _historical_timestamp_snapshot_select_sql, _quote_sql_string, _render_ansi_typed_scalar, _render_typed_value_list, @@ -32,6 +30,9 @@ HistoricalSnapshotSql, historical_insert_validity_sql, ) +from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( + HistoricalTimestampSnapshotSql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder @@ -1305,7 +1306,7 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = _historical_timestamp_snapshot_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1329,7 +1330,7 @@ def render_create_initial_historical_timestamp_changes_destination( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - historical_sql: str = _historical_timestamp_changes_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.changes_initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index 80bbad19ff..9401e6b8c4 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -21,12 +21,14 @@ HistoricalCheckSnapshotSql, ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( - HistoricalSnapshotSql, historical_insert_validity_sql, ) from sqlbuild.adapter.contract.classes.historical_snapshot_statement_sql import ( HistoricalSnapshotStatementSql, ) +from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( + HistoricalTimestampSnapshotSql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder @@ -1375,7 +1377,7 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = self._historical_timestamp_snapshot_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1401,7 +1403,7 @@ def render_create_initial_historical_timestamp_changes_destination( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - historical_sql: str = self._historical_timestamp_changes_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.changes_initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1426,7 +1428,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_new_changes_cte_sql( + new_changes_sql: str = HistoricalTimestampSnapshotSql.new_changes_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1485,7 +1487,7 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_changes_new_records_cte_sql( + new_changes_sql: str = HistoricalTimestampSnapshotSql.changes_new_records_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -3324,139 +3326,3 @@ def _snapshot_hard_delete_close_sql( f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" f")" ) - - @classmethod - def _historical_timestamp_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - updated_at_column=updated_at_column, - ).initial_select_sql(output_columns=output_columns) - - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - @classmethod - def _historical_timestamp_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - updated_at_column=updated_at_column, - ).new_changes_ctes_sql(destination=destination) - - partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - ")" - ) - - @classmethod - def _historical_timestamp_changes_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - f"FROM {origin}" - ) - - @classmethod - def _historical_timestamp_changes_new_records_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_to_column: str, - ) -> str: - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__source", right_alias="__latest", unique_key=unique_key - ) - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - return ( - "__latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column} DESC" - ") = 1" - "), __new_changes AS (" - f"SELECT __source.* FROM {origin} AS __source " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __latest.{updated_at_column}" - ")" - ) diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py index cf3cb0fefd..2546c5b7e8 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py @@ -15,7 +15,6 @@ _build_names_filter, _build_schemas_filter, _encode_typed_json, - _historical_timestamp_changes_new_records_cte_sql, _quote_sql_string, _render_ansi_typed_scalar, _render_typed_value_list, @@ -26,6 +25,9 @@ HistoricalSnapshotSql, historical_insert_validity_sql, ) +from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( + HistoricalTimestampSnapshotSql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder @@ -1922,7 +1924,7 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = _historical_timestamp_changes_new_records_cte_sql( + new_changes_sql: str = HistoricalTimestampSnapshotSql.changes_new_records_ctes_sql( destination=destination, origin=origin, unique_key=unique_key, From 96347aaa3427fb306f19aebb1592ac8ae38ff7bb Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:56:50 +0100 Subject: [PATCH 47/57] refactor(snapshots): share exact-copy current-state snapshot helpers Move _snapshot_initial_valid_from_expr and the standard _snapshot_hard_delete_close_sql, copied into the base, DuckDB, BigQuery, Snowflake and Databricks adapters, into SnapshotSql. Postgres and SQL Server use the shared initial valid_from expression; SQL Server keeps its T-SQL hard-delete close until the parameterised current-state slice. Rendered SQL is unchanged for every adapter. (cherry picked from commit 77ea82547865f536ffb158776c66c25d5a1b5e5e) --- .../adapter/contract/classes/base_adapter.py | 55 ++---------------- .../contract/classes/duckdb_backed_adapter.py | 56 ++----------------- .../adapter/contract/classes/snapshot_sql.py | 50 +++++++++++++++++ .../bigquery/classes/bigquery_adapter.py | 56 ++----------------- .../databricks/classes/databricks_adapter.py | 56 ++----------------- .../postgres/classes/postgres_adapter.py | 12 ++-- .../snowflake/classes/snowflake_adapter.py | 56 ++----------------- .../sqlserver/classes/sqlserver_adapter.py | 7 +-- 8 files changed, 83 insertions(+), 265 deletions(-) diff --git a/src/sqlbuild/adapter/contract/classes/base_adapter.py b/src/sqlbuild/adapter/contract/classes/base_adapter.py index bbfc635571..ff821e8c3e 100644 --- a/src/sqlbuild/adapter/contract/classes/base_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/base_adapter.py @@ -65,7 +65,6 @@ render_insert_source_freshness_records_sql, ) from sqlbuild.compiler.compile.types import FunctionLanguage -from sqlbuild.compiler.planner.types import InitialValidFrom, SnapshotStrategy from sqlbuild.compiler.source_freshness.models import SourceFreshnessRecord from sqlbuild.spec.contracts.constants import DEFAULT_SEED_CSV_SETTINGS from sqlbuild.spec.contracts.models import SeedCsvSettings @@ -858,7 +857,7 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - valid_from_expr: str = _snapshot_initial_valid_from_expr( + valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy=snapshot_strategy, updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -889,7 +888,7 @@ def render_apply_timestamp_snapshot_changes( invalidate_hard_deletes: bool, ) -> tuple[str, ...]: current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = _snapshot_initial_valid_from_expr( + initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy="timestamp", updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -949,7 +948,7 @@ def render_apply_timestamp_snapshot_changes( if invalidate_hard_deletes: statements = ( *statements, - _snapshot_hard_delete_close_sql( + SnapshotSql.hard_delete_close_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -976,7 +975,7 @@ def render_apply_check_snapshot_changes( valid_to_column: str = target.valid_to_column output_columns: tuple[str, ...] = target.output_columns current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = _snapshot_initial_valid_from_expr( + initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy="check", updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -1023,7 +1022,7 @@ def render_apply_check_snapshot_changes( if invalidate_hard_deletes: statements = ( *statements, - _snapshot_hard_delete_close_sql( + SnapshotSql.hard_delete_close_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -2585,50 +2584,6 @@ def _quote_sql_string(value: str) -> str: return "'" + value.replace("'", "''") + "'" -def _snapshot_initial_valid_from_expr( - *, - snapshot_strategy: str | None, - updated_at_column: str | None, - observed_at_column: str | None, - initial_valid_from: str | None, - source_alias: str | None, - current_timestamp: str, -) -> str: - prefix: str = f"{source_alias}." if source_alias is not None else "" - if initial_valid_from == InitialValidFrom.EXECUTION_TIME: - return current_timestamp - if initial_valid_from == InitialValidFrom.OBSERVED_AT and observed_at_column is not None: - return f"{prefix}{observed_at_column}" - if initial_valid_from == InitialValidFrom.UPDATED_AT and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - if snapshot_strategy == SnapshotStrategy.TIMESTAMP and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - return current_timestamp - - -def _snapshot_hard_delete_close_sql( - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - valid_to_column: str, - current_timestamp: str, -) -> str: - missing_key_condition: str = SnapshotSql.key_condition( - left_alias="__source", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __source " - f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" - f")" - ) - - def _historical_snapshot_combined_close_sql( *, destination: str, diff --git a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py index 8f3c39291c..6b359217fe 100644 --- a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py @@ -78,7 +78,6 @@ from sqlbuild.adapter.type_system.main.normalize_numeric_family import normalize_numeric_family from sqlbuild.adapter.type_system.main.types_equal import types_equal from sqlbuild.compiler.compile.types import FunctionLanguage -from sqlbuild.compiler.planner.types import InitialValidFrom, SnapshotStrategy from sqlbuild.compiler.source_freshness.models import SourceFreshnessRecord from sqlbuild.diagnostics.main.log_sql import log_sql from sqlbuild.spec.contracts.constants import DEFAULT_SEED_CSV_SETTINGS @@ -411,7 +410,7 @@ def render_create_initial_snapshot_destination( initial_valid_from: str | None, ) -> tuple[str, ...]: current_timestamp: str = self.render_current_timestamp() - valid_from_expr: str = self._snapshot_initial_valid_from_expr( + valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy=snapshot_strategy, updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -442,7 +441,7 @@ def render_apply_timestamp_snapshot_changes( invalidate_hard_deletes: bool, ) -> tuple[str, ...]: current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( + initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy="timestamp", updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -502,7 +501,7 @@ def render_apply_timestamp_snapshot_changes( if invalidate_hard_deletes: statements = ( *statements, - self._snapshot_hard_delete_close_sql( + SnapshotSql.hard_delete_close_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -697,7 +696,7 @@ def render_apply_check_snapshot_changes( valid_to_column: str = target.valid_to_column output_columns: tuple[str, ...] = target.output_columns current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( + initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy="check", updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -744,7 +743,7 @@ def render_apply_check_snapshot_changes( if invalidate_hard_deletes: statements = ( *statements, - self._snapshot_hard_delete_close_sql( + SnapshotSql.hard_delete_close_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -2590,51 +2589,6 @@ def _render_seed_csv_options(self, csv_settings: SeedCsvSettings) -> str: def _duckdb_string_literal(self, value: str) -> str: return value.replace("'", "''") - @staticmethod - def _snapshot_initial_valid_from_expr( - *, - snapshot_strategy: str | None, - updated_at_column: str | None, - observed_at_column: str | None, - initial_valid_from: str | None, - source_alias: str | None, - current_timestamp: str, - ) -> str: - prefix: str = f"{source_alias}." if source_alias is not None else "" - if initial_valid_from == InitialValidFrom.EXECUTION_TIME: - return current_timestamp - if initial_valid_from == InitialValidFrom.OBSERVED_AT and observed_at_column is not None: - return f"{prefix}{observed_at_column}" - if initial_valid_from == InitialValidFrom.UPDATED_AT and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - if snapshot_strategy == SnapshotStrategy.TIMESTAMP and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - return current_timestamp - - @classmethod - def _snapshot_hard_delete_close_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - valid_to_column: str, - current_timestamp: str, - ) -> str: - missing_key_condition: str = SnapshotSql.key_condition( - left_alias="__source", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __source " - f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" - f")" - ) - @classmethod def _historical_snapshot_combined_close_sql( cls, diff --git a/src/sqlbuild/adapter/contract/classes/snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/snapshot_sql.py index aea802d795..80f6c18132 100644 --- a/src/sqlbuild/adapter/contract/classes/snapshot_sql.py +++ b/src/sqlbuild/adapter/contract/classes/snapshot_sql.py @@ -2,6 +2,8 @@ from __future__ import annotations +from sqlbuild.compiler.planner.types import InitialValidFrom, SnapshotStrategy + class SnapshotSql: """Render SQL fragments shared by snapshot materializations.""" @@ -13,3 +15,51 @@ def key_condition(*, left_alias: str, right_alias: str, unique_key: tuple[str, . return " AND ".join( f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key ) + + @staticmethod + def initial_valid_from_expr( + *, + snapshot_strategy: str | None, + updated_at_column: str | None, + observed_at_column: str | None, + initial_valid_from: str | None, + source_alias: str | None, + current_timestamp: str, + ) -> str: + """Render the valid_from expression for a key's first snapshot version.""" + + prefix: str = f"{source_alias}." if source_alias is not None else "" + if initial_valid_from == InitialValidFrom.EXECUTION_TIME: + return current_timestamp + if initial_valid_from == InitialValidFrom.OBSERVED_AT and observed_at_column is not None: + return f"{prefix}{observed_at_column}" + if initial_valid_from == InitialValidFrom.UPDATED_AT and updated_at_column is not None: + return f"{prefix}{updated_at_column}" + if snapshot_strategy == SnapshotStrategy.TIMESTAMP and updated_at_column is not None: + return f"{prefix}{updated_at_column}" + return current_timestamp + + @staticmethod + def hard_delete_close_sql( + *, + destination: str, + origin: str, + unique_key: tuple[str, ...], + valid_to_column: str, + current_timestamp: str, + ) -> str: + """Render the statement that closes active versions whose key is missing from the source.""" + + missing_key_condition: str = SnapshotSql.key_condition( + left_alias="__source", right_alias="__target", unique_key=unique_key + ) + first_key: str = unique_key[0] + return ( + f"UPDATE {destination} AS __target " + f"SET {valid_to_column} = {current_timestamp} " + f"WHERE __target.{valid_to_column} IS NULL " + f"AND NOT EXISTS (" + f"SELECT 1 FROM {origin} AS __source " + f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" + f")" + ) diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py index c748e754dc..0d2f12ef1b 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py @@ -110,7 +110,6 @@ TIMESTAMP_TYPE_TOKEN, ) from sqlbuild.compiler.compile.types import FunctionLanguage -from sqlbuild.compiler.planner.types import InitialValidFrom, SnapshotStrategy from sqlbuild.compiler.source_freshness.models import SourceFreshnessRecord from sqlbuild.diagnostics.main.log_sql import log_sql from sqlbuild.runtime.observability.classes.statement_lifecycle import StatementLifecycle @@ -1009,7 +1008,7 @@ def render_create_initial_snapshot_destination( initial_valid_from: str | None, ) -> tuple[str, ...]: current_timestamp: str = self.render_current_timestamp() - valid_from_expr: str = self._snapshot_initial_valid_from_expr( + valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy=snapshot_strategy, updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -1040,7 +1039,7 @@ def render_apply_timestamp_snapshot_changes( invalidate_hard_deletes: bool, ) -> tuple[str, ...]: current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( + initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy="timestamp", updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -1100,7 +1099,7 @@ def render_apply_timestamp_snapshot_changes( if invalidate_hard_deletes: statements = ( *statements, - self._snapshot_hard_delete_close_sql( + SnapshotSql.hard_delete_close_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1281,7 +1280,7 @@ def render_apply_check_snapshot_changes( valid_to_column: str = target.valid_to_column output_columns: tuple[str, ...] = target.output_columns current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( + initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy="check", updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -1328,7 +1327,7 @@ def render_apply_check_snapshot_changes( if invalidate_hard_deletes: statements = ( *statements, - self._snapshot_hard_delete_close_sql( + SnapshotSql.hard_delete_close_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -2932,48 +2931,3 @@ def _render_cursor_bound_string(value: str) -> str: except ValueError: return f"TIMESTAMP '{value}'" return value - - @staticmethod - def _snapshot_initial_valid_from_expr( - *, - snapshot_strategy: str | None, - updated_at_column: str | None, - observed_at_column: str | None, - initial_valid_from: str | None, - source_alias: str | None, - current_timestamp: str, - ) -> str: - prefix: str = f"{source_alias}." if source_alias is not None else "" - if initial_valid_from == InitialValidFrom.EXECUTION_TIME: - return current_timestamp - if initial_valid_from == InitialValidFrom.OBSERVED_AT and observed_at_column is not None: - return f"{prefix}{observed_at_column}" - if initial_valid_from == InitialValidFrom.UPDATED_AT and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - if snapshot_strategy == SnapshotStrategy.TIMESTAMP and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - return current_timestamp - - @classmethod - def _snapshot_hard_delete_close_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - valid_to_column: str, - current_timestamp: str, - ) -> str: - missing_key_condition: str = SnapshotSql.key_condition( - left_alias="__source", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __source " - f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" - f")" - ) diff --git a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py index 86f7d5bb72..6a2d489506 100644 --- a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py +++ b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py @@ -95,7 +95,6 @@ TABLE_RELATION_METADATA_TYPES, ) from sqlbuild.compiler.compile.types import FunctionLanguage -from sqlbuild.compiler.planner.types import InitialValidFrom, SnapshotStrategy from sqlbuild.compiler.source_freshness.models import SourceFreshnessRecord from sqlbuild.diagnostics.main.log_sql import log_sql from sqlbuild.spec.contracts.constants import DEFAULT_SEED_CSV_SETTINGS @@ -856,7 +855,7 @@ def render_create_initial_snapshot_destination( initial_valid_from: str | None, ) -> tuple[str, ...]: current_timestamp: str = self.render_current_timestamp() - valid_from_expr: str = self._snapshot_initial_valid_from_expr( + valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy=snapshot_strategy, updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -887,7 +886,7 @@ def render_apply_timestamp_snapshot_changes( invalidate_hard_deletes: bool, ) -> tuple[str, ...]: current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( + initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy="timestamp", updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -947,7 +946,7 @@ def render_apply_timestamp_snapshot_changes( if invalidate_hard_deletes: statements = ( *statements, - self._snapshot_hard_delete_close_sql( + SnapshotSql.hard_delete_close_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1138,7 +1137,7 @@ def render_apply_check_snapshot_changes( valid_to_column: str = target.valid_to_column output_columns: tuple[str, ...] = target.output_columns current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( + initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy="check", updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -1185,7 +1184,7 @@ def render_apply_check_snapshot_changes( if invalidate_hard_deletes: statements = ( *statements, - self._snapshot_hard_delete_close_sql( + SnapshotSql.hard_delete_close_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -3067,51 +3066,6 @@ def _to_int(value: object) -> int: return value return int(str(value)) - @staticmethod - def _snapshot_initial_valid_from_expr( - *, - snapshot_strategy: str | None, - updated_at_column: str | None, - observed_at_column: str | None, - initial_valid_from: str | None, - source_alias: str | None, - current_timestamp: str, - ) -> str: - prefix: str = f"{source_alias}." if source_alias is not None else "" - if initial_valid_from == InitialValidFrom.EXECUTION_TIME: - return current_timestamp - if initial_valid_from == InitialValidFrom.OBSERVED_AT and observed_at_column is not None: - return f"{prefix}{observed_at_column}" - if initial_valid_from == InitialValidFrom.UPDATED_AT and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - if snapshot_strategy == SnapshotStrategy.TIMESTAMP and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - return current_timestamp - - @classmethod - def _snapshot_hard_delete_close_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - valid_to_column: str, - current_timestamp: str, - ) -> str: - missing_key_condition: str = SnapshotSql.key_condition( - left_alias="__source", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __source " - f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" - f")" - ) - @classmethod def _historical_snapshot_combined_close_sql( cls, diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py index 65f336f680..1e0acac905 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py @@ -19,8 +19,6 @@ _quote_sql_string, _render_ansi_typed_scalar, _render_typed_value_list, - _snapshot_hard_delete_close_sql, - _snapshot_initial_valid_from_expr, _validate_rectangular_typed_array, ) from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( @@ -1117,7 +1115,7 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - valid_from_expr: str = _snapshot_initial_valid_from_expr( + valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy=snapshot_strategy, updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -1148,7 +1146,7 @@ def render_apply_timestamp_snapshot_changes( invalidate_hard_deletes: bool, ) -> tuple[str, ...]: current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = _snapshot_initial_valid_from_expr( + initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy="timestamp", updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -1208,7 +1206,7 @@ def render_apply_timestamp_snapshot_changes( if invalidate_hard_deletes: statements = ( *statements, - _snapshot_hard_delete_close_sql( + SnapshotSql.hard_delete_close_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1235,7 +1233,7 @@ def render_apply_check_snapshot_changes( valid_to_column: str = target.valid_to_column output_columns: tuple[str, ...] = target.output_columns current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = _snapshot_initial_valid_from_expr( + initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy="check", updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -1282,7 +1280,7 @@ def render_apply_check_snapshot_changes( if invalidate_hard_deletes: statements = ( *statements, - _snapshot_hard_delete_close_sql( + SnapshotSql.hard_delete_close_sql( destination=destination, origin=origin, unique_key=unique_key, diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index 9401e6b8c4..e779865d81 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -100,7 +100,6 @@ VIEW_RELATION_TYPE_TOKEN, ) from sqlbuild.compiler.compile.types import FunctionLanguage -from sqlbuild.compiler.planner.types import InitialValidFrom, SnapshotStrategy from sqlbuild.compiler.source_freshness.models import SourceFreshnessRecord from sqlbuild.cost.main.collection import collect_snowflake_cost from sqlbuild.cost.models import RunCostSummary @@ -1261,7 +1260,7 @@ def render_create_initial_snapshot_destination( initial_valid_from: str | None, ) -> tuple[str, ...]: current_timestamp: str = self.render_current_timestamp() - valid_from_expr: str = self._snapshot_initial_valid_from_expr( + valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy=snapshot_strategy, updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -1293,7 +1292,7 @@ def render_apply_timestamp_snapshot_changes( invalidate_hard_deletes: bool, ) -> tuple[str, ...]: current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( + initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy="timestamp", updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -1353,7 +1352,7 @@ def render_apply_timestamp_snapshot_changes( if invalidate_hard_deletes: statements = ( *statements, - self._snapshot_hard_delete_close_sql( + SnapshotSql.hard_delete_close_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1538,7 +1537,7 @@ def render_apply_check_snapshot_changes( valid_to_column: str = target.valid_to_column output_columns: tuple[str, ...] = target.output_columns current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( + initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy="check", updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -1585,7 +1584,7 @@ def render_apply_check_snapshot_changes( if invalidate_hard_deletes: statements = ( *statements, - self._snapshot_hard_delete_close_sql( + SnapshotSql.hard_delete_close_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -3281,48 +3280,3 @@ def schema_exists( return cursor.fetchone() is not None finally: cursor.close() - - @staticmethod - def _snapshot_initial_valid_from_expr( - *, - snapshot_strategy: str | None, - updated_at_column: str | None, - observed_at_column: str | None, - initial_valid_from: str | None, - source_alias: str | None, - current_timestamp: str, - ) -> str: - prefix: str = f"{source_alias}." if source_alias is not None else "" - if initial_valid_from == InitialValidFrom.EXECUTION_TIME: - return current_timestamp - if initial_valid_from == InitialValidFrom.OBSERVED_AT and observed_at_column is not None: - return f"{prefix}{observed_at_column}" - if initial_valid_from == InitialValidFrom.UPDATED_AT and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - if snapshot_strategy == SnapshotStrategy.TIMESTAMP and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - return current_timestamp - - @classmethod - def _snapshot_hard_delete_close_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - valid_to_column: str, - current_timestamp: str, - ) -> str: - missing_key_condition: str = SnapshotSql.key_condition( - left_alias="__source", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __source " - f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" - f")" - ) diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py index 2546c5b7e8..6951cf9cff 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py @@ -18,7 +18,6 @@ _quote_sql_string, _render_ansi_typed_scalar, _render_typed_value_list, - _snapshot_initial_valid_from_expr, _typed_scalar_payload, ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( @@ -1787,7 +1786,7 @@ def render_apply_check_snapshot_changes( valid_to_column: str = target.valid_to_column output_columns: tuple[str, ...] = target.output_columns current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = _snapshot_initial_valid_from_expr( + initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy="check", updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -2046,7 +2045,7 @@ def render_apply_timestamp_snapshot_changes( invalidate_hard_deletes: bool, ) -> tuple[str, ...]: current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = _snapshot_initial_valid_from_expr( + initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy="timestamp", updated_at_column=updated_at_column, observed_at_column=observed_at_column, @@ -2336,7 +2335,7 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - valid_from_expr: str = _snapshot_initial_valid_from_expr( + valid_from_expr: str = SnapshotSql.initial_valid_from_expr( snapshot_strategy=snapshot_strategy, updated_at_column=updated_at_column, observed_at_column=observed_at_column, From 23583ad08cdef683727e29f7a9a3422a1f4a2f02 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:02:53 +0100 Subject: [PATCH 48/57] refactor(snapshots): share parameterised historical check snapshot SQL Add SnapshotSqlDialect, which names the places where adapters really render snapshot SQL differently: the null timestamp type, the null-safe inequality, how the latest stored version is selected (QUALIFY, derived table or ordered CTE), how historical versions are closed (correlated UPDATE, UPDATE ... FROM, Databricks MERGE for hard-delete closes, or T-SQL UPDATE ... FROM), where insert CTEs go, and how current-state versions are closed. Each adapter declares its dialect once. HistoricalCheckSnapshotSql now renders the complete incremental historical check statements for every adapter from that dialect, and HistoricalSnapshotStatementSql owns the shared latest-version, close and insert statements. render_apply_historical_check_snapshot_changes delegates on every adapter, and the Postgres and SQL Server copies of the incremental CTE builder are removed. BigQuery and Snowflake historical timestamp closes and inserts use the same statement builder. Rendered SQL is unchanged for every adapter. (cherry picked from commit 5098cd2838103532c199fd2e5cfcfc04df181231) --- .../adapter/contract/classes/base_adapter.py | 64 ++--- .../contract/classes/duckdb_backed_adapter.py | 48 +--- .../classes/historical_check_snapshot_sql.py | 143 ++++++---- .../historical_snapshot_statement_sql.py | 258 ++++++++++++++++-- src/sqlbuild/adapter/contract/models.py | 18 +- src/sqlbuild/adapter/contract/types.py | 33 +++ .../bigquery/classes/bigquery_adapter.py | 125 ++------- .../databricks/classes/databricks_adapter.py | 64 ++--- .../postgres/classes/postgres_adapter.py | 132 ++------- .../snowflake/classes/snowflake_adapter.py | 125 ++------- .../sqlserver/classes/sqlserver_adapter.py | 148 ++-------- 11 files changed, 524 insertions(+), 634 deletions(-) diff --git a/src/sqlbuild/adapter/contract/classes/base_adapter.py b/src/sqlbuild/adapter/contract/classes/base_adapter.py index ff821e8c3e..5872aa3e07 100644 --- a/src/sqlbuild/adapter/contract/classes/base_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/base_adapter.py @@ -14,6 +14,7 @@ ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( historical_insert_validity_sql, + render_is_distinct_from, ) from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( HistoricalTimestampSnapshotSql, @@ -48,14 +49,19 @@ RowDiffTolerances, SchemaDiffResult, SnapshotChangeTarget, + SnapshotSqlDialect, TableFreshnessMetadata, TableFreshnessRequest, ) from sqlbuild.adapter.contract.types import ( CursorKind, FrameworkType, + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, LoaderLogicalType, PromotionStrategy, + SnapshotLatestVersionStyle, + SnapshotUpdateStyle, TablePromotionMode, ) from sqlbuild.adapter.relations.main.get_columns_for_relations import ( @@ -96,6 +102,14 @@ class BaseAdapter(RetentionAdapterMixin, StrictAdapter): state_tables_transient: ClassVar[bool] = False allows_implicit_managed_write_schema: ClassVar[bool] = False execution_duration_limit_seconds: ClassVar[int | None] = None + _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( + timestamp_type="TIMESTAMP", + distinct_condition=render_is_distinct_from, + update_style=SnapshotUpdateStyle.UPDATE_FROM, + latest_version=SnapshotLatestVersionStyle.QUALIFY, + historical_close=HistoricalSnapshotCloseStyle.CORRELATED, + historical_insert=HistoricalSnapshotInsertStyle.WITH_INSERT, + ) def supports_zero_copy_clone(self) -> bool: return False @@ -1213,7 +1227,9 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = HistoricalCheckSnapshotSql.initial_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -1238,7 +1254,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalCheckSnapshotSql.new_changes_ctes_sql( + return HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1246,51 +1262,9 @@ def render_apply_historical_check_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = _historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=SNAPSHOT_VERSION_START_COLUMN, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - version_columns_sql: str = ( - f"__new_changes.{observed_at_column}, LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column})" - ) - if invalidate_hard_deletes: - version_columns_sql = historical_insert_validity_sql() - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_columns_sql} " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def create_table_as( self, diff --git a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py index 6b359217fe..af870264c5 100644 --- a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py @@ -769,7 +769,7 @@ def render_create_initial_historical_check_snapshot_destination( ) -> tuple[str, ...]: return self.render_create_table_as( destination=destination, - sql=HistoricalCheckSnapshotSql.initial_select_sql( + sql=HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -794,7 +794,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalCheckSnapshotSql.new_changes_ctes_sql( + return HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -802,51 +802,9 @@ def render_apply_historical_check_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=SNAPSHOT_VERSION_START_COLUMN, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - version_columns_sql: str = ( - f"__new_changes.{observed_at_column}, LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column})" - ) - if invalidate_hard_deletes: - version_columns_sql = historical_insert_validity_sql() - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_columns_sql} " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def recommended_max_sql_length(self) -> int | None: """DuckDB uses the framework default recommendation for lightweight unit-test SQL.""" diff --git a/src/sqlbuild/adapter/contract/classes/historical_check_snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/historical_check_snapshot_sql.py index 220053ada3..468015894f 100644 --- a/src/sqlbuild/adapter/contract/classes/historical_check_snapshot_sql.py +++ b/src/sqlbuild/adapter/contract/classes/historical_check_snapshot_sql.py @@ -3,14 +3,21 @@ from __future__ import annotations from sqlbuild.adapter.contract.classes.historical_snapshot_sql import HistoricalSnapshotSql +from sqlbuild.adapter.contract.classes.historical_snapshot_statement_sql import ( + HistoricalSnapshotStatementSql, +) from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql +from sqlbuild.adapter.contract.models import SnapshotSqlDialect class HistoricalCheckSnapshotSql: """Render SQL for historical check snapshots.""" - @staticmethod + def __init__(self, *, dialect: SnapshotSqlDialect) -> None: + self.dialect: SnapshotSqlDialect = dialect + def initial_select_sql( + self, *, origin: str, unique_key: tuple[str, ...], @@ -31,26 +38,23 @@ def initial_select_sql( valid_from_column=valid_from_column, valid_to_column=valid_to_column, check_columns=check_columns, + distinct_condition=self.dialect.distinct_condition, ).initial_select_sql(output_columns=output_columns) partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" + change_condition: str = " OR ".join( + self.dialect.distinct_condition(left=column, right=f"__prev_{column}") for column in check_columns ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns + output_select_sql: str = ", ".join(output_columns) + ordered_sql: str = self._ordered_select_sql( + origin=origin, + unique_key=unique_key, + check_columns=check_columns, + observed_at_column=observed_at_column, ) - output_select_sql: str = ", ".join(column for column in output_columns) return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __changes AS (" + f"WITH __ordered AS ({ordered_sql}), __changes AS (" f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" ") " f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " @@ -59,8 +63,8 @@ def initial_select_sql( "FROM __changes" ) - @staticmethod - def new_changes_ctes_sql( + def apply_sql( + self, *, destination: str, origin: str, @@ -69,60 +73,107 @@ def new_changes_ctes_sql( observed_at_column: str, valid_from_column: str, valid_to_column: str, + output_columns: tuple[str, ...], invalidate_hard_deletes: bool, - ) -> str: - """Render incremental historical check CTEs using ``QUALIFY`` for the latest versions.""" + ) -> tuple[str, ...]: + """Render the incremental statements for a historical check snapshot.""" + statements: HistoricalSnapshotStatementSql = HistoricalSnapshotStatementSql( + dialect=self.dialect + ) + new_changes_sql: str if invalidate_hard_deletes: - return HistoricalSnapshotSql( + new_changes_sql = HistoricalSnapshotSql( origin=origin, unique_key=unique_key, observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, check_columns=check_columns, + distinct_condition=self.dialect.distinct_condition, ).new_changes_ctes_sql(destination=destination) - - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns + else: + new_changes_sql = self._new_changes_ctes_sql( + statements=statements, + destination=destination, + origin=origin, + unique_key=unique_key, + check_columns=check_columns, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + ) + return statements.apply_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + change_time_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + invalidate_hard_deletes=invalidate_hard_deletes, ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" + + def _new_changes_ctes_sql( + self, + *, + statements: HistoricalSnapshotStatementSql, + destination: str, + origin: str, + unique_key: tuple[str, ...], + check_columns: tuple[str, ...], + observed_at_column: str, + valid_from_column: str, + ) -> str: delta_change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns + self.dialect.distinct_condition(left=column, right=f"__prev_{column}") + for column in check_columns ) latest_join_condition: str = SnapshotSql.key_condition( left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key ) latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" + self.dialect.distinct_condition( + left=f"__delta_changes.{column}", right=f"__latest.{column}" + ) for column in check_columns ) - first_key: str = unique_key[0] - changed_or_first_sql: str = ( + latest_sql: str = statements.latest_version_cte_sql( + destination=destination, unique_key=unique_key, order_column=valid_from_column + ) + ordered_sql: str = self._ordered_select_sql( + origin=origin, + unique_key=unique_key, + check_columns=check_columns, + observed_at_column=observed_at_column, + ) + return ( + f"__ordered AS ({ordered_sql}), __delta_changes AS (" "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " f"OR ({delta_change_condition})" + f"), {latest_sql}, __new_changes AS (" + "SELECT __delta_changes.* FROM __delta_changes " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{unique_key[0]} IS NULL OR (" + f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " + f"AND ({latest_change_condition})))" + ) + + @staticmethod + def _ordered_select_sql( + *, + origin: str, + unique_key: tuple[str, ...], + check_columns: tuple[str, ...], + observed_at_column: str, + ) -> str: + partition_sql: str = ", ".join(unique_key) + previous_columns_sql: str = "".join( + f", LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " + f"AS __prev_{column}" + for column in check_columns ) return ( - "__ordered AS (" f"SELECT *, LAG({observed_at_column}) OVER (" f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition})" - ")" - ")" ) diff --git a/src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py b/src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py index af27517117..bfb307c66b 100644 --- a/src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py +++ b/src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py @@ -2,14 +2,98 @@ from __future__ import annotations +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + historical_insert_validity_sql, +) from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql +from sqlbuild.adapter.contract.constants import SNAPSHOT_VERSION_START_COLUMN +from sqlbuild.adapter.contract.models import SnapshotSqlDialect +from sqlbuild.adapter.contract.types import ( + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, + SnapshotLatestVersionStyle, +) class HistoricalSnapshotStatementSql: - """Render close and insert statements shared by historical snapshots.""" + """Render latest-version, close and insert statements shared by historical snapshots.""" - @staticmethod - def grouped_combined_close_sql( + def __init__(self, *, dialect: SnapshotSqlDialect) -> None: + self.dialect: SnapshotSqlDialect = dialect + + def latest_version_cte_sql( + self, *, destination: str, unique_key: tuple[str, ...], order_column: str + ) -> str: + """Render CTEs ending in ``__latest``, the newest stored version of each key.""" + + partition_sql: str = ", ".join(unique_key) + window_sql: str = ( + f"ROW_NUMBER() OVER (PARTITION BY {partition_sql} ORDER BY {order_column} DESC)" + ) + if self.dialect.latest_version == SnapshotLatestVersionStyle.QUALIFY: + return f"__latest AS (SELECT * FROM {destination} QUALIFY {window_sql} = 1)" + if self.dialect.latest_version == SnapshotLatestVersionStyle.DERIVED_TABLE: + return ( + f"__latest AS (SELECT * FROM (SELECT *, {window_sql} AS __rn FROM {destination}) " + "AS __q WHERE __rn = 1)" + ) + return ( + f"__latest_ordered AS (SELECT *, {window_sql} AS __rn FROM {destination}), " + "__latest AS (SELECT * FROM __latest_ordered WHERE __rn = 1)" + ) + + def apply_sql( + self, + *, + destination: str, + new_changes_sql: str, + unique_key: tuple[str, ...], + change_time_column: str, + valid_from_column: str, + valid_to_column: str, + output_columns: tuple[str, ...], + invalidate_hard_deletes: bool, + ) -> tuple[str, ...]: + """Render the close and insert statements that apply ``__new_changes``.""" + + close_sql: str + if invalidate_hard_deletes: + close_sql = self._combined_close_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, + ) + else: + close_sql = self._change_close_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + change_time_column=change_time_column, + ) + insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) + output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) + partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{change_time_column}, LEAD(__new_changes.{change_time_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{change_time_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() + insert_sql: str = self._insert_sql( + destination=destination, + insert_column_sql=insert_column_sql, + new_changes_sql=new_changes_sql, + select_sql=f"SELECT {output_select_sql}, {version_columns_sql} FROM __new_changes", + ) + return (close_sql, insert_sql) + + def _change_close_sql( + self, *, destination: str, new_changes_sql: str, @@ -18,26 +102,110 @@ def grouped_combined_close_sql( valid_to_column: str, change_time_column: str, ) -> str: - """Render an ``UPDATE ... FROM`` close over new versions and hard deletes.""" + if self.dialect.historical_close == HistoricalSnapshotCloseStyle.UPDATE_FROM: + return self._grouped_close_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + close_candidates_sql=( + f"SELECT {', '.join(unique_key)}, {change_time_column} AS __close_at " + "FROM __new_changes" + ), + ) + key_condition: str = SnapshotSql.key_condition( + left_alias="__target", right_alias="__new_changes", unique_key=unique_key + ) + return self._correlated_close_sql( + ctes_sql=new_changes_sql, + destination=destination, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + close_at_sql=( + f"SELECT MIN(__new_changes.{change_time_column}) " + f"FROM __new_changes WHERE {key_condition}" + ), + exists_sql=f"SELECT 1 FROM __new_changes WHERE {key_condition}", + ) + def _combined_close_sql( + self, + *, + destination: str, + new_changes_sql: str, + unique_key: tuple[str, ...], + valid_from_column: str, + valid_to_column: str, + change_time_column: str, + ) -> str: candidate_key_sql: str = ", ".join(unique_key) - return HistoricalSnapshotStatementSql.grouped_close_sql( + close_candidates_sql: str = ( + f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at FROM __new_changes " + "UNION ALL " + f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " + "WHERE __close_at IS NOT NULL" + ) + if self.dialect.historical_close == HistoricalSnapshotCloseStyle.UPDATE_FROM: + return self._grouped_close_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + close_candidates_sql=close_candidates_sql, + ) + if self.dialect.historical_close == HistoricalSnapshotCloseStyle.MERGE_HARD_DELETES: + return self._merge_close_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + close_candidates_sql=close_candidates_sql, + ) + close_candidate_condition: str = SnapshotSql.key_condition( + left_alias="__close_candidates", right_alias="__target", unique_key=unique_key + ) + return self._correlated_close_sql( + ctes_sql=f"{new_changes_sql}, __close_candidates AS ({close_candidates_sql})", destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at " - "FROM __new_changes " - "UNION ALL " - f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " - "WHERE __close_at IS NOT NULL" + close_at_sql=( + "SELECT MIN(__close_candidates.__close_at) FROM __close_candidates " + f"WHERE {close_candidate_condition}" ), + exists_sql=f"SELECT 1 FROM __close_candidates WHERE {close_candidate_condition}", ) - @staticmethod - def grouped_close_sql( + def _correlated_close_sql( + self, + *, + ctes_sql: str, + destination: str, + valid_from_column: str, + valid_to_column: str, + close_at_sql: str, + exists_sql: str, + ) -> str: + update_sql: str = ( + f"WITH {ctes_sql} UPDATE {destination} AS __target " + f"SET {valid_to_column} = ({close_at_sql}) " + ) + if self.dialect.historical_close == HistoricalSnapshotCloseStyle.TSQL: + update_sql = ( + f";WITH {ctes_sql} UPDATE __target SET {valid_to_column} = ({close_at_sql}) " + f"FROM {destination} AS __target " + ) + return ( + f"{update_sql}WHERE __target.{valid_to_column} IS NULL " + f"AND __target.{valid_from_column} < ({close_at_sql}) " + f"AND EXISTS ({exists_sql})" + ) + + def _grouped_close_sql( + self, *, destination: str, new_changes_sql: str, @@ -46,17 +214,13 @@ def grouped_close_sql( valid_to_column: str, close_candidates_sql: str, ) -> str: - """Render an ``UPDATE ... FROM`` close that uses the earliest candidate per key.""" - close_candidate_condition: str = SnapshotSql.key_condition( left_alias="__close_candidates", right_alias="__target", unique_key=unique_key ) - candidate_key_sql: str = ", ".join(unique_key) - close_candidates_query: str = ( - f"WITH {new_changes_sql}, __close_candidates AS ({close_candidates_sql}) " - f"SELECT {candidate_key_sql}, MIN(__close_at) AS __close_at " - "FROM __close_candidates GROUP BY " - f"{candidate_key_sql}" + close_candidates_query: str = self._earliest_close_candidates_sql( + new_changes_sql=new_changes_sql, + unique_key=unique_key, + close_candidates_sql=close_candidates_sql, ) return ( f"UPDATE {destination} AS __target " @@ -67,12 +231,50 @@ def grouped_close_sql( f"AND {close_candidate_condition}" ) - @staticmethod - def insert_with_cte_sql( - *, destination: str, insert_column_sql: str, new_changes_sql: str, select_sql: str + def _merge_close_sql( + self, + *, + destination: str, + new_changes_sql: str, + unique_key: tuple[str, ...], + valid_from_column: str, + valid_to_column: str, + close_candidates_sql: str, ) -> str: - """Render an insert whose CTEs follow the ``INSERT INTO`` column list.""" + close_candidate_condition: str = SnapshotSql.key_condition( + left_alias="__target", right_alias="__close_candidates", unique_key=unique_key + ) + close_candidates_query: str = self._earliest_close_candidates_sql( + new_changes_sql=new_changes_sql, + unique_key=unique_key, + close_candidates_sql=close_candidates_sql, + ) + return ( + f"MERGE INTO {destination} AS __target " + f"USING ({close_candidates_query}) AS __close_candidates " + f"ON {close_candidate_condition} " + f"AND __target.{valid_to_column} IS NULL " + f"AND __target.{valid_from_column} < __close_candidates.__close_at " + f"WHEN MATCHED THEN UPDATE SET {valid_to_column} = __close_candidates.__close_at" + ) + @staticmethod + def _earliest_close_candidates_sql( + *, new_changes_sql: str, unique_key: tuple[str, ...], close_candidates_sql: str + ) -> str: + candidate_key_sql: str = ", ".join(unique_key) return ( - f"INSERT INTO {destination} ({insert_column_sql}) WITH {new_changes_sql} {select_sql}" + f"WITH {new_changes_sql}, __close_candidates AS ({close_candidates_sql}) " + f"SELECT {candidate_key_sql}, MIN(__close_at) AS __close_at " + f"FROM __close_candidates GROUP BY {candidate_key_sql}" ) + + def _insert_sql( + self, *, destination: str, insert_column_sql: str, new_changes_sql: str, select_sql: str + ) -> str: + insert_sql: str = f"INSERT INTO {destination} ({insert_column_sql})" + if self.dialect.historical_insert == HistoricalSnapshotInsertStyle.INSERT_WITH: + return f"{insert_sql} WITH {new_changes_sql} {select_sql}" + if self.dialect.historical_insert == HistoricalSnapshotInsertStyle.TSQL: + return f";WITH {new_changes_sql} {insert_sql} {select_sql}" + return f"WITH {new_changes_sql} {insert_sql} {select_sql}" diff --git a/src/sqlbuild/adapter/contract/models.py b/src/sqlbuild/adapter/contract/models.py index cc9aead38d..d376cbebc0 100644 --- a/src/sqlbuild/adapter/contract/models.py +++ b/src/sqlbuild/adapter/contract/models.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass, field from datetime import datetime from decimal import Decimal @@ -12,9 +12,13 @@ from sqlbuild.adapter.contract.types import ( CursorKind, FunctionNullabilityRule, + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, LifeCycleEventKind, RetentionChangePhase, RetentionScope, + SnapshotLatestVersionStyle, + SnapshotUpdateStyle, TypeFamily, ) from sqlbuild.compiler.compile.types import FunctionLanguage @@ -202,6 +206,18 @@ class SnapshotChangeTarget: output_columns: tuple[str, ...] +@dataclass(frozen=True) +class SnapshotSqlDialect: + """Dialect differences in adapter-rendered snapshot SQL.""" + + timestamp_type: str + distinct_condition: Callable[..., str] + update_style: SnapshotUpdateStyle + latest_version: SnapshotLatestVersionStyle + historical_close: HistoricalSnapshotCloseStyle + historical_insert: HistoricalSnapshotInsertStyle + + @dataclass(frozen=True) class FunctionDefinition: """Warehouse function definition inputs for create/render adapter methods.""" diff --git a/src/sqlbuild/adapter/contract/types.py b/src/sqlbuild/adapter/contract/types.py index 4fea9e8db9..56b0f4de32 100644 --- a/src/sqlbuild/adapter/contract/types.py +++ b/src/sqlbuild/adapter/contract/types.py @@ -100,3 +100,36 @@ class TypeDialect(StrEnum): DATABRICKS = "databricks" POSTGRES = "postgres" TSQL = "tsql" + + +class SnapshotLatestVersionStyle(StrEnum): + """How incremental snapshot SQL selects each key's latest stored version.""" + + QUALIFY = "qualify" + DERIVED_TABLE = "derived_table" + ORDERED_CTE = "ordered_cte" + + +class SnapshotUpdateStyle(StrEnum): + """How current-state snapshot SQL closes active versions.""" + + UPDATE_FROM = "update_from" + MERGE = "merge" + TSQL = "tsql" + + +class HistoricalSnapshotCloseStyle(StrEnum): + """How historical snapshot SQL closes versions superseded by new observations.""" + + CORRELATED = "correlated" + UPDATE_FROM = "update_from" + MERGE_HARD_DELETES = "merge_hard_deletes" + TSQL = "tsql" + + +class HistoricalSnapshotInsertStyle(StrEnum): + """Where historical snapshot inserts place their common table expressions.""" + + WITH_INSERT = "with_insert" + INSERT_WITH = "insert_with" + TSQL = "tsql" diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py index 0d2f12ef1b..77f765e279 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py @@ -22,7 +22,7 @@ HistoricalCheckSnapshotSql, ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( - historical_insert_validity_sql, + render_is_distinct_from, ) from sqlbuild.adapter.contract.classes.historical_snapshot_statement_sql import ( HistoricalSnapshotStatementSql, @@ -37,7 +37,6 @@ from sqlbuild.adapter.contract.constants import ( DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE, - SNAPSHOT_VERSION_START_COLUMN, ) from sqlbuild.adapter.contract.exceptions import ( AdapterUserError, @@ -65,6 +64,7 @@ RowDiffTolerances, SchemaDiffResult, SnapshotChangeTarget, + SnapshotSqlDialect, TableFreshnessMetadata, TableFreshnessRequest, ) @@ -72,10 +72,14 @@ BuiltinAdapter, CursorKind, FrameworkType, + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, LoaderLogicalType, PromotionStrategy, RetentionChangePhase, RetentionScope, + SnapshotLatestVersionStyle, + SnapshotUpdateStyle, TablePromotionMode, ) from sqlbuild.adapter.relations.main.get_columns_for_relations import ( @@ -135,6 +139,14 @@ def get_columns_for_relations( adapter_name: ClassVar[str] = BuiltinAdapter.BIGQUERY.value sql_analysis_dialect_name: ClassVar[str | None] = "bigquery" max_identifier_length: ClassVar[int] = 1024 + _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( + timestamp_type="TIMESTAMP", + distinct_condition=render_is_distinct_from, + update_style=SnapshotUpdateStyle.UPDATE_FROM, + latest_version=SnapshotLatestVersionStyle.QUALIFY, + historical_close=HistoricalSnapshotCloseStyle.UPDATE_FROM, + historical_insert=HistoricalSnapshotInsertStyle.INSERT_WITH, + ) def inspect_retention( self, *, connection: _BigQueryConnection, request: RetentionRequest @@ -1180,43 +1192,16 @@ def render_apply_historical_timestamp_snapshot_changes( valid_to_column=valid_to_column, invalidate_hard_deletes=invalidate_hard_deletes, ) - if invalidate_hard_deletes: - close_sql: str = HistoricalSnapshotStatementSql.grouped_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=SNAPSHOT_VERSION_START_COLUMN, - ) - else: - close_sql = HistoricalSnapshotStatementSql.grouped_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {', '.join(unique_key)}, {updated_at_column} AS __close_at " - "FROM __new_changes" - ), - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - version_columns_sql: str = ( - f"__new_changes.{updated_at_column}, LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column})" - ) - if invalidate_hard_deletes: - version_columns_sql = historical_insert_validity_sql() - insert_sql: str = HistoricalSnapshotStatementSql.insert_with_cte_sql( + return HistoricalSnapshotStatementSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, - insert_column_sql=insert_column_sql, new_changes_sql=new_changes_sql, - select_sql=(f"SELECT {output_select_sql}, {version_columns_sql} FROM __new_changes"), + unique_key=unique_key, + change_time_column=updated_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + invalidate_hard_deletes=invalidate_hard_deletes, ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_changes( self, @@ -1236,32 +1221,16 @@ def render_apply_historical_timestamp_changes( updated_at_column=updated_at_column, valid_to_column=valid_to_column, ) - close_sql: str = HistoricalSnapshotStatementSql.grouped_close_sql( + return HistoricalSnapshotStatementSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, new_changes_sql=new_changes_sql, unique_key=unique_key, + change_time_column=updated_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {', '.join(unique_key)}, {updated_at_column} AS __close_at " - "FROM __new_changes" - ), - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = HistoricalSnapshotStatementSql.insert_with_cte_sql( - destination=destination, - insert_column_sql=insert_column_sql, - new_changes_sql=new_changes_sql, - select_sql=( - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") FROM __new_changes" - ), + output_columns=output_columns, + invalidate_hard_deletes=False, ) - return (close_sql, insert_sql) def render_apply_check_snapshot_changes( self, @@ -1351,7 +1320,9 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = HistoricalCheckSnapshotSql.initial_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -1376,7 +1347,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalCheckSnapshotSql.new_changes_ctes_sql( + return HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1384,45 +1355,9 @@ def render_apply_historical_check_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - if invalidate_hard_deletes: - close_sql: str = HistoricalSnapshotStatementSql.grouped_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=SNAPSHOT_VERSION_START_COLUMN, - ) - else: - close_sql = HistoricalSnapshotStatementSql.grouped_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {', '.join(unique_key)}, {observed_at_column} AS __close_at " - "FROM __new_changes" - ), - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - version_columns_sql: str = ( - f"__new_changes.{observed_at_column}, LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column})" - ) - if invalidate_hard_deletes: - version_columns_sql = historical_insert_validity_sql() - insert_sql: str = HistoricalSnapshotStatementSql.insert_with_cte_sql( - destination=destination, - insert_column_sql=insert_column_sql, - new_changes_sql=new_changes_sql, - select_sql=(f"SELECT {output_select_sql}, {version_columns_sql} FROM __new_changes"), - ) - return (close_sql, insert_sql) def __init__(self) -> None: self._location: str | None = None diff --git a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py index 6a2d489506..5776009646 100644 --- a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py +++ b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py @@ -23,6 +23,7 @@ ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( historical_insert_validity_sql, + render_is_distinct_from, ) from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( HistoricalTimestampSnapshotSql, @@ -60,6 +61,7 @@ RowDiffTolerances, SchemaDiffResult, SnapshotChangeTarget, + SnapshotSqlDialect, TableFreshnessMetadata, TableFreshnessRequest, ) @@ -67,11 +69,15 @@ BuiltinAdapter, CursorKind, FrameworkType, + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, LoaderLogicalType, PromotionStrategy, RelationType, RetentionChangePhase, RetentionScope, + SnapshotLatestVersionStyle, + SnapshotUpdateStyle, TablePromotionMode, ) from sqlbuild.adapter.relations.main.get_columns_for_relations import ( @@ -118,6 +124,14 @@ def get_columns_for_relations( adapter_name: ClassVar[str] = BuiltinAdapter.DATABRICKS.value sql_analysis_dialect_name: ClassVar[str | None] = "databricks" max_identifier_length: ClassVar[int] = 255 + _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( + timestamp_type="TIMESTAMP", + distinct_condition=render_is_distinct_from, + update_style=SnapshotUpdateStyle.MERGE, + latest_version=SnapshotLatestVersionStyle.QUALIFY, + historical_close=HistoricalSnapshotCloseStyle.MERGE_HARD_DELETES, + historical_insert=HistoricalSnapshotInsertStyle.WITH_INSERT, + ) def inspect_retention(self, *, connection: Any, request: RetentionRequest) -> RetentionState: target: str = self._retention_relation_target(request=request) @@ -1208,7 +1222,9 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = HistoricalCheckSnapshotSql.initial_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -1233,7 +1249,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalCheckSnapshotSql.new_changes_ctes_sql( + return HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1241,51 +1257,9 @@ def render_apply_historical_check_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=SNAPSHOT_VERSION_START_COLUMN, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - version_columns_sql: str = ( - f"__new_changes.{observed_at_column}, LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column})" - ) - if invalidate_hard_deletes: - version_columns_sql = historical_insert_validity_sql() - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_columns_sql} " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def supports_zero_copy_clone(self) -> bool: return True diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py index 1e0acac905..b1912cee98 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py @@ -27,6 +27,7 @@ from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( HistoricalSnapshotSql, historical_insert_validity_sql, + render_is_distinct_from, ) from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( HistoricalTimestampSnapshotSql, @@ -62,6 +63,7 @@ RowDiffTolerances, SchemaDiffResult, SnapshotChangeTarget, + SnapshotSqlDialect, TableFreshnessMetadata, TableFreshnessRequest, ) @@ -69,8 +71,12 @@ BuiltinAdapter, CursorKind, FrameworkType, + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, LoaderLogicalType, PromotionStrategy, + SnapshotLatestVersionStyle, + SnapshotUpdateStyle, TablePromotionMode, ) from sqlbuild.adapter.relations.main.get_columns_for_relations import ( @@ -107,6 +113,14 @@ def get_columns_for_relations( adapter_name: ClassVar[str] = BuiltinAdapter.POSTGRES.value sql_analysis_dialect_name: ClassVar[str | None] = "postgres" max_identifier_length: ClassVar[int] = 63 + _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( + timestamp_type="TIMESTAMP", + distinct_condition=render_is_distinct_from, + update_style=SnapshotUpdateStyle.UPDATE_FROM, + latest_version=SnapshotLatestVersionStyle.DERIVED_TABLE, + historical_close=HistoricalSnapshotCloseStyle.CORRELATED, + historical_insert=HistoricalSnapshotInsertStyle.WITH_INSERT, + ) def supports_zero_copy_clone(self) -> bool: return False @@ -1352,7 +1366,9 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = HistoricalCheckSnapshotSql.initial_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -2529,7 +2545,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._pg_historical_check_new_changes_cte_sql( + return HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -2537,51 +2553,9 @@ def render_apply_historical_check_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = _historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=SNAPSHOT_VERSION_START_COLUMN, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - version_columns_sql: str = ( - f"__new_changes.{observed_at_column}, LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column})" - ) - if invalidate_hard_deletes: - version_columns_sql = historical_insert_validity_sql() - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_columns_sql} " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) @staticmethod def _pg_historical_timestamp_new_changes_cte_sql( @@ -2659,71 +2633,3 @@ def _pg_historical_timestamp_changes_new_records_cte_sql( f"OR __source.{updated_at_column} > __latest.{updated_at_column}" ")" ) - - @staticmethod - def _pg_historical_check_new_changes_cte_sql( - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - check_columns=check_columns, - ).new_changes_ctes_sql(destination=destination) - - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - delta_change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - first_key: str = unique_key[0] - changed_or_first_sql: str = ( - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL" - f" OR ({delta_change_condition})" - ) - latest_sql: str = ( - "__latest AS (SELECT * FROM (" - f"SELECT *, ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - f") AS __rn FROM {destination}) AS __q WHERE __rn = 1)" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - f"), {latest_sql}, __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition})" - ")" - ")" - ) diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index e779865d81..e2c52081d7 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -21,7 +21,7 @@ HistoricalCheckSnapshotSql, ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( - historical_insert_validity_sql, + render_is_distinct_from, ) from sqlbuild.adapter.contract.classes.historical_snapshot_statement_sql import ( HistoricalSnapshotStatementSql, @@ -36,7 +36,6 @@ from sqlbuild.adapter.contract.constants import ( DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE, - SNAPSHOT_VERSION_START_COLUMN, ) from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapter.contract.main.normalize_seed_csv_value import normalize_seed_csv_value @@ -61,6 +60,7 @@ RowDiffTolerances, SchemaDiffResult, SnapshotChangeTarget, + SnapshotSqlDialect, TableFreshnessMetadata, TableFreshnessRequest, ) @@ -68,10 +68,14 @@ BuiltinAdapter, CursorKind, FrameworkType, + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, LoaderLogicalType, PromotionStrategy, RetentionChangePhase, RetentionScope, + SnapshotLatestVersionStyle, + SnapshotUpdateStyle, TablePromotionMode, ) from sqlbuild.adapter.state_sql.main.render_insert_source_freshness_records_sql import ( @@ -126,6 +130,14 @@ class SnowflakeAdapter(MicrobatchMixin, UnkeyedDiffMixin, BaseAdapter): connection_routing_keys: ClassVar[frozenset[str]] = frozenset( {"source", "profile", "target", "project_dir", "profiles_dir"} ) + _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( + timestamp_type="TIMESTAMP", + distinct_condition=render_is_distinct_from, + update_style=SnapshotUpdateStyle.UPDATE_FROM, + latest_version=SnapshotLatestVersionStyle.QUALIFY, + historical_close=HistoricalSnapshotCloseStyle.UPDATE_FROM, + historical_insert=HistoricalSnapshotInsertStyle.INSERT_WITH, + ) def inspect_retention(self, *, connection: Any, request: RetentionRequest) -> RetentionState: self._validate_relation_retention_request(request=request) @@ -1437,43 +1449,16 @@ def render_apply_historical_timestamp_snapshot_changes( valid_to_column=valid_to_column, invalidate_hard_deletes=invalidate_hard_deletes, ) - if invalidate_hard_deletes: - close_sql: str = HistoricalSnapshotStatementSql.grouped_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=SNAPSHOT_VERSION_START_COLUMN, - ) - else: - close_sql = HistoricalSnapshotStatementSql.grouped_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {', '.join(unique_key)}, {updated_at_column} AS __close_at " - "FROM __new_changes" - ), - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - version_columns_sql: str = ( - f"__new_changes.{updated_at_column}, LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column})" - ) - if invalidate_hard_deletes: - version_columns_sql = historical_insert_validity_sql() - insert_sql: str = HistoricalSnapshotStatementSql.insert_with_cte_sql( + return HistoricalSnapshotStatementSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, - insert_column_sql=insert_column_sql, new_changes_sql=new_changes_sql, - select_sql=(f"SELECT {output_select_sql}, {version_columns_sql} FROM __new_changes"), + unique_key=unique_key, + change_time_column=updated_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + invalidate_hard_deletes=invalidate_hard_deletes, ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_changes( self, @@ -1493,32 +1478,16 @@ def render_apply_historical_timestamp_changes( updated_at_column=updated_at_column, valid_to_column=valid_to_column, ) - close_sql: str = HistoricalSnapshotStatementSql.grouped_close_sql( + return HistoricalSnapshotStatementSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, new_changes_sql=new_changes_sql, unique_key=unique_key, + change_time_column=updated_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {', '.join(unique_key)}, {updated_at_column} AS __close_at " - "FROM __new_changes" - ), - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = HistoricalSnapshotStatementSql.insert_with_cte_sql( - destination=destination, - insert_column_sql=insert_column_sql, - new_changes_sql=new_changes_sql, - select_sql=( - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") FROM __new_changes" - ), + output_columns=output_columns, + invalidate_hard_deletes=False, ) - return (close_sql, insert_sql) def render_apply_check_snapshot_changes( self, @@ -1608,7 +1577,9 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = HistoricalCheckSnapshotSql.initial_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -1635,7 +1606,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalCheckSnapshotSql.new_changes_ctes_sql( + return HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1643,45 +1614,9 @@ def render_apply_historical_check_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - if invalidate_hard_deletes: - close_sql: str = HistoricalSnapshotStatementSql.grouped_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=SNAPSHOT_VERSION_START_COLUMN, - ) - else: - close_sql = HistoricalSnapshotStatementSql.grouped_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {', '.join(unique_key)}, {observed_at_column} AS __close_at " - "FROM __new_changes" - ), - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - version_columns_sql: str = ( - f"__new_changes.{observed_at_column}, LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column})" - ) - if invalidate_hard_deletes: - version_columns_sql = historical_insert_validity_sql() - insert_sql: str = HistoricalSnapshotStatementSql.insert_with_cte_sql( - destination=destination, - insert_column_sql=insert_column_sql, - new_changes_sql=new_changes_sql, - select_sql=(f"SELECT {output_select_sql}, {version_columns_sql} FROM __new_changes"), - ) - return (close_sql, insert_sql) def supports_python_functions(self) -> bool: return True diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py index 6951cf9cff..1fd258adbb 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py @@ -20,6 +20,9 @@ _render_typed_value_list, _typed_scalar_payload, ) +from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( + HistoricalCheckSnapshotSql, +) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( HistoricalSnapshotSql, historical_insert_validity_sql, @@ -60,6 +63,7 @@ RowDiffTolerances, SchemaDiffResult, SnapshotChangeTarget, + SnapshotSqlDialect, TableFreshnessMetadata, TableFreshnessRequest, ) @@ -67,8 +71,12 @@ BuiltinAdapter, CursorKind, FrameworkType, + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, LoaderLogicalType, PromotionStrategy, + SnapshotLatestVersionStyle, + SnapshotUpdateStyle, TablePromotionMode, ) from sqlbuild.adapter.relations.main.get_columns_for_relations import ( @@ -96,6 +104,14 @@ from sqlbuild.sql_values.types import SqlValueKind +def _render_tsql_distinct_condition(*, left: str, right: str) -> str: + return ( + f"({left} <> {right} " + f"OR ({left} IS NULL AND {right} IS NOT NULL) " + f"OR ({left} IS NOT NULL AND {right} IS NULL))" + ) + + class SqlServerAdapter(MicrobatchMixin, UnkeyedDiffMixin, BaseAdapter): """Microsoft SQL Server adapter backed by pymssql.""" @@ -112,6 +128,14 @@ def get_columns_for_relations( adapter_name: ClassVar[str] = BuiltinAdapter.SQLSERVER.value sql_analysis_dialect_name: ClassVar[str | None] = "tsql" max_identifier_length: ClassVar[int] = 128 + _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( + timestamp_type="DATETIME2", + distinct_condition=_render_tsql_distinct_condition, + update_style=SnapshotUpdateStyle.TSQL, + latest_version=SnapshotLatestVersionStyle.ORDERED_CTE, + historical_close=HistoricalSnapshotCloseStyle.TSQL, + historical_insert=HistoricalSnapshotInsertStyle.TSQL, + ) def supports_table_freshness_metadata(self) -> bool: return False @@ -139,11 +163,7 @@ def get_tables_freshness_metadata( ) def _distinct_condition(self, *, left: str, right: str) -> str: - return ( - f"({left} <> {right} " - f"OR ({left} IS NULL AND {right} IS NOT NULL) " - f"OR ({left} IS NOT NULL AND {right} IS NULL))" - ) + return _render_tsql_distinct_condition(left=left, right=right) def _lag_expr(self, *, column: str, partition_sql: str, observed_at_column: str) -> str: return ( @@ -261,77 +281,6 @@ def _historical_timestamp_new_changes_cte_sql( ")" ) - def _historical_check_new_changes_cte_sql( - self, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - check_columns=check_columns, - distinct_condition=self._distinct_condition, - ).new_changes_ctes_sql(destination=destination) - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - self._lag_expr( - column=column, - partition_sql=partition_sql, - observed_at_column=observed_at_column, - ) - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - delta_change_condition: str = " OR ".join( - self._distinct_condition(left=column, right=f"__prev_{column}") - for column in check_columns - ) - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - self._distinct_condition(left=f"__delta_changes.{column}", right=f"__latest.{column}") - for column in check_columns - ) - first_key: str = unique_key[0] - changed_or_new_sql: str = ( - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - ) - changed_or_first_sql: str = ( - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({delta_change_condition})" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest_ordered AS (" - f"SELECT *, ROW_NUMBER() OVER (PARTITION BY {partition_sql} " - f"ORDER BY {valid_from_column} DESC) AS __rn FROM {destination}" - "), __latest AS (" - "SELECT * FROM __latest_ordered WHERE __rn = 1" - f"), __new_changes AS ({changed_or_new_sql})" - ) - def connect(self, config: dict[str, Any]) -> _SqlServerConnection: try: import pymssql @@ -1857,7 +1806,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_check_new_changes_cte_sql( + return HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1865,52 +1814,9 @@ def render_apply_historical_check_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=SNAPSHOT_VERSION_START_COLUMN, - ) - else: - close_sql = ( - f";WITH {new_changes_sql} " - f"UPDATE __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - version_columns_sql: str = ( - f"__new_changes.{observed_at_column}, LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column})" - ) - if invalidate_hard_deletes: - version_columns_sql = historical_insert_validity_sql() - insert_sql: str = ( - f";WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_columns_sql} " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_changes( self, From e94f8abaf1301ae53fd53e8da989b7a1d7dd485d Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:04:42 +0100 Subject: [PATCH 49/57] refactor(snapshots): share parameterised historical timestamp snapshot SQL HistoricalTimestampSnapshotSql now renders the incremental historical timestamp and historical changes statements for every adapter from its SnapshotSqlDialect. render_apply_historical_timestamp_snapshot_changes and render_apply_historical_timestamp_changes delegate on every adapter, and the Postgres, SQL Server, DuckDB, Databricks and base copies of the incremental CTE and hard-delete close builders are removed. SQL Server's historical changes path keeps rendering QUALIFY for the latest stored version, exactly as before; it passes that explicitly instead of inheriting it from the base module. Rendered SQL is unchanged for every adapter. (cherry picked from commit 6749396c958f75747b8cab78ec51e2cf9e0ba897) --- .../adapter/contract/classes/base_adapter.py | 121 +----------- .../contract/classes/duckdb_backed_adapter.py | 123 +----------- .../historical_timestamp_snapshot_sql.py | 165 ++++++++++------ .../bigquery/classes/bigquery_adapter.py | 28 +-- .../databricks/classes/databricks_adapter.py | 121 +----------- .../postgres/classes/postgres_adapter.py | 165 +--------------- .../snowflake/classes/snowflake_adapter.py | 28 +-- .../sqlserver/classes/sqlserver_adapter.py | 178 ++---------------- 8 files changed, 163 insertions(+), 766 deletions(-) diff --git a/src/sqlbuild/adapter/contract/classes/base_adapter.py b/src/sqlbuild/adapter/contract/classes/base_adapter.py index 5872aa3e07..f862f8ad86 100644 --- a/src/sqlbuild/adapter/contract/classes/base_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/base_adapter.py @@ -13,7 +13,6 @@ HistoricalCheckSnapshotSql, ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( - historical_insert_validity_sql, render_is_distinct_from, ) from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( @@ -25,7 +24,6 @@ from sqlbuild.adapter.contract.classes.strict_adapter import StrictAdapter from sqlbuild.adapter.contract.constants import ( INTEGER_TYPE_TOKEN, - SNAPSHOT_VERSION_START_COLUMN, TYPED_OBJECT_ENTRY_PART_COUNT, ) from sqlbuild.adapter.contract.exceptions import ( @@ -1063,7 +1061,9 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = HistoricalTimestampSnapshotSql.initial_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1110,59 +1110,17 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalTimestampSnapshotSql.new_changes_ctes_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, observed_at_column=observed_at_column, - valid_to_column=valid_to_column, valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = _historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=SNAPSHOT_VERSION_START_COLUMN, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - version_columns_sql: str = ( - f"__new_changes.{updated_at_column}, LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column})" - ) - if invalidate_hard_deletes: - version_columns_sql = historical_insert_validity_sql() - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_columns_sql} " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_changes( self, @@ -1175,43 +1133,15 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalTimestampSnapshotSql.changes_new_records_ctes_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, + valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - close_sql: str = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_create_initial_historical_check_snapshot_destination( self, @@ -2556,36 +2486,3 @@ def dimensions(node: SqlValue) -> tuple[int, ...] | None: def _quote_sql_string(value: str) -> str: return "'" + value.replace("'", "''") + "'" - - -def _historical_snapshot_combined_close_sql( - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - change_time_column: str, -) -> str: - close_candidate_condition: str = SnapshotSql.key_condition( - left_alias="__close_candidates", right_alias="__target", unique_key=unique_key - ) - candidate_key_sql: str = ", ".join(unique_key) - return ( - f"WITH {new_changes_sql}, __close_candidates AS (" - f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at FROM __new_changes " - "UNION ALL " - f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes WHERE __close_at IS NOT NULL" - ") " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - "SELECT MIN(__close_candidates.__close_at) FROM __close_candidates " - f"WHERE {close_candidate_condition}" - ") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - "SELECT MIN(__close_candidates.__close_at) FROM __close_candidates " - f"WHERE {close_candidate_condition}" - ") " - f"AND EXISTS (SELECT 1 FROM __close_candidates WHERE {close_candidate_condition})" - ) diff --git a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py index af870264c5..00cfa5662c 100644 --- a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py @@ -22,9 +22,6 @@ from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( HistoricalCheckSnapshotSql, ) -from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( - historical_insert_validity_sql, -) from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( HistoricalTimestampSnapshotSql, ) @@ -36,7 +33,6 @@ DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE, QUALIFIED_NAME_SEPARATOR, - SNAPSHOT_VERSION_START_COLUMN, ) from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapter.contract.models import ( @@ -527,7 +523,9 @@ def render_create_initial_historical_timestamp_snapshot_destination( ) -> tuple[str, ...]: return self.render_create_table_as( destination=destination, - sql=HistoricalTimestampSnapshotSql.initial_select_sql( + sql=HistoricalTimestampSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -576,7 +574,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalTimestampSnapshotSql.new_changes_ctes_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -584,51 +582,9 @@ def render_apply_historical_timestamp_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=SNAPSHOT_VERSION_START_COLUMN, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - version_columns_sql: str = ( - f"__new_changes.{updated_at_column}, LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column})" - ) - if invalidate_hard_deletes: - version_columns_sql = historical_insert_validity_sql() - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_columns_sql} " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_changes( self, @@ -641,43 +597,15 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalTimestampSnapshotSql.changes_new_records_ctes_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, + valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - close_sql: str = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_check_snapshot_changes( self, @@ -2547,41 +2475,6 @@ def _render_seed_csv_options(self, csv_settings: SeedCsvSettings) -> str: def _duckdb_string_literal(self, value: str) -> str: return value.replace("'", "''") - @classmethod - def _historical_snapshot_combined_close_sql( - cls, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - change_time_column: str, - ) -> str: - close_candidate_condition: str = SnapshotSql.key_condition( - left_alias="__close_candidates", right_alias="__target", unique_key=unique_key - ) - candidate_key_sql: str = ", ".join(unique_key) - return ( - f"WITH {new_changes_sql}, __close_candidates AS (" - f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at FROM __new_changes " - "UNION ALL " - f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " - "WHERE __close_at IS NOT NULL" - ") " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - "SELECT MIN(__close_candidates.__close_at) FROM __close_candidates " - f"WHERE {close_candidate_condition}" - ") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - "SELECT MIN(__close_candidates.__close_at) FROM __close_candidates " - f"WHERE {close_candidate_condition}" - ") " - f"AND EXISTS (SELECT 1 FROM __close_candidates WHERE {close_candidate_condition})" - ) - def _duckdb_string_literal(value: str) -> str: return "'" + value.replace("'", "''") + "'" diff --git a/src/sqlbuild/adapter/contract/classes/historical_timestamp_snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/historical_timestamp_snapshot_sql.py index a2eeb8d532..4b9f0157bd 100644 --- a/src/sqlbuild/adapter/contract/classes/historical_timestamp_snapshot_sql.py +++ b/src/sqlbuild/adapter/contract/classes/historical_timestamp_snapshot_sql.py @@ -3,14 +3,21 @@ from __future__ import annotations from sqlbuild.adapter.contract.classes.historical_snapshot_sql import HistoricalSnapshotSql +from sqlbuild.adapter.contract.classes.historical_snapshot_statement_sql import ( + HistoricalSnapshotStatementSql, +) from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql +from sqlbuild.adapter.contract.models import SnapshotSqlDialect class HistoricalTimestampSnapshotSql: - """Render SQL for historical timestamp snapshots.""" + """Render SQL for historical timestamp and historical changes snapshots.""" + + def __init__(self, *, dialect: SnapshotSqlDialect) -> None: + self.dialect: SnapshotSqlDialect = dialect - @staticmethod def initial_select_sql( + self, *, origin: str, unique_key: tuple[str, ...], @@ -31,27 +38,28 @@ def initial_select_sql( valid_from_column=valid_from_column, valid_to_column=valid_to_column, updated_at_column=updated_at_column, + distinct_condition=self.dialect.distinct_condition, ).initial_select_sql(output_columns=output_columns) partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) + output_select_sql: str = ", ".join(output_columns) + ordered_ctes_sql: str = self._ordered_ctes_sql( + origin=origin, + unique_key=unique_key, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + changes_cte="__changes", + ) return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - ") " + f"WITH {ordered_ctes_sql} " f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " f"ORDER BY {updated_at_column}) AS {valid_to_column} " "FROM __changes" ) - @staticmethod - def new_changes_ctes_sql( + def apply_sql( + self, *, destination: str, origin: str, @@ -60,43 +68,55 @@ def new_changes_ctes_sql( observed_at_column: str, valid_from_column: str, valid_to_column: str, + output_columns: tuple[str, ...], invalidate_hard_deletes: bool, - ) -> str: - """Render incremental historical timestamp CTEs using QUALIFY for the latest versions.""" + ) -> tuple[str, ...]: + """Render the incremental statements for a historical timestamp snapshot.""" + statements: HistoricalSnapshotStatementSql = HistoricalSnapshotStatementSql( + dialect=self.dialect + ) + new_changes_sql: str if invalidate_hard_deletes: - return HistoricalSnapshotSql( + new_changes_sql = HistoricalSnapshotSql( origin=origin, unique_key=unique_key, observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, updated_at_column=updated_at_column, + distinct_condition=self.dialect.distinct_condition, ).new_changes_ctes_sql(destination=destination) - - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - ")" + else: + ordered_ctes_sql: str = self._ordered_ctes_sql( + origin=origin, + unique_key=unique_key, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + changes_cte="__delta_changes", + ) + latest_sql: str = statements.latest_version_cte_sql( + destination=destination, unique_key=unique_key, order_column=valid_from_column + ) + latest_join_condition: str = SnapshotSql.key_condition( + left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key + ) + new_changes_sql = ( + f"{ordered_ctes_sql}, {latest_sql}, __new_changes AS (" + "SELECT __delta_changes.* FROM __delta_changes " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{unique_key[0]} IS NULL " + f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column})" + ) + return statements.apply_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + change_time_column=updated_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + invalidate_hard_deletes=invalidate_hard_deletes, ) @staticmethod @@ -112,7 +132,7 @@ def changes_initial_select_sql( """Render the initial-build query for a historical changes snapshot.""" partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) + output_select_sql: str = ", ".join(output_columns) return ( f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " @@ -120,31 +140,64 @@ def changes_initial_select_sql( f"FROM {origin}" ) - @staticmethod - def changes_new_records_ctes_sql( + def changes_apply_sql( + self, *, destination: str, origin: str, unique_key: tuple[str, ...], updated_at_column: str, + valid_from_column: str, valid_to_column: str, - ) -> str: - """Render incremental historical changes CTEs using QUALIFY for the latest versions.""" + output_columns: tuple[str, ...], + ) -> tuple[str, ...]: + """Render the incremental statements for a historical changes snapshot.""" + statements: HistoricalSnapshotStatementSql = HistoricalSnapshotStatementSql( + dialect=self.dialect + ) + latest_sql: str = statements.latest_version_cte_sql( + destination=destination, unique_key=unique_key, order_column=updated_at_column + ) latest_join_condition: str = SnapshotSql.key_condition( left_alias="__source", right_alias="__latest", unique_key=unique_key ) - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - return ( - "__latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column} DESC" - ") = 1" - "), __new_changes AS (" + new_changes_sql: str = ( + f"{latest_sql}, __new_changes AS (" f"SELECT __source.* FROM {origin} AS __source " f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __latest.{updated_at_column}" - ")" + f"WHERE __latest.{unique_key[0]} IS NULL " + f"OR __source.{updated_at_column} > __latest.{updated_at_column})" + ) + return statements.apply_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + change_time_column=updated_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + invalidate_hard_deletes=False, + ) + + def _ordered_ctes_sql( + self, + *, + origin: str, + unique_key: tuple[str, ...], + updated_at_column: str, + observed_at_column: str, + changes_cte: str, + ) -> str: + partition_sql: str = ", ".join(unique_key) + changed_sql: str = self.dialect.distinct_condition( + left=updated_at_column, right="__prev_updated_at" + ) + return ( + "__ordered AS (" + f"SELECT *, LAG({updated_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" + f") AS __prev_updated_at FROM {origin}" + f"), {changes_cte} AS (" + f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL OR {changed_sql})" ) diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py index 77f765e279..922ecdb3fd 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py @@ -24,9 +24,6 @@ from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( render_is_distinct_from, ) -from sqlbuild.adapter.contract.classes.historical_snapshot_statement_sql import ( - HistoricalSnapshotStatementSql, -) from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( HistoricalTimestampSnapshotSql, ) @@ -1135,7 +1132,9 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = HistoricalTimestampSnapshotSql.initial_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1182,7 +1181,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalTimestampSnapshotSql.new_changes_ctes_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1190,15 +1189,6 @@ def render_apply_historical_timestamp_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - invalidate_hard_deletes=invalidate_hard_deletes, - ) - return HistoricalSnapshotStatementSql(dialect=self._snapshot_sql_dialect).apply_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - change_time_column=updated_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) @@ -1214,22 +1204,14 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalTimestampSnapshotSql.changes_new_records_ctes_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, - valid_to_column=valid_to_column, - ) - return HistoricalSnapshotStatementSql(dialect=self._snapshot_sql_dialect).apply_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - change_time_column=updated_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, output_columns=output_columns, - invalidate_hard_deletes=False, ) def render_apply_check_snapshot_changes( diff --git a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py index 5776009646..dcbfe506a6 100644 --- a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py +++ b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py @@ -22,7 +22,6 @@ HistoricalCheckSnapshotSql, ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( - historical_insert_validity_sql, render_is_distinct_from, ) from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( @@ -35,7 +34,6 @@ from sqlbuild.adapter.contract.constants import ( DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE, - SNAPSHOT_VERSION_START_COLUMN, ) from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapter.contract.main.normalize_seed_csv_value import normalize_seed_csv_value @@ -984,7 +982,9 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = HistoricalTimestampSnapshotSql.initial_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1031,7 +1031,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalTimestampSnapshotSql.new_changes_ctes_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1039,51 +1039,9 @@ def render_apply_historical_timestamp_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=SNAPSHOT_VERSION_START_COLUMN, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - version_columns_sql: str = ( - f"__new_changes.{updated_at_column}, LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column})" - ) - if invalidate_hard_deletes: - version_columns_sql = historical_insert_validity_sql() - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_columns_sql} " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_changes( self, @@ -1096,43 +1054,15 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalTimestampSnapshotSql.changes_new_records_ctes_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, + valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - close_sql: str = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_check_snapshot_changes( self, @@ -3039,38 +2969,3 @@ def _to_int(value: object) -> int: if isinstance(value, int): return value return int(str(value)) - - @classmethod - def _historical_snapshot_combined_close_sql( - cls, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - change_time_column: str, - ) -> str: - close_candidate_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__close_candidates", unique_key=unique_key - ) - candidate_key_sql: str = ", ".join(unique_key) - close_candidates_query: str = ( - f"WITH {new_changes_sql}, __close_candidates AS (" - f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at FROM __new_changes " - "UNION ALL " - f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " - "WHERE __close_at IS NOT NULL" - ") " - f"SELECT {candidate_key_sql}, MIN(__close_at) AS __close_at " - "FROM __close_candidates GROUP BY " - f"{candidate_key_sql}" - ) - return ( - f"MERGE INTO {destination} AS __target " - f"USING ({close_candidates_query}) AS __close_candidates " - f"ON {close_candidate_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < __close_candidates.__close_at " - f"WHEN MATCHED THEN UPDATE SET {valid_to_column} = __close_candidates.__close_at" - ) diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py index b1912cee98..f1045689e0 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py @@ -15,7 +15,6 @@ _build_names_filter, _build_schemas_filter, _encode_typed_json, - _historical_snapshot_combined_close_sql, _quote_sql_string, _render_ansi_typed_scalar, _render_typed_value_list, @@ -25,8 +24,6 @@ HistoricalCheckSnapshotSql, ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( - HistoricalSnapshotSql, - historical_insert_validity_sql, render_is_distinct_from, ) from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( @@ -40,7 +37,6 @@ DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE, QUALIFIED_NAME_SEPARATOR, - SNAPSHOT_VERSION_START_COLUMN, ) from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapter.contract.main.normalize_seed_csv_value import normalize_seed_csv_value @@ -1318,7 +1314,9 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = HistoricalTimestampSnapshotSql.initial_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -2429,7 +2427,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._pg_historical_timestamp_new_changes_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -2437,51 +2435,9 @@ def render_apply_historical_timestamp_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = _historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=SNAPSHOT_VERSION_START_COLUMN, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - version_columns_sql: str = ( - f"__new_changes.{updated_at_column}, LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column})" - ) - if invalidate_hard_deletes: - version_columns_sql = historical_insert_validity_sql() - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_columns_sql} " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_changes( self, @@ -2494,43 +2450,15 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = self._pg_historical_timestamp_changes_new_records_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, + valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - close_sql: str = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_historical_check_snapshot_changes( self, @@ -2556,80 +2484,3 @@ def render_apply_historical_check_snapshot_changes( output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - - @staticmethod - def _pg_historical_timestamp_new_changes_cte_sql( - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - updated_at_column=updated_at_column, - ).new_changes_ctes_sql(destination=destination) - - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - latest_sql: str = ( - "__latest AS (SELECT * FROM (" - f"SELECT *, ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - f") AS __rn FROM {destination}) AS __q WHERE __rn = 1)" - ) - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - f"), {latest_sql}, __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - ")" - ) - - @staticmethod - def _pg_historical_timestamp_changes_new_records_cte_sql( - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_to_column: str, - ) -> str: - del valid_to_column - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__source", right_alias="__latest", unique_key=unique_key - ) - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - return ( - "__latest AS (SELECT * FROM (" - f"SELECT *, ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column} DESC" - f") AS __rn FROM {destination}) AS __q WHERE __rn = 1" - "), __new_changes AS (" - f"SELECT __source.* FROM {origin} AS __source " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __latest.{updated_at_column}" - ")" - ) diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index e2c52081d7..0116dacfe4 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -23,9 +23,6 @@ from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( render_is_distinct_from, ) -from sqlbuild.adapter.contract.classes.historical_snapshot_statement_sql import ( - HistoricalSnapshotStatementSql, -) from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( HistoricalTimestampSnapshotSql, ) @@ -1388,7 +1385,9 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = HistoricalTimestampSnapshotSql.initial_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1439,7 +1438,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalTimestampSnapshotSql.new_changes_ctes_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1447,15 +1446,6 @@ def render_apply_historical_timestamp_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - invalidate_hard_deletes=invalidate_hard_deletes, - ) - return HistoricalSnapshotStatementSql(dialect=self._snapshot_sql_dialect).apply_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - change_time_column=updated_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) @@ -1471,22 +1461,14 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalTimestampSnapshotSql.changes_new_records_ctes_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, - valid_to_column=valid_to_column, - ) - return HistoricalSnapshotStatementSql(dialect=self._snapshot_sql_dialect).apply_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - change_time_column=updated_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, output_columns=output_columns, - invalidate_hard_deletes=False, ) def render_apply_check_snapshot_changes( diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py index 1fd258adbb..af39b242e0 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py @@ -5,6 +5,7 @@ import json import logging import re +from dataclasses import replace from datetime import date, datetime from decimal import Decimal from pathlib import Path @@ -25,7 +26,6 @@ ) from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( HistoricalSnapshotSql, - historical_insert_validity_sql, ) from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( HistoricalTimestampSnapshotSql, @@ -37,7 +37,6 @@ from sqlbuild.adapter.contract.constants import ( DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE, - SNAPSHOT_VERSION_START_COLUMN, ) from sqlbuild.adapter.contract.exceptions import ( AdapterUserError, @@ -171,43 +170,6 @@ def _lag_expr(self, *, column: str, partition_sql: str, observed_at_column: str) f"ORDER BY {observed_at_column}) AS __prev_{column}" ) - def _historical_snapshot_combined_close_sql( - self, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - change_time_column: str, - ) -> str: - close_candidate_condition: str = SnapshotSql.key_condition( - left_alias="__close_candidates", right_alias="__target", unique_key=unique_key - ) - candidate_key_sql: str = ", ".join(unique_key) - hard_delete_select_sql: str = ( - f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " - "WHERE __close_at IS NOT NULL" - ) - return ( - f";WITH {new_changes_sql}, __close_candidates AS (" - f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at FROM __new_changes " - "UNION ALL " - f"{hard_delete_select_sql}" - ") " - f"UPDATE __target SET {valid_to_column} = (" - "SELECT MIN(__close_candidates.__close_at) FROM __close_candidates " - f"WHERE {close_candidate_condition}" - ") " - f"FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - "SELECT MIN(__close_candidates.__close_at) FROM __close_candidates " - f"WHERE {close_candidate_condition}" - ") " - f"AND EXISTS (SELECT 1 FROM __close_candidates WHERE {close_candidate_condition})" - ) - def _snapshot_hard_delete_close_sql( self, *, @@ -231,56 +193,6 @@ def _snapshot_hard_delete_close_sql( ")" ) - def _historical_timestamp_new_changes_cte_sql( - self, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - if invalidate_hard_deletes: - return HistoricalSnapshotSql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - updated_at_column=updated_at_column, - distinct_condition=self._distinct_condition, - ).new_changes_ctes_sql(destination=destination) - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - latest_join_condition: str = SnapshotSql.key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - updated_changed: str = self._distinct_condition( - left=f"{updated_at_column}", right="__prev_updated_at" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL OR {updated_changed}" - "), __latest_ordered AS (" - f"SELECT *, ROW_NUMBER() OVER (PARTITION BY {partition_sql} " - f"ORDER BY {valid_from_column} DESC) AS __rn FROM {destination}" - "), __latest AS (" - "SELECT * FROM __latest_ordered WHERE __rn = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - ")" - ) - def connect(self, config: dict[str, Any]) -> _SqlServerConnection: try: import pymssql @@ -1829,44 +1741,19 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = HistoricalTimestampSnapshotSql.changes_new_records_ctes_sql( + return HistoricalTimestampSnapshotSql( + dialect=replace( + self._snapshot_sql_dialect, latest_version=SnapshotLatestVersionStyle.QUALIFY + ) + ).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, + valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - close_sql: str = ( - f";WITH {new_changes_sql} " - f"UPDATE __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f";WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_snapshot_changes( self, @@ -1881,60 +1768,17 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_new_changes_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, observed_at_column=observed_at_column, - valid_to_column=valid_to_column, valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=SNAPSHOT_VERSION_START_COLUMN, - ) - else: - close_sql = ( - f";WITH {new_changes_sql} " - f"UPDATE __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - version_columns_sql: str = ( - f"__new_changes.{updated_at_column}, LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column})" - ) - if invalidate_hard_deletes: - version_columns_sql = historical_insert_validity_sql() - insert_sql: str = ( - f";WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_columns_sql} " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_timestamp_snapshot_changes( self, From ed92cc1b8de2397cf9c1bc4f0c7c0c4b8d062bd8 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:06:36 +0100 Subject: [PATCH 50/57] refactor(snapshots): share parameterised current-state snapshot SQL SnapshotSql now renders the current-state initial query and the timestamp and check change statements for every adapter from its SnapshotSqlDialect: UPDATE ... FROM, Databricks MERGE, or T-SQL UPDATE ... FROM ... JOIN for closes, the T-SQL hard-delete close, the dialect's null timestamp type and null-safe inequality. render_create_initial_snapshot_destination, render_apply_timestamp_snapshot_changes and render_apply_check_snapshot_changes delegate on every adapter, and the SQL Server hard-delete close copy is removed. Rendered SQL is unchanged for every adapter. (cherry picked from commit 9cfef9b60c5ed68ea684550e0b4da4896cd5f34c) --- .../adapter/contract/classes/base_adapter.py | 165 ++---------- .../contract/classes/duckdb_backed_adapter.py | 166 ++----------- .../adapter/contract/classes/snapshot_sql.py | 235 ++++++++++++++++-- .../bigquery/classes/bigquery_adapter.py | 166 ++----------- .../databricks/classes/databricks_adapter.py | 164 ++---------- .../postgres/classes/postgres_adapter.py | 163 ++---------- .../snowflake/classes/snowflake_adapter.py | 166 ++----------- .../sqlserver/classes/sqlserver_adapter.py | 186 ++------------ 8 files changed, 404 insertions(+), 1007 deletions(-) diff --git a/src/sqlbuild/adapter/contract/classes/base_adapter.py b/src/sqlbuild/adapter/contract/classes/base_adapter.py index f862f8ad86..f08a4befba 100644 --- a/src/sqlbuild/adapter/contract/classes/base_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/base_adapter.py @@ -869,19 +869,17 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy=snapshot_strategy, - updated_at_column=updated_at_column, - observed_at_column=observed_at_column, - initial_valid_from=initial_valid_from, - source_alias=None, - current_timestamp=self.render_current_timestamp(), - ) return self.render_create_table_as( destination=destination, - sql=( - f"SELECT *, {valid_from_expr} AS {valid_from_column}, " - f"CAST(NULL AS TIMESTAMP) AS {valid_to_column} FROM {origin}" + sql=SnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( + origin=origin, + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + initial_valid_from=initial_valid_from, + current_timestamp=self.render_current_timestamp(), ), ) @@ -899,76 +897,21 @@ def render_apply_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy="timestamp", + return SnapshotSql(dialect=self._snapshot_sql_dialect).timestamp_changes_sql( + target=SnapshotChangeTarget( + destination=destination, + origin=origin, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + ), updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = __source.{updated_at_column} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __source.{updated_at_column} > __target.{updated_at_column}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = SnapshotSql.key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - history_join_sql: str = "" - if invalidate_hard_deletes: - key_sql: str = ", ".join(unique_key) - history_condition: str = " AND ".join( - f"__history.{column} = __source.{column}" for column in unique_key - ) - version_valid_from_expr = ( - f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " - f"AND __history.__closed_at <> __source.{updated_at_column} " - f"THEN {current_timestamp} " - f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - history_join_sql = ( - f"LEFT JOIN (SELECT {key_sql}, MAX({valid_to_column}) AS __closed_at " - f"FROM {destination} GROUP BY {key_sql}) AS __history ON {history_condition} " - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"{history_join_sql}" - f"WHERE __active.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __active.{updated_at_column}" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - SnapshotSql.hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_apply_check_snapshot_changes( self, @@ -980,69 +923,15 @@ def render_apply_check_snapshot_changes( initial_valid_from: str | None, invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - destination: str = target.destination - origin: str = target.origin - unique_key: tuple[str, ...] = target.unique_key - valid_from_column: str = target.valid_from_column - valid_to_column: str = target.valid_to_column - output_columns: tuple[str, ...] = target.output_columns - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy="check", + return SnapshotSql(dialect=self._snapshot_sql_dialect).check_changes_sql( + target=target, + check_columns=check_columns, updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __target.{column}" for column in check_columns - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND ({change_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = SnapshotSql.key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - active_change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __active.{column}" for column in check_columns - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE {current_timestamp} END" - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL OR ({active_change_condition})" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - SnapshotSql.hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_current_timestamp(self) -> str: return "CURRENT_TIMESTAMP" diff --git a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py index 00cfa5662c..9fb8597a0a 100644 --- a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py @@ -405,20 +405,17 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy=snapshot_strategy, - updated_at_column=updated_at_column, - observed_at_column=observed_at_column, - initial_valid_from=initial_valid_from, - source_alias=None, - current_timestamp=current_timestamp, - ) return self.render_create_table_as( destination=destination, - sql=( - f"SELECT *, {valid_from_expr} AS {valid_from_column}, " - f"CAST(NULL AS TIMESTAMP) AS {valid_to_column} FROM {origin}" + sql=SnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( + origin=origin, + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + initial_valid_from=initial_valid_from, + current_timestamp=self.render_current_timestamp(), ), ) @@ -436,76 +433,21 @@ def render_apply_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy="timestamp", + return SnapshotSql(dialect=self._snapshot_sql_dialect).timestamp_changes_sql( + target=SnapshotChangeTarget( + destination=destination, + origin=origin, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + ), updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = __source.{updated_at_column} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __source.{updated_at_column} > __target.{updated_at_column}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = SnapshotSql.key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - history_join_sql: str = "" - if invalidate_hard_deletes: - key_sql: str = ", ".join(unique_key) - history_condition: str = " AND ".join( - f"__history.{column} = __source.{column}" for column in unique_key - ) - version_valid_from_expr = ( - f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " - f"AND __history.__closed_at <> __source.{updated_at_column} " - f"THEN {current_timestamp} " - f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - history_join_sql = ( - f"LEFT JOIN (SELECT {key_sql}, MAX({valid_to_column}) AS __closed_at " - f"FROM {destination} GROUP BY {key_sql}) AS __history ON {history_condition} " - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"{history_join_sql}" - f"WHERE __active.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __active.{updated_at_column}" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - SnapshotSql.hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_create_initial_historical_timestamp_snapshot_destination( self, @@ -617,69 +559,15 @@ def render_apply_check_snapshot_changes( initial_valid_from: str | None, invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - destination: str = target.destination - origin: str = target.origin - unique_key: tuple[str, ...] = target.unique_key - valid_from_column: str = target.valid_from_column - valid_to_column: str = target.valid_to_column - output_columns: tuple[str, ...] = target.output_columns - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy="check", + return SnapshotSql(dialect=self._snapshot_sql_dialect).check_changes_sql( + target=target, + check_columns=check_columns, updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __target.{column}" for column in check_columns - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND ({change_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = SnapshotSql.key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - active_change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __active.{column}" for column in check_columns - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE {current_timestamp} END" - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL OR ({active_change_condition})" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - SnapshotSql.hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_create_initial_historical_check_snapshot_destination( self, diff --git a/src/sqlbuild/adapter/contract/classes/snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/snapshot_sql.py index 80f6c18132..18870f2434 100644 --- a/src/sqlbuild/adapter/contract/classes/snapshot_sql.py +++ b/src/sqlbuild/adapter/contract/classes/snapshot_sql.py @@ -1,12 +1,17 @@ -"""Shared SQL fragments for snapshot materializations.""" +"""Shared SQL for current-state snapshots and fragments common to every snapshot.""" from __future__ import annotations +from sqlbuild.adapter.contract.models import SnapshotChangeTarget, SnapshotSqlDialect +from sqlbuild.adapter.contract.types import SnapshotUpdateStyle from sqlbuild.compiler.planner.types import InitialValidFrom, SnapshotStrategy class SnapshotSql: - """Render SQL fragments shared by snapshot materializations.""" + """Render current-state snapshot statements and shared snapshot fragments.""" + + def __init__(self, *, dialect: SnapshotSqlDialect) -> None: + self.dialect: SnapshotSqlDialect = dialect @staticmethod def key_condition(*, left_alias: str, right_alias: str, unique_key: tuple[str, ...]) -> str: @@ -39,27 +44,223 @@ def initial_valid_from_expr( return f"{prefix}{updated_at_column}" return current_timestamp - @staticmethod - def hard_delete_close_sql( + def initial_select_sql( + self, *, - destination: str, origin: str, - unique_key: tuple[str, ...], + snapshot_strategy: str | None, + updated_at_column: str | None, + observed_at_column: str | None, + valid_from_column: str, valid_to_column: str, + initial_valid_from: str | None, current_timestamp: str, ) -> str: - """Render the statement that closes active versions whose key is missing from the source.""" + """Render the initial-build query for a current-state snapshot.""" - missing_key_condition: str = SnapshotSql.key_condition( - left_alias="__source", right_alias="__target", unique_key=unique_key + valid_from_expr: str = self.initial_valid_from_expr( + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + initial_valid_from=initial_valid_from, + source_alias=None, + current_timestamp=current_timestamp, ) - first_key: str = unique_key[0] return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __source " - f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" - f")" + f"SELECT *, {valid_from_expr} AS {valid_from_column}, " + f"CAST(NULL AS {self.dialect.timestamp_type}) AS {valid_to_column} FROM {origin}" + ) + + def timestamp_changes_sql( + self, + *, + target: SnapshotChangeTarget, + updated_at_column: str, + observed_at_column: str | None, + initial_valid_from: str | None, + invalidate_hard_deletes: bool, + current_timestamp: str, + ) -> tuple[str, ...]: + """Render the incremental statements for a current-state timestamp snapshot.""" + + initial_valid_from_expr: str = self.initial_valid_from_expr( + snapshot_strategy="timestamp", + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + initial_valid_from=initial_valid_from, + source_alias="__source", + current_timestamp=current_timestamp, + ) + first_key: str = target.unique_key[0] + close_sql: str = self._keyed_close_sql( + target=target, + valid_to_sql=f"__source.{updated_at_column}", + changed_sql=f"__source.{updated_at_column} > __target.{updated_at_column}", + ) + version_valid_from_expr: str = ( + f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " + f"ELSE __source.{updated_at_column} END" + ) + history_join_sql: str = "" + if invalidate_hard_deletes: + key_sql: str = ", ".join(target.unique_key) + history_condition: str = self.key_condition( + left_alias="__history", right_alias="__source", unique_key=target.unique_key + ) + version_valid_from_expr = ( + f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " + f"AND __history.__closed_at <> __source.{updated_at_column} " + f"THEN {current_timestamp} " + f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " + f"ELSE __source.{updated_at_column} END" + ) + history_join_sql = ( + f"LEFT JOIN (SELECT {key_sql}, MAX({target.valid_to_column}) AS __closed_at " + f"FROM {target.destination} GROUP BY {key_sql}) AS __history " + f"ON {history_condition} " + ) + insert_sql: str = self._insert_sql( + target=target, + valid_from_sql=version_valid_from_expr, + joins_sql=history_join_sql, + changed_sql=f"__source.{updated_at_column} > __active.{updated_at_column}", + ) + return self._with_hard_delete_close( + statements=(close_sql, insert_sql), + target=target, + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=current_timestamp, + ) + + def check_changes_sql( + self, + *, + target: SnapshotChangeTarget, + check_columns: tuple[str, ...], + updated_at_column: str | None, + observed_at_column: str | None, + initial_valid_from: str | None, + invalidate_hard_deletes: bool, + current_timestamp: str, + ) -> tuple[str, ...]: + """Render the incremental statements for a current-state check snapshot.""" + + initial_valid_from_expr: str = self.initial_valid_from_expr( + snapshot_strategy="check", + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + initial_valid_from=initial_valid_from, + source_alias="__source", + current_timestamp=current_timestamp, + ) + change_condition: str = " OR ".join( + self.dialect.distinct_condition(left=f"__source.{column}", right=f"__target.{column}") + for column in check_columns + ) + close_sql: str = self._keyed_close_sql( + target=target, valid_to_sql=current_timestamp, changed_sql=f"({change_condition})" + ) + active_change_condition: str = " OR ".join( + self.dialect.distinct_condition(left=f"__source.{column}", right=f"__active.{column}") + for column in check_columns + ) + insert_sql: str = self._insert_sql( + target=target, + valid_from_sql=( + f"CASE WHEN __active.{target.unique_key[0]} IS NULL " + f"THEN {initial_valid_from_expr} ELSE {current_timestamp} END" + ), + joins_sql="", + changed_sql=f"({active_change_condition})", + ) + return self._with_hard_delete_close( + statements=(close_sql, insert_sql), + target=target, + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=current_timestamp, + ) + + def _keyed_close_sql( + self, *, target: SnapshotChangeTarget, valid_to_sql: str, changed_sql: str + ) -> str: + key_condition: str = self.key_condition( + left_alias="__target", right_alias="__source", unique_key=target.unique_key + ) + valid_to_column: str = target.valid_to_column + if self.dialect.update_style == SnapshotUpdateStyle.MERGE: + return ( + f"MERGE INTO {target.destination} AS __target " + f"USING {target.origin} AS __source " + f"ON {key_condition} " + f"AND __target.{valid_to_column} IS NULL " + f"AND {changed_sql} " + f"WHEN MATCHED THEN UPDATE SET {valid_to_column} = {valid_to_sql}" + ) + if self.dialect.update_style == SnapshotUpdateStyle.TSQL: + return ( + f"UPDATE __target SET {valid_to_column} = {valid_to_sql} " + f"FROM {target.destination} AS __target " + f"JOIN {target.origin} AS __source ON {key_condition} " + f"WHERE __target.{valid_to_column} IS NULL " + f"AND {changed_sql}" + ) + return ( + f"UPDATE {target.destination} AS __target " + f"SET {valid_to_column} = {valid_to_sql} " + f"FROM {target.origin} AS __source " + f"WHERE {key_condition} " + f"AND __target.{valid_to_column} IS NULL " + f"AND {changed_sql}" + ) + + def _insert_sql( + self, *, target: SnapshotChangeTarget, valid_from_sql: str, joins_sql: str, changed_sql: str + ) -> str: + insert_column_sql: str = ", ".join( + (*target.output_columns, target.valid_from_column, target.valid_to_column) + ) + output_select_sql: str = ", ".join(f"__source.{column}" for column in target.output_columns) + active_join_condition: str = self.key_condition( + left_alias="__active", right_alias="__source", unique_key=target.unique_key + ) + return ( + f"INSERT INTO {target.destination} ({insert_column_sql}) " + f"SELECT {output_select_sql}, {valid_from_sql}, " + f"CAST(NULL AS {self.dialect.timestamp_type}) " + f"FROM {target.origin} AS __source " + f"LEFT JOIN {target.destination} AS __active " + f"ON {active_join_condition} AND __active.{target.valid_to_column} IS NULL " + f"{joins_sql}" + f"WHERE __active.{target.unique_key[0]} IS NULL OR {changed_sql}" + ) + + def _with_hard_delete_close( + self, + *, + statements: tuple[str, ...], + target: SnapshotChangeTarget, + invalidate_hard_deletes: bool, + current_timestamp: str, + ) -> tuple[str, ...]: + if not invalidate_hard_deletes: + return statements + missing_key_condition: str = self.key_condition( + left_alias="__source", right_alias="__target", unique_key=target.unique_key + ) + update_sql: str = ( + f"UPDATE {target.destination} AS __target " + f"SET {target.valid_to_column} = {current_timestamp} " + ) + if self.dialect.update_style == SnapshotUpdateStyle.TSQL: + update_sql = ( + f"UPDATE __target SET {target.valid_to_column} = {current_timestamp} " + f"FROM {target.destination} AS __target " + ) + hard_delete_sql: str = ( + f"{update_sql}WHERE __target.{target.valid_to_column} IS NULL " + "AND NOT EXISTS (" + f"SELECT 1 FROM {target.origin} AS __source " + f"WHERE {missing_key_condition} AND __source.{target.unique_key[0]} IS NOT NULL" + ")" ) + return (*statements, hard_delete_sql) diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py index 922ecdb3fd..b8fcd68f76 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py @@ -1016,20 +1016,17 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy=snapshot_strategy, - updated_at_column=updated_at_column, - observed_at_column=observed_at_column, - initial_valid_from=initial_valid_from, - source_alias=None, - current_timestamp=current_timestamp, - ) return self.render_create_table_as( destination=destination, - sql=( - f"SELECT *, {valid_from_expr} AS {valid_from_column}, " - f"CAST(NULL AS TIMESTAMP) AS {valid_to_column} FROM {origin}" + sql=SnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( + origin=origin, + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + initial_valid_from=initial_valid_from, + current_timestamp=self.render_current_timestamp(), ), ) @@ -1047,76 +1044,21 @@ def render_apply_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy="timestamp", + return SnapshotSql(dialect=self._snapshot_sql_dialect).timestamp_changes_sql( + target=SnapshotChangeTarget( + destination=destination, + origin=origin, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + ), updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = __source.{updated_at_column} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __source.{updated_at_column} > __target.{updated_at_column}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = SnapshotSql.key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - history_join_sql: str = "" - if invalidate_hard_deletes: - key_sql: str = ", ".join(unique_key) - history_condition: str = " AND ".join( - f"__history.{column} = __source.{column}" for column in unique_key - ) - version_valid_from_expr = ( - f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " - f"AND __history.__closed_at <> __source.{updated_at_column} " - f"THEN {current_timestamp} " - f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - history_join_sql = ( - f"LEFT JOIN (SELECT {key_sql}, MAX({valid_to_column}) AS __closed_at " - f"FROM {destination} GROUP BY {key_sql}) AS __history ON {history_condition} " - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"{history_join_sql}" - f"WHERE __active.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __active.{updated_at_column}" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - SnapshotSql.hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_create_initial_historical_timestamp_snapshot_destination( self, @@ -1224,69 +1166,15 @@ def render_apply_check_snapshot_changes( initial_valid_from: str | None, invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - destination: str = target.destination - origin: str = target.origin - unique_key: tuple[str, ...] = target.unique_key - valid_from_column: str = target.valid_from_column - valid_to_column: str = target.valid_to_column - output_columns: tuple[str, ...] = target.output_columns - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy="check", + return SnapshotSql(dialect=self._snapshot_sql_dialect).check_changes_sql( + target=target, + check_columns=check_columns, updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __target.{column}" for column in check_columns - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND ({change_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = SnapshotSql.key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - active_change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __active.{column}" for column in check_columns - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE {current_timestamp} END" - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL OR ({active_change_condition})" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - SnapshotSql.hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_create_initial_historical_check_snapshot_destination( self, diff --git a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py index dcbfe506a6..1525cd5881 100644 --- a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py +++ b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py @@ -866,20 +866,17 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy=snapshot_strategy, - updated_at_column=updated_at_column, - observed_at_column=observed_at_column, - initial_valid_from=initial_valid_from, - source_alias=None, - current_timestamp=current_timestamp, - ) return self.render_create_table_as( destination=destination, - sql=( - f"SELECT *, {valid_from_expr} AS {valid_from_column}, " - f"CAST(NULL AS TIMESTAMP) AS {valid_to_column} FROM {origin}" + sql=SnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( + origin=origin, + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + initial_valid_from=initial_valid_from, + current_timestamp=self.render_current_timestamp(), ), ) @@ -897,76 +894,21 @@ def render_apply_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy="timestamp", + return SnapshotSql(dialect=self._snapshot_sql_dialect).timestamp_changes_sql( + target=SnapshotChangeTarget( + destination=destination, + origin=origin, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + ), updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - close_sql: str = ( - f"MERGE INTO {destination} AS __target " - f"USING {origin} AS __source " - f"ON {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __source.{updated_at_column} > __target.{updated_at_column} " - f"WHEN MATCHED THEN UPDATE SET {valid_to_column} = __source.{updated_at_column}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = SnapshotSql.key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - history_join_sql: str = "" - if invalidate_hard_deletes: - key_sql: str = ", ".join(unique_key) - history_condition: str = " AND ".join( - f"__history.{column} = __source.{column}" for column in unique_key - ) - version_valid_from_expr = ( - f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " - f"AND __history.__closed_at <> __source.{updated_at_column} " - f"THEN {current_timestamp} " - f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - history_join_sql = ( - f"LEFT JOIN (SELECT {key_sql}, MAX({valid_to_column}) AS __closed_at " - f"FROM {destination} GROUP BY {key_sql}) AS __history ON {history_condition} " - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"{history_join_sql}" - f"WHERE __active.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __active.{updated_at_column}" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - SnapshotSql.hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_create_initial_historical_timestamp_snapshot_destination( self, @@ -1074,69 +1016,15 @@ def render_apply_check_snapshot_changes( initial_valid_from: str | None, invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - destination: str = target.destination - origin: str = target.origin - unique_key: tuple[str, ...] = target.unique_key - valid_from_column: str = target.valid_from_column - valid_to_column: str = target.valid_to_column - output_columns: tuple[str, ...] = target.output_columns - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy="check", + return SnapshotSql(dialect=self._snapshot_sql_dialect).check_changes_sql( + target=target, + check_columns=check_columns, updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __target.{column}" for column in check_columns - ) - close_sql: str = ( - f"MERGE INTO {destination} AS __target " - f"USING {origin} AS __source " - f"ON {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND ({change_condition}) " - f"WHEN MATCHED THEN UPDATE SET {valid_to_column} = {current_timestamp}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = SnapshotSql.key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - active_change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __active.{column}" for column in check_columns - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE {current_timestamp} END" + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL OR ({active_change_condition})" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - SnapshotSql.hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements def render_create_initial_historical_check_snapshot_destination( self, diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py index f1045689e0..f926bfa888 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py @@ -1125,19 +1125,17 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy=snapshot_strategy, - updated_at_column=updated_at_column, - observed_at_column=observed_at_column, - initial_valid_from=initial_valid_from, - source_alias=None, - current_timestamp=self.render_current_timestamp(), - ) return self.render_create_table_as( destination=destination, - sql=( - f"SELECT *, {valid_from_expr} AS {valid_from_column}, " - f"CAST(NULL AS TIMESTAMP) AS {valid_to_column} FROM {origin}" + sql=SnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( + origin=origin, + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + initial_valid_from=initial_valid_from, + current_timestamp=self.render_current_timestamp(), ), ) @@ -1155,76 +1153,21 @@ def render_apply_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy="timestamp", + return SnapshotSql(dialect=self._snapshot_sql_dialect).timestamp_changes_sql( + target=SnapshotChangeTarget( + destination=destination, + origin=origin, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + ), updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = __source.{updated_at_column} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __source.{updated_at_column} > __target.{updated_at_column}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = SnapshotSql.key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - history_join_sql: str = "" - if invalidate_hard_deletes: - key_sql: str = ", ".join(unique_key) - history_condition: str = " AND ".join( - f"__history.{column} = __source.{column}" for column in unique_key - ) - version_valid_from_expr = ( - f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " - f"AND __history.__closed_at <> __source.{updated_at_column} " - f"THEN {current_timestamp} " - f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - history_join_sql = ( - f"LEFT JOIN (SELECT {key_sql}, MAX({valid_to_column}) AS __closed_at " - f"FROM {destination} GROUP BY {key_sql}) AS __history ON {history_condition} " - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"{history_join_sql}" - f"WHERE __active.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __active.{updated_at_column}" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - SnapshotSql.hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_apply_check_snapshot_changes( self, @@ -1236,69 +1179,15 @@ def render_apply_check_snapshot_changes( initial_valid_from: str | None, invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - destination: str = target.destination - origin: str = target.origin - unique_key: tuple[str, ...] = target.unique_key - valid_from_column: str = target.valid_from_column - valid_to_column: str = target.valid_to_column - output_columns: tuple[str, ...] = target.output_columns - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy="check", + return SnapshotSql(dialect=self._snapshot_sql_dialect).check_changes_sql( + target=target, + check_columns=check_columns, updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __target.{column}" for column in check_columns - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND ({change_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = SnapshotSql.key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - active_change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __active.{column}" for column in check_columns - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE {current_timestamp} END" + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL OR ({active_change_condition})" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - SnapshotSql.hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements def render_create_initial_historical_timestamp_snapshot_destination( self, diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index 0116dacfe4..fd5bb0a32a 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -1268,21 +1268,18 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy=snapshot_strategy, - updated_at_column=updated_at_column, - observed_at_column=observed_at_column, - initial_valid_from=initial_valid_from, - source_alias=None, - current_timestamp=current_timestamp, - ) return self._render_create_table_as_type( - destination=destination, table_type=table_type, - sql=( - f"SELECT *, {valid_from_expr} AS {valid_from_column}, " - f"CAST(NULL AS TIMESTAMP) AS {valid_to_column} FROM {origin}" + destination=destination, + sql=SnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( + origin=origin, + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + initial_valid_from=initial_valid_from, + current_timestamp=self.render_current_timestamp(), ), ) @@ -1300,76 +1297,21 @@ def render_apply_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy="timestamp", + return SnapshotSql(dialect=self._snapshot_sql_dialect).timestamp_changes_sql( + target=SnapshotChangeTarget( + destination=destination, + origin=origin, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + ), updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = __source.{updated_at_column} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __source.{updated_at_column} > __target.{updated_at_column}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = SnapshotSql.key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - history_join_sql: str = "" - if invalidate_hard_deletes: - key_sql: str = ", ".join(unique_key) - history_condition: str = " AND ".join( - f"__history.{column} = __source.{column}" for column in unique_key - ) - version_valid_from_expr = ( - f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " - f"AND __history.__closed_at <> __source.{updated_at_column} " - f"THEN {current_timestamp} " - f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - history_join_sql = ( - f"LEFT JOIN (SELECT {key_sql}, MAX({valid_to_column}) AS __closed_at " - f"FROM {destination} GROUP BY {key_sql}) AS __history ON {history_condition} " - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"{history_join_sql}" - f"WHERE __active.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __active.{updated_at_column}" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - SnapshotSql.hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_create_initial_historical_timestamp_snapshot_destination( self, @@ -1481,69 +1423,15 @@ def render_apply_check_snapshot_changes( initial_valid_from: str | None, invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - destination: str = target.destination - origin: str = target.origin - unique_key: tuple[str, ...] = target.unique_key - valid_from_column: str = target.valid_from_column - valid_to_column: str = target.valid_to_column - output_columns: tuple[str, ...] = target.output_columns - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy="check", + return SnapshotSql(dialect=self._snapshot_sql_dialect).check_changes_sql( + target=target, + check_columns=check_columns, updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __target.{column}" for column in check_columns - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND ({change_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = SnapshotSql.key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - active_change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __active.{column}" for column in check_columns - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE {current_timestamp} END" + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL OR ({active_change_condition})" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - SnapshotSql.hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements def render_create_initial_historical_check_snapshot_destination( self, diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py index af39b242e0..4a3d4682bc 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py @@ -170,29 +170,6 @@ def _lag_expr(self, *, column: str, partition_sql: str, observed_at_column: str) f"ORDER BY {observed_at_column}) AS __prev_{column}" ) - def _snapshot_hard_delete_close_sql( - self, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - valid_to_column: str, - current_timestamp: str, - ) -> str: - missing_key_condition: str = SnapshotSql.key_condition( - left_alias="__source", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"UPDATE __target SET {valid_to_column} = {current_timestamp} " - f"FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL " - "AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __source " - f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" - ")" - ) - def connect(self, config: dict[str, Any]) -> _SqlServerConnection: try: import pymssql @@ -1640,70 +1617,15 @@ def render_apply_check_snapshot_changes( initial_valid_from: str | None, invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - destination: str = target.destination - origin: str = target.origin - unique_key: tuple[str, ...] = target.unique_key - valid_from_column: str = target.valid_from_column - valid_to_column: str = target.valid_to_column - output_columns: tuple[str, ...] = target.output_columns - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy="check", + return SnapshotSql(dialect=self._snapshot_sql_dialect).check_changes_sql( + target=target, + check_columns=check_columns, updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - change_condition: str = " OR ".join( - self._distinct_condition(left=f"__source.{column}", right=f"__target.{column}") - for column in check_columns - ) - close_sql: str = ( - f"UPDATE __target SET {valid_to_column} = {current_timestamp} " - f"FROM {destination} AS __target " - f"JOIN {origin} AS __source ON {key_condition} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND ({change_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = SnapshotSql.key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - active_change_condition: str = " OR ".join( - self._distinct_condition(left=f"__source.{column}", right=f"__active.{column}") - for column in check_columns - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE {current_timestamp} END" + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS DATETIME2) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL OR ({active_change_condition})" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - self._snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements def render_apply_historical_check_snapshot_changes( self, @@ -1794,75 +1716,21 @@ def render_apply_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy="timestamp", + return SnapshotSql(dialect=self._snapshot_sql_dialect).timestamp_changes_sql( + target=SnapshotChangeTarget( + destination=destination, + origin=origin, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + ), updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - close_sql: str = ( - f"UPDATE __target SET {valid_to_column} = __source.{updated_at_column} " - f"FROM {destination} AS __target " - f"JOIN {origin} AS __source ON {key_condition} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __source.{updated_at_column} > __target.{updated_at_column}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = SnapshotSql.key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - history_join_sql: str = "" - if invalidate_hard_deletes: - key_sql: str = ", ".join(unique_key) - history_condition: str = " AND ".join( - f"__history.{column} = __source.{column}" for column in unique_key - ) - version_valid_from_expr = ( - f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " - f"AND __history.__closed_at <> __source.{updated_at_column} " - f"THEN {current_timestamp} " - f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - history_join_sql = ( - f"LEFT JOIN (SELECT {key_sql}, MAX({valid_to_column}) AS __closed_at " - f"FROM {destination} GROUP BY {key_sql}) AS __history ON {history_condition} " - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS DATETIME2) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"{history_join_sql}" - f"WHERE __active.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __active.{updated_at_column}" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - self._snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_clone( self, @@ -2085,19 +1953,17 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - valid_from_expr: str = SnapshotSql.initial_valid_from_expr( - snapshot_strategy=snapshot_strategy, - updated_at_column=updated_at_column, - observed_at_column=observed_at_column, - initial_valid_from=initial_valid_from, - source_alias=None, - current_timestamp=self.render_current_timestamp(), - ) return self.render_create_table_as( destination=destination, - sql=( - f"SELECT *, {valid_from_expr} AS {valid_from_column}, " - f"CAST(NULL AS DATETIME2) AS {valid_to_column} FROM {origin}" + sql=SnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( + origin=origin, + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + initial_valid_from=initial_valid_from, + current_timestamp=self.render_current_timestamp(), ), ) From c5c39bb75fc81cdda4d585fff5f2fb56f49b1825 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:08:19 +0100 Subject: [PATCH 51/57] refactor(snapshots): render grouped and MERGE historical closes from one builder The BigQuery/Snowflake UPDATE ... FROM close and the Databricks MERGE close shared the same earliest-candidate query and close condition. Render both from one method. Rendered SQL is unchanged for every adapter. (cherry picked from commit 11fecb7ec9409c864bdf47b5ba80ae2df8865cf3) --- .../historical_snapshot_statement_sql.py | 83 ++++++------------- 1 file changed, 26 insertions(+), 57 deletions(-) diff --git a/src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py b/src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py index bfb307c66b..e0424356e1 100644 --- a/src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py +++ b/src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py @@ -146,7 +146,10 @@ def _combined_close_sql( f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " "WHERE __close_at IS NOT NULL" ) - if self.dialect.historical_close == HistoricalSnapshotCloseStyle.UPDATE_FROM: + if self.dialect.historical_close in { + HistoricalSnapshotCloseStyle.UPDATE_FROM, + HistoricalSnapshotCloseStyle.MERGE_HARD_DELETES, + }: return self._grouped_close_sql( destination=destination, new_changes_sql=new_changes_sql, @@ -155,15 +158,6 @@ def _combined_close_sql( valid_to_column=valid_to_column, close_candidates_sql=close_candidates_sql, ) - if self.dialect.historical_close == HistoricalSnapshotCloseStyle.MERGE_HARD_DELETES: - return self._merge_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - close_candidates_sql=close_candidates_sql, - ) close_candidate_condition: str = SnapshotSql.key_condition( left_alias="__close_candidates", right_alias="__target", unique_key=unique_key ) @@ -214,59 +208,34 @@ def _grouped_close_sql( valid_to_column: str, close_candidates_sql: str, ) -> str: - close_candidate_condition: str = SnapshotSql.key_condition( - left_alias="__close_candidates", right_alias="__target", unique_key=unique_key + candidate_key_sql: str = ", ".join(unique_key) + close_candidates_query: str = ( + f"WITH {new_changes_sql}, __close_candidates AS ({close_candidates_sql}) " + f"SELECT {candidate_key_sql}, MIN(__close_at) AS __close_at " + f"FROM __close_candidates GROUP BY {candidate_key_sql}" ) - close_candidates_query: str = self._earliest_close_candidates_sql( - new_changes_sql=new_changes_sql, - unique_key=unique_key, - close_candidates_sql=close_candidates_sql, + close_condition: str = ( + f"__target.{valid_to_column} IS NULL " + f"AND __target.{valid_from_column} < __close_candidates.__close_at" + ) + if self.dialect.historical_close == HistoricalSnapshotCloseStyle.MERGE_HARD_DELETES: + merge_key_condition: str = SnapshotSql.key_condition( + left_alias="__target", right_alias="__close_candidates", unique_key=unique_key + ) + return ( + f"MERGE INTO {destination} AS __target " + f"USING ({close_candidates_query}) AS __close_candidates " + f"ON {merge_key_condition} AND {close_condition} " + f"WHEN MATCHED THEN UPDATE SET {valid_to_column} = __close_candidates.__close_at" + ) + update_key_condition: str = SnapshotSql.key_condition( + left_alias="__close_candidates", right_alias="__target", unique_key=unique_key ) return ( f"UPDATE {destination} AS __target " f"SET {valid_to_column} = __close_candidates.__close_at " f"FROM ({close_candidates_query}) AS __close_candidates " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < __close_candidates.__close_at " - f"AND {close_candidate_condition}" - ) - - def _merge_close_sql( - self, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - close_candidates_sql: str, - ) -> str: - close_candidate_condition: str = SnapshotSql.key_condition( - left_alias="__target", right_alias="__close_candidates", unique_key=unique_key - ) - close_candidates_query: str = self._earliest_close_candidates_sql( - new_changes_sql=new_changes_sql, - unique_key=unique_key, - close_candidates_sql=close_candidates_sql, - ) - return ( - f"MERGE INTO {destination} AS __target " - f"USING ({close_candidates_query}) AS __close_candidates " - f"ON {close_candidate_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < __close_candidates.__close_at " - f"WHEN MATCHED THEN UPDATE SET {valid_to_column} = __close_candidates.__close_at" - ) - - @staticmethod - def _earliest_close_candidates_sql( - *, new_changes_sql: str, unique_key: tuple[str, ...], close_candidates_sql: str - ) -> str: - candidate_key_sql: str = ", ".join(unique_key) - return ( - f"WITH {new_changes_sql}, __close_candidates AS ({close_candidates_sql}) " - f"SELECT {candidate_key_sql}, MIN(__close_at) AS __close_at " - f"FROM __close_candidates GROUP BY {candidate_key_sql}" + f"WHERE {close_condition} AND {update_key_condition}" ) def _insert_sql( From 1d9446a2698038dcffc287372f4b83844e7f66a6 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:08:29 +0100 Subject: [PATCH 52/57] docs(agents): share private helpers behind duplicated adapter contract methods (cherry picked from commit 2ef529ac564303a7b993f2a4915254ae322848a6) --- AGENTS.md | 1 + 1 file changed, 1 insertion(+) diff --git a/AGENTS.md b/AGENTS.md index c74814585b..6bb5192c5e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -120,6 +120,7 @@ - Treat a genuine duplicate as evidence of possible wider drift, not only a cleanup task: a missing shared owner, parallel implementations of one concept, logic on the wrong side of a boundary, or copies that have already diverged (use `--diff`; a fix present in only one copy may be a bug in the others). Diagnose and record that wider smell before consolidating, because removing the copies erases the only deterministic signal of it. - Consolidate genuine duplication the change introduces or touches. Report unrelated findings as follow-ups instead of refactoring them opportunistically. - Members marked `[forced]` are adapter contract overrides required by `test_strict_adapter.py`. Record other intentional mirrors in `[dupes]` in `fensu.toml` with a reason. +- Public adapter contract methods are duplicated per adapter on purpose: `test_strict_adapter.py` requires every first-class adapter to define them, and `[dupes]` exempts them. The private helpers behind them must not be copied. Keep them in shared builders, such as the snapshot SQL builders in `src/sqlbuild/adapter/contract/classes/`, and let each contract method delegate. Express real dialect differences as explicit parameters, such as a `SnapshotSqlDialect` field, not as a forked copy of the helper. ## Bounded Review Process From a8c1f874bf323e160d8c6440d602ef3555873f3b Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:15:22 +0100 Subject: [PATCH 53/57] fix(snapshots): stop rendering QUALIFY for SQL Server historical changes snapshots Incremental builds of SQL Server historical timestamp snapshots with historical_input changes selected each key's latest stored version with QUALIFY, which SQL Server does not support, so those builds failed. The path now uses the same ordered-CTE form as the other SQL Server snapshot paths: ROW_NUMBER() in a __latest_ordered CTE filtered to the first row in __latest, inside the ;WITH statement. No other adapter or snapshot configuration changes. (cherry picked from commit 9b72b1f98e08a3188f1af09bf89fe69bf26f588c) --- .../sqlserver/classes/sqlserver_adapter.py | 7 +- .../snapshot_sql_rendering/_test_types.py | 10 + .../snapshot_sql_rendering/helpers.py | 191 ++++++++++++++++++ .../snapshot_sql_rendering/test_rendering.py | 55 +++++ 4 files changed, 257 insertions(+), 6 deletions(-) diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py index 4a3d4682bc..69ecbd7016 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py @@ -5,7 +5,6 @@ import json import logging import re -from dataclasses import replace from datetime import date, datetime from decimal import Decimal from pathlib import Path @@ -1663,11 +1662,7 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - return HistoricalTimestampSnapshotSql( - dialect=replace( - self._snapshot_sql_dialect, latest_version=SnapshotLatestVersionStyle.QUALIFY - ) - ).changes_apply_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py index 9e5a812187..0d45ec37fd 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py @@ -38,3 +38,13 @@ class SnapshotExecutionTestCase: history_sql: str builds: tuple[tuple[tuple[object, ...], ...], ...] expected_history: tuple[tuple[object, ...], ...] + + +@dataclass(frozen=True) +class SqlServerSnapshotRenderingTestCase: + """SQL Server snapshot renders that must stay valid T-SQL.""" + + description: str + unexpected_fragments: tuple[str, ...] + expected_changes_apply_prefix: str + expected_changes_apply_fragments: tuple[str, ...] diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py index cc43637595..84063d8205 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py @@ -380,3 +380,194 @@ def _run_build( statement: str for statement in render(adapter): connection.execute(test_case.normalize_sql(statement)) + + +def render_snapshot_sql_matrix(adapter: BaseAdapter) -> dict[str, str]: + """Render every snapshot contract method over keys, strategies, hard deletes and paths.""" + + statements_by_label: dict[str, str] = {} + unique_key: tuple[str, ...] + for unique_key in (("customer_id",), ("customer_id", "region")): + description: str + statements: tuple[str, ...] + for description, statements in ( + *_current_state_renders(adapter=adapter, unique_key=unique_key), + *_historical_renders(adapter=adapter, unique_key=unique_key), + ): + index: int + statement: str + for index, statement in enumerate(statements): + statements_by_label[f"{description} key {unique_key} #{index}"] = statement + return statements_by_label + + +def _matrix_output_columns(unique_key: tuple[str, ...]) -> tuple[str, ...]: + return (*unique_key, "plan", "status", "updated_at", "observed_at") + + +def _current_state_renders( + *, adapter: BaseAdapter, unique_key: tuple[str, ...] +) -> list[tuple[str, tuple[str, ...]]]: + renders: list[tuple[str, tuple[str, ...]]] = [] + target: SnapshotChangeTarget = SnapshotChangeTarget( + destination="target_table", + origin="source_table", + unique_key=unique_key, + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=_matrix_output_columns(unique_key), + ) + strategy: str + for strategy in ("timestamp", "check"): + renders.append( + ( + f"current-state {strategy} initial", + adapter.render_create_initial_snapshot_destination( + table_type=TableType.PERMANENT, + destination="target_table", + origin="source_table", + snapshot_strategy=strategy, + updated_at_column="updated_at", + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + initial_valid_from=None, + ), + ) + ) + hard_deletes: bool + for hard_deletes in (False, True): + renders.append( + ( + f"current-state timestamp apply (hard deletes {hard_deletes})", + adapter.render_apply_timestamp_snapshot_changes( + destination="target_table", + origin="source_table", + unique_key=unique_key, + updated_at_column="updated_at", + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + initial_valid_from=None, + output_columns=target.output_columns, + invalidate_hard_deletes=hard_deletes, + ), + ) + ) + renders.append( + ( + f"current-state check apply (hard deletes {hard_deletes})", + adapter.render_apply_check_snapshot_changes( + target=target, + check_columns=("plan", "status"), + updated_at_column="updated_at", + observed_at_column="observed_at", + initial_valid_from=None, + invalidate_hard_deletes=hard_deletes, + ), + ) + ) + return renders + + +def _historical_renders( + *, adapter: BaseAdapter, unique_key: tuple[str, ...] +) -> list[tuple[str, tuple[str, ...]]]: + output_columns: tuple[str, ...] = _matrix_output_columns(unique_key) + renders: list[tuple[str, tuple[str, ...]]] = [ + ( + "historical changes initial", + adapter.render_create_initial_historical_timestamp_changes_destination( + table_type=TableType.PERMANENT, + destination="target_table", + origin="source_table", + unique_key=unique_key, + updated_at_column="updated_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=output_columns, + ), + ), + ( + "historical changes apply", + adapter.render_apply_historical_timestamp_changes( + destination="target_table", + origin="source_table", + unique_key=unique_key, + updated_at_column="updated_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=output_columns, + ), + ), + ] + hard_deletes: bool + for hard_deletes in (False, True): + suffix: str = f"(hard deletes {hard_deletes})" + renders.append( + ( + f"historical timestamp initial {suffix}", + adapter.render_create_initial_historical_timestamp_snapshot_destination( + table_type=TableType.PERMANENT, + destination="target_table", + origin="source_table", + unique_key=unique_key, + updated_at_column="updated_at", + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=output_columns, + invalidate_hard_deletes=hard_deletes, + ), + ) + ) + renders.append( + ( + f"historical timestamp apply {suffix}", + adapter.render_apply_historical_timestamp_snapshot_changes( + destination="target_table", + origin="source_table", + unique_key=unique_key, + updated_at_column="updated_at", + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=output_columns, + invalidate_hard_deletes=hard_deletes, + ), + ) + ) + renders.append( + ( + f"historical check initial {suffix}", + adapter.render_create_initial_historical_check_snapshot_destination( + table_type=TableType.PERMANENT, + destination="target_table", + origin="source_table", + unique_key=unique_key, + check_columns=("plan", "status"), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=output_columns, + invalidate_hard_deletes=hard_deletes, + ), + ) + ) + renders.append( + ( + f"historical check apply {suffix}", + adapter.render_apply_historical_check_snapshot_changes( + destination="target_table", + origin="source_table", + unique_key=unique_key, + check_columns=("plan", "status"), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=output_columns, + invalidate_hard_deletes=hard_deletes, + ), + ) + ) + return renders diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py index 1d859dabfa..b9c2a92d6b 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py @@ -14,6 +14,10 @@ from tests.integration.src.sqlbuild.adapters.snapshot_sql_rendering._test_types import ( SnapshotReappearanceRenderingTestCase, SnapshotSqlRenderingAdapterTestCase, + SqlServerSnapshotRenderingTestCase, +) +from tests.integration.src.sqlbuild.adapters.snapshot_sql_rendering.helpers import ( + render_snapshot_sql_matrix, ) _REAPPEARANCE_INITIAL_FRAGMENTS: tuple[str, ...] = ( @@ -462,3 +466,54 @@ def test_given_hard_deletes_when_rendering_historical_check_snapshot_then_reopen assert "__prev_group_observed_at" not in without_hard_deletes_sql assert "__reappearances" not in without_hard_deletes_sql assert "__next_absence_at" not in without_hard_deletes_sql + + +@pytest.mark.parametrize( + "test_case", + [ + SqlServerSnapshotRenderingTestCase( + description="sqlserver renders every snapshot statement without QUALIFY", + unexpected_fragments=("QUALIFY",), + expected_changes_apply_prefix=( + ";WITH __latest_ordered AS (SELECT *, ROW_NUMBER() OVER (" + "PARTITION BY customer_id ORDER BY updated_at DESC) AS __rn FROM target_table), " + "__latest AS (SELECT * FROM __latest_ordered WHERE __rn = 1), __new_changes AS (" + ), + expected_changes_apply_fragments=( + "UPDATE __target SET valid_to = (", + "FROM target_table AS __target WHERE __target.valid_to IS NULL", + "INSERT INTO target_table (", + ), + ) + ], + ids=lambda case: case.description, +) +def test_given_sqlserver_when_rendering_snapshot_matrix_then_sql_is_tsql_compatible( + test_case: SqlServerSnapshotRenderingTestCase, +) -> None: + statements_by_label: dict[str, str] = render_snapshot_sql_matrix(SqlServerAdapter()) + changes_apply_statements: tuple[str, ...] = ( + SqlServerAdapter().render_apply_historical_timestamp_changes( + destination="target_table", + origin="source_table", + unique_key=("customer_id",), + updated_at_column="updated_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "plan", "updated_at"), + ) + ) + + fragment: str + for fragment in test_case.unexpected_fragments: + assert { + label: fragment in statement for label, statement in statements_by_label.items() + } == dict.fromkeys(statements_by_label, False) + assert len(changes_apply_statements) == 2 + statement: str + for statement in changes_apply_statements: + assert statement.startswith(test_case.expected_changes_apply_prefix) + changes_apply_sql: str = "\n".join(changes_apply_statements) + expected_fragment: str + for expected_fragment in test_case.expected_changes_apply_fragments: + assert expected_fragment in changes_apply_sql From cc9c86f540b5348d0c5ca23a1e8b55f9573d877f Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:52:48 +0100 Subject: [PATCH 54/57] fix(snapshots): make reappearance replay and interrupted changes idempotent --- .../classes/historical_snapshot_sql.py | 2 +- .../adapter/contract/classes/snapshot_sql.py | 31 ++++++++++++++++--- .../snapshot_sql_rendering/_test_types.py | 1 + .../snapshot_sql_rendering/helpers.py | 7 ++++- .../test_duckdb_execution.py | 25 ++++++++++++++- .../snapshot_sql_rendering/test_rendering.py | 10 +++--- 6 files changed, 63 insertions(+), 13 deletions(-) diff --git a/src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py index e2fe4878dc..a20f2933db 100644 --- a/src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py +++ b/src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py @@ -230,7 +230,7 @@ def _incremental_start_condition_sql(self) -> str: return ( "(__has_latest = 0 AND (__is_change = 1 OR __is_reappearance = 1)) " "OR (__has_latest = 1 AND __is_reappearance = 1 " - f"AND {self.observed_at_column} > __latest_valid_from) " + f"AND {self._incremental_version_start_sql()} > __latest_valid_from) " "OR (__has_latest = 1 AND __is_reappearance = 0 AND __is_change = 1 " f"AND {self._latest_change_condition_sql()})" ) diff --git a/src/sqlbuild/adapter/contract/classes/snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/snapshot_sql.py index 18870f2434..7187fe778e 100644 --- a/src/sqlbuild/adapter/contract/classes/snapshot_sql.py +++ b/src/sqlbuild/adapter/contract/classes/snapshot_sql.py @@ -81,7 +81,7 @@ def timestamp_changes_sql( invalidate_hard_deletes: bool, current_timestamp: str, ) -> tuple[str, ...]: - """Render the incremental statements for a current-state timestamp snapshot.""" + """Render changes, inserting before closing when hard deletes need active-row evidence.""" initial_valid_from_expr: str = self.initial_valid_from_expr( snapshot_strategy="timestamp", @@ -102,6 +102,7 @@ def timestamp_changes_sql( f"ELSE __source.{updated_at_column} END" ) history_join_sql: str = "" + insert_guard_sql: str | None = None if invalidate_hard_deletes: key_sql: str = ", ".join(target.unique_key) history_condition: str = self.key_condition( @@ -109,7 +110,6 @@ def timestamp_changes_sql( ) version_valid_from_expr = ( f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " - f"AND __history.__closed_at <> __source.{updated_at_column} " f"THEN {current_timestamp} " f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " f"ELSE __source.{updated_at_column} END" @@ -119,14 +119,26 @@ def timestamp_changes_sql( f"FROM {target.destination} GROUP BY {key_sql}) AS __history " f"ON {history_condition} " ) + inserted_key_condition: str = self.key_condition( + left_alias="__inserted", right_alias="__source", unique_key=target.unique_key + ) + insert_guard_sql = ( + f"NOT EXISTS (SELECT 1 FROM {target.destination} AS __inserted " + f"WHERE {inserted_key_condition} " + f"AND __inserted.{target.valid_to_column} IS NULL " + f"AND __inserted.{updated_at_column} = __source.{updated_at_column})" + ) insert_sql: str = self._insert_sql( target=target, valid_from_sql=version_valid_from_expr, joins_sql=history_join_sql, changed_sql=f"__source.{updated_at_column} > __active.{updated_at_column}", + insert_guard_sql=insert_guard_sql, ) return self._with_hard_delete_close( - statements=(close_sql, insert_sql), + statements=(insert_sql, close_sql) + if invalidate_hard_deletes + else (close_sql, insert_sql), target=target, invalidate_hard_deletes=invalidate_hard_deletes, current_timestamp=current_timestamp, @@ -214,7 +226,13 @@ def _keyed_close_sql( ) def _insert_sql( - self, *, target: SnapshotChangeTarget, valid_from_sql: str, joins_sql: str, changed_sql: str + self, + *, + target: SnapshotChangeTarget, + valid_from_sql: str, + joins_sql: str, + changed_sql: str, + insert_guard_sql: str | None = None, ) -> str: insert_column_sql: str = ", ".join( (*target.output_columns, target.valid_from_column, target.valid_to_column) @@ -223,6 +241,9 @@ def _insert_sql( active_join_condition: str = self.key_condition( left_alias="__active", right_alias="__source", unique_key=target.unique_key ) + insert_condition: str = f"__active.{target.unique_key[0]} IS NULL OR {changed_sql}" + if insert_guard_sql is not None: + insert_condition = f"({insert_condition}) AND {insert_guard_sql}" return ( f"INSERT INTO {target.destination} ({insert_column_sql}) " f"SELECT {output_select_sql}, {valid_from_sql}, " @@ -231,7 +252,7 @@ def _insert_sql( f"LEFT JOIN {target.destination} AS __active " f"ON {active_join_condition} AND __active.{target.valid_to_column} IS NULL " f"{joins_sql}" - f"WHERE __active.{target.unique_key[0]} IS NULL OR {changed_sql}" + f"WHERE {insert_condition}" ) def _with_hard_delete_close( diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py index 0d45ec37fd..4e5bfe9e98 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py @@ -38,6 +38,7 @@ class SnapshotExecutionTestCase: history_sql: str builds: tuple[tuple[tuple[object, ...], ...], ...] expected_history: tuple[tuple[object, ...], ...] + retry_after_insert: bool = False @dataclass(frozen=True) diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py index 84063d8205..9d8fb97b3f 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py @@ -347,6 +347,7 @@ def run_snapshot_builds( rows=rows, render=test_case.render_apply, day=FIRST_CLOCK_DAY + build_index, + retry_after_insert=test_case.retry_after_insert, ) history: tuple[tuple[object, ...], ...] = tuple( tuple(row) for row in connection.execute(test_case.history_sql).fetchall() @@ -373,12 +374,16 @@ def _run_build( rows: tuple[tuple[object, ...], ...], render: Callable[[BaseAdapter], tuple[str, ...]], day: int, + retry_after_insert: bool = False, ) -> None: source_sql: str = " UNION ALL ".join(test_case.source_select_sql(row) for row in rows) connection.execute(f"CREATE OR REPLACE TABLE {_SOURCE} AS {source_sql}") adapter: BaseAdapter = _pinned_clock_adapter(test_case.adapter_type, day) + statements: tuple[str, ...] = render(adapter) + if retry_after_insert: + connection.execute(test_case.normalize_sql(statements[0])) statement: str - for statement in render(adapter): + for statement in statements: connection.execute(test_case.normalize_sql(statement)) diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py index 8ec71fccc4..9c7b223ce2 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py @@ -207,6 +207,7 @@ _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2], _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:3], (*_TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:3], (1, "basic", 1, 4), (2, "team", 3, 4)), + (*_TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:3], (1, "basic", 1, 4), (2, "team", 3, 4)), ), expected_history=((1, "basic", 1, None), (2, "pro", 1, 2), (2, "team", 3, None)), ), @@ -240,6 +241,21 @@ ) _CURRENT_STATE_SCENARIOS: tuple[SnapshotExecutionScenario, ...] = ( + SnapshotExecutionScenario( + description="current timestamp reappearance updated_at equals deletion time", + kind=CURRENT_TIMESTAMP_HARD_DELETES, + builds=( + ((1, "basic", 1), (2, "pro", 1)), + ((1, "basic", 1),), + ((1, "basic", 1), (2, "team", FIRST_CLOCK_DAY + 1)), + ((1, "basic", 1), (2, "team", FIRST_CLOCK_DAY + 1)), + ), + expected_history=( + (1, "basic", 1, None), + (2, "pro", 1, FIRST_CLOCK_DAY + 1), + (2, "team", FIRST_CLOCK_DAY + 2, None), + ), + ), SnapshotExecutionScenario( description="current timestamp without hard deletes records updated_at changes", kind=CURRENT_TIMESTAMP, @@ -306,7 +322,10 @@ "test_case", [ SnapshotExecutionTestCase( - description=f"{run.adapter.name}: {run.scenario.description} ({run.path})", + description=( + f"{run.adapter.name}: {run.scenario.description} ({run.path})" + f" retry_after_insert={retry_after_insert}" + ), adapter_type=run.adapter.adapter_type, normalize_sql=run.adapter.normalize_sql, source_select_sql=run.scenario.kind.source_select_sql, @@ -315,8 +334,12 @@ history_sql=run.scenario.kind.history_sql, builds=run.builds, expected_history=run.scenario.expected_history, + retry_after_insert=retry_after_insert, ) for run in _RUNS + for retry_after_insert in ( + (False, True) if run.scenario.kind is CURRENT_TIMESTAMP_HARD_DELETES else (False,) + ) ], ids=lambda case: case.description, ) diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py index b9c2a92d6b..f297d819fb 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py @@ -61,7 +61,7 @@ "WHERE __target.effective_to IS NULL AND NOT EXISTS", "MAX(effective_to) AS __closed_at FROM target_table " "GROUP BY customer_id, region) AS __history", - "AND __history.__closed_at <> __source.updated_at THEN", + "AND __history.__closed_at IS NOT NULL THEN", ), expected_historical_check_initial_hard_delete_fragments=( "__next_absence_at", @@ -99,7 +99,7 @@ "WHERE __target.effective_to IS NULL AND NOT EXISTS", "MAX(effective_to) AS __closed_at FROM target_table " "GROUP BY customer_id, region) AS __history", - "AND __history.__closed_at <> __source.updated_at THEN", + "AND __history.__closed_at IS NOT NULL THEN", ), expected_historical_check_initial_hard_delete_fragments=( "__next_absence_at", @@ -137,7 +137,7 @@ "WHERE __target.effective_to IS NULL AND NOT EXISTS", "MAX(effective_to) AS __closed_at FROM target_table " "GROUP BY customer_id, region) AS __history", - "AND __history.__closed_at <> __source.updated_at THEN", + "AND __history.__closed_at IS NOT NULL THEN", ), expected_historical_check_initial_hard_delete_fragments=( "__next_absence_at", @@ -175,7 +175,7 @@ "WHEN MATCHED THEN UPDATE SET effective_to = __source.updated_at", "MAX(effective_to) AS __closed_at FROM target_table " "GROUP BY customer_id, region) AS __history", - "AND __history.__closed_at <> __source.updated_at THEN", + "AND __history.__closed_at IS NOT NULL THEN", ), expected_historical_check_initial_hard_delete_fragments=( "__next_absence_at", @@ -213,7 +213,7 @@ "WHERE __target.effective_to IS NULL AND NOT EXISTS", "MAX(effective_to) AS __closed_at FROM target_table " "GROUP BY customer_id, region) AS __history", - "AND __history.__closed_at <> __source.updated_at THEN", + "AND __history.__closed_at IS NOT NULL THEN", ), expected_historical_check_initial_hard_delete_fragments=( "__next_absence_at", From 91a05a900c1b5a32ec9d3d32ea8144a85a7c46b4 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:55:14 +0100 Subject: [PATCH 55/57] test(snapshots): expose interrupted retry cases to architecture checks --- .../snapshot_sql_rendering/helpers.py | 4 +- .../test_duckdb_execution.py | 38 +++++++++++++++---- 2 files changed, 31 insertions(+), 11 deletions(-) diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py index 9d8fb97b3f..a138efe157 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py @@ -380,10 +380,8 @@ def _run_build( connection.execute(f"CREATE OR REPLACE TABLE {_SOURCE} AS {source_sql}") adapter: BaseAdapter = _pinned_clock_adapter(test_case.adapter_type, day) statements: tuple[str, ...] = render(adapter) - if retry_after_insert: - connection.execute(test_case.normalize_sql(statements[0])) statement: str - for statement in statements: + for statement in (*statements[: int(retry_after_insert)], *statements): connection.execute(test_case.normalize_sql(statement)) diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py index 9c7b223ce2..a1bd95630e 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py @@ -322,10 +322,7 @@ "test_case", [ SnapshotExecutionTestCase( - description=( - f"{run.adapter.name}: {run.scenario.description} ({run.path})" - f" retry_after_insert={retry_after_insert}" - ), + description=f"{run.adapter.name}: {run.scenario.description} ({run.path})", adapter_type=run.adapter.adapter_type, normalize_sql=run.adapter.normalize_sql, source_select_sql=run.scenario.kind.source_select_sql, @@ -334,12 +331,8 @@ history_sql=run.scenario.kind.history_sql, builds=run.builds, expected_history=run.scenario.expected_history, - retry_after_insert=retry_after_insert, ) for run in _RUNS - for retry_after_insert in ( - (False, True) if run.scenario.kind is CURRENT_TIMESTAMP_HARD_DELETES else (False,) - ) ], ids=lambda case: case.description, ) @@ -352,3 +345,32 @@ def test_given_snapshot_scenario_when_executing_adapter_sql_on_duckdb_then_histo assert history == test_case.expected_history assert violations == (0, 0, 0) + + +@pytest.mark.parametrize( + "test_case", + [ + SnapshotExecutionTestCase( + description=f"{run.adapter.name}: {run.scenario.description} interrupted", + adapter_type=run.adapter.adapter_type, + normalize_sql=run.adapter.normalize_sql, + source_select_sql=run.scenario.kind.source_select_sql, + render_initial=run.scenario.kind.render_initial, + render_apply=run.scenario.kind.render_apply, + history_sql=run.scenario.kind.history_sql, + builds=run.builds, + expected_history=run.scenario.expected_history, + retry_after_insert=True, + ) + for run in _RUNS + if run.scenario.kind is CURRENT_TIMESTAMP_HARD_DELETES + ], + ids=lambda case: case.description, +) +def test_given_interrupted_timestamp_insert_when_retrying_then_history_is_idempotent( + test_case: SnapshotExecutionTestCase, +) -> None: + history, violations = run_snapshot_builds(test_case) + + assert history == test_case.expected_history + assert violations == (0, 0, 0) From 649264e8a6f6ca04a714726cb8f4d2724d973260 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:55:14 +0100 Subject: [PATCH 56/57] fix(scenarios): preserve scenario labels for unclosed comments --- .../compile/_helpers/scenarios/core.py | 2 +- .../compile/_helpers/sql_tests/core.py | 38 ++++++++++--------- .../commands/main/scenario/test_scenario.py | 30 +++++++++++++++ .../compile/_helpers/test_scenarios.py | 20 ++++++++++ 4 files changed, 71 insertions(+), 19 deletions(-) diff --git a/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py b/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py index 1fc69eab28..111313f063 100644 --- a/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py +++ b/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py @@ -60,7 +60,7 @@ def extract_sql_scenario_ctes(*, sql: str, file_label: str) -> CompileSqlScenari def extract_sql_scenario_expected_model_names(*, sql: str, file_label: str) -> tuple[str, ...]: """Extract explicit expected-model relationships without inspecting CTE bodies.""" - start: int = _skip_ignorable(sql=sql, start=0) + start: int = _skip_ignorable(sql=sql, start=0, context_label=_CONTEXT) if _try_consume_keyword(sql=sql, start=start, keyword=SQL_WITH_KEYWORD) is None: return () ctes: tuple[CompileSqlScenarioCte, ...] = extract_top_level_ctes_with_scanner( diff --git a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py index f1ca638c9d..9887de4642 100644 --- a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py +++ b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py @@ -215,16 +215,18 @@ def extract_top_level_ctes_with_scanner[CteT]( """Scan top-level `WITH` CTEs followed by the ceremonial `SELECT 1` of a test file.""" with_end: int | None = _try_consume_keyword( - sql=sql, start=_skip_ignorable(sql=sql, start=0), keyword=SQL_WITH_KEYWORD + sql=sql, + start=_skip_ignorable(sql=sql, start=0, context_label=context_label), + keyword=SQL_WITH_KEYWORD, ) if with_end is None: raise CompileInputError( f"{context_label} '{file_label}' must declare {with_requirement} before `SELECT 1`" ) - index: int = _skip_ignorable(sql=sql, start=with_end) + index: int = _skip_ignorable(sql=sql, start=with_end, context_label=context_label) recursive_end: int | None = _try_consume_keyword(sql=sql, start=index, keyword="RECURSIVE") if recursive_end is not None: - index = _skip_ignorable(sql=sql, start=recursive_end) + index = _skip_ignorable(sql=sql, start=recursive_end, context_label=context_label) ctes: list[CteT] = [] seen_cte_names: set[str] = set() @@ -238,10 +240,10 @@ def extract_top_level_ctes_with_scanner[CteT]( ) seen_cte_names.add(cte_name) - index = _skip_ignorable(sql=sql, start=index) + index = _skip_ignorable(sql=sql, start=index, context_label=context_label) if index < len(sql) and sql[index] == SQL_OPEN_PAREN_TOKEN: index = find_matching_paren(sql=sql, open_paren_index=index, context=context_label) + 1 - index = _skip_ignorable(sql=sql, start=index) + index = _skip_ignorable(sql=sql, start=index, context_label=context_label) index = _consume_keyword( sql=sql, start=index, @@ -249,7 +251,7 @@ def extract_top_level_ctes_with_scanner[CteT]( file_label=file_label, context_label=context_label, ) - index = _skip_ignorable(sql=sql, start=index) + index = _skip_ignorable(sql=sql, start=index, context_label=context_label) if index >= len(sql) or sql[index] != SQL_OPEN_PAREN_TOKEN: raise CompileInputError( f"{context_label} '{file_label}' CTE '{cte_name}' must use AS (...)" @@ -259,9 +261,9 @@ def extract_top_level_ctes_with_scanner[CteT]( sql=sql, open_paren_index=index, context=context_label ) ctes.append(cte_type(name=cte_name, sql_body=sql[cte_body_start:cte_body_end].strip())) - index = _skip_ignorable(sql=sql, start=cte_body_end + 1) + index = _skip_ignorable(sql=sql, start=cte_body_end + 1, context_label=context_label) if index < len(sql) and sql[index] == SQL_ARGUMENT_SEPARATOR_TOKEN: - index = _skip_ignorable(sql=sql, start=index + 1) + index = _skip_ignorable(sql=sql, start=index + 1, context_label=context_label) continue break @@ -857,7 +859,7 @@ def _validate_ceremonial_select( file_label: str, context_label: str = _CONTEXT, ) -> None: - if _is_ceremonial_select_statement(sql=sql, start=start): + if _is_ceremonial_select_statement(sql=sql, start=start, context_label=context_label): return raise CompileInputError( f"{context_label} '{file_label}' must end with a ceremonial top-level `SELECT 1` " @@ -865,22 +867,22 @@ def _validate_ceremonial_select( ) -def _is_ceremonial_select_statement(*, sql: str, start: int) -> bool: - index: int = _skip_ignorable(sql=sql, start=start) +def _is_ceremonial_select_statement(*, sql: str, start: int, context_label: str = _CONTEXT) -> bool: + index: int = _skip_ignorable(sql=sql, start=start, context_label=context_label) select_end: int | None = _try_consume_keyword(sql=sql, start=index, keyword="SELECT") if select_end is None: return False - index = _skip_ignorable(sql=sql, start=select_end) + index = _skip_ignorable(sql=sql, start=select_end, context_label=context_label) if index >= len(sql) or sql[index] != SQL_CEREMONIAL_SELECT_VALUE: return False - index = _skip_ignorable(sql=sql, start=index + 1) - return _is_statement_end(sql=sql, start=index) + index = _skip_ignorable(sql=sql, start=index + 1, context_label=context_label) + return _is_statement_end(sql=sql, start=index, context_label=context_label) -def _is_statement_end(*, sql: str, start: int) -> bool: +def _is_statement_end(*, sql: str, start: int, context_label: str = _CONTEXT) -> bool: index: int = start if index < len(sql) and sql[index] == SQL_STATEMENT_TERMINATOR_TOKEN: - index = _skip_ignorable(sql=sql, start=index + 1) + index = _skip_ignorable(sql=sql, start=index + 1, context_label=context_label) return index == len(sql) @@ -1097,7 +1099,7 @@ def _read_identifier( return sql[start:index], index -def _skip_ignorable(*, sql: str, start: int) -> int: +def _skip_ignorable(*, sql: str, start: int, context_label: str = _CONTEXT) -> int: index: int = start while index < len(sql): if sql[index].isspace(): @@ -1107,7 +1109,7 @@ def _skip_ignorable(*, sql: str, start: int) -> int: index = skip_line_comment(sql=sql, start=index) continue if sql.startswith("/*", index): - index = skip_block_comment(sql=sql, start=index, context=_CONTEXT) + index = skip_block_comment(sql=sql, start=index, context=context_label) continue return index return index diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/scenario/test_scenario.py b/tests/e2e/src/sqlbuild/cli/commands/main/scenario/test_scenario.py index 07793fe235..2907149c44 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/scenario/test_scenario.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/scenario/test_scenario.py @@ -42,6 +42,36 @@ ) +@pytest.mark.parametrize( + "test_case", + ( + ScenarioCliE2ETestCase( + description="unclosed scenario comment names the scenario", + command=("--no-color", "scenario", "test"), + expected_exit_code=1, + expected_stderr_fragments=("SQL scenario contains an unclosed block comment",), + ), + ), + ids=lambda case: case.description, +) +def test_given_unclosed_scenario_comment_when_running_cli_then_error_names_scenario( + test_case: ScenarioCliE2ETestCase, + tmp_path: Path, +) -> None: + repo_files: dict[str, str] = build_scenario_project_files() + repo_files["tests/scenarios/order_totals_pass.sql"] = "SCENARIO ();\nWITH /* unfinished" + project_dir: Path = prepare_inline_project( + tmp_path=tmp_path, project_name="scenario_comment", repo_files=repo_files + ) + result: subprocess.CompletedProcess[str] = run_sqb( + command=test_case.command, project_dir=project_dir + ) + + assert result.returncode == test_case.expected_exit_code + for fragment in test_case.expected_stderr_fragments: + assert fragment in result.stderr + + @pytest.mark.parametrize( "test_case", ( diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_scenarios.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_scenarios.py index 6031c6786d..3b191b7ba2 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_scenarios.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_scenarios.py @@ -457,6 +457,16 @@ def test_given_invalid_scenario_source_refs_when_building_inputs_then_it_raises_ "one __expected__ or __assert__ CTE before `SELECT 1`" ), ), + CteScannerMessageTestCase( + description="unclosed comment before a fixture uses scenario wording", + sql="WITH /* unfinished", + expected_message="SQL scenario contains an unclosed block comment", + ), + CteScannerMessageTestCase( + description="unclosed trailing comment uses scenario wording", + sql="WITH __expected__orders AS (SELECT 1 AS order_id) SELECT 1; /* unfinished", + expected_message="SQL scenario contains an unclosed block comment", + ), CteScannerMessageTestCase( description="scenario with an invalid cte name uses scenario wording", sql="WITH 1orders AS (SELECT 1) SELECT 1", @@ -505,6 +515,16 @@ def test_given_malformed_scenario_ctes_when_extracting_then_messages_name_the_sc sql="WITH __source__raw_orders (SELECT 1 AS order_id) SELECT 1", expected_message="SQL scenario 'tests/scenarios/orders.sql' expected keyword AS", ), + CteScannerMessageTestCase( + description="expected model scan unclosed comment uses scenario wording", + sql="WITH /* unfinished", + expected_message="SQL scenario contains an unclosed block comment", + ), + CteScannerMessageTestCase( + description="expected model scan leading unclosed comment uses scenario wording", + sql="/* unfinished", + expected_message="SQL scenario contains an unclosed block comment", + ), CteScannerMessageTestCase( description="expected model scan without a target uses scenario wording", sql=( From ebb867203ecf2008b50b9be7583bfa9d74d5f63a Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:55:15 +0100 Subject: [PATCH 57/57] test(ci): include hook identity regression in build shard --- Makefile | 1 + 1 file changed, 1 insertion(+) diff --git a/Makefile b/Makefile index 701ed84de0..1043a3ee74 100644 --- a/Makefile +++ b/Makefile @@ -73,6 +73,7 @@ E2E_DUCKDB_BUILD_CORE_PATHS := \ tests/e2e/src/sqlbuild/cli/commands/main/build/test_no_tests_no_audits_flags.py \ tests/e2e/src/sqlbuild/cli/commands/main/build/test_plan_command_surface.py \ tests/e2e/src/sqlbuild/cli/commands/main/build/test_python_hooks.py \ + tests/e2e/src/sqlbuild/cli/commands/main/build/test_sql_hook_identity.py \ tests/e2e/src/sqlbuild/cli/commands/main/build/test_query_change_tracking.py \ tests/e2e/src/sqlbuild/cli/commands/main/build/test_query_propagation.py \ tests/e2e/src/sqlbuild/cli/commands/main/build/test_remove_column_semantics.py \