Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ def compile_build_plan(
resolve_python_run_selectors=(request.include_python or invocation.should_load_sources),
max_microbatches=request.max_microbatches,
selection_diagnostics=request.selection_diagnostics,
plan_sql_tests=request.run_tests,
),
hooks=ConnectionHooks(
on_progress=invocation.planning_progress.on_progress,
Expand Down
1 change: 1 addition & 0 deletions src/sqlbuild/compiler/pipeline/main/compile.py
Original file line number Diff line number Diff line change
Expand Up @@ -232,6 +232,7 @@ def _build_result(
policies=PlannerPolicies(
auto_load_sources=options.auto_load_sources,
selection_diagnostics=options.selection_diagnostics,
plan_sql_tests=options.plan_sql_tests,
),
on_progress=on_progress,
project_config=discovered_inputs.project_config,
Expand Down
1 change: 1 addition & 0 deletions src/sqlbuild/compiler/pipeline/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ class CompilePipelineOptions:
no_cache: bool = False
max_microbatches: int | None = None
selection_diagnostics: bool = False
plan_sql_tests: bool = True


@dataclass(frozen=True)
Expand Down
33 changes: 26 additions & 7 deletions src/sqlbuild/compiler/planner/_helpers/output/plan_output.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@
FunctionChangeResult,
FunctionPlanEntry,
ModelPlanEntry,
PlannedSqlTests,
PlannerChangeResults,
PlannerModelEntryResults,
PlannerRelationsContext,
Expand Down Expand Up @@ -133,13 +134,13 @@ def build_plan_output(
relations=relations,
model_materializations=model_materializations,
)
test_entries: list[SqlTestPlanEntry]
test_warnings: list[PlanWarning]
test_entries, test_warnings = build_selected_test_entries(
project=project,
adapter=adapter,
selected_keys=scope.selected_keys,
)
planned_tests: PlannedSqlTests | None = resolved_extras.planned_sql_tests
if planned_tests is None or planned_tests.selected_keys != scope.selected_keys:
planned_tests = plan_selected_sql_tests(
project=project, adapter=adapter, selected_keys=scope.selected_keys
)
test_entries: list[SqlTestPlanEntry] = list(planned_tests.entries)
test_warnings: list[PlanWarning] = list(planned_tests.warnings)
selected_test_keys: frozenset[CompiledObjectKey] = frozenset(
entry.key for entry in test_entries
)
Expand Down Expand Up @@ -439,6 +440,24 @@ def build_selected_audit_entries(
return entries


def plan_selected_sql_tests(
*,
project: CompiledProject,
adapter: BaseAdapter,
selected_keys: frozenset[CompiledObjectKey],
) -> PlannedSqlTests:
"""Plan SQL tests overlapping one selection into a reusable result."""

entries: list[SqlTestPlanEntry]
warnings: list[PlanWarning]
entries, warnings = build_selected_test_entries(
project=project, adapter=adapter, selected_keys=selected_keys
)
return PlannedSqlTests(
selected_keys=selected_keys, entries=tuple(entries), warnings=tuple(warnings)
)


def build_selected_test_entries(
*,
project: CompiledProject,
Expand Down
24 changes: 24 additions & 0 deletions src/sqlbuild/compiler/planner/_helpers/planning/output_assembly.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,12 @@

from sqlbuild.compiler.compile.types import CompiledResourceType
from sqlbuild.compiler.planner._helpers.output.plan_output import build_plan_output
from sqlbuild.compiler.planner._helpers.planning.retention import plan_retention, plan_table_types
from sqlbuild.compiler.planner._helpers.pruning.selection_staleness import (
build_stale_out_of_selection_warnings,
)
from sqlbuild.compiler.planner.models import (
PlannedSqlTests,
PlannerChangeReconciliation,
PlannerChangeResults,
PlannerEntryResults,
Expand Down Expand Up @@ -39,6 +41,7 @@ def assemble_base_plan_output(
reconciliation: PlannerChangeReconciliation,
entries: PlannerEntryResults,
source_freshness: DirectSourceFreshnessPlanningResult,
planned_sql_tests: PlannedSqlTests | None = None,
) -> PlanOutput:
"""Assemble the base plan output with freshness and pruning metadata attached."""

Expand All @@ -54,6 +57,7 @@ def assemble_base_plan_output(
extras=PlanOutputExtras(
seed_version_hashes=identities.version_identities.seed_version_hashes,
seed_metadata_jsons=identities.version_identities.seed_metadata_jsons,
planned_sql_tests=planned_sql_tests,
),
)
plan_output = replace(plan_output, source_freshness=source_freshness)
Expand All @@ -68,6 +72,26 @@ def assemble_base_plan_output(
return plan_output


def with_storage_policies(
*,
plan_output: PlanOutput,
runtime: PlannerRuntime,
warehouse: PlannerWarehouseState,
scopes: PlannerScopeResolution,
) -> PlanOutput:
"""Attach selected Snowflake table-type and retention work to the plan."""

return replace(
plan_output,
table_type_entries=plan_table_types(
runtime=runtime, warehouse=warehouse, scope=scopes.selected_scope
),
retention_entries=plan_retention(
runtime=runtime, warehouse=warehouse, scope=scopes.selected_scope
),
)


def with_plan_warnings(
*,
runtime: PlannerRuntime,
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
"""Plan SQL tests on a worker thread while the planner waits on the warehouse."""

from __future__ import annotations

import threading
from concurrent.futures import Future
from types import TracebackType

from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter
from sqlbuild.compiler.compile.models import CompiledObjectKey, CompiledProject
from sqlbuild.compiler.planner._helpers.output.plan_output import plan_selected_sql_tests
from sqlbuild.compiler.planner.models import PlannedSqlTests


class BackgroundSqlTestPlanning:
"""Start warehouse-independent SQL test planning early and join it before assembly."""

def __init__(
self,
*,
project: CompiledProject,
adapter: BaseAdapter,
selected_keys: frozenset[CompiledObjectKey],
enabled: bool,
) -> None:
self._future: Future[PlannedSqlTests] | None = None
if enabled:
future: Future[PlannedSqlTests] = Future()
self._future = future
threading.Thread(
target=_plan_into,
kwargs={
"future": future,
"project": project,
"adapter": adapter,
"selected_keys": selected_keys,
},
name="sqlbuild-test-planning",
daemon=True,
).start()
self._selected_keys: frozenset[CompiledObjectKey] = selected_keys

def __enter__(self) -> BackgroundSqlTestPlanning:
return self

def __exit__(
self,
exception_type: type[BaseException] | None,
exception: BaseException | None,
traceback: TracebackType | None,
) -> None:
del exception_type, exception, traceback

def result(self) -> PlannedSqlTests:
"""Return planned tests, raising any planning error at the original assembly point."""

if self._future is None:
return PlannedSqlTests(selected_keys=self._selected_keys)
return self._future.result()


def _plan_into(
*,
future: Future[PlannedSqlTests],
project: CompiledProject,
adapter: BaseAdapter,
selected_keys: frozenset[CompiledObjectKey],
) -> None:
"""Run planning on a daemon thread so a failed plan never waits for it at exit."""

try:
future.set_result(
plan_selected_sql_tests(project=project, adapter=adapter, selected_keys=selected_keys)
)
except BaseException as error: # noqa: BLE001 - re-raised on the joining thread
future.set_exception(error)
Loading
Loading