Fix: Preserve execution observation context and propagate worker cancellation

Signed-off-by: mday-io <mdaytn@gmail.com>

Author: mday-io

docs/reference/execution_observations.md

@@ -67,8 +67,13 @@ Ordinary observer `Exception` failures are suppressed and do not replace the
 workload result or error. This guarantee does **not** cover observer-raised
 `BaseException` subclasses such as `KeyboardInterrupt` or `SystemExit`: a failure
 in start can prevent work and finish, and a failure in finish can replace a
-workload error. Blocking callbacks and process termination are also outside the
-guarantee. SQLMesh does not recursively log observer failures.
+workload error. A worker's `BaseException` can propagate to the caller instead
+of leaving parallel completion pending. Once a DAG run records a terminal error,
+it stops admitting further nodes and raises that error after already-running
+workers exit. The parallel value helper records a worker's `BaseException` in its
+result and joins the pool before resolving results. Blocking callbacks and
+process termination are also outside the guarantee. SQLMesh does not recursively
+log observer failures.
 
 While a callback is being delivered, observations triggered by that callback are
 suppressed, and internal metadata updates are ignored. This guard applies only
@@ -117,16 +122,18 @@ consumer associate existing logs with the invocation.
 ## Worker context
 
 When observation is active, the instrumented scheduler and snapshot DAG dispatch
-paths copy the admitting context at execution time. Each node runs in its own
-copy, including serial DAG execution. Consumer-owned `ContextVar` values can
-therefore carry parent identities into workers; successors inherit the admission
-context rather than the preceding worker's changes. Reusing a DAG executor
-captures a fresh context for each run.
+paths copy the admitting context at execution time. Snapshot schema creation and
+data-object listing also copy the context for each admitted value. Each node or
+value runs in its own copy, including serial execution. Consumer-owned
+`ContextVar` values can therefore carry parent identities into workers;
+successors inherit the admission context rather than the preceding worker's
+changes. Reusing a DAG executor captures a fresh context for each run.
 
 Context copies are shallow. Mutable objects stored in ContextVars can still be
 shared; consumers should use immutable values or replace values instead of
-mutating shared state. This behavior applies to the instrumented DAG paths, not
-to arbitrary consumer threads or every SQLMesh concurrency helper.
+mutating shared state. This behavior applies to these instrumented scheduler,
+DAG, schema-creation, and data-object-listing paths, not to arbitrary consumer
+threads or every SQLMesh concurrency helper.
 
-With no observer registered, these call sites do not opt into context copying
-and retain their existing serial and worker execution behavior.
+With no observer registered, these call sites do not opt into context copying;
+ContextVars keep the concurrency helper's existing serial and worker behavior.

sqlmesh/core/scheduler.py

@@ -1513,6 +1513,7 @@ def _create_schema(
                 list(unique_schemas),
                 lambda item: _create_schema(item[0], item[1], item[2]),
                 self.ddl_concurrent_tasks,
+                context_factory=execution_context_factory(),
             )
 
     def get_adapter(self, gateway: t.Optional[str] = None) -> EngineAdapter:
@@ -1687,6 +1688,7 @@ def _get_data_objects_in_schema(
                         schema=s, object_names=tables_by_schema.get(s), gateway=gateway
                     ),
                     self.ddl_concurrent_tasks,
+                    context_factory=execution_context_factory(),
                 )
 
                 for schema, objs in zip(schema_list, results):

sqlmesh/core/snapshot/evaluator.py

@@ -76,26 +76,40 @@ def run(self) -> t.Tuple[t.List[NodeExecutionFailedError[H]], t.List[H]]:
             self._context = None
 
     def _process_node(self, node: H, executor: Executor) -> None:
+        with self._unprocessed_nodes_lock:
+            if self._finished_future.done():
+                return
+
         try:
             self.fn(node)
 
             with self._unprocessed_nodes_lock:
+                if self._finished_future.done():
+                    return
                 self._unprocessed_nodes_num -= 1
                 self._submit_next_nodes(executor, node)
         except Exception as ex:
             error = NodeExecutionFailedError(node)
             error.__cause__ = ex
 
-            if self.raise_on_error:
-                self._finished_future.set_exception(error)
-                return
-
             with self._unprocessed_nodes_lock:
+                if self._finished_future.done():
+                    return
+                if self.raise_on_error:
+                    self._finished_future.set_exception(error)
+                    return
                 self._unprocessed_nodes_num -= 1
                 self._node_errors.append(error)
                 self._skip_next_nodes(node)
+        except BaseException as ex:
+            with self._unprocessed_nodes_lock:
+                if not self._finished_future.done():
+                    self._finished_future.set_exception(ex)
 
     def _submit_next_nodes(self, executor: Executor, processed_node: t.Optional[H] = None) -> None:
+        if self._finished_future.done():
+            return
+
         if not self._unprocessed_nodes_num:
             self._finished_future.set_result(None)
             return
@@ -117,6 +131,9 @@ def _submit_next_nodes(self, executor: Executor, processed_node: t.Optional[H] =
                 )
 
     def _skip_next_nodes(self, parent: H) -> None:
+        if self._finished_future.done():
+            return
+
         if not self._unprocessed_nodes_num:
             self._finished_future.set_result(None)
             return
@@ -280,30 +297,43 @@ def concurrent_apply_to_values(
     values: t.Sequence[A],
     fn: t.Callable[[A], R],
     tasks_num: int,
+    context_factory: t.Optional[t.Callable[[], Context]] = None,
 ) -> t.List[R]:
     """Applies a function to the given collection of values concurrently.
 
     Args:
         values: Target values.
         fn: The function that will be applied concurrently to each value.
         tasks_num: The number of concurrent tasks.
+        context_factory: Optional factory called once when values are admitted. Each value runs
+            in a separate copy of that context, including serial execution. By default no
+            context is captured or copied.
 
     Returns:
         A list of results.
     """
+    context = context_factory() if context_factory else None
+
     if tasks_num == 1:
-        return [fn(value) for value in values]
+        if context is None:
+            return [fn(value) for value in values]
+        return [context.copy().run(fn, value) for value in values]
 
     futures: t.List[Future] = [Future() for _ in values]
 
     def _process_value(value: A, index: int) -> None:
         try:
-            futures[index].set_result(fn(value))
-        except Exception as ex:
+            result = fn(value)
+        except BaseException as ex:
             futures[index].set_exception(ex)
+        else:
+            futures[index].set_result(result)
 
     with ThreadPoolExecutor(max_workers=tasks_num) as pool:
         for index, value in enumerate(values):
-            pool.submit(_process_value, value, index)
+            if context is None:
+                pool.submit(_process_value, value, index)
+            else:
+                pool.submit(context.copy().run, _process_value, value, index)
 
     return [f.result() for f in futures]

sqlmesh/utils/concurrency.py

@@ -401,3 +401,110 @@ def run_scheduler(*args):
         ]
     else:
         assert observer.events == []
+
+
+@pytest.mark.parametrize("registration", ["console", "disabled"])
+def test_standalone_scheduler_observes_environment_statement_queries(
+    duck_conn, mocker, registration
+):
+    from contextlib import nullcontext
+
+    from sqlmesh.core.engine_adapter import DuckDBEngineAdapter
+    from sqlmesh.core.environment import EnvironmentNamingInfo, EnvironmentStatements
+    from sqlmesh.core.scheduler import Scheduler
+    from sqlmesh.core.snapshot import DeployabilityIndex, SnapshotEvaluator
+
+    observer = Observer()
+    console = mocker.MagicMock()
+    getter = mocker.patch.object(console, "get_execution_observer", return_value=observer)
+    state_sync = mocker.MagicMock()
+    state_sync.get_environment_statements.return_value = [
+        EnvironmentStatements(
+            before_all=["CREATE TABLE observed_environment (stage VARCHAR)"],
+            after_all=["INSERT INTO observed_environment VALUES ('after')"],
+            python_env={},
+        )
+    ]
+    scheduler = Scheduler(
+        snapshots=[],
+        snapshot_evaluator=SnapshotEvaluator(DuckDBEngineAdapter(lambda: duck_conn)),
+        state_sync=state_sync,
+        default_catalog=None,
+        console=console,
+    )
+
+    scope = observer_scope(None) if registration == "disabled" else nullcontext()
+    with scope:
+        errors, skipped = scheduler.run_merged_intervals(
+            merged_intervals={},
+            deployability_index=DeployabilityIndex.all_deployable(),
+            environment_naming_info=EnvironmentNamingInfo(name="observed"),
+            run_environment_statements=True,
+        )
+
+    assert not errors and not skipped
+    assert duck_conn.execute("SELECT stage FROM observed_environment").fetchall() == [("after",)]
+    if registration == "console":
+        assert [(phase, facts["kind"], facts["action"]) for phase, facts, _ in observer.events] == [
+            ("start", "query", "execute"),
+            ("finish", "query", "execute"),
+            ("start", "query", "execute"),
+            ("finish", "query", "execute"),
+        ]
+        assert all(error is None for phase, _, error in observer.events if phase == "finish")
+        getter.assert_called_once_with()
+    else:
+        assert observer.events == []
+        getter.assert_not_called()
+
+
+@pytest.mark.parametrize("failing_stage", ["before_all", "after_all"])
+def test_standalone_scheduler_restores_observer_scope_after_environment_query_error(
+    duck_conn, mocker, failing_stage
+):
+    from sqlmesh.core.engine_adapter import DuckDBEngineAdapter
+    from sqlmesh.core.environment import EnvironmentNamingInfo, EnvironmentStatements
+    from sqlmesh.core.scheduler import Scheduler
+    from sqlmesh.core.snapshot import DeployabilityIndex, SnapshotEvaluator
+    from sqlmesh.utils.errors import SQLMeshError
+
+    observer = Observer()
+    console = mocker.MagicMock()
+    getter = mocker.patch.object(console, "get_execution_observer", return_value=observer)
+    state_sync = mocker.MagicMock()
+    state_sync.get_environment_statements.return_value = [
+        EnvironmentStatements(
+            before_all=(
+                ["SELECT * FROM missing_before_all"]
+                if failing_stage == "before_all"
+                else ["SELECT 1"]
+            ),
+            after_all=(["SELECT * FROM missing_after_all"] if failing_stage == "after_all" else []),
+            python_env={},
+        )
+    ]
+    scheduler = Scheduler(
+        snapshots=[],
+        snapshot_evaluator=SnapshotEvaluator(DuckDBEngineAdapter(lambda: duck_conn)),
+        state_sync=state_sync,
+        default_catalog=None,
+        console=console,
+    )
+
+    with pytest.raises(SQLMeshError):
+        scheduler.run_merged_intervals(
+            merged_intervals={},
+            deployability_index=DeployabilityIndex.all_deployable(),
+            environment_naming_info=EnvironmentNamingInfo(name="observed"),
+            run_environment_statements=True,
+        )
+
+    query_events = [event for event in observer.events if event[1]["kind"] == "query"]
+    assert len(query_events) == (2 if failing_stage == "before_all" else 4)
+    assert query_events[-1][0] == "finish"
+    assert query_events[-1][2] is not None
+    getter.assert_called_once_with()
+    state_sync.recycle.assert_not_called()
+
+    with action("query", "outside") as facts:
+        assert facts is None

tests/core/test_execution_observation.py

@@ -248,6 +248,122 @@ def x(evaluator, y=None) -> None:
     )
 
 
+def test_evaluate_reports_materialization_through_transaction_exit(
+    duck_conn, mocker: MockerFixture, make_snapshot
+):
+    from sqlmesh.core.engine_adapter import DuckDBEngineAdapter
+    from sqlmesh.core.execution_observation import observer_scope
+
+    model = SqlModel(
+        name="test_schema.test_model",
+        kind=FullKind(),
+        query=parse_one("SELECT 1 AS value"),
+    )
+    snapshot = make_snapshot(model)
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+    duck_conn.execute("CREATE SCHEMA sqlmesh__test_schema")
+    evaluator = SnapshotEvaluator(DuckDBEngineAdapter(lambda: duck_conn))
+    observations: t.List[t.Tuple[str, t.Dict[str, t.Any], t.Optional[BaseException]]] = []
+    materialize_active = False
+
+    def observer(phase, facts, error):
+        nonlocal materialize_active
+        observations.append((phase, dict(facts), error))
+        if facts["action"] == "materialize":
+            materialize_active = phase == "start"
+
+    target = snapshot.table_name()
+    transaction_failure = RuntimeError("transaction exit failed")
+    transaction_depth = 0
+    transaction = evaluator.adapter.transaction
+
+    @contextlib.contextmanager
+    def fail_outer_transaction(*args: t.Any, **kwargs: t.Any) -> t.Iterator[None]:
+        nonlocal transaction_depth
+        outermost = transaction_depth == 0
+        fail_on_exit = outermost and materialize_active
+        transaction_depth += 1
+        try:
+            with transaction(*args, **kwargs):
+                yield
+            if fail_on_exit:
+                raise transaction_failure
+        finally:
+            transaction_depth -= 1
+
+    mocker.patch.object(evaluator.adapter, "transaction", new=fail_outer_transaction)
+    with observer_scope(observer), pytest.raises(RuntimeError) as raised:
+        evaluator.evaluate(
+            snapshot,
+            start="2020-01-01",
+            end="2020-01-02",
+            execution_time="2020-01-02",
+            snapshots={},
+            target_table_exists=False,
+        )
+
+    assert raised.value is transaction_failure
+    assert evaluator.adapter.table_exists(target)
+    materialize_events = [event for event in observations if event[1]["action"] == "materialize"]
+    assert [phase for phase, _, _ in materialize_events] == ["start", "finish"]
+    assert all(facts["kind"] == "physical" for _, facts, _ in materialize_events)
+    assert all(facts["target"] == target for _, facts, _ in materialize_events)
+    assert all(facts["creating"] is True for _, facts, _ in materialize_events)
+    assert materialize_events[0][2] is None
+    assert materialize_events[1][2] is transaction_failure
+
+
+@pytest.mark.parametrize("workers", [1, 2])
+def test_create_physical_schemas_preserves_observation_context(duck_conn, make_snapshot, workers):
+    from contextvars import ContextVar
+
+    from sqlmesh.core.engine_adapter import DuckDBEngineAdapter
+    from sqlmesh.core.execution_observation import observer_scope
+
+    snapshots = []
+    for index in ("one", "two"):
+        snapshot = make_snapshot(
+            SqlModel(
+                name=f"schema_{index}.model",
+                kind=FullKind(),
+                query=parse_one("SELECT 1 AS value"),
+            )
+        )
+        snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+        snapshots.append(snapshot)
+    evaluator = SnapshotEvaluator(
+        DuckDBEngineAdapter(lambda: duck_conn), ddl_concurrent_tasks=workers
+    )
+    consumer = ContextVar("schema_consumer", default=None)
+    observed = []
+
+    def observer(phase, facts, error):
+        observed.append((phase, dict(facts), error, consumer.get()))
+
+    token = consumer.set("parent")
+    try:
+        with observer_scope(observer):
+            evaluator.create_physical_schemas(snapshots, DeployabilityIndex.all_deployable())
+    finally:
+        consumer.reset(token)
+
+    query_events = [event for event in observed if event[1]["kind"] == "query"]
+    assert len(query_events) == 4
+    assert sum(phase == "start" for phase, _, _, _ in query_events) == 2
+    assert sum(phase == "finish" for phase, _, _, _ in query_events) == 2
+    assert all(facts["action"] == "execute" for _, facts, _, _ in query_events)
+    assert all(context == "parent" for _, _, _, context in query_events)
+    assert all(error is None for phase, _, error, _ in query_events if phase == "finish")
+
+    created_schemas = {
+        row[0]
+        for row in duck_conn.execute(
+            "SELECT schema_name FROM information_schema.schemata"
+        ).fetchall()
+    }
+    assert {"sqlmesh__schema_one", "sqlmesh__schema_two"} <= created_schemas
+
+
 def test_runtime_stages(capsys, mocker, adapter_mock, make_snapshot):
     evaluator = SnapshotEvaluator(adapter_mock)