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
1 change: 1 addition & 0 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,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_python_node_selection.py \
tests/e2e/src/sqlbuild/cli/commands/main/build/test_python_source_deferral.py \
tests/e2e/src/sqlbuild/cli/commands/main/build/test_sql_hook_reads.py \
tests/e2e/src/sqlbuild/cli/commands/main/build/test_python_node_identity.py \
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,12 @@ Python nodes run in two phases relative to SQL:
- **Ingress (pre-SQL):** loaders, and tasks/assets that feed sources, run before SQL models are built.
- **Read-side (post-SQL):** tasks/assets that read SQL run after their SQL dependencies are built.

Read-side tasks and assets follow the same selection rule as SQL models. A SQL dependency that is
part of the run gates the node: it waits for that model, seed, or source load, and is skipped if the
dependency fails or is skipped. A dependency that is not part of the run (not selected, or an
unmanaged source) is read as it currently exists, so `sqb build --select task:export_orders` runs
the task against the relations already in the warehouse.

The commands differ in what they include by default:

| Command | SQL | Loaders / tasks / assets | Checks | Audits |
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter
from sqlbuild.cli.commands._helpers.build_python_nodes.python_lifecycle_selection import (
python_node_result_names,
scheduled_sql_resource_names,
task_asset_python_node_names,
)
from sqlbuild.cli.commands._helpers.build_python_nodes.python_node_output import (
Expand Down Expand Up @@ -177,6 +178,8 @@ def prepare_direct_python_lifecycle(
read_side_tracker = create_read_side_python_execution_tracker(
python_graph=python_graph,
selected_python_names=read_side_names,
scheduled_sql_names=scheduled_sql_resource_names(plan_output)
| frozenset(result.source_name for result in ingress_load_results),
runtime=PythonNodeRuntime(
adapter=adapter,
connection_config=connection_config,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,17 @@ def task_asset_python_node_names(
return frozenset(name for name in selected_names if name in task_asset_names)


def scheduled_sql_resource_names(plan: PlanOutput) -> frozenset[str]:
"""Return the models, seeds, and source loads this run executes and reports results for."""

scheduled_keys: frozenset[CompiledObjectKey] = frozenset(plan.execution_order)
return frozenset(
entry.name
for entry in (*plan.model_entries, *plan.seed_entries, *plan.source_load_entries)
if entry.key in scheduled_keys
)


def sql_loader_functions_for_lifecycle_handoff(
*,
discovered_inputs: DiscoveredProjectInputs,
Expand Down
23 changes: 1 addition & 22 deletions src/sqlbuild/cli/commands/_helpers/check/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@
PythonSqlRunLifecyclePlan,
)
from sqlbuild.compiler.python_nodes.types import PythonNodeKind, PythonNodeStatus
from sqlbuild.executor.build.models import SeedExecutionResult
from sqlbuild.executor.build.types import ExecutionStatus
from sqlbuild.executor.load.models import LoadExecutionResult
from sqlbuild.executor.python_nodes.main.read_side import create_read_side_python_execution_tracker
Expand All @@ -46,12 +45,10 @@
PythonNodeRunState,
PythonNodeRuntime,
)
from sqlbuild.executor.run.models import ModelExecutionResult
from sqlbuild.presentation.classes.cli_style import CliStyle
from sqlbuild.provider.main.runtime import ProviderContainer
from sqlbuild.python_nodes.models import SqlResourceRef
from sqlbuild.python_nodes.types import SqlResourceRefKind
from sqlbuild.runtime.contracts.types import ExecutionResourceKind
from sqlbuild.spec.contracts.models import SourceEntry

_IGNORED_EXCLUDE_SELECTOR_ERROR_CODES: frozenset[str] = frozenset({"S007", "S008", "S009"})
Expand Down Expand Up @@ -234,6 +231,7 @@ def run_check_read_side_dependencies(
tracker: Any = create_read_side_python_execution_tracker(
python_graph=python_graph,
selected_python_names=read_side_names,
scheduled_sql_names=frozenset(),
runtime=PythonNodeRuntime(
adapter=adapter,
connection_config=connection_config,
Expand All @@ -249,25 +247,6 @@ def run_check_read_side_dependencies(
providers=providers,
),
)
for sql_ref in sorted(read_side_refs, key=_sql_ref_sort_key):
if sql_ref.kind == SqlResourceRefKind.MODEL:
tracker.record_sql_result(
ModelExecutionResult(model_name=sql_ref.name, status=ExecutionStatus.SUCCESS)
)
elif sql_ref.kind == SqlResourceRefKind.SEED:
tracker.record_sql_result(
SeedExecutionResult(seed_name=sql_ref.name, status=ExecutionStatus.SUCCESS)
)
elif sql_ref.kind == SqlResourceRefKind.SOURCE:
tracker.record_sql_result(
LoadExecutionResult(
source_name=sql_ref.name,
loader_name=sql_ref.name,
status=ExecutionStatus.SUCCESS,
target=relation_targets[sql_ref],
resource_kind=ExecutionResourceKind.SOURCE,
)
)
tracker.dispatch_ready_python_nodes()
tracker.finalize_unrun_python_nodes()
return tracker.results
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@ def _sql_result_name(result: object) -> str | None:


def _sql_result_failed(result: object) -> bool:
if isinstance(result, ModelExecutionResult | LoadExecutionResult | SeedExecutionResult):
if isinstance(result, ModelExecutionResult):
return result.status in {ExecutionStatus.FAILED, ExecutionStatus.SKIPPED}
if isinstance(result, LoadExecutionResult | SeedExecutionResult):
return result.status == ExecutionStatus.FAILED
return False
Original file line number Diff line number Diff line change
Expand Up @@ -31,18 +31,20 @@


class ReadSidePythonExecutionTracker:
"""Dispatch selected read-only Python nodes as SQL dependencies complete."""
"""Dispatch selected read-only Python nodes as SQL dependencies scheduled in this run finish."""

def __init__(
self,
*,
python_graph: PythonNodeGraph,
selected_python_names: frozenset[str],
scheduled_sql_names: frozenset[str],
runtime: PythonNodeRuntime,
identity_recorder: PythonIdentityRecorder | None = None,
) -> None:
self._python_graph: PythonNodeGraph = python_graph
self._selected_python_names: frozenset[str] = selected_python_names
self._scheduled_sql_names: frozenset[str] = scheduled_sql_names
self._runtime: PythonNodeRuntime = runtime
self._identity_recorder: PythonIdentityRecorder | None = identity_recorder
self._result_store: Any | None = (
Expand Down Expand Up @@ -133,7 +135,10 @@ def _sql_dependencies_satisfied(self, node: DiscoveredPythonNode) -> bool:
for sql_dep_name in (dependency.name for dependency in node.sql_deps):
if sql_dep_name in self._failed_sql_names:
return False
if sql_dep_name not in self._completed_sql_names:
if (
sql_dep_name in self._scheduled_sql_names
and sql_dep_name not in self._completed_sql_names
):
return False
return True

Expand Down Expand Up @@ -191,7 +196,10 @@ def _unrun_reason(self, node: DiscoveredPythonNode) -> str:
for sql_dep_name in (dependency.name for dependency in node.sql_deps):
if sql_dep_name in self._failed_sql_names:
return f"Upstream SQL resource did not succeed: {sql_dep_name}"
if sql_dep_name not in self._completed_sql_names:
if (
sql_dep_name in self._scheduled_sql_names
and sql_dep_name not in self._completed_sql_names
):
return f"Upstream SQL resource did not complete: {sql_dep_name}"
upstream_name: str
for upstream_name in self._python_graph.upstream_deps.get(node.name, ()):
Expand Down
2 changes: 2 additions & 0 deletions src/sqlbuild/executor/python_nodes/main/read_side.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ def create_read_side_python_execution_tracker(
*,
python_graph: PythonNodeGraph,
selected_python_names: frozenset[str],
scheduled_sql_names: frozenset[str],
runtime: PythonNodeRuntime,
identity_recorder: PythonIdentityRecorder | None = None,
) -> ReadSidePythonExecutionTracker:
Expand All @@ -22,6 +23,7 @@ def create_read_side_python_execution_tracker(
return ReadSidePythonExecutionTracker(
python_graph=python_graph,
selected_python_names=selected_python_names,
scheduled_sql_names=scheduled_sql_names,
runtime=runtime,
identity_recorder=identity_recorder,
)
17 changes: 16 additions & 1 deletion tests/e2e/src/sqlbuild/cli/commands/main/build/_test_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -1104,7 +1104,7 @@ class PythonSourceTaskOnlyE2ETestCase:

description: str
target_config: str
expected_task_pattern: str
expected_task_count: int


@dataclass(frozen=True)
Expand Down Expand Up @@ -1142,3 +1142,18 @@ class AwareCursorDateWatermarkE2ETestCase:
plan_args: tuple[str, ...]
expected_plan_fragment: str
expected_event_ids: tuple[int, ...]


@dataclass(frozen=True)
class PythonNodeSelectionE2ETestCase:
"""Python nodes selected with or without their SQL upstreams."""

description: str
select: tuple[str, ...]
expected_exit_code: int
expected_output_patterns: tuple[str, ...]
expected_row_counts: dict[str, int]
expected_node_counts: dict[str, int]
expected_missing_tables: tuple[str, ...] = ()
orders_sql: str = 'MODEL (materialized table);\nSELECT order_id FROM __source("raw_orders")\n'
existing_sql: str = ""
65 changes: 65 additions & 0 deletions tests/e2e/src/sqlbuild/cli/commands/main/build/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -751,3 +751,68 @@ def explicit_reference_literal_loader(*, table: str) -> str:
f' ctx.query("SELECT count(*) FROM {table}")\n'
" return [{'id': 1}]\n"
)


PYTHON_NODE_SELECTION_DATABASE: str = "warehouse.duckdb"


def _counting_python_node(*, decorator: str, name: str, ref: str, returned: str) -> str:
return (
f"@{decorator}(depends_on={ref})\n"
f"def {name}(ctx):\n"
f" relation = ctx.relation({ref})\n"
f' ctx.execute_sql(f"CREATE OR REPLACE TABLE main.{name}_result AS '
f'SELECT count(*) AS n FROM {{relation}}")\n'
f" return {returned}\n\n\n"
)


def python_node_selection_project_files(*, orders_sql: str) -> dict[str, str]:
"""Return a DuckDB project whose tasks and assets count managed, unmanaged, and model rows."""

task_nodes: str = "".join(
_counting_python_node(decorator="task", name=name, ref=ref, returned="None")
for name, ref in (
("count_customers", 'source("raw_customers")'),
("count_raw_orders", 'source("raw_orders")'),
("count_orders", 'model("orders")'),
("count_order_summary", 'model("order_summary")'),
)
)
asset_nodes: str = "".join(
_counting_python_node(
decorator="asset", name=name, ref=ref, returned="ctx.result(materialized=True)"
)
for name, ref in (
("customers_extract", 'source("raw_customers")'),
("orders_extract", 'model("orders")'),
)
)
return {
"sqlbuild_project.toml": (
'name = "order_nodes"\nadapter = "duckdb"\n\n'
f'[connection]\ndatabase = "{PYTHON_NODE_SELECTION_DATABASE}"\n'
),
"sources/raw.yml": (
"sources:\n"
" - name: raw_orders\n schema: main\n table: raw_orders\n"
" - name: raw_customers\n managed: true\n write_strategy: table\n"
" columns:\n - name: customer_id\n type: INTEGER\n"
),
"python/loaders/customers.py": (
"from sqlbuild.loaders import loader\n\n\n"
"@loader\ndef raw_customers(ctx):\n return [{'customer_id': 99}]\n"
),
"python/tasks/counts.py": (
"from sqlbuild.refs import model, source\nfrom sqlbuild.tasks import task\n\n\n"
+ task_nodes
),
"python/assets/extracts.py": (
"from sqlbuild.assets import asset\nfrom sqlbuild.refs import model, source\n\n\n"
+ asset_nodes
),
"models/orders.sql": orders_sql,
"models/order_summary.sql": (
'MODEL (materialized table);\nSELECT count(*) AS order_count FROM __ref("orders")\n'
),
}
Loading
Loading