fix: run selected Python nodes against upstreams outside the run

Read-side tasks and assets now wait only for SQL resources scheduled in the current build; unselected upstreams and unmanaged sources are read as they exist, like unselected upstreams of a selected SQL model. A skipped upstream model now skips its dependent Python nodes, as it already blocks dependent models.

Author: kvlonge

Makefile

@@ -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 \

src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/python-nodes/overview.md

@@ -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 |

src/sqlbuild/cli/commands/_helpers/build_python_nodes/python_lifecycle.py

@@ -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 (
@@ -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,

src/sqlbuild/cli/commands/_helpers/build_python_nodes/python_lifecycle_selection.py

@@ -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,

src/sqlbuild/cli/commands/_helpers/check/core.py

@@ -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
@@ -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"})
@@ -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,
@@ -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

src/sqlbuild/executor/python_nodes/_helpers/read_side_execution.py

@@ -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

src/sqlbuild/executor/python_nodes/classes/read_side_python_execution_tracker.py

@@ -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 = (
@@ -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
 
@@ -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, ()):

src/sqlbuild/executor/python_nodes/main/read_side.py

@@ -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:
@@ -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,
     )

tests/e2e/src/sqlbuild/cli/commands/main/build/_test_types.py

@@ -1104,7 +1104,7 @@ class PythonSourceTaskOnlyE2ETestCase:
 
     description: str
     target_config: str
-    expected_task_pattern: str
+    expected_task_count: int
 
 
 @dataclass(frozen=True)
@@ -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 = ""

tests/e2e/src/sqlbuild/cli/commands/main/build/helpers.py

@@ -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'
+        ),
+    }