Feat: Add optional execution observations

Extract native lifecycle callbacks with isolated observer dispatch and opt-in worker context propagation.

Co-authored-by: Pavan Mylavarapu <pavan.mylavarapu@atoms.co>
Signed-off-by: mday-io <mdaytn@gmail.com>

Author: mday-io

docs/reference/execution_observations.md

@@ -0,0 +1,132 @@
+# Execution observations
+
+Execution observations let Python integrations observe native SQLMesh work through
+an optional callback. They do not replace SQL generation, execution, retries,
+transactions, or Console progress reporting.
+
+## Register an observer
+
+Override `Console.get_execution_observer()` to return a callback accepting
+`phase`, `facts`, and `error`. Its default return value is `None`.
+Install the console before constructing the SQLMesh `Context`:
+
+```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=["path/to/project"])
+```
+
+Built-in plan-stage evaluation and scheduler dispatch resolve the getter at the
+outer observation boundary. Nested dispatch reuses the active registration.
+An ordinary `Exception` from the getter disables observation at that boundary;
+native work continues.
+
+For an explicitly bounded registration, use `observer_scope(observer)` from
+`sqlmesh.core.execution_observation`. Nested scopes restore the preceding
+registration on exit. `observer_scope(None)` explicitly disables observation
+inside the scope, including automatic Console registration.
+
+## Callback lifecycle
+
+`start` runs before the observed work, with `error=None`. Once the start callback
+returns, `finish` runs from `finally` after that work, with its original exception
+object or `None`. Both callbacks run synchronously on the invoking thread.
+
+SQLMesh passes the **same read-only mapping object** to start and finish for one
+invocation. Distinct invocations have distinct mappings. Consumers may retain the
+mapping until its matching finish callback to pair the calls by object identity;
+release it afterward. SQLMesh may add facts discovered during execution, so this
+mapping is a live view, not an immutable snapshot of the start facts.
+
+Nested workload calls produce nested observations. Callbacks for concurrent
+workers can interleave, and there is no global callback ordering. Observers must
+support concurrent calls and return promptly. SQLMesh does not allocate execution
+IDs, timestamps, durations, or parent IDs; consumers can maintain that information.
+
+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.
+
+While a callback is being delivered, observations triggered by that callback are
+suppressed, and internal metadata updates are ignored. This guard applies only
+during callback delivery: a normal model operation can still contain observed
+physical operations, audits, and queries.
+
+## Native facts and boundaries
+
+Every mapping contains `kind` and `action`. Other fields depend on the boundary
+and may be absent, particularly when work exits early or fails.
+
+| Kind | Action | Native facts |
+| --- | --- | --- |
+| `stage` | Native plan-stage 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` and `AuditResult` values are borrowed native objects. Treat them and
+other nested values as read-only, and do not retain them beyond finish. The
+mapping's read-only wrapper does not freeze these objects. Normalize any values
+needed for later processing during the callback; facts are not a JSON event
+schema. SQLMesh does not interpret audit results as consumer statuses or inherit
+model facts into child mappings.
+
+Audit observations finish before later audits or WAP publication can fail. An
+audit query failure can produce a finish callback without an `AuditResult`.
+Physical materialization includes the surrounding transaction/session exit.
+An invocation can perform no work, multiple statements, or internal driver
+retries; observations are not an expected-work manifest or a retry inventory.
+
+The query boundary surrounds the base adapter's existing SQL logging and
+`_execute` call. It does not include a surrounding transaction's commit or later
+fetch/result consumption. Adapter paths that bypass this boundary, including
+BigQuery session queries and dataframe loads and ClickHouse dataframe inserts,
+do not necessarily produce query observations. These hooks therefore do not
+provide a complete warehouse activity history or a transaction-success guarantee.
+
+SQL logs retain their existing behavior, including VALUES redaction for
+non-query expressions. The callback does not receive SQL text, adapter objects,
+or warehouse query IDs. Starting the observation before SQL logging lets a
+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.
+
+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.
+
+With no observer registered, these call sites do not opt into context copying
+and retain their existing serial and worker execution behavior.

mkdocs.yml

@@ -105,6 +105,7 @@ nav:
       - reference/cli.md
       - reference/notebook.md
       - reference/python.md
+      - reference/execution_observations.md
     - Configuration:
       - reference/configuration.md
       - reference/model_configuration.md

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."""

sqlmesh/core/engine_adapter/base.py

@@ -2624,12 +2624,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:

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

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(

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
@@ -518,6 +523,21 @@ def run_merged_intervals(
         )
 
         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):
@@ -596,12 +616,13 @@ def run_node(node: SchedulingUnit) -> None:
                 )
 
         try:
-            with self.snapshot_evaluator.concurrent_context():
+            with console_observer_scope(self.console), 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)