From edfb13b16a5741b4de1f766d1041fbfb0f3f56cc Mon Sep 17 00:00:00 2001 From: mday-io Date: Tue, 13 Oct 2026 19:18:22 -0400 Subject: [PATCH 1/3] Fix: Propagate terminal worker failures and stop queued DAG work Signed-off-by: mday-io --- sqlmesh/utils/concurrency.py | 31 ++++- tests/utils/test_concurrency.py | 215 ++++++++++++++++++++++++++++++++ 2 files changed, 240 insertions(+), 6 deletions(-) diff --git a/sqlmesh/utils/concurrency.py b/sqlmesh/utils/concurrency.py index c5f78645f6..1a9a5e2134 100644 --- a/sqlmesh/utils/concurrency.py +++ b/sqlmesh/utils/concurrency.py @@ -65,26 +65,40 @@ def run(self) -> t.Tuple[t.List[NodeExecutionFailedError[H]], t.List[H]]: return self._node_errors, self._skipped_nodes 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 @@ -101,6 +115,9 @@ def _submit_next_nodes(self, executor: Executor, processed_node: t.Optional[H] = executor.submit(self._process_node, submitted_node, executor) 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 @@ -267,9 +284,11 @@ def concurrent_apply_to_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): diff --git a/tests/utils/test_concurrency.py b/tests/utils/test_concurrency.py index 5e1e4326f7..89f9789725 100644 --- a/tests/utils/test_concurrency.py +++ b/tests/utils/test_concurrency.py @@ -1,3 +1,9 @@ +import os +import subprocess +import sys +import textwrap +from pathlib import Path + import pytest from pytest_mock.plugin import MockerFixture @@ -147,3 +153,212 @@ def test_concurrent_apply(tasks_num: int): values = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] results = concurrent_apply_to_values(values, lambda x: x * 2, tasks_num) assert results == [x * 2 for x in values] + + +def _run_bounded_python(program: str) -> None: + repo_root = Path(__file__).resolve().parents[2] + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join( + path for path in (str(repo_root), env.get("PYTHONPATH")) if path + ) + result = subprocess.run( + [sys.executable, "-c", textwrap.dedent(program)], + cwd=repo_root, + env=env, + capture_output=True, + text=True, + timeout=8, + ) + assert result.returncode == 0, f"{result.stdout}\n{result.stderr}" + + +def test_concurrent_apply_to_values_propagates_base_exception(): + _run_bounded_python( + """ + from sqlmesh.utils.concurrency import concurrent_apply_to_values + + error = SystemExit("cancel value execution") + + def process(value): + if value == 0: + raise error + return value + + try: + concurrent_apply_to_values([0, 1], process, 2) + except SystemExit as raised: + assert raised is error + else: + raise AssertionError("the value worker cancellation was not propagated") + """ + ) + + +@pytest.mark.parametrize("raise_on_error", [True, False]) +def test_concurrent_dag_propagates_base_exception_drains_workers_and_can_be_reused( + raise_on_error: bool, +): + _run_bounded_python( + f""" + from threading import Event, Thread + + from sqlmesh.utils.concurrency import ConcurrentDAGExecutor + from sqlmesh.utils.dag import DAG + + dag = DAG() + dag.add("cancel") + dag.add("sibling") + dag.add("dependent", ["sibling"]) + error = KeyboardInterrupt("cancel execution") + release_sibling = Event() + sibling_started = Event() + called = [] + + def execute(node): + if node == "cancel": + assert sibling_started.wait(timeout=3) + raise error + if node == "sibling": + sibling_started.set() + assert release_sibling.wait(timeout=6) + called.append(node) + + executor = ConcurrentDAGExecutor(dag, execute, 2, {raise_on_error!r}) + run_errors = [] + + def run(): + try: + executor.run() + except BaseException as ex: + run_errors.append(ex) + + runner = Thread(target=run) + runner.start() + try: + assert sibling_started.wait(timeout=3) + try: + executor._finished_future.result(timeout=3) + except KeyboardInterrupt as raised: + assert raised is error + else: + raise AssertionError("the worker cancellation was not propagated") + assert runner.is_alive(), "run returned before the active worker drained" + finally: + release_sibling.set() + + runner.join(timeout=3) + assert not runner.is_alive() + assert run_errors == [error] + assert called == ["sibling"] + + second_error = SystemExit("cancel reused executor") + + def cancel_again(node): + raise second_error + + executor.fn = cancel_again + try: + executor.run() + except SystemExit as raised: + assert raised is second_error + else: + raise AssertionError("the reused executor did not propagate cancellation") + """ + ) + + +def test_concurrent_dag_fail_fast_stops_queued_work_while_aggregate_mode_continues(): + _run_bounded_python( + """ + from threading import Event, Thread + + from sqlmesh.utils.concurrency import ConcurrentDAGExecutor + from sqlmesh.utils.dag import DAG + + def make_dag(): + dag = DAG() + dag.add("fail") + dag.add("running") + dag.add("queued") + dag.add("failed_child", ["fail"]) + dag.add("running_child", ["running"]) + dag.add("queued_child", ["queued"]) + return dag + + failure = RuntimeError("fail node") + release_running = Event() + running_started = Event() + queued_started = Event() + called = [] + + def execute(node): + if node == "fail": + assert running_started.wait(timeout=3) + raise failure + if node == "running": + running_started.set() + assert release_running.wait(timeout=6) + if node == "queued": + queued_started.set() + called.append(node) + + executor = ConcurrentDAGExecutor(make_dag(), execute, 2, True) + run_errors = [] + + def run_fail_fast(): + try: + executor.run() + except BaseException as ex: + run_errors.append(ex) + + runner = Thread(target=run_fail_fast) + runner.start() + try: + assert running_started.wait(timeout=3) + try: + executor._finished_future.result(timeout=3) + except Exception as raised: + assert raised.node == "fail" + assert raised.__cause__ is failure + else: + raise AssertionError("fail-fast execution did not raise") + assert runner.is_alive(), "run returned before the active worker drained" + assert not queued_started.is_set(), "queued work started after fail-fast" + finally: + release_running.set() + + runner.join(timeout=3) + assert not runner.is_alive() + assert len(run_errors) == 1 + assert run_errors[0].node == "fail" + assert run_errors[0].__cause__ is failure + assert called == ["running"] + + release_running = Event() + running_started = Event() + queued_started = Event() + called = [] + executor = ConcurrentDAGExecutor(make_dag(), execute, 2, False) + run_results = [] + + def run_aggregate(): + run_results.append(executor.run()) + + runner = Thread(target=run_aggregate) + runner.start() + try: + assert running_started.wait(timeout=3) + assert queued_started.wait(timeout=3), "aggregate mode stopped independent work" + finally: + release_running.set() + + runner.join(timeout=3) + assert not runner.is_alive() + errors, skipped = run_results[0] + assert len(errors) == 1 + assert errors[0].node == "fail" + assert errors[0].__cause__ is failure + assert skipped == ["failed_child"] + assert set(called) == {"running", "queued", "running_child", "queued_child"} + """ + ) From 1766e524255f4d7c826f197f6fa6c0d1db1aa0ca Mon Sep 17 00:00:00 2001 From: mday-io Date: Tue, 13 Oct 2026 19:18:33 -0400 Subject: [PATCH 2/3] Feat: Add optional per-task context propagation to concurrency helpers Signed-off-by: mday-io --- sqlmesh/utils/concurrency.py | 62 +++++++++++++++++++++++++++------ tests/utils/test_concurrency.py | 50 ++++++++++++++++++++++++++ 2 files changed, 102 insertions(+), 10 deletions(-) diff --git a/sqlmesh/utils/concurrency.py b/sqlmesh/utils/concurrency.py index 1a9a5e2134..a4271f8072 100644 --- a/sqlmesh/utils/concurrency.py +++ b/sqlmesh/utils/concurrency.py @@ -1,6 +1,7 @@ import typing as t from concurrent.futures import Executor, Future, ThreadPoolExecutor from threading import Lock +from contextvars import Context from sqlmesh.core.snapshot import SnapshotId, SnapshotInfoLike from sqlmesh.utils.dag import DAG @@ -30,6 +31,9 @@ class ConcurrentDAGExecutor(t.Generic[H]): raise_on_error: If set to True raises an exception on a first encountered error, otherwises returns a tuple which contains a list of failed nodes and a list of skipped nodes. + context_factory: Optional factory called at each run admission. Each node runs + in a separate copy of that context, including serial execution. By default + no context is captured or copied. """ def __init__( @@ -38,11 +42,14 @@ def __init__( fn: t.Callable[[H], None], tasks_num: int, raise_on_error: bool, + context_factory: t.Optional[t.Callable[[], Context]] = None, ): self.dag = dag self.fn = fn self.tasks_num = tasks_num self.raise_on_error = raise_on_error + self.context_factory = context_factory + self._context: t.Optional[Context] = None self._init_state() @@ -55,14 +62,18 @@ def run(self) -> t.Tuple[t.List[NodeExecutionFailedError[H]], t.List[H]]: Returns: A pair which contains a list of node errors and a list of skipped nodes. """ + self._context = self.context_factory() if self.context_factory else None if self._finished_future.done(): self._init_state() - with ThreadPoolExecutor(max_workers=self.tasks_num) as pool: - with self._unprocessed_nodes_lock: - self._submit_next_nodes(pool) - self._finished_future.result() - return self._node_errors, self._skipped_nodes + try: + with ThreadPoolExecutor(max_workers=self.tasks_num) as pool: + with self._unprocessed_nodes_lock: + self._submit_next_nodes(pool) + self._finished_future.result() + return self._node_errors, self._skipped_nodes + finally: + self._context = None def _process_node(self, node: H, executor: Executor) -> None: with self._unprocessed_nodes_lock: @@ -112,7 +123,12 @@ def _submit_next_nodes(self, executor: Executor, processed_node: t.Optional[H] = for submitted_node in submitted_nodes: self._unprocessed_nodes.pop(submitted_node) - executor.submit(self._process_node, submitted_node, executor) + if self._context is None: + executor.submit(self._process_node, submitted_node, executor) + else: + executor.submit( + self._context.copy().run, self._process_node, submitted_node, executor + ) def _skip_next_nodes(self, parent: H) -> None: if self._finished_future.done(): @@ -156,6 +172,7 @@ def concurrent_apply_to_snapshots( tasks_num: int, reverse_order: bool = False, raise_on_error: bool = True, + context_factory: t.Optional[t.Callable[[], Context]] = None, ) -> t.Tuple[t.List[NodeExecutionFailedError[SnapshotId]], t.List[SnapshotId]]: """Applies a function to the given collection of snapshots concurrently while preserving the topological order between snapshots. @@ -168,6 +185,9 @@ def concurrent_apply_to_snapshots( raise_on_error: If set to True raises an exception on a first encountered error, otherwises returns a tuple which contains a list of failed nodes and a list of skipped nodes. + context_factory: Optional factory called at each run admission. Each node runs + in a separate copy of that context, including serial execution. By default + no context is captured or copied. Raises: NodeExecutionFailedError if `raise_on_error` is set to True and execution fails for any snapshot. @@ -189,6 +209,7 @@ def concurrent_apply_to_snapshots( lambda s_id: fn(snapshots_by_id[s_id]), tasks_num, raise_on_error=raise_on_error, + context_factory=context_factory, ) @@ -197,6 +218,7 @@ def concurrent_apply_to_dag( fn: t.Callable[[H], None], tasks_num: int, raise_on_error: bool = True, + context_factory: t.Optional[t.Callable[[], Context]] = None, ) -> t.Tuple[t.List[NodeExecutionFailedError[H]], t.List[H]]: """Applies a function to the given DAG concurrently while preserving the topological order between snapshots. @@ -208,6 +230,9 @@ def concurrent_apply_to_dag( raise_on_error: If set to True raises an exception on a first encountered error, otherwises returns a tuple which contains a list of failed nodes and a list of skipped nodes. + context_factory: Optional factory called at each run admission. Each node runs + in a separate copy of that context, including serial execution. By default + no context is captured or copied. Raises: NodeExecutionFailedError if `raise_on_error` is set to True and execution fails for any snapshot. @@ -219,13 +244,14 @@ def concurrent_apply_to_dag( raise ConfigError(f"Invalid number of concurrent tasks {tasks_num}") if tasks_num == 1: - return sequential_apply_to_dag(dag, fn, raise_on_error) + return sequential_apply_to_dag(dag, fn, raise_on_error, context_factory) return ConcurrentDAGExecutor( dag, fn, tasks_num, raise_on_error, + context_factory, ).run() @@ -233,7 +259,9 @@ def sequential_apply_to_dag( dag: DAG[H], fn: t.Callable[[H], None], raise_on_error: bool = True, + context_factory: t.Optional[t.Callable[[], Context]] = None, ) -> t.Tuple[t.List[NodeExecutionFailedError[H]], t.List[H]]: + context = context_factory() if context_factory else None dependencies = dag.graph node_errors: t.List[NodeExecutionFailedError[H]] = [] @@ -248,7 +276,10 @@ def sequential_apply_to_dag( continue try: - fn(node) + if context is None: + fn(node) + else: + context.copy().run(fn, node) except Exception as ex: error = NodeExecutionFailedError(node) error.__cause__ = ex @@ -266,6 +297,7 @@ 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. @@ -273,12 +305,19 @@ def concurrent_apply_to_values( 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] @@ -292,6 +331,9 @@ def _process_value(value: A, index: int) -> None: 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] diff --git a/tests/utils/test_concurrency.py b/tests/utils/test_concurrency.py index 89f9789725..083897fe56 100644 --- a/tests/utils/test_concurrency.py +++ b/tests/utils/test_concurrency.py @@ -2,6 +2,7 @@ import subprocess import sys import textwrap +from contextvars import ContextVar, copy_context from pathlib import Path import pytest @@ -362,3 +363,52 @@ def run_aggregate(): assert set(called) == {"running", "queued", "running_child", "queued_child"} """ ) + + +@pytest.mark.parametrize("tasks_num", [1, 3]) +def test_concurrent_apply_to_values_copies_admission_context(tasks_num: int): + value = ContextVar("admission", default="worker") + token = value.set("admitted") + + def process(item: int) -> str: + assert value.get() == "admitted" + value.set(str(item)) + return value.get() + + try: + assert concurrent_apply_to_values( + [1, 2, 3], process, tasks_num, context_factory=copy_context + ) == ["1", "2", "3"] + assert value.get() == "admitted" + default_results = concurrent_apply_to_values([1, 2, 3], lambda _: value.get(), tasks_num) + expected = "admitted" if tasks_num == 1 else "worker" + assert default_results == [expected, expected, expected] + finally: + value.reset(token) + + +@pytest.mark.parametrize("tasks_num", [1, 2]) +def test_context_admission_and_sibling_isolation(tasks_num: int): + from sqlmesh.utils.concurrency import ConcurrentDAGExecutor, concurrent_apply_to_dag + from sqlmesh.utils.dag import DAG + + consumer = ContextVar("admission", default="construction") + dag: DAG[str] = DAG() + dag.add("second", {"first"}) + seen = [] + + def execute(node): + seen.append(consumer.get()) + consumer.set(node) + + executor = ConcurrentDAGExecutor(dag, execute, tasks_num, False, copy_context) + token = consumer.set("first admission") + try: + executor.run() + consumer.set("second admission") + executor.run() + concurrent_apply_to_dag(dag, execute, tasks_num, context_factory=copy_context) + assert consumer.get() == "second admission" + finally: + consumer.reset(token) + assert seen == ["first admission"] * 2 + ["second admission"] * 4 From 4ef4829dd39002d622ec38a4cf2c047a30a19bdb Mon Sep 17 00:00:00 2001 From: mday-io Date: Tue, 13 Oct 2026 19:18:34 -0400 Subject: [PATCH 3/3] Feat: Add optional execution observations Signed-off-by: mday-io --- docs/reference/execution_observations.md | 170 ++++++++ mkdocs.yml | 1 + sqlmesh/core/console.py | 14 + sqlmesh/core/engine_adapter/base.py | 17 +- sqlmesh/core/execution_observation.py | 132 +++++++ sqlmesh/core/plan/evaluator.py | 7 +- sqlmesh/core/scheduler.py | 353 +++++++++-------- sqlmesh/core/snapshot/evaluator.py | 44 ++- tests/core/test_execution_observation.py | 482 +++++++++++++++++++++++ tests/core/test_snapshot_evaluator.py | 379 +++++++++++++++++- tests/utils/test_concurrency.py | 77 ++++ 11 files changed, 1487 insertions(+), 189 deletions(-) create mode 100644 docs/reference/execution_observations.md create mode 100644 sqlmesh/core/execution_observation.py create mode 100644 tests/core/test_execution_observation.py diff --git a/docs/reference/execution_observations.md b/docs/reference/execution_observations.md new file mode 100644 index 0000000000..316445d53b --- /dev/null +++ b/docs/reference/execution_observations.md @@ -0,0 +1,170 @@ +# Execution observations + +Execution observations expose native SQLMesh work to Python integrations through +an optional callback. They do not replace SQL generation, execution, retries, +transactions, or `Console` progress reporting. + +## Console observer registration + +`Console.get_execution_observer()` returns `None` by default. A `Console` +subclass can return an `Observer` callback from `sqlmesh.core.execution_observation`. +The callback receives a +phase, a read-only mapping of facts, and the original exception or `None`. +For general Python integration details, see the [Python API reference](./python.md). + +The following console prints the kind, action, and finish error for each observed +operation. The example uses `examples/sushi` relative to the SQLMesh repository root: + +```python +from typing import Any, Literal, Mapping, Optional + +from sqlmesh.core.console import TerminalConsole, set_console +from sqlmesh.core.context import Context +from sqlmesh.core.execution_observation import Observer + + +def observe( + phase: Literal["start", "finish"], + facts: Mapping[str, Any], + error: Optional[BaseException], +) -> None: + if phase == "finish": + print(f"{facts['kind']}/{facts['action']}: error={error!r}") + + +class ObservedConsole(TerminalConsole): + def get_execution_observer(self) -> Optional[Observer]: + return observe + + +set_console(ObservedConsole()) +context = Context(paths=["examples/sushi"]) +``` + +A `Context` uses the `Console` that was current when the context was created. +Unless an explicit `observer_scope()` is active, SQLMesh resolves +`get_execution_observer()` once for each outer plan evaluation or scheduler run. +Nested calls reuse that observer. If the getter raises an ordinary `Exception`, +SQLMesh disables observation for that plan evaluation or scheduler run and +continues native work. + +## Observer scopes + +`observer_scope(observer)` from `sqlmesh.core.execution_observation` sets the +observer for the duration of the scope. A nested scope restores the previously +active observer when it exits. `observer_scope(None)` disables observation inside +that scope, including automatic Console registration. + +## Callback lifecycle + +SQLMesh calls `start` before observed work with `error=None`. When `start` +returns or raises an ordinary `Exception`, SQLMesh continues the work and calls +`finish` from a `finally` block afterward. The `finish` callback receives the +original workload exception object or `None`. Both callbacks run synchronously +on the thread that performs the observed work. Callbacks for work dispatched to +a worker run on that worker's thread. + +Every invocation that reaches `finish` passes the same read-only mapping object +to `start` and `finish`. Distinct invocations use distinct mapping objects. +Consumers can pair the phases by mapping identity. SQLMesh may add facts while +the work runs, so values in the mapping at `finish` can differ from its values +at `start`. The mapping's supported lifetime ends when the matching `finish` +callback returns. + +Ordinary observer `Exception` failures are suppressed and do not replace the +workload result or error. SQLMesh does not suppress observer-raised +`BaseException` subclasses that are not `Exception` subclasses, such as +`KeyboardInterrupt` and `SystemExit`. If `start` raises a `BaseException` that +is not an `Exception`, the work and `finish` callback do not run. If `finish` +raises a `BaseException` that is not an `Exception`, that failure can replace +the workload error. + +If a worker raises a `BaseException`, the error can reach the caller after its +worker pool exits. If a DAG run records a terminal error, SQLMesh stops +scheduling nodes and waits for running workers to exit before it raises that +error. Parallel value operations also wait for their workers to finish before +raising a worker's `BaseException`. Blocking callbacks and process termination +are outside the observer failure guarantee. SQLMesh does not log observer +failures. + +Nested workload calls produce nested observations. Callbacks for concurrent +workers can interleave, and SQLMesh provides no global callback order. Observers +must support concurrent calls and return promptly. SQLMesh allocates no +execution IDs, timestamps, durations, or parent IDs. Consumers supply those +values when needed. + +While SQLMesh delivers a callback, it suppresses observations triggered by that +callback and ignores calls to `update_execution()`. The guard covers callback +delivery only. Normal model work can still contain observations for physical +operations, audits, and queries. + +## Native facts and limits + +Every facts mapping contains `kind` and `action`. Other fields depend on the +observed operation and can be absent when work exits early or fails. + +| Kind | Action | Native facts | +| --- | --- | --- | +| `stage` | Plan-stage class name | `native_plan_id` | +| `model` | `evaluate`, `audit_only` | `snapshot`, `interval`, `batch_index`, `execution_time`, `audit_only` | +| `physical` | `create`, `materialize`, `schema_migration` | `snapshot`; `target` when known; `creating` for materialization; `existed` for schema migration when known | +| `virtual` | `promote`, `demote` | `snapshot`; `target` when known; `source_table` for promotion when known | +| `audit` | `audit` | `audit_name`; `audit_result` at finish if produced | +| `query` | `execute` | `engine` dialect string | + +[`Snapshot`](../concepts/architecture/snapshots.md) and `AuditResult` values are +borrowed native objects. Consumers must +treat these objects and other nested values as read-only and must not retain them +beyond the matching `finish` callback. The mapping's read-only wrapper does not +freeze nested objects. Facts are not a JSON event schema. SQLMesh does not +interpret audit results as consumer statuses or inherit model facts into child +mappings. Consumers that need facts after `finish` copy and normalize those +values during the callback. + +SQLMesh calls `finish` for an audit before it runs later audits or completes +write-audit-publish (WAP). A later audit or WAP publication can still fail. An audit +query failure can produce a `finish` callback without an `AuditResult`. The +[audits guide](../concepts/audits.md) describes audit behavior in plans and +runs. + +For `physical/materialize`, `target` is the logical destination table from the +snapshot. During WAP, SQLMesh writes to a staging destination, but the `target` +fact remains the logical destination. The `finish` callback runs after the +materialization's transaction and session exit. It does not report the result +of later audits or WAP publication. + +An observed operation can do no work, issue multiple statements, or trigger +internal driver retries. The observation does not predict how many statements +the operation will issue or report each retry. + +The query observation surrounds the base adapter's existing SQL logging and +`_execute` call. It ends before a surrounding transaction commits and before +later fetch or result consumption. Some adapter operations do not call this +query hook. BigQuery session queries, BigQuery dataframe loads, and ClickHouse +dataframe inserts may not produce query observations. These hooks do not provide +a complete warehouse activity history or guarantee transaction success. + +SQL logging retains its existing behavior, including VALUES redaction for +non-query expressions. The callback receives no SQL text, adapter objects, or +warehouse query IDs. The `start` callback runs before SQL logging. Consumers can +use that ordering to associate existing log entries with the invocation. + +## Worker context + +When an observer is active outside callback delivery, SQLMesh copies the current +Python context for each task in scheduler and snapshot DAG work, snapshot schema creation, and +data-object listing. That context carries Python's [context variables](https://docs.python.org/3/library/contextvars.html). +Each task gets its own context copy, including during serial execution. DAG +successors use the context captured for the run. They do not inherit changes +made by earlier workers. Reusing a DAG executor captures a fresh context for +each run. + +Context copies are shallow. Mutable objects stored in `ContextVar` values remain +shared across copies. A parent identity stored in a consumer-owned `ContextVar` +is available to each worker. SQLMesh uses these copies in the scheduler and +snapshot DAG, schema creation, and data-object listing. Arbitrary consumer +threads and other concurrency helpers do not receive these copies. + +With no observer registered, these operations do not copy contexts. The concurrency +helpers retain their existing `ContextVar` behavior for serial and worker +execution. diff --git a/mkdocs.yml b/mkdocs.yml index 49c4b9163b..1a1a34e624 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -106,6 +106,7 @@ nav: - reference/cli.md - reference/notebook.md - reference/python.md + - reference/execution_observations.md - Configuration: - reference/configuration.md - reference/model_configuration.md diff --git a/sqlmesh/core/console.py b/sqlmesh/core/console.py index ca11b6e607..8ae3be9083 100644 --- a/sqlmesh/core/console.py +++ b/sqlmesh/core/console.py @@ -407,6 +407,20 @@ class Console( INDIRECTLY_MODIFIED_DISPLAY_THRESHOLD = 10 + def get_execution_observer( + self, + ) -> t.Optional[ + t.Callable[ + [t.Literal["start", "finish"], t.Mapping[str, t.Any], t.Optional[BaseException]], None + ] + ]: + """Optional start/finish hook receiving native facts and original errors. + + May run concurrently. No IDs, clocks or product status are supplied. + Implementations must return promptly; ordinary hook errors are isolated. + """ + return None + @abc.abstractmethod def start_plan_evaluation(self, plan: EvaluatablePlan) -> None: """Indicates that a new evaluation has begun.""" diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py index 930cdf7cd4..825159d2f9 100644 --- a/sqlmesh/core/engine_adapter/base.py +++ b/sqlmesh/core/engine_adapter/base.py @@ -2636,12 +2636,17 @@ def execute( sql = self._attach_correlation_id(sql) - self._log_sql( - sql, - expression=e if isinstance(e, exp.Expr) else None, - quote_identifiers=quote_identifiers, - ) - self._execute(sql, track_rows_processed, **kwargs) + # The rendered statement and correlation comment are finalized + # here; observe the exact invocation, including driver errors. + from sqlmesh.core.execution_observation import action + + with action("query", "execute", engine=self.dialect): + self._log_sql( + sql, + expression=e if isinstance(e, exp.Expr) else None, + quote_identifiers=quote_identifiers, + ) + self._execute(sql, track_rows_processed, **kwargs) def _attach_correlation_id(self, sql: str) -> str: if self.ATTACH_CORRELATION_ID and self.correlation_id: diff --git a/sqlmesh/core/execution_observation.py b/sqlmesh/core/execution_observation.py new file mode 100644 index 0000000000..36f426c90e --- /dev/null +++ b/sqlmesh/core/execution_observation.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Optional native lifecycle hooks, independent of telemetry storage or policy. + +Call sites report start before native work and finish in finally, including the +original error. Consumers must return promptly. Ordinary observer exceptions never +replace workload results. The scope only routes hooks and late native metadata; +execution IDs, timing, parent correlation and status mapping belong to consumers. +""" + +from __future__ import annotations + +import contextlib +import contextvars +import functools +import typing as t +from types import MappingProxyType + +Observer = t.Callable[ + [t.Literal["start", "finish"], t.Mapping[str, t.Any], t.Optional[BaseException]], None +] +_dispatching: contextvars.ContextVar[bool] = contextvars.ContextVar( + "execution_observer_dispatch", default=False +) +_scoped: contextvars.ContextVar[bool] = contextvars.ContextVar( + "execution_observer_scoped", default=False +) +_observer: contextvars.ContextVar[t.Optional[Observer]] = contextvars.ContextVar( + "execution_observer", default=None +) +_metadata: contextvars.ContextVar[t.Optional[t.Dict[str, t.Any]]] = contextvars.ContextVar( + "execution_observation_metadata", default=None +) + + +def update_execution(**metadata: t.Any) -> None: + """Supply native facts discovered inside the current hook boundary.""" + current = _metadata.get() + if current is not None and not _dispatching.get(): + current.update(metadata) + + +@contextlib.contextmanager +def observer_scope(observer: t.Optional[Observer]) -> t.Iterator[None]: + token = _observer.set(observer) + scope_token = _scoped.set(True) + metadata_token = _metadata.set(None) + try: + yield + finally: + _metadata.reset(metadata_token) + _scoped.reset(scope_token) + _observer.reset(token) + + +def _notify( + observer: Observer, + phase: t.Literal["start", "finish"], + metadata: t.Mapping[str, t.Any], + error: t.Optional[BaseException] = None, +) -> None: + token = _dispatching.set(True) + try: + observer(phase, metadata, error) + except Exception: + # Hooks are observational: do not alter execution or recursively log a + # telemetry failure. Consumers own their capture-loss reporting. + pass + finally: + _dispatching.reset(token) + + +@contextlib.contextmanager +def action(kind: str, name: str, **metadata: t.Any) -> t.Iterator[t.Optional[t.Dict[str, t.Any]]]: + observer = _observer.get() + if observer is None or _dispatching.get(): + yield None + return + data = {**metadata, "kind": kind, "action": name} + view = MappingProxyType(data) + token = _metadata.set(data) + try: + _notify(observer, "start", view) + error: t.Optional[BaseException] = None + try: + yield data + except BaseException as ex: + error = ex + raise + finally: + _notify(observer, "finish", view, error) + finally: + _metadata.reset(token) + + +def observe_snapshot(kind: str, name: str) -> t.Callable: + """Hook evaluator methods whose first argument after self is snapshot. + + Verified callers pass Snapshot as that argument (positional or keyword). + Pass the native object unchanged; no reflection, identity formatting, skip + classification or snapshot interpretation is performed by this hook. + """ + + def decorate(function: t.Callable) -> t.Callable: + @functools.wraps(function) + def wrapped(self: t.Any, snapshot: t.Any, *args: t.Any, **kwargs: t.Any) -> t.Any: + with action(kind, name, snapshot=snapshot): + return function(self, snapshot, *args, **kwargs) + + return wrapped + + return decorate + + +def execution_context_factory() -> t.Optional[t.Callable[[], contextvars.Context]]: + """Opt in to worker context propagation only during active observation.""" + return ( + contextvars.copy_context if _observer.get() is not None and not _dispatching.get() else None + ) + + +@contextlib.contextmanager +def console_observer_scope(console: t.Any) -> t.Iterator[None]: + """Resolve once at the outer boundary; an explicit None scope disables hooks.""" + if _scoped.get(): + yield + return + try: + observer = console.get_execution_observer() + except Exception: + observer = None + with observer_scope(observer): + yield diff --git a/sqlmesh/core/plan/evaluator.py b/sqlmesh/core/plan/evaluator.py index f2f432a97e..71e5285d63 100644 --- a/sqlmesh/core/plan/evaluator.py +++ b/sqlmesh/core/plan/evaluator.py @@ -18,6 +18,7 @@ import logging import typing as t from sqlmesh.core import analytics +from sqlmesh.core.execution_observation import action, console_observer_scope from sqlmesh.core import constants as c from sqlmesh.core.console import Console, get_console from sqlmesh.core.environment import EnvironmentNamingInfo, execute_environment_statements @@ -102,7 +103,8 @@ def evaluate( try: plan_stages = stages.build_plan_stages(plan, self.state_sync, self.default_catalog) - self._evaluate_stages(plan_stages, plan) + with console_observer_scope(self.console): + self._evaluate_stages(plan_stages, plan) except Exception as e: analytics.collector.on_plan_apply_end(plan_id=plan.plan_id, error=e) raise @@ -122,7 +124,8 @@ def _evaluate_stages( raise SQLMeshError(f"Unexpected plan stage: {stage_name}") logger.info("Evaluating plan stage %s", stage_name) handler = getattr(self, handler_name) - handler(stage, plan) + with action("stage", stage_name, native_plan_id=plan.plan_id): + handler(stage, plan) def visit_before_all_stage(self, stage: stages.BeforeAllStage, plan: EvaluatablePlan) -> None: execute_environment_statements( diff --git a/sqlmesh/core/scheduler.py b/sqlmesh/core/scheduler.py index 5eb0ff40ff..7345d92888 100644 --- a/sqlmesh/core/scheduler.py +++ b/sqlmesh/core/scheduler.py @@ -8,6 +8,11 @@ from sqlglot import exp from sqlmesh.core import constants as c from sqlmesh.core.console import Console, get_console +from sqlmesh.core.execution_observation import ( + action, + console_observer_scope, + execution_context_factory, +) from sqlmesh.core.environment import EnvironmentNamingInfo, execute_environment_statements from sqlmesh.core.macros import RuntimeStage from sqlmesh.core.model.definition import AuditResult @@ -454,178 +459,43 @@ def run_merged_intervals( Returns: A tuple of errors and skipped intervals. """ - execution_time = execution_time or now_timestamp() + with console_observer_scope(self.console): + execution_time = execution_time or now_timestamp() - selected_snapshots = [self.snapshots[sid] for sid in (selected_snapshot_ids or set())] - if not selected_snapshots: - selected_snapshots = list(merged_intervals) + selected_snapshots = [self.snapshots[sid] for sid in (selected_snapshot_ids or set())] + if not selected_snapshots: + selected_snapshots = list(merged_intervals) - # Build the full DAG from all snapshots to preserve transitive dependencies - full_dag = snapshots_to_dag(self.snapshots.values()) + # Build the full DAG from all snapshots to preserve transitive dependencies + full_dag = snapshots_to_dag(self.snapshots.values()) - # Create a subdag that includes the selected snapshots and all their upstream dependencies - # This ensures that transitive dependencies are preserved even when intermediate nodes are not selected - selected_snapshot_ids_set = {s.snapshot_id for s in selected_snapshots} - snapshot_dag = full_dag.subdag(*selected_snapshot_ids_set) - - batched_intervals = self.batch_intervals( - merged_intervals, - deployability_index, - environment_naming_info, - dag=snapshot_dag, - is_restatement=is_restatement, - ) - self.console.start_evaluation_progress( - batched_intervals, - environment_naming_info, - self.default_catalog, - audit_only=audit_only, - ) + # Create a subdag that includes the selected snapshots and all their upstream dependencies + # This ensures that transitive dependencies are preserved even when intermediate nodes are not selected + selected_snapshot_ids_set = {s.snapshot_id for s in selected_snapshots} + snapshot_dag = full_dag.subdag(*selected_snapshot_ids_set) - if run_environment_statements: - environment_statements = self.state_sync.get_environment_statements( - environment_naming_info.name - ) - execute_environment_statements( - adapter=self.snapshot_evaluator.adapter, - environment_statements=environment_statements, - runtime_stage=RuntimeStage.BEFORE_ALL, - environment_naming_info=environment_naming_info, - default_catalog=self.default_catalog, - snapshots=self.snapshots_by_name, - start=start, - end=end, - execution_time=execution_time, - selected_models=selected_models, + batched_intervals = self.batch_intervals( + merged_intervals, + deployability_index, + environment_naming_info, + dag=snapshot_dag, + is_restatement=is_restatement, ) - - # We only need to create physical tables if the snapshot is not representative or if it - # needs backfill - snapshots_to_create_candidates = [ - s - for s in selected_snapshots - if not deployability_index.is_representative(s) or s in batched_intervals - ] - snapshots_to_create = { - s.snapshot_id - for s in self.snapshot_evaluator.get_snapshots_to_create( - snapshots_to_create_candidates, deployability_index + self.console.start_evaluation_progress( + batched_intervals, + environment_naming_info, + self.default_catalog, + audit_only=audit_only, ) - } - - dag = self._dag( - batched_intervals, snapshot_dag=snapshot_dag, snapshots_to_create=snapshots_to_create - ) - def run_node(node: SchedulingUnit) -> None: - if circuit_breaker and circuit_breaker(): - raise CircuitBreakerError() - if isinstance(node, DummyNode): - return - - snapshot = self.snapshots_by_name[node.snapshot_name] - - if isinstance(node, EvaluateNode): - self.console.start_snapshot_evaluation_progress(snapshot) - execution_start_ts = now_timestamp() - evaluation_duration_ms: t.Optional[int] = None - start, end = node.interval - - audit_results: t.List[AuditResult] = [] - try: - assert execution_time # mypy - assert deployability_index # mypy - - if audit_only: - audit_results = self._audit_snapshot( - snapshot=snapshot, - environment_naming_info=environment_naming_info, - deployability_index=deployability_index, - snapshots=self.snapshots_by_name, - start=start, - end=end, - execution_time=execution_time, - ) - else: - # If batch_index > 0, then the target table must exist since the first batch would have created it - target_table_exists = ( - snapshot.snapshot_id not in snapshots_to_create or node.batch_index > 0 - ) - audit_results = self.evaluate( - snapshot=snapshot, - environment_naming_info=environment_naming_info, - start=start, - end=end, - execution_time=execution_time, - deployability_index=deployability_index, - batch_index=node.batch_index, - allow_destructive_snapshots=allow_destructive_snapshots, - allow_additive_snapshots=allow_additive_snapshots, - target_table_exists=target_table_exists, - selected_models=selected_models, - ) - - evaluation_duration_ms = now_timestamp() - execution_start_ts - finally: - num_audits = len(audit_results) - num_audits_failed = sum(1 for result in audit_results if result.count) - - execution_stats = self.snapshot_evaluator.execution_tracker.get_execution_stats( - SnapshotIdBatch(snapshot_id=snapshot.snapshot_id, batch_id=node.batch_index) - ) - - self.console.update_snapshot_evaluation_progress( - snapshot, - batched_intervals[snapshot][node.batch_index], - node.batch_index, - evaluation_duration_ms, - num_audits - num_audits_failed, - num_audits_failed, - execution_stats=execution_stats, - auto_restatement_triggers=auto_restatement_triggers.get( - snapshot.snapshot_id - ), - ) - elif isinstance(node, CreateNode): - self.snapshot_evaluator.create_snapshot( - snapshot=snapshot, - snapshots=self.snapshots_by_name, - deployability_index=deployability_index, - allow_destructive_snapshots=allow_destructive_snapshots or set(), - allow_additive_snapshots=allow_additive_snapshots or set(), - ) - - try: - with self.snapshot_evaluator.concurrent_context(): - errors, skipped_intervals = concurrent_apply_to_dag( - dag, - run_node, - self.max_workers, - raise_on_error=False, - ) - self.console.stop_evaluation_progress(success=not errors) - - skipped_snapshots = { - i.snapshot_name for i in skipped_intervals if isinstance(i, EvaluateNode) - } - self.console.log_skipped_models(skipped_snapshots) - for skipped in skipped_snapshots: - logger.info(f"SKIPPED snapshot {skipped}\n") - - for error in errors: - if isinstance(error.__cause__, CircuitBreakerError): - raise error.__cause__ - logger.info(str(error), exc_info=error) - - self.console.log_failed_models(errors) - - return errors, skipped_intervals - finally: if run_environment_statements: + environment_statements = self.state_sync.get_environment_statements( + environment_naming_info.name + ) execute_environment_statements( adapter=self.snapshot_evaluator.adapter, environment_statements=environment_statements, - runtime_stage=RuntimeStage.AFTER_ALL, + runtime_stage=RuntimeStage.BEFORE_ALL, environment_naming_info=environment_naming_info, default_catalog=self.default_catalog, snapshots=self.snapshots_by_name, @@ -635,7 +505,166 @@ def run_node(node: SchedulingUnit) -> None: selected_models=selected_models, ) - self.state_sync.recycle() + # We only need to create physical tables if the snapshot is not representative or if it + # needs backfill + snapshots_to_create_candidates = [ + s + for s in selected_snapshots + if not deployability_index.is_representative(s) or s in batched_intervals + ] + snapshots_to_create = { + s.snapshot_id + for s in self.snapshot_evaluator.get_snapshots_to_create( + snapshots_to_create_candidates, deployability_index + ) + } + + dag = self._dag( + batched_intervals, + snapshot_dag=snapshot_dag, + snapshots_to_create=snapshots_to_create, + ) + + def run_node(node: SchedulingUnit) -> None: + if isinstance(node, EvaluateNode): + with action( + "model", + "audit_only" if audit_only else "evaluate", + snapshot=self.snapshots_by_name[node.snapshot_name], + interval=node.interval, + batch_index=node.batch_index, + execution_time=execution_time, + audit_only=audit_only, + ): + execute_node(node) + else: + execute_node(node) + + def execute_node(node: SchedulingUnit) -> None: + if circuit_breaker and circuit_breaker(): + raise CircuitBreakerError() + if isinstance(node, DummyNode): + return + + snapshot = self.snapshots_by_name[node.snapshot_name] + + if isinstance(node, EvaluateNode): + self.console.start_snapshot_evaluation_progress(snapshot) + execution_start_ts = now_timestamp() + evaluation_duration_ms: t.Optional[int] = None + start, end = node.interval + + audit_results: t.List[AuditResult] = [] + try: + assert execution_time # mypy + assert deployability_index # mypy + + if audit_only: + audit_results = self._audit_snapshot( + snapshot=snapshot, + environment_naming_info=environment_naming_info, + deployability_index=deployability_index, + snapshots=self.snapshots_by_name, + start=start, + end=end, + execution_time=execution_time, + ) + else: + # If batch_index > 0, then the target table must exist since the first batch would have created it + target_table_exists = ( + snapshot.snapshot_id not in snapshots_to_create + or node.batch_index > 0 + ) + audit_results = self.evaluate( + snapshot=snapshot, + environment_naming_info=environment_naming_info, + start=start, + end=end, + execution_time=execution_time, + deployability_index=deployability_index, + batch_index=node.batch_index, + allow_destructive_snapshots=allow_destructive_snapshots, + allow_additive_snapshots=allow_additive_snapshots, + target_table_exists=target_table_exists, + selected_models=selected_models, + ) + + evaluation_duration_ms = now_timestamp() - execution_start_ts + finally: + num_audits = len(audit_results) + num_audits_failed = sum(1 for result in audit_results if result.count) + + execution_stats = ( + self.snapshot_evaluator.execution_tracker.get_execution_stats( + SnapshotIdBatch( + snapshot_id=snapshot.snapshot_id, batch_id=node.batch_index + ) + ) + ) + + self.console.update_snapshot_evaluation_progress( + snapshot, + batched_intervals[snapshot][node.batch_index], + node.batch_index, + evaluation_duration_ms, + num_audits - num_audits_failed, + num_audits_failed, + execution_stats=execution_stats, + auto_restatement_triggers=auto_restatement_triggers.get( + snapshot.snapshot_id + ), + ) + elif isinstance(node, CreateNode): + self.snapshot_evaluator.create_snapshot( + snapshot=snapshot, + snapshots=self.snapshots_by_name, + deployability_index=deployability_index, + allow_destructive_snapshots=allow_destructive_snapshots or set(), + allow_additive_snapshots=allow_additive_snapshots or set(), + ) + + try: + with self.snapshot_evaluator.concurrent_context(): + errors, skipped_intervals = concurrent_apply_to_dag( + dag, + run_node, + self.max_workers, + raise_on_error=False, + context_factory=execution_context_factory(), + ) + self.console.stop_evaluation_progress(success=not errors) + + skipped_snapshots = { + i.snapshot_name for i in skipped_intervals if isinstance(i, EvaluateNode) + } + self.console.log_skipped_models(skipped_snapshots) + for skipped in skipped_snapshots: + logger.info(f"SKIPPED snapshot {skipped}\n") + + for error in errors: + if isinstance(error.__cause__, CircuitBreakerError): + raise error.__cause__ + logger.info(str(error), exc_info=error) + + self.console.log_failed_models(errors) + + return errors, skipped_intervals + finally: + if run_environment_statements: + execute_environment_statements( + adapter=self.snapshot_evaluator.adapter, + environment_statements=environment_statements, + runtime_stage=RuntimeStage.AFTER_ALL, + environment_naming_info=environment_naming_info, + default_catalog=self.default_catalog, + snapshots=self.snapshots_by_name, + start=start, + end=end, + execution_time=execution_time, + selected_models=selected_models, + ) + + self.state_sync.recycle() def _dag( self, diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py index 9d9ecf7672..345f8816e6 100644 --- a/sqlmesh/core/snapshot/evaluator.py +++ b/sqlmesh/core/snapshot/evaluator.py @@ -29,6 +29,12 @@ from collections import defaultdict from contextlib import contextmanager from functools import reduce +from sqlmesh.core.execution_observation import ( + action, + observe_snapshot, + update_execution, + execution_context_factory, +) from sqlglot import exp, select from sqlglot.executor import execute @@ -329,6 +335,7 @@ def promote( on_complete=on_complete, ), self.ddl_concurrent_tasks, + context_factory=execution_context_factory(), ) def demote( @@ -357,6 +364,7 @@ def demote( table_mapping=table_mapping, ), self.ddl_concurrent_tasks, + context_factory=execution_context_factory(), ) def create( @@ -467,6 +475,7 @@ def _create_snapshots( on_complete=on_complete, ), self.ddl_concurrent_tasks, + context_factory=execution_context_factory(), raise_on_error=False, ) if errors: @@ -514,6 +523,7 @@ def migrate( deployability_index, ), self.ddl_concurrent_tasks, + context_factory=execution_context_factory(), ) def cleanup( @@ -558,6 +568,7 @@ def cleanup( on_complete, ), self.ddl_concurrent_tasks, + context_factory=execution_context_factory(), reverse_order=True, raise_on_error=False, ) @@ -636,8 +647,11 @@ def audit( # so that we can fall back to the audit's setting, which we override to blocking: False audit = audit.model_copy(update={"blocking": False}) - results.append( - self._audit( + # Observe before calling _audit: query exceptions need a terminal + # event even when no AuditResult exists. Preserve completed audits + # before later aggregate handling or WAP publication can fail. + with action("audit", "audit", audit_name=audit.name) as observation: + audit_result = self._audit( audit=audit, audit_args=audit_args, snapshot=snapshot, @@ -648,7 +662,9 @@ def audit( deployability_index=deployability_index, **kwargs, ) - ) + if observation is not None: + observation["audit_result"] = audit_result + results.append(audit_result) if wap_id is not None: logger.info( @@ -778,6 +794,16 @@ def _evaluate_snapshot( ) with ( + # The first batch may create via CTAS instead of create_snapshot. + # Observe the materialization, including its transaction exit, so + # lazy creation is not missing from physical-layer history. + action( + "physical", + "materialize", + target=target_table_name, + creating=not target_table_exists, + snapshot=snapshot, + ), adapter.transaction(), adapter.session(snapshot.model.render_session_properties(**render_statements_kwargs)), ): @@ -866,6 +892,7 @@ def _evaluate_snapshot( return wap_id + @observe_snapshot("physical", "create") def create_snapshot( self, snapshot: Snapshot, @@ -889,6 +916,9 @@ def create_snapshot( return logger.info("Creating a physical table for snapshot %s", snapshot.snapshot_id) + update_execution( + target=snapshot.table_name(is_deployable=deployability_index.is_deployable(snapshot)), + ) adapter = self.get_adapter(snapshot.model.gateway) create_render_kwargs: t.Dict[str, t.Any] = dict( @@ -1124,6 +1154,7 @@ def _clone_snapshot_in_dev( adapter.drop_table(target_table_name) raise + @observe_snapshot("physical", "schema_migration") def _migrate_snapshot( self, snapshot: Snapshot, @@ -1145,6 +1176,7 @@ def _migrate_snapshot( deployability_index=deployability_index, ) target_table_name = snapshot.table_name() + update_execution(target=target_table_name, existed=target_data_object is not None) evaluation_strategy = _evaluation_strategy(snapshot, adapter) evaluation_strategy.run_pre_statements( @@ -1253,6 +1285,7 @@ def _migrate_target_table( if snapshot.is_materialized: adapter.drop_table(tmp_table_name) + @observe_snapshot("virtual", "promote") def _promote_snapshot( self, snapshot: Snapshot, @@ -1277,6 +1310,7 @@ def _promote_snapshot( view_name = snapshot.qualified_view_name.for_environment( environment_naming_info, dialect=adapter.dialect ) + update_execution(target=view_name, source_table=table_name) render_kwargs: t.Dict[str, t.Any] = dict( start=start, end=end, @@ -1308,6 +1342,7 @@ def _promote_snapshot( if on_complete is not None: on_complete(snapshot) + @observe_snapshot("virtual", "demote") def _demote_snapshot( self, snapshot: Snapshot, @@ -1327,6 +1362,7 @@ def _demote_snapshot( view_name = snapshot.qualified_view_name.for_environment( environment_naming_info, dialect=adapter.dialect ) + update_execution(target=view_name) with ( adapter.transaction(), adapter.session( @@ -1477,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: @@ -1651,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): diff --git a/tests/core/test_execution_observation.py b/tests/core/test_execution_observation.py new file mode 100644 index 0000000000..626eae3d98 --- /dev/null +++ b/tests/core/test_execution_observation.py @@ -0,0 +1,482 @@ +# SPDX-License-Identifier: Apache-2.0 +import logging +from dataclasses import dataclass, field +from contextvars import ContextVar + +import pytest + +from sqlmesh.core.execution_observation import ( + action, + observer_scope, + update_execution, + execution_context_factory, +) + + +@dataclass +class Observer: + events: list = field(default_factory=list) + + def __call__(self, phase, facts, error): + self.events.append((phase, dict(facts), error)) + + +def test_native_facts_have_no_telemetry_policy(): + observer = Observer() + error = ValueError("view failed") + with observer_scope(observer), pytest.raises(ValueError) as raised: + with action("virtual", "promote", target="view"): + update_execution(source_table="table") + raise error + assert raised.value is error + assert observer.events == [ + ("start", {"kind": "virtual", "action": "promote", "target": "view"}, None), + ( + "finish", + {"kind": "virtual", "action": "promote", "target": "view", "source_table": "table"}, + error, + ), + ] + + +@pytest.mark.parametrize("phase", ["start", "finish"]) +@pytest.mark.parametrize("failure", [None, ValueError("native failure"), KeyboardInterrupt()]) +def test_observer_errors_cannot_change_workload_outcome(phase, failure): + calls = [] + + def observer(event, facts, error): + calls.append(event) + if event == phase: + raise RuntimeError("observer failure") + + def workload(): + with observer_scope(observer), action("query", "execute"): + if failure is not None: + raise failure + return 42 + + if failure is None: + assert workload() == 42 + else: + with pytest.raises(type(failure)) as raised: + workload() + assert raised.value is failure + assert calls == ["start", "finish"] + + +def test_disabled_observer_does_not_create_metadata(): + from sqlmesh.core.console import TerminalConsole + + assert TerminalConsole().get_execution_observer() is None + with action("query", "execute") as metadata: + assert metadata is None + + +def test_real_adapter_preserves_sql_and_errors(duck_conn, caplog): + from sqlmesh.core.engine_adapter import DuckDBEngineAdapter + + adapter = DuckDBEngineAdapter(lambda: duck_conn) + observer = Observer() + sql = "SELECT\n 1 AS observed_value" + with caplog.at_level(logging.DEBUG), observer_scope(observer): + adapter.execute(sql) + with pytest.raises(Exception): + adapter.execute("SELECT * FROM missing_observed_table") + queries = [event for event in observer.events if event[0] == "finish"] + assert len(queries) == 2 + assert queries[0][2] is None + assert queries[1][2] is not None + assert all( + event[1] == {"kind": "query", "action": "execute", "engine": "duckdb"} for event in queries + ) + assert f"Executing SQL: {sql}" in [record.getMessage() for record in caplog.records] + assert "Executing SQL: SELECT * FROM missing_observed_table" in [ + record.getMessage() for record in caplog.records + ] + + +def test_native_dag_copies_consumer_context_not_predecessor_context(): + from sqlmesh.utils.concurrency import ConcurrentDAGExecutor + from sqlmesh.utils.dag import DAG + + consumer_context = ContextVar("consumer_context", default=None) + observer = Observer() + dag = DAG() + dag.add("second", {"first"}) + + def execute(name): + assert consumer_context.get() == "stage" + consumer_context.set(name) + with action("physical", "create", target=name): + pass + + with observer_scope(observer): + token = consumer_context.set("stage") + try: + errors, skipped = ConcurrentDAGExecutor( + dag, execute, 2, False, execution_context_factory() + ).run() + finally: + consumer_context.reset(token) + assert not errors and not skipped + assert len(observer.events) == 4 + + +def test_concurrent_model_hooks_keep_batch_context_and_original_failures(): + """Real DAG workers retain independent batch scopes even when one fails.""" + from threading import Barrier, Lock, get_ident + + from sqlmesh.utils.concurrency import ConcurrentDAGExecutor + from sqlmesh.utils.dag import DAG + + consumer_context = ContextVar("batch_context", default="plan") + barrier = Barrier(2) + lock = Lock() + events = [] + failure = RuntimeError("model evaluation failed") + dag = DAG() + dag.add(0) + dag.add(1) + + def observer(phase, facts, error): + with lock: + events.append((phase, dict(facts), error, get_ident(), consumer_context.get())) + + def execute(batch): + assert consumer_context.get() == "plan" + token = consumer_context.set(batch) + try: + with action("model", "evaluate", batch_index=batch, interval=(batch, batch + 1)): + barrier.wait(timeout=5) + if batch == 1: + raise failure + finally: + consumer_context.reset(token) + + with observer_scope(observer): + errors, skipped = ConcurrentDAGExecutor( + dag, execute, 2, False, execution_context_factory() + ).run() + assert len(errors) == 1 and errors[0].__cause__ is failure + assert not skipped + assert len({event[3] for event in events}) == 2 + for batch in (0, 1): + batch_events = [event for event in events if event[1]["batch_index"] == batch] + assert [event[0] for event in batch_events] == ["start", "finish"] + assert all(event[4] == batch for event in batch_events) + assert batch_events[-1][2] is (failure if batch == 1 else None) + assert consumer_context.get() == "plan" + + +def test_read_only_identity_and_callback_reentry(): + events = [] + native = object() + + def observer(phase, facts, error): + events.append((phase, facts, error)) + with pytest.raises(TypeError): + facts["kind"] = "corrupt" + update_execution(corrupt=True) + with action("query", "observer-query"): + update_execution(corrupt=True) + + with observer_scope(observer), action("model", "evaluate", snapshot=native): + with action("query", "execute"): + update_execution(late=True) + update_execution(target="late table") + assert [event[0] for event in events] == ["start", "start", "finish", "finish"] + assert events[0][1] is events[-1][1] + assert events[1][1] is events[2][1] + assert events[-1][1]["snapshot"] is native + assert events[-1][1]["target"] == "late table" + assert "late" not in events[-1][1] + assert all("corrupt" not in event[1] for event in events) + + +def test_explicit_disabled_scope_isolates_parent_metadata(): + observer = Observer() + with observer_scope(observer), action("model", "evaluate"): + with observer_scope(None), action("query", "execute"): + update_execution(corrupt=True) + update_execution(target="parent") + assert len(observer.events) == 2 + assert observer.events[-1][1] == {"kind": "model", "action": "evaluate", "target": "parent"} + + +def test_console_scope_getter_failure_and_nested_reuse(): + from sqlmesh.core.execution_observation import console_observer_scope + from unittest.mock import Mock + + console = Mock() + console.get_execution_observer.side_effect = ValueError("getter failed") + with console_observer_scope(console), action("query", "execute") as facts: + assert facts is None + observer = Observer() + console.reset_mock(side_effect=True) + console.get_execution_observer.return_value = observer + with console_observer_scope(console): + with console_observer_scope(console), action("model", "evaluate"): + pass + console.get_execution_observer.assert_called_once_with() + assert len(observer.events) == 2 + console.reset_mock() + with observer_scope(None), console_observer_scope(console): + with action("query", "execute") as facts: + assert facts is None + console.get_execution_observer.assert_not_called() + + +def test_unobserved_context_factory_is_disabled(): + assert execution_context_factory() is None + with observer_scope(None): + assert execution_context_factory() is None + + +@pytest.mark.parametrize("phase", ["start", "finish"]) +def test_callback_base_exception_propagates_and_scope_is_restored(phase): + interrupt = KeyboardInterrupt() + events = [] + + def observer(event, facts, error): + events.append(event) + if event == phase: + raise interrupt + + with pytest.raises(KeyboardInterrupt) as raised: + with observer_scope(observer), action("query", "execute"): + pass + assert raised.value is interrupt + assert events == (["start"] if phase == "start" else ["start", "finish"]) + with action("query", "execute") as facts: + assert facts is None + observed = Observer() + with observer_scope(observed), action("model", "evaluate"): + update_execution(target="clean") + assert len(observed.events) == 2 + assert observed.events[-1][1]["target"] == "clean" + + +def test_disabled_and_observed_adapter_sql_and_side_effects_match(duck_conn, mocker): + from sqlmesh.core.engine_adapter import DuckDBEngineAdapter + + adapter = DuckDBEngineAdapter(lambda: duck_conn) + execute = mocker.spy(adapter, "_execute") + sql_log = mocker.spy(adapter, "_log_sql") + sql = "CREATE OR REPLACE TABLE observed_effect AS SELECT 42 AS value" + calls = [] + for observer in (None, Observer()): + execute.reset_mock() + sql_log.reset_mock() + with observer_scope(observer): + adapter.execute(sql) + calls.append((execute.call_args_list, sql_log.call_args_list)) + assert duck_conn.execute("SELECT value FROM observed_effect").fetchall() == [(42,)] + assert calls[0] == calls[1] + assert len(calls[0][0]) == 1 + + +def test_default_dag_does_not_capture_context(mocker): + from sqlmesh.utils.concurrency import ConcurrentDAGExecutor, concurrent_apply_to_dag + from sqlmesh.utils.dag import DAG + + copy = mocker.patch("contextvars.copy_context", side_effect=AssertionError("unexpected copy")) + dag = DAG() + dag.add("node") + for workers in (1, 2): + concurrent_apply_to_dag(dag, lambda _: None, workers) + executor = ConcurrentDAGExecutor(dag, lambda _: None, workers, False) + executor.run() + assert executor._context is None + copy.assert_not_called() + + +@pytest.mark.parametrize("workers", [1, 2]) +@pytest.mark.parametrize("registration", ["console", "explicit", "disabled", "broken"]) +def test_plan_and_scheduler_registration(mocker, make_snapshot, workers, registration): + from contextlib import nullcontext + from sqlglot import parse_one + from sqlmesh.core.console import TerminalConsole + from sqlmesh.core.environment import EnvironmentNamingInfo + from sqlmesh.core.model import SqlModel, FullKind + from sqlmesh.core.plan.evaluator import BuiltInPlanEvaluator + from sqlmesh.core.plan import stages + from sqlmesh.core.scheduler import Scheduler + from sqlmesh.core.snapshot import DeployabilityIndex, SnapshotChangeCategory + from sqlmesh.utils.date import to_timestamp + + snapshot = make_snapshot( + SqlModel(name="test.model", kind=FullKind(), query=parse_one("SELECT 1 AS a")) + ) + snapshot.categorize_as(SnapshotChangeCategory.BREAKING) + console = TerminalConsole() + observer = Observer() + getter = mocker.patch.object(console, "get_execution_observer", return_value=observer) + if registration == "broken": + getter.side_effect = RuntimeError("getter failure") + native = mocker.MagicMock() + native.execution_tracker.get_execution_stats.return_value = None + native.get_snapshots_to_create.return_value = [] + native.set_correlation_id.return_value = native + scheduler = Scheduler( + [snapshot], native, mocker.MagicMock(), None, max_workers=workers, console=console + ) + workload = [] + consumer = ContextVar("plan_consumer", default=None) + + def evaluate(*args, **kwargs): + if registration in ("console", "explicit"): + assert consumer.get() == "stage" + with action("query", "execute"): + workload.append(True) + return [] + + mocker.patch.object(scheduler, "evaluate", side_effect=evaluate) + evaluator = BuiltInPlanEvaluator( + mocker.MagicMock(), native, mocker.Mock(), None, console=console + ) + stage = stages.BeforeAllStage(statements=[], all_snapshots=[]) + mocker.patch("sqlmesh.core.plan.evaluator.stages.build_plan_stages", return_value=[stage]) + mocker.patch("sqlmesh.core.plan.evaluator.analytics.collector") + + def run_scheduler(*args): + token = consumer.set("stage") + try: + errors, skipped = scheduler.run_merged_intervals( + merged_intervals={ + snapshot: [(to_timestamp("2024-01-01"), to_timestamp("2024-01-02"))] + }, + deployability_index=DeployabilityIndex.all_deployable(), + environment_naming_info=EnvironmentNamingInfo(), + ) + assert not errors and not skipped + finally: + consumer.reset(token) + + mocker.patch.object(evaluator, "visit_before_all_stage", side_effect=run_scheduler) + scope = ( + observer_scope(observer if registration == "explicit" else None) + if registration in ("explicit", "disabled") + else nullcontext() + ) + with scope: + evaluator.evaluate(mocker.Mock(plan_id="observed-plan")) + assert workload == [True] + assert getter.call_count == (1 if registration in ("console", "broken") else 0) + if registration in ("console", "explicit"): + assert [(phase, facts["kind"]) for phase, facts, _ in observer.events] == [ + ("start", "stage"), + ("start", "model"), + ("start", "query"), + ("finish", "query"), + ("finish", "model"), + ("finish", "stage"), + ] + 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 diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py index a9ca86496a..0955af331d 100644 --- a/tests/core/test_snapshot_evaluator.py +++ b/tests/core/test_snapshot_evaluator.py @@ -28,6 +28,7 @@ from sqlmesh.core.environment import EnvironmentNamingInfo from sqlmesh.core.macros import RuntimeStage, macro, MacroEvaluator, MacroFunc from sqlmesh.core.model import ( + AuditResult, Model, FullKind, IncrementalByTimeRangeKind, @@ -247,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) @@ -3285,6 +3402,26 @@ def test_standalone_audit(mocker: MockerFixture, adapter_mock, make_snapshot): adapter_mock.session.assert_not_called() +def _audit_observed( + evaluator: SnapshotEvaluator, + snapshot: Snapshot, + *, + observe_result: t.Callable[[AuditResult], None], + **kwargs: t.Any, +) -> t.List[AuditResult]: + """Exercise the generic hook at the real per-audit evaluator boundary.""" + from sqlmesh.core.execution_observation import observer_scope + + def observer( + phase: str, facts: t.Mapping[str, t.Any], error: t.Optional[BaseException] + ) -> None: + if phase == "finish" and facts["kind"] == "audit" and "audit_result" in facts: + observe_result(facts["audit_result"]) + + with observer_scope(observer): + return evaluator.audit(snapshot, **kwargs) + + def test_audit_wap(adapter_mock: Mock, make_snapshot: t.Callable[..., Snapshot]) -> None: evaluator = SnapshotEvaluator(adapter_mock) @@ -3310,8 +3447,15 @@ def test_audit_wap(adapter_mock: Mock, make_snapshot: t.Callable[..., Snapshot]) expected_table_name = f"spark_catalog.test_schema.test_table.branch_wap_{wap_id}" adapter_mock.wap_table_name.return_value = expected_table_name adapter_mock.fetchone.return_value = (0,) + observed_results: t.List[AuditResult] = [] - evaluator.audit(snapshot, snapshots={}, wap_id=wap_id) + _audit_observed( + evaluator, + snapshot, + snapshots={}, + wap_id=wap_id, + observe_result=observed_results.append, + ) call_args = adapter_mock.fetchone.call_args_list assert len(call_args) == 2 @@ -3330,6 +3474,123 @@ def test_audit_wap(adapter_mock: Mock, make_snapshot: t.Callable[..., Snapshot]) adapter_mock.wap_table_name.assert_called_once_with(snapshot.table_name(), wap_id) adapter_mock.wap_publish.assert_called_once_with(snapshot.table_name(), wap_id) + assert [result.audit.name for result in observed_results] == ["not_null", "test_audit"] + + +def test_audit_observer_preserves_completed_results_when_a_later_audit_fails( + adapter_mock: Mock, make_snapshot: t.Callable[..., Snapshot] +) -> None: + evaluator = SnapshotEvaluator(adapter_mock) + model = SqlModel( + name="test_schema.test_table", + kind=FullKind(), + query=parse_one("SELECT a::int FROM tbl"), + audits=[ + ("not_null", {"columns": exp.to_column("a")}), + ("unique_values", {"columns": exp.convert(["a"])}), + ], + ) + snapshot = make_snapshot(model) + snapshot.categorize_as(SnapshotChangeCategory.BREAKING) + observed_results: t.List[AuditResult] = [] + fetch_count = 0 + + def fetch_audit_result(*args: t.Any, **kwargs: t.Any) -> t.Tuple[int]: + nonlocal fetch_count + fetch_count += 1 + if fetch_count == 2: + assert [result.audit.name for result in observed_results] == ["not_null"] + raise RuntimeError("second audit failed") + return (0,) + + adapter_mock.fetchone.side_effect = fetch_audit_result + + with pytest.raises(RuntimeError, match="second audit failed"): + _audit_observed( + evaluator, + snapshot, + snapshots={}, + observe_result=observed_results.append, + ) + + assert [result.audit.name for result in observed_results] == ["not_null"] + + +def test_audit_observer_failure_does_not_abort_remaining_audits( + adapter_mock: Mock, make_snapshot: t.Callable[..., Snapshot] +) -> None: + evaluator = SnapshotEvaluator(adapter_mock) + model = SqlModel( + name="test_schema.test_table", + kind=FullKind(), + query=parse_one("SELECT a::int, b::int FROM tbl"), + audits=[ + ("not_null", {"columns": exp.to_column("a")}), + ("not_null", {"columns": exp.to_column("b")}), + ], + ) + snapshot = make_snapshot(model) + snapshot.categorize_as(SnapshotChangeCategory.BREAKING) + adapter_mock.fetchone.return_value = (0,) + + def fail_observer(_result: AuditResult) -> None: + raise RuntimeError("observer unavailable") + + results = _audit_observed( + evaluator, + snapshot, + snapshots={}, + observe_result=fail_observer, + ) + + assert len(results) == adapter_mock.fetchone.call_count == 2 + + +def test_audit_observer_reports_failed_skipped_and_repeated_results_once_in_order( + adapter_mock: Mock, make_snapshot: t.Callable[..., Snapshot] +) -> None: + evaluator = SnapshotEvaluator(adapter_mock) + skipped_audit = ModelAudit( + name="skipped_audit", + query="SELECT * FROM test_schema.test_table", + skip=True, + ) + model = SqlModel( + name="test_schema.test_table", + kind=FullKind(), + query=parse_one("SELECT a::int, b::int FROM tbl"), + audits=[ + ("not_null", {"columns": exp.to_column("a")}), + ("skipped_audit", {}), + ("not_null", {"columns": exp.to_column("b")}), + ], + audit_definitions={skipped_audit.name: skipped_audit}, + ) + snapshot = make_snapshot(model) + snapshot.categorize_as(SnapshotChangeCategory.BREAKING) + adapter_mock.fetchone.side_effect = [(2,), (0,)] + observed_results: t.List[AuditResult] = [] + + returned_results = _audit_observed( + evaluator, + snapshot, + snapshots={}, + observe_result=observed_results.append, + ) + + assert observed_results == returned_results + assert [result.audit.name for result in observed_results] == [ + "not_null", + "skipped_audit", + "not_null", + ] + assert [result.count for result in observed_results] == [2, None, 0] + assert [result.skipped for result in observed_results] == [False, True, False] + assert [observed_results[index].audit_args["columns"].sql() for index in (0, 2)] == [ + "a", + "b", + ] + assert adapter_mock.fetchone.call_count == 2 def test_audit_with_datetime_macros(adapter_mock, make_snapshot): @@ -4810,9 +5071,15 @@ def test_multiple_engine_virtual_layer(snapshot: Snapshot, adapters, make_snapsh ) +@pytest.mark.parametrize("observer_enabled", [False, True]) def test_wap_basic( - adapter_mock: Mock, make_snapshot: t.Callable[..., Snapshot], mocker: MockerFixture + adapter_mock: Mock, + make_snapshot: t.Callable[..., Snapshot], + mocker: MockerFixture, + observer_enabled: bool, ) -> None: + from sqlmesh.core.execution_observation import observer_scope + evaluator = SnapshotEvaluator(adapter_mock) model = SqlModel( @@ -4828,19 +5095,25 @@ def test_wap_basic( expected_wap_table = "test_schema.test_table.branch_wap_12345678" adapter_mock.wap_prepare.return_value = expected_wap_table - wap_id = evaluator.evaluate( - snapshot, - start="2020-01-01", - end="2020-01-01", - execution_time="2020-01-01", - snapshots={}, - target_table_exists=True, # Use parameter to control table existence - ) + observations = [] + + def observer(phase, facts, error): + observations.append((phase, dict(facts), error)) + + expected_table_name = snapshot.table_name() + with observer_scope(observer if observer_enabled else None): + wap_id = evaluator.evaluate( + snapshot, + start="2020-01-01", + end="2020-01-01", + execution_time="2020-01-01", + snapshots={}, + target_table_exists=True, # Use parameter to control table existence + ) assert wap_id is not None assert len(wap_id) == 8 - expected_table_name = snapshot.table_name() adapter_mock.wap_prepare.assert_called_once_with(expected_table_name, wap_id) adapter_mock.replace_query.assert_called_once_with( expected_wap_table, @@ -4857,6 +5130,21 @@ def test_wap_basic( source_columns=mocker.ANY, ) + materialization_observations = [ + (phase, facts, error) + for phase, facts, error in observations + if facts.get("kind") == "physical" and facts.get("action") == "materialize" + ] + if observer_enabled: + assert [ + (phase, facts["target"], error) for phase, facts, error in materialization_observations + ] == [ + ("start", expected_table_name, None), + ("finish", expected_table_name, None), + ] + else: + assert observations == [] + def test_wap_model_wap_supported( adapter_mock: Mock, make_snapshot: t.Callable[..., Snapshot], mocker: MockerFixture @@ -5052,12 +5340,24 @@ def test_wap_publish_failure(adapter_mock: Mock, make_snapshot: t.Callable[..., adapter_mock.wap_table_name.return_value = expected_wap_table_name adapter_mock.fetchone.return_value = (0,) - # Mock WAP publish to raise an exception - adapter_mock.wap_publish.side_effect = Exception("WAP publish failed") + observed_results: t.List[AuditResult] = [] + + def fail_wap_publish(*args: t.Any, **kwargs: t.Any) -> None: + assert [result.audit.name for result in observed_results] == ["not_null"] + raise RuntimeError("WAP publish failed") + + adapter_mock.wap_publish.side_effect = fail_wap_publish + + with pytest.raises(RuntimeError, match="WAP publish failed"): + _audit_observed( + evaluator, + snapshot, + snapshots={}, + wap_id=wap_id, + observe_result=observed_results.append, + ) - # Execute audit with WAP ID and expect it to raise the exception - with pytest.raises(Exception, match="WAP publish failed"): - evaluator.audit(snapshot, snapshots={}, wap_id=wap_id) + assert [result.audit.name for result in observed_results] == ["not_null"] def test_properties_are_preserved_in_both_create_statements( @@ -5766,3 +6066,50 @@ def test_promote_resolves_this_model_with_single_mapping_entry( assert spy.call_count == 20 for call in spy.call_args_list: assert len(call.args[1]) == 1 + + +@pytest.mark.parametrize("workers", [1, 2]) +@pytest.mark.parametrize("operation", ["create", "promote", "demote", "migrate"]) +def test_observation_propagates_through_snapshot_dispatch( + mocker, adapter_mock, snapshot, workers, operation +): + from contextvars import ContextVar + from sqlmesh.core.execution_observation import observer_scope + + evaluator = SnapshotEvaluator(adapter_mock, ddl_concurrent_tasks=workers) + consumer = ContextVar("snapshot_dispatch_consumer", default=None) + events = [] + + def observer(phase, facts, error): + assert consumer.get() == "admission" + assert facts["snapshot"] is snapshot + events.append((phase, facts["action"], error)) + + naming = EnvironmentNamingInfo(name="test") + mocker.patch.object(evaluator, "_get_virtual_data_objects", return_value={}) + mocker.patch.object( + evaluator, "_get_physical_data_objects", return_value={snapshot.snapshot_id: mocker.Mock()} + ) + mocker.patch.object(evaluator, "_create_schemas") + token = consumer.set("admission") + try: + with observer_scope(observer): + if operation == "create": + evaluator._create_snapshots( + [snapshot], + {snapshot.name: snapshot}, + DeployabilityIndex.all_deployable(), + None, + set(), + set(), + ) + elif operation == "promote": + evaluator.promote([snapshot], naming) + elif operation == "demote": + evaluator.demote([snapshot], naming) + else: + evaluator.migrate([snapshot], {snapshot.snapshot_id: snapshot}) + finally: + consumer.reset(token) + action = "schema_migration" if operation == "migrate" else operation + assert events == [("start", action, None), ("finish", action, None)] diff --git a/tests/utils/test_concurrency.py b/tests/utils/test_concurrency.py index 083897fe56..0d45a375f1 100644 --- a/tests/utils/test_concurrency.py +++ b/tests/utils/test_concurrency.py @@ -412,3 +412,80 @@ def execute(node): finally: consumer.reset(token) assert seen == ["first admission"] * 2 + ["second admission"] * 4 + + +@pytest.mark.parametrize("phase", ["start", "finish"]) +def test_concurrent_dag_propagates_observer_base_exception_and_can_be_reused(phase: str): + _run_bounded_python( + f""" + from threading import Event, Thread + + from sqlmesh.core.execution_observation import action, observer_scope + from sqlmesh.utils.concurrency import ConcurrentDAGExecutor + from sqlmesh.utils.dag import DAG + + dag = DAG() + dag.add("cancel") + dag.add("sibling") + dag.add("dependent", ["sibling"]) + error = KeyboardInterrupt("cancel execution") + release_sibling = Event() + sibling_started = Event() + called = [] + failing_phase = {phase!r} + + def observe(current_phase, facts, callback_error): + if current_phase == failing_phase: + raise error + + def execute(node): + if node == "cancel": + assert sibling_started.wait(timeout=3) + with observer_scope(observe), action("query", "execute"): + pass + if node == "sibling": + sibling_started.set() + assert release_sibling.wait(timeout=6) + called.append(node) + + executor = ConcurrentDAGExecutor(dag, execute, 2, True) + run_errors = [] + + def run(): + try: + executor.run() + except BaseException as ex: + run_errors.append(ex) + + runner = Thread(target=run) + runner.start() + assert sibling_started.wait(timeout=3) + try: + try: + executor._finished_future.result(timeout=3) + except KeyboardInterrupt as raised: + assert raised is error + else: + raise AssertionError("the worker cancellation was not propagated") + finally: + release_sibling.set() + + runner.join(timeout=3) + assert not runner.is_alive() + assert run_errors == [error] + assert "dependent" not in called + + second_error = SystemExit("cancel reused executor") + + def cancel_again(node): + raise second_error + + executor.fn = cancel_again + try: + executor.run() + except SystemExit as raised: + assert raised is second_error + else: + raise AssertionError("the reused executor did not propagate cancellation") + """ + )