From f554a1a1117ab414b31ea0853729e86205493301 Mon Sep 17 00:00:00 2001 From: mday-io Date: Sat, 31 Oct 2026 14:01:22 -0400 Subject: [PATCH 1/4] feat: add query execution context provider transport Signed-off-by: mday-io --- docs/concepts/query_execution_context.md | 143 ++++++ docs/examples/query_execution_context.md | 22 + docs/reference/configuration.md | 1 + examples/query_execution_context/README.md | 24 + .../models/aux_events.sql | 13 + .../query_execution_context/models/events.sql | 13 + examples/query_execution_context/observer.py | 119 +++++ examples/query_execution_context/run_local.py | 110 +++++ mkdocs.yml | 3 + sqlmesh/core/config/connection.py | 13 +- sqlmesh/core/config/scheduler.py | 18 +- sqlmesh/core/engine_adapter/base.py | 89 ++++ sqlmesh/core/query_execution.py | 301 ++++++++++++ .../test_query_execution_context.py | 454 ++++++++++++++++++ .../test_query_execution_context.py | 152 ++++++ tests/core/test_config.py | 2 + tests/core/test_query_execution.py | 151 ++++++ .../test_query_execution_context_config.py | 107 +++++ tests/integrations/jupyter/test_magics.py | 8 +- 19 files changed, 1735 insertions(+), 8 deletions(-) create mode 100644 docs/concepts/query_execution_context.md create mode 100644 docs/examples/query_execution_context.md create mode 100644 examples/query_execution_context/README.md create mode 100644 examples/query_execution_context/models/aux_events.sql create mode 100644 examples/query_execution_context/models/events.sql create mode 100644 examples/query_execution_context/observer.py create mode 100644 examples/query_execution_context/run_local.py create mode 100644 sqlmesh/core/query_execution.py create mode 100644 tests/core/engine_adapter/test_query_execution_context.py create mode 100644 tests/core/integration/test_query_execution_context.py create mode 100644 tests/core/test_query_execution.py create mode 100644 tests/core/test_query_execution_context_config.py diff --git a/docs/concepts/query_execution_context.md b/docs/concepts/query_execution_context.md new file mode 100644 index 0000000000..74e3c33c82 --- /dev/null +++ b/docs/concepts/query_execution_context.md @@ -0,0 +1,143 @@ +# Query execution context + +SQLMesh can add structured metadata to warehouse statements so a query-history +tool can connect a statement to the plan or run, model, and evaluation that issued +it. The metadata is a versioned JSON object in a `QUERY_CONTEXT` SQL comment. + +Enable `query_execution_context` on a connection and give its engine adapter a +provider that returns a `QueryExecutionContext`. The flag defaults to `false`. +For YAML configuration, set the option under the gateway's `connection`: + +```yaml +gateways: + local: + connection: + type: duckdb + query_execution_context: true +``` + +The configuration flag enables transport. It does not create a model-context +provider. Without a provider, SQLMesh adds parent metadata only when a plan or run +correlation ID is available. SQLMesh calls the provider after rendering each +statement and before executing it. The provider sees the current worker's context, +so concurrent model evaluations can supply different IDs to a shared adapter. + +This provider API does not add automatic model tagging to a stock project. +For a provider that reads native observer events, see the +[query execution context example](../examples/query_execution_context.md). + +## Configure a provider + +For a Python `Context`, configure its warehouse adapters before loading the +project. The `Context` clones these adapters when it creates its snapshot +evaluator, so the provider and flag carry to those clones. + +```python +from sqlmesh.core.context import Context +from sqlmesh.core.execution_observation import observer_scope +from examples.query_execution_context.observer import QueryContextObserver + +context = Context(paths=".", load=False) +observer = QueryContextObserver(context.console.get_execution_observer()) + +for adapter in context.engine_adapters.values(): + adapter.query_execution_context = True + adapter.query_execution_context_provider = observer.current_query_context + +context.load() +with observer_scope(observer): + plan = context.plan("dev", no_prompts=True) + context.apply(plan) + context.run("dev") +``` + +The provider returns the immutable context currently active in that worker, or +`None` when no context applies. An observer can set an evaluation ID when +SQLMesh reports a model batch, then restore its prior context when that +action ends. The sample observer composes with the observer returned by the +console. It preserves that observer's facts and events. + +Use the adapter instance that the evaluator will use. Adapters created with +`with_settings` inherit the flag and provider. Passing +`query_execution_context_provider=None` to `with_settings` clears the inherited +provider. Setting the flag to `false` disables provider calls and comment output. + +Connection settings also provide defaults to other adapters. Configure the +warehouse adapters where you want model context in query history. The built-in +scheduler can derive its state adapter from the warehouse settings, but it +disables query context on that fallback. To annotate state SQL, configure the +`state_connection` explicitly and provide a context provider for its adapter. + +## Read the payload + +The version-one payload uses the SQLMesh model's canonical name and opaque +snapshot identifier. Operation intervals use UTC half-open boundaries. Fields +that are unavailable are omitted. + +```json +{ + "schema_version": 1, + "emitter": "sqlmesh", + "parent": {"system": "sqlmesh", "type": "run", "id": "run-123"}, + "workload": { + "system": "sqlmesh", + "type": "model", + "name": "analytics.orders", + "version": "snapshot-456" + }, + "operation": { + "id": "evaluation-789", + "type": "evaluate", + "interval_start": "2026-10-01T00:00:00+00:00", + "interval_end": "2026-10-02T00:00:00+00:00", + "batch_index": 0 + }, + "query": {"phase": "materialize", "role": "query"} +} +``` + +The operation ID groups statements for one model evaluation, including nested +audits. A later evaluation gets a new ID. The provider owns the ID; SQLMesh does +not create another operation ID at the adapter boundary. Consumers can compare +the value they supplied with records from their execution observer. + +The adapter's own correlation ID sets the payload's `parent`. SQLMesh replaces +a different provider-supplied parent with the adapter's value. When the adapter +has no correlation ID, SQLMesh clears a parent from the provided context. + +Consumers must accept absent fields and unknown future phase or role values. +The `workload` identity is the tuple `(system, type, name)`. Its `version` is +opaque and does not change that identity. + +## SQL comment behavior + +SQLMesh prefixes the context comment after rendering. When a plan or run marker +is present, SQLMesh keeps that legacy marker first. The submitted statement may +look like this: + +```sql +/* SQLMESH_RUN: run-123 */ /* QUERY_CONTEXT: {"schema_version":1,"emitter":"sqlmesh","parent":{"system":"sqlmesh","type":"run","id":"run-123"},"query":{"phase":"unknown","role":"query"}} */ SELECT 1 +``` + +The generic comment is bounded to 8,192 UTF-8 bytes. Each identifier is bounded +to 1,024 UTF-8 bytes before JSON escaping. SQLMesh escapes metadata so it cannot +end the comment. In JSON, SQLMesh also escapes slash characters so an identifier +cannot form a nested comment delimiter. If an identifier is invalid or either +limit is exceeded, SQLMesh +omits the complete `QUERY_CONTEXT` comment and logs a warning. It never truncates +the JSON or fails the query for an omitted annotation. + +The submitted SQL changes, while rendered model definitions and model fingerprints +do not. If a downstream query-history system hashes raw SQL text, it must account +for the additional comment. + +An adapter that does not support generic SQL comments keeps its SQL unchanged +and logs one warning per adapter when the option is enabled. It does not call the +provider. Athena does not support this transport. Engines can override the adapter +hook to add native metadata transport in a separate implementation. + +If the provider raises an ordinary exception or returns an invalid value, SQLMesh +omits the rich comment, logs a bounded warning without provider data, and executes +the statement with its legacy correlation marker. It tries the provider again +for the next statement. `KeyboardInterrupt`, `SystemExit`, and other cancellation +exceptions keep their normal behavior and propagate. diff --git a/docs/examples/query_execution_context.md b/docs/examples/query_execution_context.md new file mode 100644 index 0000000000..e767a6d0a3 --- /dev/null +++ b/docs/examples/query_execution_context.md @@ -0,0 +1,22 @@ +# Query execution context example + +This example runs two incremental models with daily batches and `not_null` +audits in a temporary DuckDB project. It applies a plan to `dev`, then runs a +later interval. + +Run it from the SQLMesh repository root with the development environment: + +```bash +.venv/bin/python examples/query_execution_context/run_local.py +``` + +The script attaches `QueryContextObserver.current_query_context` to the warehouse +adapter providers. The observer builds one context for each model batch and +delegates events to the observer supplied by the console. The run checks that four +initial batches and two later batches use the same IDs in the model events, +materialization SQL, and audits. It also checks that the resulting model rows are +correct and that the query-context scope is restored after the run. + +The example uses a temporary DuckDB database and makes no external connections. +See [query execution context](../concepts/query_execution_context.md) for the +payload, provider contract, and transport limits. diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index ae036a7774..c817583c67 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -212,6 +212,7 @@ Most parameters are specific to the connection engine `type` - see [below](#engi | `register_comments` | Whether SQLMesh should register model comments with the SQL engine (if the engine supports it). (Default: `true`.) | bool | N | | `pre_ping` | Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive. This can only be enabled for engines with transaction support. | bool | N | | `pretty_sql` | If SQL should be formatted before being executed, not recommended in a production setting. (Default: `false`.) | bool | N | +| `query_execution_context` | Whether an opted-in engine adapter adds bounded execution metadata to SQL comments. (Default: `false`.) See [query execution context](../concepts/query_execution_context.md). | bool | N | #### Engine-specific diff --git a/examples/query_execution_context/README.md b/examples/query_execution_context/README.md new file mode 100644 index 0000000000..e04c12608d --- /dev/null +++ b/examples/query_execution_context/README.md @@ -0,0 +1,24 @@ +# Local query execution context example + +This example runs two incremental models with daily batches and `not_null` +audits in a temporary DuckDB database. It applies a plan to `dev`, then runs a +later interval. + +From the SQLMesh repository root, run: + +```bash +.venv/bin/python examples/query_execution_context/run_local.py +``` + +The script registers the query-context provider on each warehouse adapter before +it loads the project. `QueryContextObserver` builds model and batch metadata from +SQLMesh's native observation events. Its provider returns the active immutable +context in each worker. The existing console observer still receives its events. + +The example checks four batches during the plan and two during the later run. It +confirms that the observer events, materialization SQL, and audits share each +batch's operation ID. It also verifies the resulting rows and scope cleanup. The +temporary project and DuckDB database are removed when the script exits. + +The example does not connect to an external database. For configuration and +payload details, see [Query execution context](../../docs/concepts/query_execution_context.md). diff --git a/examples/query_execution_context/models/aux_events.sql b/examples/query_execution_context/models/aux_events.sql new file mode 100644 index 0000000000..30abb0bde7 --- /dev/null +++ b/examples/query_execution_context/models/aux_events.sql @@ -0,0 +1,13 @@ +-- SPDX-License-Identifier: Apache-2.0 +MODEL ( + name e2e.aux_events, + kind INCREMENTAL_BY_TIME_RANGE ( + time_column ds, + batch_size 1 + ), + start '2024-01-01', + cron '@daily', + audits (not_null(columns := (id))) +); + +SELECT CAST(@start_ds AS DATE) AS ds, 2 AS id diff --git a/examples/query_execution_context/models/events.sql b/examples/query_execution_context/models/events.sql new file mode 100644 index 0000000000..cb00a913f8 --- /dev/null +++ b/examples/query_execution_context/models/events.sql @@ -0,0 +1,13 @@ +-- SPDX-License-Identifier: Apache-2.0 +MODEL ( + name e2e.events, + kind INCREMENTAL_BY_TIME_RANGE ( + time_column ds, + batch_size 1 + ), + start '2024-01-01', + cron '@daily', + audits (not_null(columns := (id))) +); + +SELECT CAST(@start_ds AS DATE) AS ds, 1 AS id diff --git a/examples/query_execution_context/observer.py b/examples/query_execution_context/observer.py new file mode 100644 index 0000000000..780a3ff978 --- /dev/null +++ b/examples/query_execution_context/observer.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Use native observation events to supply query-context IDs.""" + +from __future__ import annotations + +import contextvars +import dataclasses +import typing as t + +from sqlmesh.core.execution_observation import Observer +from sqlmesh.core.query_execution import ( + ModelExecutionContext, + QueryExecutionContext, + QueryExecutionOperation, + QueryExecutionPhase, + QueryExecutionRole, +) +from sqlmesh.utils import random_id +from sqlmesh.utils.date import to_datetime + + +@dataclasses.dataclass(frozen=True) +class _Frame: + facts: t.Mapping[str, t.Any] + token: contextvars.Token[t.Optional[QueryExecutionContext]] + + +class QueryContextObserver: + def __init__(self, observer: t.Optional[Observer] = None): + self._observer = observer + self._context: contextvars.ContextVar[t.Optional[QueryExecutionContext]] = ( + contextvars.ContextVar("example_query_context", default=None) + ) + self._frames: contextvars.ContextVar[t.Tuple[_Frame, ...]] = contextvars.ContextVar( + "example_query_context_frames", default=() + ) + + def current_query_context(self) -> t.Optional[QueryExecutionContext]: + return self._context.get() + + def __call__( + self, + event: t.Literal["start", "finish"], + facts: t.Mapping[str, t.Any], + error: t.Optional[BaseException], + ) -> None: + if event == "start": + context = self._derive_context(facts) + token = self._context.set(context) + self._frames.set(self._frames.get() + (_Frame(facts, token),)) + try: + self._notify(event, facts, error) + except BaseException: + self._restore(facts) + raise + else: + try: + self._notify(event, facts, error) + finally: + self._restore(facts) + + def _notify( + self, + event: t.Literal["start", "finish"], + facts: t.Mapping[str, t.Any], + error: t.Optional[BaseException], + ) -> None: + if self._observer is not None: + try: + self._observer(event, facts, error) + except Exception: + pass + + def _restore(self, facts: t.Mapping[str, t.Any]) -> None: + frames = self._frames.get() + assert frames[-1].facts is facts + self._context.reset(frames[-1].token) + self._frames.set(frames[:-1]) + + def _derive_context(self, facts: t.Mapping[str, t.Any]) -> t.Optional[QueryExecutionContext]: + current = self._context.get() + kind, action = facts["kind"], facts["action"] + if kind == "model": + snapshot = facts["snapshot"] + if not snapshot.is_model: + return None + start, end = facts["interval"] + return QueryExecutionContext( + evaluation=ModelExecutionContext( + evaluation_id=random_id(), + model_name=snapshot.name, + snapshot_id=snapshot.snapshot_id.identifier, + interval_start=to_datetime(start), + interval_end=to_datetime(end), + batch_index=facts["batch_index"], + operation_type=QueryExecutionOperation.AUDIT + if action == "audit_only" + else QueryExecutionOperation.EVALUATE, + ), + phase=QueryExecutionPhase.AUDIT + if action == "audit_only" + else QueryExecutionPhase.MATERIALIZE, + role=QueryExecutionRole.AUDIT + if action == "audit_only" + else QueryExecutionRole.QUERY, + ) + if current is not None and kind == "audit": + return dataclasses.replace( + current, phase=QueryExecutionPhase.AUDIT, role=QueryExecutionRole.AUDIT + ) + if current is not None and kind in ("physical", "virtual"): + return dataclasses.replace( + current, + phase=QueryExecutionPhase.MATERIALIZE + if action == "materialize" + else QueryExecutionPhase.DDL, + role=QueryExecutionRole.QUERY, + ) + return current diff --git a/examples/query_execution_context/run_local.py b/examples/query_execution_context/run_local.py new file mode 100644 index 0000000000..d170997e39 --- /dev/null +++ b/examples/query_execution_context/run_local.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run a local plan and scheduler with observer-owned query context.""" + +from __future__ import annotations + +import datetime as dt +import json +import logging +import re +import shutil +import tempfile +import typing as t +from pathlib import Path + +from examples.query_execution_context.observer import QueryContextObserver +from sqlmesh.core.config import Config, ModelDefaultsConfig +from sqlmesh.core.config.connection import DuckDBConnectionConfig +from sqlmesh.core.config.gateway import GatewayConfig +from sqlmesh.core.console import NoopConsole +from sqlmesh.core.context import Context +from sqlmesh.core.execution_observation import observer_scope + + +class _QueryLog(logging.Handler): + def __init__(self) -> None: + super().__init__() + self.statements: t.List[str] = [] + + def emit(self, record: logging.LogRecord) -> None: + message = record.getMessage() + if message.startswith("Executing SQL: "): + self.statements.append(message.removeprefix("Executing SQL: ")) + + +def main() -> None: + with tempfile.TemporaryDirectory(prefix="sqlmesh-query-context-") as directory: + root = Path(directory) + shutil.copytree(Path(__file__).parent / "models", root / "models") + context = Context( + paths=root, + config=Config( + gateways={ + "local": GatewayConfig( + connection=DuckDBConnectionConfig( + database=str(root / "warehouse.duckdb"), + concurrent_tasks=2, + query_execution_context=True, + ) + ) + }, + default_gateway="local", + model_defaults=ModelDefaultsConfig(dialect="duckdb"), + ), + load=False, + ) + context.console = NoopConsole() + ids: t.Set[str] = set() + + def record( + event: t.Literal["start", "finish"], + facts: t.Mapping[str, t.Any], + error: t.Optional[BaseException], + ) -> None: + if event == "finish" and facts["kind"] == "model": + value = consumer.current_query_context() + assert value is not None and value.evaluation is not None + ids.add(value.evaluation.evaluation_id) + + consumer = QueryContextObserver(record) + for adapter in context.engine_adapters.values(): + adapter.query_execution_context_provider = consumer.current_query_context + context.load() + query_log = _QueryLog() + logger = logging.getLogger("sqlmesh.core.engine_adapter.base") + previous_level = logger.level + logger.setLevel(logging.INFO) + logger.addHandler(query_log) + try: + with observer_scope(consumer): + plan = context.plan( + "dev", start="2024-01-01", end="2024-01-02", no_prompts=True, skip_tests=True + ) + context.apply(plan) + assert context.run("dev", start="2024-01-03", end="2024-01-03").is_success + assert consumer.current_query_context() is None + assert len(ids) == 6 + phases: t.Dict[str, t.Set[str]] = {} + for sql in query_log.statements: + match = re.search(r"/\* QUERY_CONTEXT: (.*?) \*/", sql) + if match is None: + continue + payload = json.loads(match[1]) + operation = payload.get("operation") + if operation is not None: + phases.setdefault(operation["id"], set()).add(payload["query"]["phase"]) + assert set(phases) == ids + assert all({"materialize", "audit"} <= value for value in phases.values()) + rows = context.engine_adapter.fetchall( + f"SELECT ds, id FROM {context.resolve_table('e2e.events')} ORDER BY ds" + ) + assert rows == [(dt.date(2024, 1, day), 1) for day in (1, 2, 3)] + print("Verified six model-batch IDs and three resulting rows.") + finally: + logger.removeHandler(query_log) + logger.setLevel(previous_level) + context.close() + + +if __name__ == "__main__": + main() diff --git a/mkdocs.yml b/mkdocs.yml index 1a1a34e624..e64de773bd 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -45,6 +45,7 @@ nav: - concepts/tests.md - concepts/audits.md - concepts/state.md + - concepts/query_execution_context.md - Models: - concepts/models/overview.md - concepts/models/model_kinds.md @@ -70,6 +71,8 @@ nav: - Walkthroughs: - examples/sqlmesh_cli_crash_course.md - examples/incremental_time_full_walkthrough.md + - Projects: + - examples/query_execution_context.md - Integrations: - "Overview": integrations/overview.md - Tools: diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py index 29cb1ef467..e96f54b9b2 100644 --- a/sqlmesh/core/config/connection.py +++ b/sqlmesh/core/config/connection.py @@ -109,6 +109,7 @@ class ConnectionConfig(abc.ABC, BaseConfig): register_comments: bool pre_ping: bool pretty_sql: bool = False + query_execution_context: bool = False schema_differ_overrides: t.Optional[t.Dict[str, t.Any]] = None catalog_type_overrides: t.Optional[t.Dict[str, str]] = None @@ -183,7 +184,7 @@ def create_engine_adapter( """Returns a new instance of the Engine Adapter.""" concurrent_tasks = concurrent_tasks or self.concurrent_tasks - return self._engine_adapter( + adapter = self._engine_adapter( self._connection_factory_with_kwargs, multithreaded=concurrent_tasks > 1, default_catalog=self.get_catalog(), @@ -196,6 +197,8 @@ def create_engine_adapter( catalog_type_overrides=self.catalog_type_overrides, **self._extra_engine_config, ) + adapter.query_execution_context = self.query_execution_context + return adapter def get_catalog(self) -> t.Optional[str]: """The catalog for this connection""" @@ -506,6 +509,14 @@ def create_engine_adapter( logger.info( f"Using existing DuckDB adapter due to overlapping data file: {self._mask_sensitive_data(key)}" ) + if ( + adapter.query_execution_context != self.query_execution_context + or adapter.query_execution_context_provider is not None + ): + return adapter.with_settings( + query_execution_context=self.query_execution_context, + query_execution_context_provider=None, + ) return adapter if data_files: diff --git a/sqlmesh/core/config/scheduler.py b/sqlmesh/core/config/scheduler.py index 4cce9b0f76..6e4c19eec8 100644 --- a/sqlmesh/core/config/scheduler.py +++ b/sqlmesh/core/config/scheduler.py @@ -65,9 +65,8 @@ def state_sync_fingerprint(self, context: GenericContext) -> str: class _EngineAdapterStateSyncSchedulerConfig(SchedulerConfig): def create_state_sync(self, context: GenericContext) -> StateSync: - state_connection = ( - context.config.get_state_connection(context.gateway) or context.connection_config - ) + configured_state_connection = context.config.get_state_connection(context.gateway) + state_connection = configured_state_connection or context.connection_config warehouse_connection = context.config.get_connection(context.gateway) @@ -86,6 +85,19 @@ def create_state_sync(self, context: GenericContext) -> StateSync: state_connection.concurrent_tasks = warehouse_connection.concurrent_tasks engine_adapter = state_connection.create_engine_adapter() + state_query_execution_context = ( + configured_state_connection.query_execution_context + if configured_state_connection is not None + else False + ) + if ( + engine_adapter.query_execution_context != state_query_execution_context + or engine_adapter.query_execution_context_provider is not None + ): + engine_adapter = engine_adapter.with_settings( + query_execution_context=state_query_execution_context, + query_execution_context_provider=None, + ) if state_connection.is_forbidden_for_state_sync: raise ConfigError( f"The {engine_adapter.DIALECT.upper()} engine cannot be used to store SQLMesh state - please specify a different `state_connection` engine." diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py index 2dd234f70a..01d4e023a3 100644 --- a/sqlmesh/core/engine_adapter/base.py +++ b/sqlmesh/core/engine_adapter/base.py @@ -10,11 +10,13 @@ from __future__ import annotations import contextlib +import dataclasses import itertools import logging import sys import typing as t from functools import cached_property, partial +from threading import Lock from sqlglot import Dialect, exp from sqlglot.errors import ErrorLevel @@ -39,6 +41,12 @@ set_catalog, ) from sqlmesh.core.model.kind import TimeColumn +from sqlmesh.core.query_execution import ( + QueryExecutionContext, + _query_execution_context_transport_is_active, + _query_execution_context_transport_scope, + serialize_query_execution_context, +) from sqlmesh.core.schema_diff import SchemaDiffer, TableAlterOperation from sqlmesh.core.snapshot.execution_tracker import QueryExecutionTracker from sqlmesh.utils import ( @@ -144,6 +152,10 @@ def __init__( correlation_id: t.Optional[CorrelationId] = None, schema_differ_overrides: t.Optional[t.Dict[str, t.Any]] = None, query_execution_tracker: t.Optional[QueryExecutionTracker] = None, + query_execution_context: bool = False, + query_execution_context_provider: t.Optional[ + t.Callable[[], t.Optional[QueryExecutionContext]] + ] = None, **kwargs: t.Any, ): self.dialect = dialect.lower() or self.DIALECT @@ -169,8 +181,18 @@ def __init__( self._schema_differ_overrides = schema_differ_overrides self._query_execution_tracker = query_execution_tracker self._data_object_cache: t.Dict[str, t.Optional[DataObject]] = {} + self.query_execution_context = query_execution_context + self.query_execution_context_provider = query_execution_context_provider + self._query_execution_context_warning_categories: t.Set[str] = set() + self._query_execution_context_warning_lock = Lock() def with_settings(self, **kwargs: t.Any) -> EngineAdapter: + query_execution_context = kwargs.pop( + "query_execution_context", self.query_execution_context + ) + query_execution_context_provider = kwargs.pop( + "query_execution_context_provider", self.query_execution_context_provider + ) extra_kwargs = { "null_connection": True, "execute_log_level": kwargs.pop("execute_log_level", self._execute_log_level), @@ -193,6 +215,8 @@ def with_settings(self, **kwargs: t.Any) -> EngineAdapter: pretty_sql=self._pretty_sql, **extra_kwargs, ) + adapter.query_execution_context = query_execution_context + adapter.query_execution_context_provider = query_execution_context_provider return adapter @@ -2642,6 +2666,8 @@ def execute( else: sql = t.cast(str, e) + if self.query_execution_context: + sql = self._attach_query_execution_context(sql) sql = self._attach_correlation_id(sql) # The rendered statement and correlation comment are finalized @@ -2656,6 +2682,69 @@ def execute( ) self._execute(sql, track_rows_processed, **kwargs) + def _warn_query_execution_context_once(self, category: str, message: str, *args: t.Any) -> None: + with self._query_execution_context_warning_lock: + if category in self._query_execution_context_warning_categories: + return + self._query_execution_context_warning_categories.add(category) + logger.warning(message, *args) + + def _attach_query_execution_context(self, sql: str) -> str: + """Attach bounded query context when enabled and generic comments are supported.""" + if _query_execution_context_transport_is_active(): + return sql + + with _query_execution_context_transport_scope(): + if not self.ATTACH_CORRELATION_ID: + self._warn_query_execution_context_once( + "unsupported", + "Adapter %s does not support generic SQL comment transport for " + "query execution context; metadata will be omitted", + self.__class__.__name__, + ) + return sql + + provider = self.query_execution_context_provider + try: + context = provider() if provider is not None else None + except Exception: + self._warn_query_execution_context_once( + "provider", + "Query execution context provider failed for adapter %s; metadata will be omitted", + self.__class__.__name__, + ) + return sql + + if context is not None and not isinstance(context, QueryExecutionContext): + self._warn_query_execution_context_once( + "invalid_context", + "Query execution context provider returned an invalid context for adapter %s; " + "metadata will be omitted", + self.__class__.__name__, + ) + return sql + + try: + parent = self.correlation_id + if context is None: + if parent is None: + return sql + context = QueryExecutionContext(parent=parent) + elif context.parent != parent: + context = dataclasses.replace(context, parent=parent) + + comment = serialize_query_execution_context(context) + except Exception: + self._warn_query_execution_context_once( + "serialization", + "Query execution context normalization or serialization failed for adapter %s; " + "metadata will be omitted", + self.__class__.__name__, + ) + return sql + + return f"{comment}{sql}" if comment is not None else sql + def _attach_correlation_id(self, sql: str) -> str: if self.ATTACH_CORRELATION_ID and self.correlation_id: return f"/* {self.correlation_id} */ {sql}" diff --git a/sqlmesh/core/query_execution.py b/sqlmesh/core/query_execution.py new file mode 100644 index 0000000000..f7e6f3ade9 --- /dev/null +++ b/sqlmesh/core/query_execution.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Immutable, engine-neutral context for SQL query execution.""" + +from __future__ import annotations + +import contextlib +import contextvars +import dataclasses +import datetime as dt +import json +import logging +import typing as t +from dataclasses import dataclass +from enum import Enum +from functools import cached_property + +from sqlmesh.utils import CorrelationId, JobType + +logger = logging.getLogger(__name__) + +MAX_QUERY_EXECUTION_IDENTIFIER_BYTES = 1024 +MAX_QUERY_EXECUTION_CONTEXT_BYTES = 8192 +QUERY_EXECUTION_CONTEXT_PREFIX = "/* QUERY_CONTEXT: " + + +class QueryExecutionPhase(str, Enum): + EXECUTE = "execute" + MATERIALIZE = "materialize" + AUDIT = "audit" + DDL = "ddl" + CLEANUP = "cleanup" + METADATA = "metadata" + UNKNOWN = "unknown" + + +class QueryExecutionOperation(str, Enum): + EVALUATE = "evaluate" + CREATE = "create" + MIGRATE = "migrate" + PROMOTE = "promote" + DEMOTE = "demote" + CLEANUP = "cleanup" + AUDIT = "audit" + UNKNOWN = "unknown" + + +class QueryExecutionRole(str, Enum): + QUERY = "query" + CREATE_TABLE = "create_table" + INSERT = "insert" + MERGE = "merge" + REPLACE = "replace" + AUDIT = "audit" + DROP = "drop" + INSPECT = "inspect" + PRE_STATEMENT = "pre_statement" + POST_STATEMENT = "post_statement" + PROMOTE = "promote" + DEMOTE = "demote" + + +def _validate_identifier(name: str, value: t.Any) -> None: + if not isinstance(value, str): + raise TypeError(f"{name} must be a string") + if not value: + raise ValueError(f"{name} must not be empty") + + +def _validate_datetime(name: str, value: t.Optional[dt.datetime]) -> None: + if value is None: + return + if not isinstance(value, dt.datetime): + raise TypeError(f"{name} must be a datetime or None") + if value.tzinfo is None or value.utcoffset() is None: + raise ValueError(f"{name} must be timezone-aware") + + +@dataclass(frozen=True) +class ModelExecutionContext: + evaluation_id: str + model_name: str + snapshot_id: str + interval_start: t.Optional[dt.datetime] = None + interval_end: t.Optional[dt.datetime] = None + batch_index: t.Optional[int] = None + operation_type: QueryExecutionOperation = QueryExecutionOperation.EVALUATE + + def __post_init__(self) -> None: + _validate_identifier("evaluation_id", self.evaluation_id) + _validate_identifier("model_name", self.model_name) + _validate_identifier("snapshot_id", self.snapshot_id) + if not isinstance(self.operation_type, QueryExecutionOperation): + raise TypeError("operation_type must be a QueryExecutionOperation") + _validate_datetime("interval_start", self.interval_start) + _validate_datetime("interval_end", self.interval_end) + if self.interval_start is not None: + object.__setattr__( + self, "interval_start", self.interval_start.astimezone(dt.timezone.utc) + ) + if self.interval_end is not None: + object.__setattr__(self, "interval_end", self.interval_end.astimezone(dt.timezone.utc)) + if self.interval_start is not None and self.interval_end is not None: + if self.interval_start >= self.interval_end: + raise ValueError("interval_start must be earlier than interval_end") + if self.batch_index is not None: + if type(self.batch_index) is not int: + raise TypeError("batch_index must be an integer or None") + if self.batch_index < 0: + raise ValueError("batch_index must be nonnegative") + + +@dataclass(frozen=True) +class QueryExecutionContext: + parent: t.Optional[CorrelationId] = None + evaluation: t.Optional[ModelExecutionContext] = None + phase: QueryExecutionPhase = QueryExecutionPhase.UNKNOWN + role: QueryExecutionRole = QueryExecutionRole.QUERY + schema_version: int = 1 + + @cached_property + def _serialized(self) -> t.Optional[str]: + """Return the cached transport representation derived from this value.""" + return _serialize_query_execution_context(self) + + def __post_init__(self) -> None: + if self.parent is not None: + if not isinstance(self.parent, CorrelationId): + raise TypeError("parent must be a CorrelationId or None") + if not isinstance(self.parent.job_type, JobType): + raise TypeError("parent.job_type must be a JobType") + if self.parent.job_type not in (JobType.PLAN, JobType.RUN): + raise ValueError("parent.job_type must be a plan or run job") + _validate_identifier("parent.job_id", self.parent.job_id) + if self.evaluation is not None and not isinstance(self.evaluation, ModelExecutionContext): + raise TypeError("evaluation must be a ModelExecutionContext or None") + if not isinstance(self.phase, QueryExecutionPhase): + raise TypeError("phase must be a QueryExecutionPhase") + if not isinstance(self.role, QueryExecutionRole): + raise TypeError("role must be a QueryExecutionRole") + if type(self.schema_version) is not int: + raise TypeError("schema_version must be an integer") + if self.schema_version != 1: + raise ValueError("schema_version is fixed at 1") + + +_current_query_execution_context: contextvars.ContextVar[t.Optional[QueryExecutionContext]] = ( + contextvars.ContextVar("query_execution_context", default=None) +) +_query_execution_context_transport_active: contextvars.ContextVar[bool] = contextvars.ContextVar( + "query_execution_context_transport_active", default=False +) + + +def get_query_execution_context() -> t.Optional[QueryExecutionContext]: + """Return the query execution context active in this context, if any.""" + return _current_query_execution_context.get() + + +def _query_execution_context_transport_is_active() -> bool: + return _query_execution_context_transport_active.get() + + +@contextlib.contextmanager +def _query_execution_context_transport_scope() -> t.Iterator[None]: + token = _query_execution_context_transport_active.set(True) + try: + yield + finally: + _query_execution_context_transport_active.reset(token) + + +@contextlib.contextmanager +def query_execution_scope(context: QueryExecutionContext) -> t.Iterator[None]: + """Set a query execution context for this scope and restore it on exit.""" + if not isinstance(context, QueryExecutionContext): + raise TypeError("context must be a QueryExecutionContext") + token = _current_query_execution_context.set(context) + try: + yield + finally: + _current_query_execution_context.reset(token) + + +@contextlib.contextmanager +def query_phase_scope(phase: QueryExecutionPhase, role: QueryExecutionRole) -> t.Iterator[None]: + """Temporarily derive the active context with a new phase and query role.""" + if not isinstance(phase, QueryExecutionPhase): + raise TypeError("phase must be a QueryExecutionPhase") + if not isinstance(role, QueryExecutionRole): + raise TypeError("role must be a QueryExecutionRole") + current = get_query_execution_context() + if current is None: + yield + return + token = _current_query_execution_context.set( + dataclasses.replace(current, phase=phase, role=role) + ) + try: + yield + finally: + _current_query_execution_context.reset(token) + + +def serialize_query_execution_context( + context: QueryExecutionContext, +) -> t.Optional[str]: + """Serialize a context as one complete, bounded SQL comment. + + The transport has no truncation mode: a context with overlong identifiers + or an oversized encoded payload is omitted and logged without its contents. + """ + if not isinstance(context, QueryExecutionContext): + raise TypeError("context must be a QueryExecutionContext") + + return context._serialized + + +def _serialize_query_execution_context( + context: QueryExecutionContext, +) -> t.Optional[str]: + payload: t.Dict[str, t.Any] = {"schema_version": context.schema_version, "emitter": "sqlmesh"} + if context.parent is not None: + payload["parent"] = { + "system": "sqlmesh", + "type": "plan" if context.parent.job_type is JobType.PLAN else "run", + "id": context.parent.job_id, + } + if context.evaluation is not None: + evaluation = context.evaluation + payload["workload"] = { + "system": "sqlmesh", + "type": "model", + "name": evaluation.model_name, + "version": evaluation.snapshot_id, + } + payload["operation"] = { + "id": evaluation.evaluation_id, + "type": evaluation.operation_type.value, + } + if evaluation.interval_start is not None: + payload["operation"]["interval_start"] = evaluation.interval_start.astimezone( + dt.timezone.utc + ).isoformat() + if evaluation.interval_end is not None: + payload["operation"]["interval_end"] = evaluation.interval_end.astimezone( + dt.timezone.utc + ).isoformat() + if evaluation.batch_index is not None: + payload["operation"]["batch_index"] = evaluation.batch_index + payload["query"] = {"phase": context.phase.value, "role": context.role.value} + + identifiers = [] + if context.parent is not None: + identifiers.append(context.parent.job_id) + if context.evaluation is not None: + identifiers.extend( + ( + context.evaluation.evaluation_id, + context.evaluation.model_name, + context.evaluation.snapshot_id, + ) + ) + identifier_sizes: t.List[int] = [] + has_unencodable_identifier = False + for value in identifiers: + try: + identifier_sizes.append(len(value.encode("utf-8"))) + except UnicodeEncodeError: + has_unencodable_identifier = True + if has_unencodable_identifier or any( + size > MAX_QUERY_EXECUTION_IDENTIFIER_BYTES for size in identifier_sizes + ): + if has_unencodable_identifier: + logger.warning( + "Omitting SQLMesh query execution context: identifier cannot be encoded as UTF-8 " + "(identifier transport limit %s UTF-8 bytes)", + MAX_QUERY_EXECUTION_IDENTIFIER_BYTES, + ) + else: + logger.warning( + "Omitting SQLMesh query execution context: identifier is %s UTF-8 bytes " + "(transport limit %s UTF-8 bytes)", + max(identifier_sizes), + MAX_QUERY_EXECUTION_IDENTIFIER_BYTES, + ) + return None + + # JSON escapes controls and quotes; escaping slash prevents either comment + # delimiter in an identifier from affecting the surrounding SQL comment. + encoded = json.dumps(payload, ensure_ascii=True, separators=(",", ":")).replace("/", r"\u002f") + comment = f"{QUERY_EXECUTION_CONTEXT_PREFIX}{encoded} */ " + comment_bytes = len(comment.encode("utf-8")) + if comment_bytes > MAX_QUERY_EXECUTION_CONTEXT_BYTES: + logger.warning( + "Omitting SQLMesh query execution context: comment is %s UTF-8 bytes " + "(transport limit %s UTF-8 bytes)", + comment_bytes, + MAX_QUERY_EXECUTION_CONTEXT_BYTES, + ) + return None + return comment diff --git a/tests/core/engine_adapter/test_query_execution_context.py b/tests/core/engine_adapter/test_query_execution_context.py new file mode 100644 index 0000000000..9a00e01ee3 --- /dev/null +++ b/tests/core/engine_adapter/test_query_execution_context.py @@ -0,0 +1,454 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +from concurrent.futures import ThreadPoolExecutor +from contextvars import ContextVar +from threading import Barrier +from typing import Any, Callable, Optional + +import pytest +from sqlglot import parse_one + +from sqlmesh.core.config.connection import DuckDBConnectionConfig +from sqlmesh.core.engine_adapter import EngineAdapter +from sqlmesh.core.engine_adapter.athena import AthenaEngineAdapter +import sqlmesh.core.engine_adapter.base as engine_adapter_base +from sqlmesh.core.query_execution import ( + ModelExecutionContext, + QueryExecutionContext, + QueryExecutionPhase, + QueryExecutionRole, + get_query_execution_context, +) +from sqlmesh.utils import CorrelationId + + +def _context(evaluation_id: str = "eval-1") -> QueryExecutionContext: + return QueryExecutionContext( + evaluation=ModelExecutionContext( + evaluation_id=evaluation_id, + model_name="analytics.orders", + snapshot_id="snapshot-3", + ), + phase=QueryExecutionPhase.MATERIALIZE, + role=QueryExecutionRole.INSERT, + ) + + +def _executed_sql(adapter: Any) -> list[str]: + return [call.args[0] for call in adapter.cursor.execute.call_args_list] + + +def _query_execution_payload(sql: str) -> Optional[dict[str, Any]]: + prefix = "/* QUERY_CONTEXT: " + start = sql.find(prefix) + if start == -1: + return None + end = sql.index(" */ ", start) + return json.loads(sql[start + len(prefix) : end]) + + +def test_execute_attaches_provider_context_before_legacy_parent( + make_mocked_engine_adapter: Callable[..., EngineAdapter], +): + parent = CorrelationId.from_plan_id("plan-8") + adapter = make_mocked_engine_adapter( + EngineAdapter, + query_execution_context=True, + query_execution_context_provider=lambda: _context(), + correlation_id=parent, + ) + + adapter.execute("SELECT 1") + + assert _executed_sql(adapter) == [ + '/* SQLMESH_PLAN: plan-8 */ /* QUERY_CONTEXT: {"schema_version":1,"emitter":"sqlmesh",' + '"parent":{"system":"sqlmesh","type":"plan","id":"plan-8"},' + '"workload":{"system":"sqlmesh","type":"model","name":"analytics.orders",' + '"version":"snapshot-3"},"operation":{"id":"eval-1","type":"evaluate"},' + '"query":{"phase":"materialize","role":"insert"}} */ SELECT 1' + ] + + +def test_execute_uses_parent_only_context_when_provider_returns_none( + make_mocked_engine_adapter: Callable[..., EngineAdapter], +): + parent = CorrelationId.from_run_id("run-7") + adapter = make_mocked_engine_adapter( + EngineAdapter, + query_execution_context=True, + query_execution_context_provider=lambda: None, + correlation_id=parent, + ) + + adapter.execute("SELECT 1") + + assert _executed_sql(adapter) == [ + '/* SQLMESH_RUN: run-7 */ /* QUERY_CONTEXT: {"schema_version":1,"emitter":"sqlmesh",' + '"parent":{"system":"sqlmesh","type":"run","id":"run-7"},' + '"query":{"phase":"unknown","role":"query"}} */ SELECT 1' + ] + + +def test_adapter_correlation_marker_overrides_and_clears_provider_parent( + make_mocked_engine_adapter: Callable[..., EngineAdapter], +): + stale_parent = CorrelationId.from_plan_id("stale-plan") + first_parent = CorrelationId.from_plan_id("plan-8") + second_parent = CorrelationId.from_run_id("run-9") + context = QueryExecutionContext( + parent=stale_parent, + evaluation=_context().evaluation, + phase=QueryExecutionPhase.MATERIALIZE, + role=QueryExecutionRole.INSERT, + ) + adapter = make_mocked_engine_adapter( + EngineAdapter, + query_execution_context=True, + query_execution_context_provider=lambda: context, + correlation_id=first_parent, + ) + + adapter.execute("SELECT 1") + adapter.correlation_id = second_parent + adapter.execute("SELECT 2") + adapter.correlation_id = None + adapter.execute("SELECT 3") + + sql = _executed_sql(adapter) + assert [query.startswith("/* SQLMESH_") for query in sql] == [True, True, False] + assert sql[0].startswith("/* SQLMESH_PLAN: plan-8 */ /* QUERY_CONTEXT:") + assert sql[1].startswith("/* SQLMESH_RUN: run-9 */ /* QUERY_CONTEXT:") + payloads = [_query_execution_payload(query) for query in sql] + assert [ + payload["parent"]["id"] if payload is not None and "parent" in payload else None + for payload in payloads + ] == ["plan-8", "run-9", None] + assert [ + payload["parent"]["type"] if payload is not None and "parent" in payload else None + for payload in payloads + ] == ["plan", "run", None] + + +@pytest.mark.parametrize("has_provider", [False, True], ids=["provider-absent", "provider-valid"]) +def test_invalid_adapter_parent_omits_rich_context_and_keeps_legacy_comment( + has_provider: bool, + make_mocked_engine_adapter: Callable[..., EngineAdapter], +): + provider = (lambda: _context()) if has_provider else None + adapter = make_mocked_engine_adapter( + EngineAdapter, + query_execution_context=True, + query_execution_context_provider=provider, + correlation_id=CorrelationId.from_plan_id(""), + ) + + adapter.execute("SELECT 1") + + assert _executed_sql(adapter) == ["/* SQLMESH_PLAN: */ SELECT 1"] + + +def test_disabled_transport_never_calls_provider_or_changes_sql( + make_mocked_engine_adapter: Callable[..., EngineAdapter], +): + calls = [] + + def provider() -> Optional[QueryExecutionContext]: + calls.append(True) + raise AssertionError("disabled provider was called") + + parent = CorrelationId.from_plan_id("plan-8") + adapter = make_mocked_engine_adapter( + EngineAdapter, + query_execution_context=False, + query_execution_context_provider=provider, + correlation_id=parent, + ) + + adapter.execute("SELECT 1") + + assert _executed_sql(adapter) == [f"/* {parent} */ SELECT 1"] + assert calls == [] + + +def test_adapter_clone_inherits_provider_and_allows_explicit_clear( + make_mocked_engine_adapter: Callable[..., EngineAdapter], +): + adapter = make_mocked_engine_adapter( + EngineAdapter, + query_execution_context=True, + query_execution_context_provider=lambda: _context(), + ) + + inherited = adapter.with_settings() + cleared = adapter.with_settings(query_execution_context_provider=None) + inherited.execute("SELECT 1") + cleared.execute("SELECT 2") + + sql = _executed_sql(inherited) + assert "QUERY_CONTEXT" in sql[0] + assert sql[1] == "SELECT 2" + assert inherited.query_execution_context_provider is adapter.query_execution_context_provider + assert cleared.query_execution_context_provider is None + + +def test_connection_config_propagates_flag_without_extra_driver_config(): + adapter = DuckDBConnectionConfig(query_execution_context=True).create_engine_adapter() + + assert adapter.query_execution_context is True + assert "query_execution_context" not in adapter._extra_config + assert adapter.query_execution_context_provider is None + + +def test_legacy_adapter_constructor_accepts_connection_and_clone_settings(): + class LegacyAdapter(EngineAdapter): + def __init__(self, *args: Any, **kwargs: Any): + assert "query_execution_context" not in kwargs + assert "query_execution_context_provider" not in kwargs + super().__init__(*args, **kwargs) + + class LegacyConnectionConfig(DuckDBConnectionConfig): + @property + def _engine_adapter(self) -> type[EngineAdapter]: + return LegacyAdapter + + adapter = LegacyConnectionConfig(query_execution_context=True).create_engine_adapter() + assert adapter.query_execution_context is True + assert adapter.with_settings().query_execution_context is True + assert adapter.with_settings(query_execution_context=False).query_execution_context is False + assert adapter.with_settings( + query_execution_context_provider=lambda: _context() + ).query_execution_context_provider + + +def test_unsupported_adapter_skips_provider_and_warns_once( + caplog, + make_mocked_engine_adapter: Callable[..., EngineAdapter], +): + calls = [] + + def provider() -> QueryExecutionContext: + calls.append(True) + return _context() + + adapter = make_mocked_engine_adapter( + AthenaEngineAdapter, + query_execution_context=True, + query_execution_context_provider=provider, + ) + + with caplog.at_level("WARNING"): + adapter.execute("SELECT 1") + adapter.execute("SELECT 2") + + assert _executed_sql(adapter) == ["SELECT 1", "SELECT 2"] + assert calls == [] + assert caplog.text.count("does not support generic SQL comment transport") == 1 + + +def test_provider_failure_is_omitted_and_later_call_recovers( + caplog, + make_mocked_engine_adapter: Callable[..., EngineAdapter], +): + attempts = [] + + def provider() -> QueryExecutionContext: + attempts.append(True) + if len(attempts) == 1: + raise RuntimeError("sensitive warehouse identifier") + return _context() + + parent = CorrelationId.from_run_id("run-7") + adapter = make_mocked_engine_adapter( + EngineAdapter, + query_execution_context=True, + query_execution_context_provider=provider, + correlation_id=parent, + ) + + with caplog.at_level("WARNING"): + adapter.execute("SELECT 1") + adapter.execute("SELECT 2") + + sql = _executed_sql(adapter) + assert sql[0] == "/* SQLMESH_RUN: run-7 */ SELECT 1" + assert sql[1].startswith("/* SQLMESH_RUN: run-7 */ /* QUERY_CONTEXT:") + payload = _query_execution_payload(sql[1]) + assert payload is not None + assert payload["parent"] == {"system": "sqlmesh", "type": "run", "id": "run-7"} + assert payload["operation"]["id"] == "eval-1" + assert attempts == [True, True] + assert "provider failed" in caplog.text + assert "sensitive warehouse identifier" not in caplog.text + + +def test_serialization_failure_is_omitted_and_later_call_recovers( + caplog, + make_mocked_engine_adapter: Callable[..., EngineAdapter], + monkeypatch, +): + real_serializer = engine_adapter_base.serialize_query_execution_context + attempts = [] + + def fail_once(context: QueryExecutionContext) -> Optional[str]: + attempts.append(context) + if len(attempts) == 1: + raise RuntimeError("sensitive serializer detail") + return real_serializer(context) + + monkeypatch.setattr(engine_adapter_base, "serialize_query_execution_context", fail_once) + adapter = make_mocked_engine_adapter( + EngineAdapter, + query_execution_context=True, + query_execution_context_provider=lambda: _context(), + correlation_id=CorrelationId.from_run_id("run-7"), + ) + + with caplog.at_level("WARNING"): + adapter.execute("SELECT 1") + adapter.execute("SELECT 2") + + sql = _executed_sql(adapter) + assert sql[0] == "/* SQLMESH_RUN: run-7 */ SELECT 1" + assert sql[1].startswith("/* SQLMESH_RUN: run-7 */ /* QUERY_CONTEXT:") + payload = _query_execution_payload(sql[1]) + assert payload is not None + assert payload["operation"]["id"] == "eval-1" + assert len(attempts) == 2 + assert "serialization failed" in caplog.text + assert "sensitive serializer detail" not in caplog.text + + +def test_provider_runs_once_for_each_statement( + make_mocked_engine_adapter: Callable[..., EngineAdapter], +): + attempts: list[QueryExecutionContext] = [] + + def provider() -> QueryExecutionContext: + context = _context(f"eval-{len(attempts)}") + attempts.append(context) + return context + + adapter = make_mocked_engine_adapter( + EngineAdapter, + query_execution_context=True, + query_execution_context_provider=provider, + ) + + adapter.execute([parse_one("SELECT 1"), parse_one("SELECT 2")]) + + assert len(attempts) == 2 + payloads = [_query_execution_payload(query) for query in _executed_sql(adapter)] + assert all(payload is not None for payload in payloads) + assert [payload["operation"]["id"] for payload in payloads if payload is not None] == [ + "eval-0", + "eval-1", + ] + + +def test_provider_returning_invalid_type_omits_comment( + caplog, + make_mocked_engine_adapter: Callable[..., EngineAdapter], +): + adapter = make_mocked_engine_adapter( + EngineAdapter, + query_execution_context=True, + ) + + def invalid_provider() -> Any: + return object() + + adapter.query_execution_context_provider = invalid_provider + + with caplog.at_level("WARNING"): + adapter.execute("SELECT 1") + + assert _executed_sql(adapter) == ["SELECT 1"] + assert "invalid context" in caplog.text + + +def test_provider_cancellation_resets_transport_guard_and_scope( + make_mocked_engine_adapter: Callable[..., EngineAdapter], +): + attempts = [] + + def provider() -> QueryExecutionContext: + attempts.append(True) + if len(attempts) == 1: + raise KeyboardInterrupt() + return _context("eval-after-cancellation") + + adapter = make_mocked_engine_adapter( + EngineAdapter, + query_execution_context=True, + query_execution_context_provider=provider, + ) + with pytest.raises(KeyboardInterrupt): + adapter.execute("SELECT 1") + + adapter.execute("SELECT 2") + + assert get_query_execution_context() is None + assert attempts == [True, True] + sql = _executed_sql(adapter) + assert len(sql) == 1 + payload = _query_execution_payload(sql[0]) + assert payload is not None + assert payload["operation"]["id"] == "eval-after-cancellation" + + +def test_provider_reentry_skips_nested_rich_context_across_adapters( + make_mocked_engine_adapter: Callable[..., EngineAdapter], +): + nested = make_mocked_engine_adapter(EngineAdapter, query_execution_context=True) + attempts = [] + adapter = make_mocked_engine_adapter(EngineAdapter, query_execution_context=True) + + def provider() -> QueryExecutionContext: + attempts.append(True) + nested.execute("SELECT 2") + return _context() + + nested.query_execution_context_provider = provider + adapter.query_execution_context_provider = provider + adapter.execute("SELECT 1") + + assert _executed_sql(nested) == ["SELECT 2"] + assert "QUERY_CONTEXT" in _executed_sql(adapter)[0] + assert attempts == [True] + + +def test_concurrent_provider_calls_keep_their_worker_contexts( + make_mocked_engine_adapter: Callable[..., EngineAdapter], +): + current: ContextVar[Optional[QueryExecutionContext]] = ContextVar( + "test_query_execution_context", default=None + ) + barrier = Barrier(2) + adapter = make_mocked_engine_adapter( + EngineAdapter, + query_execution_context=True, + query_execution_context_provider=current.get, + ) + + def execute(evaluation_id: str) -> None: + context = _context(evaluation_id) + token = current.set(context) + try: + barrier.wait(timeout=5) + adapter.execute("SELECT 1") + finally: + current.reset(token) + + with ThreadPoolExecutor(max_workers=2) as pool: + futures = [pool.submit(execute, value) for value in ("eval-a", "eval-b")] + for future in futures: + future.result(timeout=10) + + sql = _executed_sql(adapter) + assert len(sql) == 2 + payloads = [_query_execution_payload(query) for query in sql] + assert all(payload is not None for payload in payloads) + assert {payload["operation"]["id"] for payload in payloads if payload is not None} == { + "eval-a", + "eval-b", + } diff --git a/tests/core/integration/test_query_execution_context.py b/tests/core/integration/test_query_execution_context.py new file mode 100644 index 0000000000..fdc0eb6c6b --- /dev/null +++ b/tests/core/integration/test_query_execution_context.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import datetime as dt +import json +import re +import shutil +import typing as t +from pathlib import Path + +from examples.query_execution_context.observer import QueryContextObserver +from sqlmesh.core.config import Config, ModelDefaultsConfig +from sqlmesh.core.config.connection import DuckDBConnectionConfig +from sqlmesh.core.config.gateway import GatewayConfig +from sqlmesh.core.context import Context +from sqlmesh.core.engine_adapter.base import EngineAdapter +from sqlmesh.core.execution_observation import observer_scope +from sqlmesh.core.state_sync import CachingStateSync, EngineAdapterStateSync + + +def _payload(sql: str) -> t.Optional[t.Dict[str, t.Any]]: + match = re.search(r"/\* QUERY_CONTEXT: (.*?) \*/", sql) + return json.loads(match[1]) if match else None + + +def _new_context(root: Path, enabled: bool) -> Context: + models = Path(__file__).parents[3] / "examples/query_execution_context/models" + shutil.copytree(models, root / "models") + config = Config( + gateways={ + "local": GatewayConfig( + connection=DuckDBConnectionConfig( + database=str(root / "warehouse.duckdb"), + concurrent_tasks=2, + query_execution_context=enabled, + ) + ) + }, + default_gateway="local", + model_defaults=ModelDefaultsConfig(dialect="duckdb"), + ) + return Context(paths=root, config=config, load=False) + + +def _batch_ids( + sqls: t.List[str], parent_type: str, phase: str +) -> t.Dict[t.Tuple[str, str, int], str]: + batches: t.Dict[t.Tuple[str, str, int], str] = {} + for sql in sqls: + payload = _payload(sql) + if not payload or not payload.get("operation") or payload["query"]["phase"] != phase: + continue + assert payload["schema_version"] == 1 + assert payload["emitter"] == "sqlmesh" + assert payload["parent"]["type"] == parent_type + assert payload["parent"]["system"] == "sqlmesh" + assert sql.startswith(f"/* SQLMESH_{parent_type.upper()}:") + assert payload["workload"]["system"] == "sqlmesh" + assert payload["workload"]["type"] == "model" + assert payload["workload"]["version"] + operation = payload["operation"] + assert operation["type"] == "evaluate" + key = ( + payload["workload"]["name"], + operation["interval_start"], + operation["batch_index"], + ) + if key in batches: + assert batches[key] == operation["id"] + batches[key] = operation["id"] + return batches + + +def test_provider_connects_native_events_to_warehouse_history(tmp_path, monkeypatch): + submitted = [] + original_execute = EngineAdapter._execute + + def capture(adapter, sql, *args, **kwargs): + submitted.append((adapter, sql)) + return original_execute(adapter, sql, *args, **kwargs) + + monkeypatch.setattr(EngineAdapter, "_execute", capture) + results = [] + for enabled in (False, True): + context = _new_context(tmp_path / str(enabled), enabled) + events = [] + + def observer(event, facts, error): + current = consumer.current_query_context() + if facts["kind"] == "model": + events.append( + (event, facts["snapshot"].name, current.evaluation.evaluation_id, error) + ) + + consumer = QueryContextObserver(observer) + for adapter in context.engine_adapters.values(): + adapter.query_execution_context_provider = consumer.current_query_context + context.load() + try: + hashes = {m.name: (m.data_hash, m.metadata_hash) for m in context.models.values()} + submitted.clear() + with observer_scope(consumer): + plan = context.plan( + "dev", start="2024-01-01", end="2024-01-02", no_prompts=True, skip_tests=True + ) + context.apply(plan) + assert consumer.current_query_context() is None + plan_sql = [sql for _, sql in submitted] + submitted.clear() + assert context.run("dev", start="2024-01-03", end="2024-01-03").is_success + assert consumer.current_query_context() is None + run_sql = [sql for _, sql in submitted] + state = context.state_sync + assert isinstance(state, CachingStateSync) + assert isinstance(state.state_sync, EngineAdapterStateSync) + state_adapter = state.state_sync.engine_adapter + assert any(adapter is state_adapter for adapter, _ in submitted) + assert all( + _payload(sql) is None for adapter, sql in submitted if adapter is state_adapter + ) + rows = context.engine_adapter.fetchall( + f"SELECT ds, id FROM {context.resolve_table('e2e.events')} ORDER BY ds" + ) + assert rows == [(dt.date(2024, 1, day), 1) for day in (1, 2, 3)] + results.append((hashes, rows)) + if enabled: + materialize = _batch_ids(plan_sql, "plan", "materialize") + assert _batch_ids(plan_sql, "plan", "audit") == materialize + assert set(materialize) == { + ('"warehouse"."e2e"."events"', "2024-01-01T00:00:00+00:00", 0), + ('"warehouse"."e2e"."events"', "2024-01-02T00:00:00+00:00", 1), + ('"warehouse"."e2e"."aux_events"', "2024-01-01T00:00:00+00:00", 0), + ('"warehouse"."e2e"."aux_events"', "2024-01-02T00:00:00+00:00", 1), + } + later = _batch_ids(run_sql, "run", "materialize") + assert _batch_ids(run_sql, "run", "audit") == later + assert set(later) == { + ('"warehouse"."e2e"."events"', "2024-01-03T00:00:00+00:00", 0), + ('"warehouse"."e2e"."aux_events"', "2024-01-03T00:00:00+00:00", 0), + } + all_ids = set(materialize.values()) | set(later.values()) + assert len(all_ids) == 6 + assert {operation_id for _, _, operation_id, _ in events} == all_ids + assert [event for event, _, _, _ in events].count("start") == 6 + assert [event for event, _, _, _ in events].count("finish") == 6 + assert all(error is None for _, _, _, error in events) + else: + assert all(_payload(sql) is None for sql in plan_sql + run_sql) + assert any(sql.startswith("/* SQLMESH_PLAN:") for sql in plan_sql) + finally: + context.close() + assert results[0] == results[1] diff --git a/tests/core/test_config.py b/tests/core/test_config.py index 12f461913a..3370b1d153 100644 --- a/tests/core/test_config.py +++ b/tests/core/test_config.py @@ -532,6 +532,7 @@ def test_connection_config_serialization(): "extensions": [], "pre_ping": False, "pretty_sql": False, + "query_execution_context": False, "shared_connection": True, "connector_config": {}, "secrets": [], @@ -545,6 +546,7 @@ def test_connection_config_serialization(): "extensions": [], "pre_ping": False, "pretty_sql": False, + "query_execution_context": False, "shared_connection": True, "connector_config": {}, "secrets": [], diff --git a/tests/core/test_query_execution.py b/tests/core/test_query_execution.py new file mode 100644 index 0000000000..b2fee57182 --- /dev/null +++ b/tests/core/test_query_execution.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +from datetime import datetime, timezone + +import pytest + +from sqlmesh.core.query_execution import ( + MAX_QUERY_EXECUTION_CONTEXT_BYTES, + MAX_QUERY_EXECUTION_IDENTIFIER_BYTES, + ModelExecutionContext, + QueryExecutionContext, + QueryExecutionOperation, + QueryExecutionPhase, + QueryExecutionRole, + get_query_execution_context, + query_execution_scope, + query_phase_scope, + serialize_query_execution_context, +) +from sqlmesh.utils import CorrelationId + + +def _context( + *, model_name: str = "analytics.orders", evaluation_id: str = "eval-1" +) -> QueryExecutionContext: + return QueryExecutionContext( + parent=CorrelationId.from_run_id("run-7"), + evaluation=ModelExecutionContext( + evaluation_id=evaluation_id, + model_name=model_name, + snapshot_id="snapshot-3", + interval_start=datetime(2026, 1, 1, tzinfo=timezone.utc), + interval_end=datetime(2026, 1, 2, tzinfo=timezone.utc), + batch_index=2, + ), + phase=QueryExecutionPhase.MATERIALIZE, + role=QueryExecutionRole.INSERT, + ) + + +def test_serializer_emits_version_one_context_with_utc_interval(): + context = _context() + + actual = serialize_query_execution_context(context) + + assert actual == ( + '/* QUERY_CONTEXT: {"schema_version":1,"emitter":"sqlmesh",' + '"parent":{"system":"sqlmesh","type":"run","id":"run-7"},' + '"workload":{"system":"sqlmesh","type":"model","name":"analytics.orders",' + '"version":"snapshot-3"},"operation":{"id":"eval-1","type":"evaluate",' + '"interval_start":"2026-01-01T00:00:00+00:00",' + '"interval_end":"2026-01-02T00:00:00+00:00","batch_index":2},' + '"query":{"phase":"materialize","role":"insert"}} */ ' + ) + + +def test_serializer_omits_unavailable_context_fields(): + actual = serialize_query_execution_context( + QueryExecutionContext(parent=CorrelationId.from_plan_id("plan-1")) + ) + + assert actual == ( + '/* QUERY_CONTEXT: {"schema_version":1,"emitter":"sqlmesh",' + '"parent":{"system":"sqlmesh","type":"plan","id":"plan-1"},' + '"query":{"phase":"unknown","role":"query"}} */ ' + ) + + +@pytest.mark.parametrize("name", ['analytics."a/*b"', 'analytics."a*/b"']) +def test_serializer_escapes_sql_comment_delimiters_and_round_trips(name: str): + import json + + comment = serialize_query_execution_context(_context(model_name=name)) + assert comment is not None + payload = comment.removeprefix("/* QUERY_CONTEXT: ").removesuffix(" */ ") + + assert "/*" not in payload + assert "*/" not in payload + assert json.loads(payload)["workload"]["name"] == name + + +def test_scope_restores_nested_and_exceptional_contexts(): + outer = _context() + inner = _context(evaluation_id="eval-2") + with query_execution_scope(outer): + assert get_query_execution_context() is outer + with query_phase_scope(QueryExecutionPhase.AUDIT, QueryExecutionRole.AUDIT): + audit = get_query_execution_context() + assert audit is not None + assert audit.evaluation == outer.evaluation + assert audit.phase is QueryExecutionPhase.AUDIT + assert get_query_execution_context() is outer + with pytest.raises(RuntimeError, match="failed"): + with query_execution_scope(inner): + raise RuntimeError("failed") + assert get_query_execution_context() is outer + + assert get_query_execution_context() is None + + +def test_serializer_accepts_identifier_at_utf8_byte_limit(): + context = _context(model_name="m" * MAX_QUERY_EXECUTION_IDENTIFIER_BYTES) + + comment = serialize_query_execution_context(context) + + assert comment is not None + assert len(comment.encode("utf-8")) <= MAX_QUERY_EXECUTION_CONTEXT_BYTES + + +def test_serializer_omits_identifier_over_utf8_byte_limit(caplog): + context = _context(model_name="m" * (MAX_QUERY_EXECUTION_IDENTIFIER_BYTES + 1)) + + assert serialize_query_execution_context(context) is None + + assert "identifier is 1025 UTF-8 bytes" in caplog.text + assert "analytics" not in caplog.text + + +def test_serializer_omits_unencodable_identifier(caplog): + context = _context(model_name="invalid-\ud800") + + assert serialize_query_execution_context(context) is None + + assert "identifier cannot be encoded as UTF-8" in caplog.text + assert "invalid-" not in caplog.text + + +def test_serializer_omits_complete_comment_over_byte_limit(caplog): + value = "😀" * (MAX_QUERY_EXECUTION_IDENTIFIER_BYTES // 4) + context = QueryExecutionContext( + evaluation=ModelExecutionContext( + evaluation_id=value, + model_name=value, + snapshot_id=value, + operation_type=QueryExecutionOperation.EVALUATE, + ) + ) + + assert serialize_query_execution_context(context) is None + + assert "comment is" in caplog.text + assert "8192 UTF-8 bytes" in caplog.text + assert "😀" not in caplog.text + + +def test_context_rejects_wrong_types_and_empty_identifiers(): + with pytest.raises(ValueError, match="must not be empty"): + ModelExecutionContext("", "model", "snapshot") + with pytest.raises(TypeError, match="must be an integer"): + ModelExecutionContext("eval", "model", "snapshot", batch_index=True) + with pytest.raises(ValueError, match="fixed at 1"): + QueryExecutionContext(schema_version=2) diff --git a/tests/core/test_query_execution_context_config.py b/tests/core/test_query_execution_context_config.py new file mode 100644 index 0000000000..09fd7afd1d --- /dev/null +++ b/tests/core/test_query_execution_context_config.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +from unittest.mock import Mock + +from sqlmesh.core.config.connection import DuckDBConnectionConfig +from sqlmesh.core.config.scheduler import BuiltInSchedulerConfig +from sqlmesh.core.query_execution import ModelExecutionContext, QueryExecutionContext + + +def test_cached_duckdb_adapter_honors_query_context_flag_in_config_order(tmp_path): + database = str(tmp_path / "shared.duckdb") + disabled = DuckDBConnectionConfig( + database=database, query_execution_context=False + ).create_engine_adapter() + enabled = DuckDBConnectionConfig( + database=database, query_execution_context=True + ).create_engine_adapter() + + assert disabled.query_execution_context is False + assert enabled.query_execution_context is True + assert enabled is not disabled + assert enabled._connection_pool is disabled._connection_pool + + +def test_cached_duckdb_adapter_does_not_leak_provider_to_default_config(tmp_path): + database = str(tmp_path / "shared.duckdb") + enabled = DuckDBConnectionConfig( + database=database, query_execution_context=True + ).create_engine_adapter() + provider = lambda: QueryExecutionContext( + evaluation=ModelExecutionContext("eval-1", "analytics.orders", "snapshot-3") + ) + enabled.query_execution_context_provider = provider + + disabled = DuckDBConnectionConfig(database=database).create_engine_adapter() + + assert enabled.query_execution_context is True + assert enabled.query_execution_context_provider is provider + assert disabled.query_execution_context is False + assert disabled.query_execution_context_provider is None + assert disabled is not enabled + assert disabled._connection_pool is enabled._connection_pool + + +def test_state_sync_clones_cached_warehouse_adapter_without_query_context(mocker, tmp_path): + warehouse = DuckDBConnectionConfig( + database=str(tmp_path / "warehouse.db"), + query_execution_context=True, + ) + warehouse_adapter = warehouse.create_engine_adapter() + warehouse_adapter.query_execution_context_provider = lambda: QueryExecutionContext( + evaluation=ModelExecutionContext("eval-1", "analytics.orders", "snapshot-3") + ) + + context = Mock() + context.config.get_state_connection.return_value = None + context.config.get_connection.return_value = warehouse + context.config.get_state_schema.return_value = None + context.connection_config = warehouse + context.cache_dir = tmp_path / "cache" + context.console = Mock() + state_sync = Mock() + state_sync_factory = mocker.patch( + "sqlmesh.core.config.scheduler.EngineAdapterStateSync", return_value=state_sync + ) + + actual = BuiltInSchedulerConfig().create_state_sync(context) + + state_adapter = state_sync_factory.call_args.args[0] + assert actual is state_sync + assert state_adapter is not warehouse_adapter + assert state_adapter.query_execution_context is False + assert state_adapter.query_execution_context_provider is None + assert warehouse_adapter.query_execution_context is True + assert warehouse_adapter.query_execution_context_provider is not None + + +def test_explicit_state_connection_can_enable_parent_query_context(mocker, tmp_path): + warehouse = DuckDBConnectionConfig( + database=str(tmp_path / "warehouse.db"), query_execution_context=True + ) + warehouse_adapter = warehouse.create_engine_adapter() + provider = lambda: QueryExecutionContext( + evaluation=ModelExecutionContext("eval-1", "analytics.orders", "snapshot-3") + ) + warehouse_adapter.query_execution_context_provider = provider + state_connection = DuckDBConnectionConfig( + database=str(tmp_path / "warehouse.db"), query_execution_context=True + ) + + context = Mock() + context.config.get_state_connection.return_value = state_connection + context.config.get_connection.return_value = warehouse + context.config.get_state_schema.return_value = None + context.connection_config = warehouse + context.cache_dir = tmp_path / "cache" + context.console = Mock() + state_sync = Mock() + state_sync_factory = mocker.patch( + "sqlmesh.core.config.scheduler.EngineAdapterStateSync", return_value=state_sync + ) + + BuiltInSchedulerConfig().create_state_sync(context) + + state_adapter = state_sync_factory.call_args.args[0] + assert state_adapter.query_execution_context is True + assert state_adapter.query_execution_context_provider is None + assert warehouse_adapter.query_execution_context_provider is provider diff --git a/tests/integrations/jupyter/test_magics.py b/tests/integrations/jupyter/test_magics.py index 8449a64b69..9fef5dacef 100644 --- a/tests/integrations/jupyter/test_magics.py +++ b/tests/integrations/jupyter/test_magics.py @@ -749,8 +749,8 @@ def test_info(notebook, sushi_context, convert_all_html_output_to_text, get_all_ "Models: 20", "Macros: 8", "", - "Connection:\n type: duckdb\n concurrent_tasks: 1\n register_comments: true\n pre_ping: false\n pretty_sql: false\n extensions: []\n connector_config: {}\n secrets: None\n filesystems: []\n shared_connection: true", - "Test Connection:\n type: duckdb\n concurrent_tasks: 1\n register_comments: true\n pre_ping: false\n pretty_sql: false\n extensions: []\n connector_config: {}\n secrets: None\n filesystems: []\n shared_connection: true", + "Connection:\n type: duckdb\n concurrent_tasks: 1\n register_comments: true\n pre_ping: false\n pretty_sql: false\n query_execution_context: false\n extensions: []\n connector_config: {}\n secrets: None\n filesystems: []\n shared_connection: true", + "Test Connection:\n type: duckdb\n concurrent_tasks: 1\n register_comments: true\n pre_ping: false\n pretty_sql: false\n query_execution_context: false\n extensions: []\n connector_config: {}\n secrets: None\n filesystems: []\n shared_connection: true", "Data warehouse connection succeeded", "State backend versions:", "Schema version: 0", @@ -761,8 +761,8 @@ def test_info(notebook, sushi_context, convert_all_html_output_to_text, get_all_ "
Models: 20
", "
Macros: 8
", "
",
-        '
Connection:  type: duckdb  concurrent_tasks: 1  register_comments: true  pre_ping: false  pretty_sql: false  extensions: []  connector_config: {}  secrets: None  filesystems: []  shared_connection: true
',
-        '
Test Connection:  type: duckdb  concurrent_tasks: 1  register_comments: true  pre_ping: false  pretty_sql: false  extensions: []  connector_config: {}  secrets: None  filesystems: []  shared_connection: true
',
+        '
Connection:  type: duckdb  concurrent_tasks: 1  register_comments: true  pre_ping: false  pretty_sql: false  query_execution_context: false  extensions: []  connector_config: {}  secrets: None  filesystems: []  shared_connection: true
',
+        '
Test Connection:  type: duckdb  concurrent_tasks: 1  register_comments: true  pre_ping: false  pretty_sql: false  query_execution_context: false  extensions: []  connector_config: {}  secrets: None  filesystems: []  shared_connection: true
',
         "
Data warehouse connection succeeded
",
         "
State backend versions:
",
         "
Schema version: 0
",

From 41f5d2741160d564d859f50b7963243584ad34c9 Mon Sep 17 00:00:00 2001
From: mday-io 
Date: Sat, 31 Oct 2026 19:20:04 -0400
Subject: [PATCH 2/4] fix: clean up cancelled query context transactions and
 config tests

Signed-off-by: mday-io 
---
 docs/concepts/query_execution_context.md      | 116 +++++++++--------
 docs/examples/query_execution_context.md      |  15 ++-
 sqlmesh/core/engine_adapter/base.py           |   4 +-
 tests/cli/test_cli.py                         |  11 +-
 .../test_query_execution_context.py           | 117 +++++++++++++++++-
 5 files changed, 185 insertions(+), 78 deletions(-)

diff --git a/docs/concepts/query_execution_context.md b/docs/concepts/query_execution_context.md
index 74e3c33c82..60ff83ff7b 100644
--- a/docs/concepts/query_execution_context.md
+++ b/docs/concepts/query_execution_context.md
@@ -4,9 +4,10 @@ SQLMesh can add structured metadata to warehouse statements so a query-history
 tool can connect a statement to the plan or run, model, and evaluation that issued
 it. The metadata is a versioned JSON object in a `QUERY_CONTEXT` SQL comment.
 
-Enable `query_execution_context` on a connection and give its engine adapter a
-provider that returns a `QueryExecutionContext`. The flag defaults to `false`.
-For YAML configuration, set the option under the gateway's `connection`:
+Enable `query_execution_context` on a connection and set a provider on its
+engine adapter. The provider returns a `QueryExecutionContext`. The flag defaults
+to `false`. For YAML configuration, set the option under the gateway's
+`connection`:
 
 ```yaml
 gateways:
@@ -16,21 +17,20 @@ gateways:
       query_execution_context: true
 ```
 
-The configuration flag enables transport. It does not create a model-context
-provider. Without a provider, SQLMesh adds parent metadata only when a plan or run
-correlation ID is available. SQLMesh calls the provider after rendering each
-statement and before executing it. The provider sees the current worker's context,
-so concurrent model evaluations can supply different IDs to a shared adapter.
+The flag enables context comments but does not infer model context. Without a
+provider, SQLMesh adds parent metadata only when a plan or run correlation ID is
+available. SQLMesh calls the provider after rendering each statement and before
+executing it. The call runs in the current worker's context, so concurrent model
+evaluations can return different IDs for a shared adapter.
 
-This provider API does not add automatic model tagging to a stock project.
-For a provider that reads native observer events, see the
+To build a provider from native observer events, see the
 [query execution context example](../examples/query_execution_context.md).
 
 ## Configure a provider
 
-For a Python `Context`, configure its warehouse adapters before loading the
-project. The `Context` clones these adapters when it creates its snapshot
-evaluator, so the provider and flag carry to those clones.
+Set the provider on warehouse adapters before loading a Python `Context`. When
+`Context` creates its snapshot evaluator, it clones those adapters and preserves
+the provider and flag.
 
 ```python
 from sqlmesh.core.context import Context
@@ -51,22 +51,19 @@ with observer_scope(observer):
     context.run("dev")
 ```
 
-The provider returns the immutable context currently active in that worker, or
-`None` when no context applies. An observer can set an evaluation ID when
-SQLMesh reports a model batch, then restore its prior context when that
-action ends. The sample observer composes with the observer returned by the
-console. It preserves that observer's facts and events.
+`QueryContextObserver` returns the immutable context active in its worker, or
+`None` when no context applies. It assigns an evaluation ID when SQLMesh reports
+a model batch and restores the previous context when that event finishes. It also
+forwards the same event, facts, and error to the observer returned by the console.
 
-Use the adapter instance that the evaluator will use. Adapters created with
-`with_settings` inherit the flag and provider. Passing
-`query_execution_context_provider=None` to `with_settings` clears the inherited
-provider. Setting the flag to `false` disables provider calls and comment output.
+Adapters created with `with_settings` inherit the flag and provider. Pass
+`query_execution_context_provider=None` to clear the inherited provider. Setting
+the flag to `false` disables provider calls and comment output.
 
-Connection settings also provide defaults to other adapters. Configure the
-warehouse adapters where you want model context in query history. The built-in
-scheduler can derive its state adapter from the warehouse settings, but it
-disables query context on that fallback. To annotate state SQL, configure the
-`state_connection` explicitly and provide a context provider for its adapter.
+When `state_connection` is omitted, the built-in scheduler uses the warehouse
+connection for its state adapter and disables query-context comments on that
+adapter. To add context to state SQL, configure `state_connection` explicitly
+and set a provider on its adapter.
 
 ## Read the payload
 
@@ -96,14 +93,13 @@ that are unavailable are omitted.
 }
 ```
 
-The operation ID groups statements for one model evaluation, including nested
-audits. A later evaluation gets a new ID. The provider owns the ID; SQLMesh does
-not create another operation ID at the adapter boundary. Consumers can compare
-the value they supplied with records from their execution observer.
+The provider assigns an operation ID to the statements for one model evaluation,
+including nested audits. A later evaluation gets a new ID. The adapter preserves
+the provider's ID. Consumers can compare it with the execution observer's record.
 
-The adapter's own correlation ID sets the payload's `parent`. SQLMesh replaces
-a different provider-supplied parent with the adapter's value. When the adapter
-has no correlation ID, SQLMesh clears a parent from the provided context.
+The adapter's correlation ID sets `parent` and replaces any parent supplied by
+the provider. If the adapter has no correlation ID, SQLMesh clears the provider's
+parent.
 
 Consumers must accept absent fields and unknown future phase or role values.
 The `workload` identity is the tuple `(system, type, name)`. Its `version` is
@@ -111,33 +107,35 @@ opaque and does not change that identity.
 
 ## SQL comment behavior
 
-SQLMesh prefixes the context comment after rendering. When a plan or run marker
-is present, SQLMesh keeps that legacy marker first. The submitted statement may
-look like this:
+SQLMesh adds the context comment after rendering. If a plan or run marker is
+present, SQLMesh keeps that marker first. The submitted statement may look like
+this:
 
 ```sql
 /* SQLMESH_RUN: run-123 */ /* QUERY_CONTEXT: {"schema_version":1,"emitter":"sqlmesh","parent":{"system":"sqlmesh","type":"run","id":"run-123"},"query":{"phase":"unknown","role":"query"}} */ SELECT 1
 ```
 
-The generic comment is bounded to 8,192 UTF-8 bytes. Each identifier is bounded
-to 1,024 UTF-8 bytes before JSON escaping. SQLMesh escapes metadata so it cannot
-end the comment. In JSON, SQLMesh also escapes slash characters so an identifier
-cannot form a nested comment delimiter. If an identifier is invalid or either
-limit is exceeded, SQLMesh
-omits the complete `QUERY_CONTEXT` comment and logs a warning. It never truncates
-the JSON or fails the query for an omitted annotation.
-
-The submitted SQL changes, while rendered model definitions and model fingerprints
-do not. If a downstream query-history system hashes raw SQL text, it must account
-for the additional comment.
-
-An adapter that does not support generic SQL comments keeps its SQL unchanged
-and logs one warning per adapter when the option is enabled. It does not call the
-provider. Athena does not support this transport. Engines can override the adapter
-hook to add native metadata transport in a separate implementation.
-
-If the provider raises an ordinary exception or returns an invalid value, SQLMesh
-omits the rich comment, logs a bounded warning without provider data, and executes
-the statement with its legacy correlation marker. It tries the provider again
-for the next statement. `KeyboardInterrupt`, `SystemExit`, and other cancellation
-exceptions keep their normal behavior and propagate.
+The comment is at most 8,192 UTF-8 bytes. Each identifier is at most 1,024 bytes
+before JSON escaping. SQLMesh JSON-escapes metadata and encodes `/` as `\u002f`,
+so values cannot end or nest the SQL comment.
+
+If an identifier cannot be encoded or a size limit is exceeded, SQLMesh omits the
+complete `QUERY_CONTEXT` comment and logs a warning. It does not truncate the JSON
+or fail the query because context was omitted.
+
+SQLMesh adds the comment after rendering. Model definitions and fingerprints
+exclude it. Query-history systems that hash submitted SQL must account for the
+added comment.
+
+Adapters that do not support generic SQL comments leave SQL unchanged and skip
+the provider. SQLMesh logs one warning per adapter when the option is enabled.
+Athena does not support this transport. Engines can override the adapter hook to
+add native metadata.
+
+If the provider raises an ordinary exception, returns an invalid value, or
+produces a context that cannot be serialized, SQLMesh omits the context comment
+and logs a bounded warning without provider data. It continues with the
+correlation marker if one is available. SQLMesh calls the provider again for the
+next statement. On transaction-capable adapters, a provider or driver
+`BaseException` rolls back any transaction SQLMesh opened for the statement and
+propagates the original exception.
diff --git a/docs/examples/query_execution_context.md b/docs/examples/query_execution_context.md
index e767a6d0a3..bff4f0c9ef 100644
--- a/docs/examples/query_execution_context.md
+++ b/docs/examples/query_execution_context.md
@@ -10,12 +10,15 @@ Run it from the SQLMesh repository root with the development environment:
 .venv/bin/python examples/query_execution_context/run_local.py
 ```
 
-The script attaches `QueryContextObserver.current_query_context` to the warehouse
-adapter providers. The observer builds one context for each model batch and
-delegates events to the observer supplied by the console. The run checks that four
-initial batches and two later batches use the same IDs in the model events,
-materialization SQL, and audits. It also checks that the resulting model rows are
-correct and that the query-context scope is restored after the run.
+The script registers `QueryContextObserver.current_query_context` as the provider
+on each warehouse adapter. `QueryContextObserver` tracks the active context from
+SQLMesh's native observation events. A callback records the model operation IDs
+used by the assertions.
+
+The plan covers four model batches, and the later run covers two. The script
+checks that model events, materialization SQL, and audits use matching operation
+IDs. It also checks the model rows and confirms that the query context is cleared
+after the run.
 
 The example uses a temporary DuckDB database and makes no external connections.
 See [query execution context](../concepts/query_execution_context.md) for the
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index 01d4e023a3..3d20b44189 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -2617,9 +2617,9 @@ def transaction(
         self._connection_pool.begin()
         try:
             yield
-        except Exception as e:
+        except BaseException:
             self._connection_pool.rollback()
-            raise e
+            raise
         else:
             self._connection_pool.commit()
 
diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py
index f1540727b1..0b5150c8e6 100644
--- a/tests/cli/test_cli.py
+++ b/tests/cli/test_cli.py
@@ -1209,6 +1209,7 @@ def test_dlt_filesystem_pipeline(tmp_path):
         "      # register_comments: False\n"
         "      # pre_ping: False\n"
         "      # pretty_sql: False\n"
+        "      # query_execution_context: False\n"
         "      # schema_differ_overrides: \n"
         "      # catalog_type_overrides: \n"
         "      # aws_access_key_id: \n"
@@ -2266,11 +2267,11 @@ def test_init_dbt_template(runner: CliRunner, tmp_path: Path):
 @time_machine.travel(FREEZE_TIME)
 def test_init_project_engine_configs(tmp_path):
     engine_type_to_config = {
-        "redshift": "# concurrent_tasks: 4\n      # register_comments: True\n      # pre_ping: False\n      # pretty_sql: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # user: \n      # password: \n      # database: \n      # host: \n      # port: \n      # source_address: \n      # unix_sock: \n      # ssl: \n      # sslmode: \n      # timeout: \n      # tcp_keepalive: \n      # application_name: \n      # preferred_role: \n      # principal_arn: \n      # credentials_provider: \n      # region: \n      # cluster_identifier: \n      # iam: \n      # db_user: \n      # is_serverless: \n      # serverless_acct_id: \n      # serverless_work_group: \n      # enable_merge: ",
-        "bigquery": "# concurrent_tasks: 1\n      # register_comments: True\n      # pre_ping: False\n      # pretty_sql: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # method: oauth\n      # project: \n      # execution_project: \n      # quota_project: \n      # location: \n      # keyfile: \n      # keyfile_json: \n      # token: \n      # refresh_token: \n      # client_id: \n      # client_secret: \n      # token_uri: \n      # scopes: \n      # impersonated_service_account: \n      # job_creation_timeout_seconds: \n      # job_execution_timeout_seconds: \n      # job_retries: 1\n      # job_retry_deadline_seconds: \n      # priority: \n      # maximum_bytes_billed: \n      # reservation: ",
-        "snowflake": "account: \n      # concurrent_tasks: 4\n      # register_comments: True\n      # pre_ping: False\n      # pretty_sql: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # user: \n      # password: \n      # warehouse: \n      # database: \n      # role: \n      # authenticator: \n      # token: \n      # host: \n      # port: \n      # application: Tobiko_SQLMesh\n      # private_key: \n      # private_key_path: \n      # private_key_passphrase: \n      # session_parameters: ",
-        "databricks": "# concurrent_tasks: 1\n      # register_comments: True\n      # pre_ping: False\n      # pretty_sql: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # server_hostname: \n      # http_path: \n      # access_token: \n      # auth_type: \n      # oauth_client_id: \n      # oauth_client_secret: \n      # catalog: \n      # http_headers: \n      # session_configuration: \n      # databricks_connect_server_hostname: \n      # databricks_connect_access_token: \n      # databricks_connect_cluster_id: \n      # databricks_connect_use_serverless: False\n      # force_databricks_connect: False\n      # disable_databricks_connect: False\n      # disable_spark_session: False",
-        "postgres": "host: \n      user: \n      password: \n      port: \n      database: \n      # concurrent_tasks: 4\n      # register_comments: True\n      # pre_ping: True\n      # pretty_sql: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # keepalives_idle: \n      # connect_timeout: 10\n      # role: \n      # sslmode: \n      # application_name: ",
+        "redshift": "# concurrent_tasks: 4\n      # register_comments: True\n      # pre_ping: False\n      # pretty_sql: False\n      # query_execution_context: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # user: \n      # password: \n      # database: \n      # host: \n      # port: \n      # source_address: \n      # unix_sock: \n      # ssl: \n      # sslmode: \n      # timeout: \n      # tcp_keepalive: \n      # application_name: \n      # preferred_role: \n      # principal_arn: \n      # credentials_provider: \n      # region: \n      # cluster_identifier: \n      # iam: \n      # db_user: \n      # is_serverless: \n      # serverless_acct_id: \n      # serverless_work_group: \n      # enable_merge: ",
+        "bigquery": "# concurrent_tasks: 1\n      # register_comments: True\n      # pre_ping: False\n      # pretty_sql: False\n      # query_execution_context: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # method: oauth\n      # project: \n      # execution_project: \n      # quota_project: \n      # location: \n      # keyfile: \n      # keyfile_json: \n      # token: \n      # refresh_token: \n      # client_id: \n      # client_secret: \n      # token_uri: \n      # scopes: \n      # impersonated_service_account: \n      # job_creation_timeout_seconds: \n      # job_execution_timeout_seconds: \n      # job_retries: 1\n      # job_retry_deadline_seconds: \n      # priority: \n      # maximum_bytes_billed: \n      # reservation: ",
+        "snowflake": "account: \n      # concurrent_tasks: 4\n      # register_comments: True\n      # pre_ping: False\n      # pretty_sql: False\n      # query_execution_context: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # user: \n      # password: \n      # warehouse: \n      # database: \n      # role: \n      # authenticator: \n      # token: \n      # host: \n      # port: \n      # application: Tobiko_SQLMesh\n      # private_key: \n      # private_key_path: \n      # private_key_passphrase: \n      # session_parameters: ",
+        "databricks": "# concurrent_tasks: 1\n      # register_comments: True\n      # pre_ping: False\n      # pretty_sql: False\n      # query_execution_context: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # server_hostname: \n      # http_path: \n      # access_token: \n      # auth_type: \n      # oauth_client_id: \n      # oauth_client_secret: \n      # catalog: \n      # http_headers: \n      # session_configuration: \n      # databricks_connect_server_hostname: \n      # databricks_connect_access_token: \n      # databricks_connect_cluster_id: \n      # databricks_connect_use_serverless: False\n      # force_databricks_connect: False\n      # disable_databricks_connect: False\n      # disable_spark_session: False",
+        "postgres": "host: \n      user: \n      password: \n      port: \n      database: \n      # concurrent_tasks: 4\n      # register_comments: True\n      # pre_ping: True\n      # pretty_sql: False\n      # query_execution_context: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # keepalives_idle: \n      # connect_timeout: 10\n      # role: \n      # sslmode: \n      # application_name: ",
     }
 
     for engine_type, expected_config in engine_type_to_config.items():
diff --git a/tests/core/engine_adapter/test_query_execution_context.py b/tests/core/engine_adapter/test_query_execution_context.py
index 9a00e01ee3..a83361f388 100644
--- a/tests/core/engine_adapter/test_query_execution_context.py
+++ b/tests/core/engine_adapter/test_query_execution_context.py
@@ -22,6 +22,10 @@
 from sqlmesh.utils import CorrelationId
 
 
+class _QueryCancellation(BaseException):
+    pass
+
+
 def _context(evaluation_id: str = "eval-1") -> QueryExecutionContext:
     return QueryExecutionContext(
         evaluation=ModelExecutionContext(
@@ -366,15 +370,28 @@ def invalid_provider() -> Any:
     assert "invalid context" in caplog.text
 
 
-def test_provider_cancellation_resets_transport_guard_and_scope(
+@pytest.mark.parametrize(
+    "exception_type",
+    [KeyboardInterrupt, SystemExit, _QueryCancellation],
+    ids=["keyboard-interrupt", "system-exit", "custom-base-exception"],
+)
+@pytest.mark.parametrize(
+    "nested",
+    [False, True],
+    ids=["adapter-owned-transaction", "nested-under-owner"],
+)
+def test_provider_base_exception_rolls_back_and_adapter_recovers(
+    exception_type: type[BaseException],
+    nested: bool,
     make_mocked_engine_adapter: Callable[..., EngineAdapter],
 ):
+    cancellation = exception_type("cancel query")
     attempts = []
 
     def provider() -> QueryExecutionContext:
         attempts.append(True)
         if len(attempts) == 1:
-            raise KeyboardInterrupt()
+            raise cancellation
         return _context("eval-after-cancellation")
 
     adapter = make_mocked_engine_adapter(
@@ -382,18 +399,106 @@ def provider() -> QueryExecutionContext:
         query_execution_context=True,
         query_execution_context_provider=provider,
     )
-    with pytest.raises(KeyboardInterrupt):
-        adapter.execute("SELECT 1")
 
-    adapter.execute("SELECT 2")
+    with pytest.raises(BaseException) as raised:
+        if nested:
+            with adapter.transaction():
+                adapter.execute("SELECT 1")
+        else:
+            adapter.execute("SELECT 1")
 
+    assert raised.value is cancellation
     assert get_query_execution_context() is None
-    assert attempts == [True, True]
+    assert adapter.cursor.execute.call_count == 0
+    assert adapter.cursor.rollback.call_count == 1
+    assert not adapter._connection_pool.is_transaction_active
+
+    adapter.execute("SELECT 2")
+
     sql = _executed_sql(adapter)
     assert len(sql) == 1
+    assert sql[0].endswith("SELECT 2")
     payload = _query_execution_payload(sql[0])
     assert payload is not None
     assert payload["operation"]["id"] == "eval-after-cancellation"
+    assert adapter.cursor.begin.call_count == 2
+    assert adapter.cursor.commit.call_count == 1
+    assert adapter.cursor.rollback.call_count == 1
+    assert not adapter._connection_pool.is_transaction_active
+    assert attempts == [True, True]
+
+
+def test_provider_cancellation_does_not_open_transaction_when_adapter_disables_them(
+    make_mocked_engine_adapter: Callable[..., EngineAdapter],
+):
+    class NonTransactionalEngineAdapter(EngineAdapter):
+        SUPPORTS_TRANSACTIONS = False
+
+    cancellation = _QueryCancellation("cancel without transaction")
+    attempts = []
+
+    def provider() -> QueryExecutionContext:
+        attempts.append(True)
+        if len(attempts) == 1:
+            raise cancellation
+        return _context("eval-after-cancellation")
+
+    adapter = make_mocked_engine_adapter(
+        NonTransactionalEngineAdapter,
+        query_execution_context=True,
+        query_execution_context_provider=provider,
+    )
+
+    with pytest.raises(_QueryCancellation) as raised:
+        adapter.execute("SELECT 1")
+
+    assert raised.value is cancellation
+    assert adapter.cursor.execute.call_count == 0
+    assert adapter.cursor.begin.call_count == 0
+    assert adapter.cursor.rollback.call_count == 0
+    assert adapter.cursor.commit.call_count == 0
+
+    adapter.execute("SELECT 2")
+
+    sql = _executed_sql(adapter)
+    assert len(sql) == 1
+    assert sql[0].endswith("SELECT 2")
+    assert adapter.cursor.begin.call_count == 0
+    assert adapter.cursor.rollback.call_count == 0
+    assert adapter.cursor.commit.call_count == 0
+    assert attempts == [True, True]
+
+
+def test_disabled_query_context_driver_cancellation_rolls_back_and_recovers(
+    make_mocked_engine_adapter: Callable[..., EngineAdapter],
+):
+    cancellation = _QueryCancellation("driver cancelled query")
+    execute_attempts = []
+
+    def execute(sql: str) -> None:
+        execute_attempts.append(sql)
+        if len(execute_attempts) == 1:
+            raise cancellation
+
+    adapter = make_mocked_engine_adapter(EngineAdapter, query_execution_context=False)
+    adapter.cursor.execute.side_effect = execute
+
+    with pytest.raises(_QueryCancellation) as raised:
+        adapter.execute("SELECT 1")
+
+    assert raised.value is cancellation
+    assert execute_attempts == ["SELECT 1"]
+    assert adapter.cursor.rollback.call_count == 1
+    assert not adapter._connection_pool.is_transaction_active
+
+    adapter.execute("SELECT 2")
+
+    assert execute_attempts == ["SELECT 1", "SELECT 2"]
+    assert _executed_sql(adapter) == ["SELECT 1", "SELECT 2"]
+    assert adapter.cursor.begin.call_count == 2
+    assert adapter.cursor.commit.call_count == 1
+    assert adapter.cursor.rollback.call_count == 1
+    assert not adapter._connection_pool.is_transaction_active
 
 
 def test_provider_reentry_skips_nested_rich_context_across_adapters(

From ca1b1e6dc3673eb50b79a22d3fcbca87e9b4103c Mon Sep 17 00:00:00 2001
From: mday-io 
Date: Sat, 31 Oct 2026 14:15:11 -0400
Subject: [PATCH 3/4] feat: add automatic model query context

Signed-off-by: mday-io 
---
 docs/concepts/query_execution_context.md      |  97 +++--
 docs/reference/configuration.md               |   2 +-
 sqlmesh/core/engine_adapter/base.py           |   9 +-
 sqlmesh/core/execution_observation.py         |  21 +-
 sqlmesh/core/plan/evaluator.py                |  10 +-
 sqlmesh/core/query_execution.py               |   8 +-
 .../core/query_execution_context_observer.py  | 184 ++++++++++
 sqlmesh/core/scheduler.py                     |  10 +-
 .../test_query_execution_context.py           |  74 +++-
 .../test_query_execution_context_observer.py  | 345 ++++++++++++++++++
 10 files changed, 714 insertions(+), 46 deletions(-)
 create mode 100644 sqlmesh/core/query_execution_context_observer.py
 create mode 100644 tests/core/test_query_execution_context_observer.py

diff --git a/docs/concepts/query_execution_context.md b/docs/concepts/query_execution_context.md
index 60ff83ff7b..37e419f306 100644
--- a/docs/concepts/query_execution_context.md
+++ b/docs/concepts/query_execution_context.md
@@ -4,10 +4,11 @@ SQLMesh can add structured metadata to warehouse statements so a query-history
 tool can connect a statement to the plan or run, model, and evaluation that issued
 it. The metadata is a versioned JSON object in a `QUERY_CONTEXT` SQL comment.
 
-Enable `query_execution_context` on a connection and set a provider on its
-engine adapter. The provider returns a `QueryExecutionContext`. The flag defaults
-to `false`. For YAML configuration, set the option under the gateway's
-`connection`:
+Set `query_execution_context` to `true` on a connection to enable bounded query
+context comments on supported adapters. The flag defaults to `false`. With this
+flag enabled, `Context.apply()` and built-in scheduler runs add model context to
+SQL on supported adapters without a custom provider. For YAML configuration,
+set the option under the gateway's `connection`:
 
 ```yaml
 gateways:
@@ -17,20 +18,59 @@ gateways:
       query_execution_context: true
 ```
 
-The flag enables context comments but does not infer model context. Without a
-provider, SQLMesh adds parent metadata only when a plan or run correlation ID is
-available. SQLMesh calls the provider after rendering each statement and before
-executing it. The call runs in the current worker's context, so concurrent model
-evaluations can return different IDs for a shared adapter.
+## Automatic model context
 
-To build a provider from native observer events, see the
-[query execution context example](../examples/query_execution_context.md).
+For each model batch, SQLMesh creates a context with a new operation ID, the
+model's canonical name, an opaque snapshot ID, and the batch interval and index
+when available. Intervals use UTC and an exclusive end.
 
-## Configure a provider
+The context stays active for SQL inside the batch, including nested audits.
+Automatic context uses the broad phases `materialize`, `audit`, and `ddl`. The
+`role` field does not identify the engine operation for each SQL statement.
 
-Set the provider on warehouse adapters before loading a Python `Context`. When
-`Context` creates its snapshot evaluator, it clones those adapters and preserves
-the provider and flag.
+Outside a model batch, SQLMesh can add parent metadata when a plan or run
+correlation ID is available. Use a custom provider when you need caller-defined
+context for other statements.
+
+## Read the context from an observer
+
+SQLMesh establishes the model context before it calls an observer's `start`
+callback. It restores the context after the `finish` callback. Both callbacks can
+read the active immutable value with `get_query_execution_context()`:
+
+```python
+from sqlmesh.core.query_execution import get_query_execution_context
+
+
+def observe(event, facts, error):
+    context = get_query_execution_context()
+    if context is not None and context.evaluation is not None:
+        print(context.evaluation.evaluation_id)
+```
+
+For console observer registration and callback behavior, see
+[Execution observations](../reference/execution_observations.md).
+
+SQLMesh forwards each observer event to the configured console observer once. If
+the console has no observer, model context still works. Nested observer scopes
+restore the outer observer and context on exit. Use `observer_scope(None)` to
+disable the built-in observer for its block.
+
+An observer callback can read the active context with
+`get_query_execution_context()`. SQL issued inside the callback does not use that
+automatic context. A custom adapter provider remains authoritative. Without one,
+an active plan or run correlation ID can still produce parent-only metadata.
+
+## Configure a custom provider
+
+Use a custom provider when you need caller-defined operation IDs or context
+outside a model batch. Set it on warehouse adapters before loading a Python
+`Context`. When `Context` creates its snapshot evaluator, it clones those
+adapters and preserves the provider and flag.
+
+SQLMesh calls the provider after rendering each statement and before executing
+it. The provider runs in the current worker's context, so concurrent workers can
+return different values for a shared adapter.
 
 ```python
 from sqlmesh.core.context import Context
@@ -53,23 +93,32 @@ with observer_scope(observer):
 
 `QueryContextObserver` returns the immutable context active in its worker, or
 `None` when no context applies. It assigns an evaluation ID when SQLMesh reports
-a model batch and restores the previous context when that event finishes. It also
+a model batch and restores the previous context when that event finishes. It
 forwards the same event, facts, and error to the observer returned by the console.
 
+A custom provider replaces the built-in provider for its adapter. If it returns
+`None`, SQLMesh does not use automatic model context for that statement. SQLMesh
+can still add parent-only metadata when the adapter has a correlation ID.
+
 Adapters created with `with_settings` inherit the flag and provider. Pass
 `query_execution_context_provider=None` to clear the inherited provider. Setting
 the flag to `false` disables provider calls and comment output.
 
 When `state_connection` is omitted, the built-in scheduler uses the warehouse
-connection for its state adapter and disables query-context comments on that
-adapter. To add context to state SQL, configure `state_connection` explicitly
-and set a provider on its adapter.
+connection for its state adapter and disables query context on that adapter. An
+explicitly configured `state_connection` uses its own `query_execution_context`
+setting. The automatic observer activates only if a supported warehouse adapter
+has `query_execution_context` enabled and no custom provider. An enabled state
+adapter without a custom provider can use that context during a model batch.
+Setting the flag only on `state_connection` does not activate the observer. Use a
+custom provider for state SQL outside a model batch.
 
 ## Read the payload
 
 The version-one payload uses the SQLMesh model's canonical name and opaque
-snapshot identifier. Operation intervals use UTC half-open boundaries. Fields
-that are unavailable are omitted.
+snapshot identifier. Operation intervals use UTC half-open boundaries, and the
+payload includes `batch_index` when available. Fields that are unavailable are
+omitted.
 
 ```json
 {
@@ -93,9 +142,9 @@ that are unavailable are omitted.
 }
 ```
 
-The provider assigns an operation ID to the statements for one model evaluation,
-including nested audits. A later evaluation gets a new ID. The adapter preserves
-the provider's ID. Consumers can compare it with the execution observer's record.
+Each model batch gets one operation ID, which nested audits reuse. The built-in
+provider creates this ID. A custom provider supplies its own. The adapter preserves
+the provider's ID when it adds the context comment.
 
 The adapter's correlation ID sets `parent` and replaces any parent supplied by
 the provider. If the adapter has no correlation ID, SQLMesh clears the provider's
diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md
index c817583c67..e34d7cdef7 100644
--- a/docs/reference/configuration.md
+++ b/docs/reference/configuration.md
@@ -212,7 +212,7 @@ Most parameters are specific to the connection engine `type` - see [below](#engi
 | `register_comments` | Whether SQLMesh should register model comments with the SQL engine (if the engine supports it). (Default: `true`.)                                                      | bool |    N     |
 | `pre_ping`          | Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive. This can only be enabled for engines with transaction support. | bool |    N     |
 | `pretty_sql`        | If SQL should be formatted before being executed, not recommended in a production setting. (Default: `false`.)                                                          | bool |    N     |
-| `query_execution_context` | Whether an opted-in engine adapter adds bounded execution metadata to SQL comments. (Default: `false`.) See [query execution context](../concepts/query_execution_context.md). | bool | N |
+| `query_execution_context` | Enables bounded query-context comments on supported adapters (Default: `false`). During `Context.apply()` and built-in scheduler runs, SQLMesh adds context for each model batch to enabled adapters without a custom provider. See [query execution context](../concepts/query_execution_context.md). | bool | N |
 
 #### Engine-specific
 
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index 3d20b44189..14eb52858e 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -2706,7 +2706,14 @@ def _attach_query_execution_context(self, sql: str) -> str:
 
             provider = self.query_execution_context_provider
             try:
-                context = provider() if provider is not None else None
+                if provider is None:
+                    from sqlmesh.core.query_execution_context_observer import (
+                        get_automatic_query_execution_context,
+                    )
+
+                    context = get_automatic_query_execution_context()
+                else:
+                    context = provider()
             except Exception:
                 self._warn_query_execution_context_once(
                     "provider",
diff --git a/sqlmesh/core/execution_observation.py b/sqlmesh/core/execution_observation.py
index 36f426c90e..c343f9b852 100644
--- a/sqlmesh/core/execution_observation.py
+++ b/sqlmesh/core/execution_observation.py
@@ -119,14 +119,33 @@ def execution_context_factory() -> t.Optional[t.Callable[[], contextvars.Context
 
 
 @contextlib.contextmanager
-def console_observer_scope(console: t.Any) -> t.Iterator[None]:
+def console_observer_scope(
+    console: t.Any, *, query_execution_context: bool = False
+) -> t.Iterator[None]:
     """Resolve once at the outer boundary; an explicit None scope disables hooks."""
     if _scoped.get():
+        observer = _observer.get()
+        if query_execution_context and observer is not None:
+            from sqlmesh.core.query_execution_context_observer import (
+                with_automatic_query_execution_observer,
+            )
+
+            automatic_observer = with_automatic_query_execution_observer(observer)
+            if automatic_observer is not observer:
+                with observer_scope(automatic_observer):
+                    yield
+                return
         yield
         return
     try:
         observer = console.get_execution_observer()
     except Exception:
         observer = None
+    if query_execution_context:
+        from sqlmesh.core.query_execution_context_observer import (
+            with_automatic_query_execution_observer,
+        )
+
+        observer = with_automatic_query_execution_observer(observer)
     with observer_scope(observer):
         yield
diff --git a/sqlmesh/core/plan/evaluator.py b/sqlmesh/core/plan/evaluator.py
index 71e5285d63..1d76025329 100644
--- a/sqlmesh/core/plan/evaluator.py
+++ b/sqlmesh/core/plan/evaluator.py
@@ -19,6 +19,9 @@
 import typing as t
 from sqlmesh.core import analytics
 from sqlmesh.core.execution_observation import action, console_observer_scope
+from sqlmesh.core.query_execution_context_observer import (
+    should_enable_automatic_query_execution_context,
+)
 from sqlmesh.core import constants as c
 from sqlmesh.core.console import Console, get_console
 from sqlmesh.core.environment import EnvironmentNamingInfo, execute_environment_statements
@@ -103,7 +106,12 @@ def evaluate(
 
         try:
             plan_stages = stages.build_plan_stages(plan, self.state_sync, self.default_catalog)
-            with console_observer_scope(self.console):
+            with console_observer_scope(
+                self.console,
+                query_execution_context=should_enable_automatic_query_execution_context(
+                    self.snapshot_evaluator.adapters.values()
+                ),
+            ):
                 self._evaluate_stages(plan_stages, plan)
         except Exception as e:
             analytics.collector.on_plan_apply_end(plan_id=plan.plan_id, error=e)
diff --git a/sqlmesh/core/query_execution.py b/sqlmesh/core/query_execution.py
index f7e6f3ade9..8d21c725e2 100644
--- a/sqlmesh/core/query_execution.py
+++ b/sqlmesh/core/query_execution.py
@@ -170,10 +170,10 @@ def _query_execution_context_transport_scope() -> t.Iterator[None]:
 
 
 @contextlib.contextmanager
-def query_execution_scope(context: QueryExecutionContext) -> t.Iterator[None]:
-    """Set a query execution context for this scope and restore it on exit."""
-    if not isinstance(context, QueryExecutionContext):
-        raise TypeError("context must be a QueryExecutionContext")
+def query_execution_scope(context: t.Optional[QueryExecutionContext]) -> t.Iterator[None]:
+    """Set or clear the query execution context for this scope."""
+    if context is not None and not isinstance(context, QueryExecutionContext):
+        raise TypeError("context must be a QueryExecutionContext or None")
     token = _current_query_execution_context.set(context)
     try:
         yield
diff --git a/sqlmesh/core/query_execution_context_observer.py b/sqlmesh/core/query_execution_context_observer.py
new file mode 100644
index 0000000000..986eb72ca6
--- /dev/null
+++ b/sqlmesh/core/query_execution_context_observer.py
@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Build model query context from the native execution observer boundary."""
+
+from __future__ import annotations
+
+import contextvars
+import dataclasses
+import typing as t
+from dataclasses import dataclass
+
+from sqlmesh.core.execution_observation import Observer, _dispatching, _observer
+from sqlmesh.core.query_execution import (
+    ModelExecutionContext,
+    QueryExecutionContext,
+    QueryExecutionOperation,
+    QueryExecutionPhase,
+    QueryExecutionRole,
+    get_query_execution_context,
+    query_execution_scope,
+)
+from sqlmesh.utils import random_id
+from sqlmesh.utils.date import to_datetime
+
+if t.TYPE_CHECKING:
+    from sqlmesh.core.engine_adapter.base import EngineAdapter
+
+
+@dataclass(frozen=True)
+class _Frame:
+    facts: t.Mapping[str, t.Any]
+    scope: t.ContextManager[None]
+
+
+class AutomaticQueryExecutionObserver:
+    def __init__(self, observer: t.Optional[Observer] = None):
+        self._observer = observer
+        self._frames: contextvars.ContextVar[t.Tuple[_Frame, ...]] = contextvars.ContextVar(
+            "automatic_query_execution_frames", default=()
+        )
+
+    def __call__(
+        self,
+        event: t.Literal["start", "finish"],
+        facts: t.Mapping[str, t.Any],
+        error: t.Optional[BaseException],
+    ) -> None:
+        if event == "start":
+            frame = self._begin(facts)
+            try:
+                self._notify(event, facts, error)
+            except BaseException as ex:
+                self._restore(frame, ex)
+                raise
+            return
+
+        matching_frame = self._matching_frame(facts)
+        try:
+            self._notify(event, facts, error)
+        finally:
+            if matching_frame is not None:
+                self._restore(matching_frame, error)
+
+    def _begin(self, facts: t.Mapping[str, t.Any]) -> _Frame:
+        try:
+            context = self._derive_context(facts)
+        except Exception:
+            context = None if facts.get("kind") == "model" else get_query_execution_context()
+
+        scope = query_execution_scope(context)
+        scope.__enter__()
+        frame = _Frame(facts=facts, scope=scope)
+        try:
+            self._frames.set(self._frames.get() + (frame,))
+        except BaseException:
+            scope.__exit__(None, None, None)
+            raise
+        return frame
+
+    def _matching_frame(self, facts: t.Mapping[str, t.Any]) -> t.Optional[_Frame]:
+        frames = self._frames.get()
+        if frames and frames[-1].facts is facts:
+            return frames[-1]
+        return None
+
+    def _restore(self, frame: _Frame, error: t.Optional[BaseException]) -> None:
+        frames = self._frames.get()
+        if not frames or frames[-1] is not frame:
+            return
+        try:
+            frame.scope.__exit__(
+                type(error) if error is not None else None,
+                error,
+                error.__traceback__ if error is not None else None,
+            )
+        finally:
+            self._frames.set(frames[:-1])
+
+    def _notify(
+        self,
+        event: t.Literal["start", "finish"],
+        facts: t.Mapping[str, t.Any],
+        error: t.Optional[BaseException],
+    ) -> None:
+        if self._observer is None:
+            return
+        try:
+            self._observer(event, facts, error)
+        except Exception:
+            pass
+
+    def _derive_context(self, facts: t.Mapping[str, t.Any]) -> t.Optional[QueryExecutionContext]:
+        current = get_query_execution_context()
+        kind = facts.get("kind")
+        action = facts.get("action")
+        if kind == "model":
+            snapshot = facts.get("snapshot")
+            if snapshot is None or not snapshot.is_model:
+                return None
+            interval = facts.get("interval")
+            interval_start, interval_end = interval if interval is not None else (None, None)
+            audit_only = action == "audit_only" or facts.get("audit_only", False)
+            return QueryExecutionContext(
+                evaluation=ModelExecutionContext(
+                    evaluation_id=random_id(),
+                    model_name=snapshot.name,
+                    snapshot_id=snapshot.snapshot_id.identifier,
+                    interval_start=to_datetime(interval_start)
+                    if interval_start is not None
+                    else None,
+                    interval_end=to_datetime(interval_end) if interval_end is not None else None,
+                    batch_index=facts.get("batch_index"),
+                    operation_type=(
+                        QueryExecutionOperation.AUDIT
+                        if audit_only
+                        else QueryExecutionOperation.EVALUATE
+                    ),
+                ),
+                phase=QueryExecutionPhase.AUDIT if audit_only else QueryExecutionPhase.MATERIALIZE,
+                role=QueryExecutionRole.AUDIT if audit_only else QueryExecutionRole.QUERY,
+            )
+        if current is None:
+            return None
+        if kind == "audit":
+            return dataclasses.replace(
+                current, phase=QueryExecutionPhase.AUDIT, role=QueryExecutionRole.AUDIT
+            )
+        if kind in ("physical", "virtual"):
+            return dataclasses.replace(
+                current,
+                phase=(
+                    QueryExecutionPhase.MATERIALIZE
+                    if action == "materialize"
+                    else QueryExecutionPhase.DDL
+                ),
+                role=QueryExecutionRole.QUERY,
+            )
+        return current
+
+
+def get_automatic_query_execution_context() -> t.Optional[QueryExecutionContext]:
+    if _dispatching.get():
+        return None
+    if isinstance(_observer.get(), AutomaticQueryExecutionObserver):
+        return get_query_execution_context()
+    return None
+
+
+def should_enable_automatic_query_execution_context(
+    adapters: t.Iterable[EngineAdapter],
+) -> bool:
+    return any(
+        adapter.query_execution_context
+        and adapter.ATTACH_CORRELATION_ID
+        and adapter.query_execution_context_provider is None
+        for adapter in adapters
+    )
+
+
+def with_automatic_query_execution_observer(
+    observer: t.Optional[Observer],
+) -> AutomaticQueryExecutionObserver:
+    if isinstance(observer, AutomaticQueryExecutionObserver):
+        return observer
+    return AutomaticQueryExecutionObserver(observer)
diff --git a/sqlmesh/core/scheduler.py b/sqlmesh/core/scheduler.py
index 7345d92888..bb5e78770b 100644
--- a/sqlmesh/core/scheduler.py
+++ b/sqlmesh/core/scheduler.py
@@ -13,6 +13,9 @@
     console_observer_scope,
     execution_context_factory,
 )
+from sqlmesh.core.query_execution_context_observer import (
+    should_enable_automatic_query_execution_context,
+)
 from sqlmesh.core.environment import EnvironmentNamingInfo, execute_environment_statements
 from sqlmesh.core.macros import RuntimeStage
 from sqlmesh.core.model.definition import AuditResult
@@ -459,7 +462,12 @@ def run_merged_intervals(
         Returns:
             A tuple of errors and skipped intervals.
         """
-        with console_observer_scope(self.console):
+        with console_observer_scope(
+            self.console,
+            query_execution_context=should_enable_automatic_query_execution_context(
+                self.snapshot_evaluator.adapters.values()
+            ),
+        ):
             execution_time = execution_time or now_timestamp()
 
             selected_snapshots = [self.snapshots[sid] for sid in (selected_snapshot_ids or set())]
diff --git a/tests/core/integration/test_query_execution_context.py b/tests/core/integration/test_query_execution_context.py
index fdc0eb6c6b..1d32ff4ce5 100644
--- a/tests/core/integration/test_query_execution_context.py
+++ b/tests/core/integration/test_query_execution_context.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
+import contextlib
 import datetime as dt
 import json
 import re
@@ -8,6 +9,8 @@
 import typing as t
 from pathlib import Path
 
+import pytest
+
 from examples.query_execution_context.observer import QueryContextObserver
 from sqlmesh.core.config import Config, ModelDefaultsConfig
 from sqlmesh.core.config.connection import DuckDBConnectionConfig
@@ -15,6 +18,7 @@
 from sqlmesh.core.context import Context
 from sqlmesh.core.engine_adapter.base import EngineAdapter
 from sqlmesh.core.execution_observation import observer_scope
+from sqlmesh.core.query_execution import get_query_execution_context
 from sqlmesh.core.state_sync import CachingStateSync, EngineAdapterStateSync
 
 
@@ -23,9 +27,18 @@ def _payload(sql: str) -> t.Optional[t.Dict[str, t.Any]]:
     return json.loads(match[1]) if match else None
 
 
-def _new_context(root: Path, enabled: bool) -> Context:
+def _new_context(root: Path, enabled: bool, yaml_config: bool = False) -> Context:
     models = Path(__file__).parents[3] / "examples/query_execution_context/models"
     shutil.copytree(models, root / "models")
+    if yaml_config:
+        (root / "config.yaml").write_text(
+            "gateways:\n  local:\n    connection:\n      type: duckdb\n"
+            f"      database: {json.dumps(str(root / 'warehouse.duckdb'))}\n"
+            "      concurrent_tasks: 2\n"
+            f"      query_execution_context: {str(enabled).lower()}\n"
+            "default_gateway: local\nmodel_defaults:\n  dialect: duckdb\n"
+        )
+        return Context(paths=root, load=False)
     config = Config(
         gateways={
             "local": GatewayConfig(
@@ -60,6 +73,11 @@ def _batch_ids(
         assert payload["workload"]["version"]
         operation = payload["operation"]
         assert operation["type"] == "evaluate"
+        assert payload["query"]["role"] == ("audit" if phase == "audit" else "query")
+        assert (
+            dt.datetime.fromisoformat(operation["interval_end"])
+            - dt.datetime.fromisoformat(operation["interval_start"])
+        ) == dt.timedelta(days=1)
         key = (
             payload["workload"]["name"],
             operation["interval_start"],
@@ -71,7 +89,8 @@ def _batch_ids(
     return batches
 
 
-def test_provider_connects_native_events_to_warehouse_history(tmp_path, monkeypatch):
+@pytest.mark.parametrize("mode", ["provider", "automatic", "automatic_default"])
+def test_context_connects_native_events_to_warehouse_history(tmp_path, monkeypatch, mode):
     submitted = []
     original_execute = EngineAdapter._execute
 
@@ -82,33 +101,54 @@ def capture(adapter, sql, *args, **kwargs):
     monkeypatch.setattr(EngineAdapter, "_execute", capture)
     results = []
     for enabled in (False, True):
-        context = _new_context(tmp_path / str(enabled), enabled)
+        context = _new_context(tmp_path / str(enabled), enabled, yaml_config=mode != "provider")
         events = []
 
         def observer(event, facts, error):
-            current = consumer.current_query_context()
+            current = (
+                consumer.current_query_context()
+                if mode == "provider"
+                else get_query_execution_context()
+            )
             if facts["kind"] == "model":
                 events.append(
-                    (event, facts["snapshot"].name, current.evaluation.evaluation_id, error)
+                    (
+                        event,
+                        facts["snapshot"].name,
+                        current.evaluation.evaluation_id
+                        if current and current.evaluation
+                        else None,
+                        error,
+                    )
                 )
 
         consumer = QueryContextObserver(observer)
         for adapter in context.engine_adapters.values():
-            adapter.query_execution_context_provider = consumer.current_query_context
+            assert adapter.query_execution_context_provider is None
+            if mode == "provider":
+                adapter.query_execution_context_provider = consumer.current_query_context
         context.load()
         try:
             hashes = {m.name: (m.data_hash, m.metadata_hash) for m in context.models.values()}
             submitted.clear()
-            with observer_scope(consumer):
+            scope = (
+                contextlib.nullcontext()
+                if mode == "automatic_default"
+                else observer_scope(consumer if mode == "provider" else observer)
+            )
+            with scope:
                 plan = context.plan(
                     "dev", start="2024-01-01", end="2024-01-02", no_prompts=True, skip_tests=True
                 )
                 context.apply(plan)
                 assert consumer.current_query_context() is None
-                plan_sql = [sql for _, sql in submitted]
+                assert get_query_execution_context() is None
+                plan_submitted = list(submitted)
+                plan_sql = [sql for _, sql in plan_submitted]
                 submitted.clear()
                 assert context.run("dev", start="2024-01-03", end="2024-01-03").is_success
                 assert consumer.current_query_context() is None
+                assert get_query_execution_context() is None
                 run_sql = [sql for _, sql in submitted]
                 state = context.state_sync
                 assert isinstance(state, CachingStateSync)
@@ -116,7 +156,9 @@ def observer(event, facts, error):
                 state_adapter = state.state_sync.engine_adapter
                 assert any(adapter is state_adapter for adapter, _ in submitted)
                 assert all(
-                    _payload(sql) is None for adapter, sql in submitted if adapter is state_adapter
+                    _payload(sql) is None
+                    for adapter, sql in plan_submitted + submitted
+                    if adapter is state_adapter
                 )
             rows = context.engine_adapter.fetchall(
                 f"SELECT ds, id FROM {context.resolve_table('e2e.events')} ORDER BY ds"
@@ -124,6 +166,11 @@ def observer(event, facts, error):
             assert rows == [(dt.date(2024, 1, day), 1) for day in (1, 2, 3)]
             results.append((hashes, rows))
             if enabled:
+                for sql in plan_sql + run_sql:
+                    payload = _payload(sql)
+                    if payload and payload.get("operation"):
+                        snapshot = context.snapshots[payload["workload"]["name"]]
+                        assert payload["workload"]["version"] == snapshot.snapshot_id.identifier
                 materialize = _batch_ids(plan_sql, "plan", "materialize")
                 assert _batch_ids(plan_sql, "plan", "audit") == materialize
                 assert set(materialize) == {
@@ -140,10 +187,11 @@ def observer(event, facts, error):
                 }
                 all_ids = set(materialize.values()) | set(later.values())
                 assert len(all_ids) == 6
-                assert {operation_id for _, _, operation_id, _ in events} == all_ids
-                assert [event for event, _, _, _ in events].count("start") == 6
-                assert [event for event, _, _, _ in events].count("finish") == 6
-                assert all(error is None for _, _, _, error in events)
+                if mode != "automatic_default":
+                    assert {operation_id for _, _, operation_id, _ in events} == all_ids
+                    assert [event for event, _, _, _ in events].count("start") == 6
+                    assert [event for event, _, _, _ in events].count("finish") == 6
+                    assert all(error is None for _, _, _, error in events)
             else:
                 assert all(_payload(sql) is None for sql in plan_sql + run_sql)
                 assert any(sql.startswith("/* SQLMESH_PLAN:") for sql in plan_sql)
diff --git a/tests/core/test_query_execution_context_observer.py b/tests/core/test_query_execution_context_observer.py
new file mode 100644
index 0000000000..378dc95aff
--- /dev/null
+++ b/tests/core/test_query_execution_context_observer.py
@@ -0,0 +1,345 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import contextvars
+import datetime as dt
+import json
+import re
+from concurrent.futures import ThreadPoolExecutor
+from threading import Barrier
+from types import SimpleNamespace
+from typing import Any, Callable, Optional
+
+import pytest
+
+from sqlmesh.core.engine_adapter import EngineAdapter
+from sqlmesh.core.engine_adapter.athena import AthenaEngineAdapter
+from sqlmesh.core.execution_observation import (
+    action,
+    console_observer_scope,
+    observer_scope,
+)
+from sqlmesh.core.query_execution import (
+    ModelExecutionContext,
+    QueryExecutionContext,
+    get_query_execution_context,
+)
+from sqlmesh.core.query_execution_context_observer import (
+    should_enable_automatic_query_execution_context,
+)
+from sqlmesh.utils import CorrelationId
+
+
+def _snapshot(name: str = "analytics.orders") -> Any:
+    return SimpleNamespace(
+        is_model=True,
+        name=name,
+        snapshot_id=SimpleNamespace(identifier=f"version-{name}"),
+    )
+
+
+def _interval(day: int) -> tuple[dt.datetime, dt.datetime]:
+    start = dt.datetime(2024, 1, day, tzinfo=dt.timezone.utc)
+    return start, start + dt.timedelta(days=1)
+
+
+def _payload(sql: str) -> Optional[dict[str, Any]]:
+    match = re.search(r"/\* QUERY_CONTEXT: (.*?) \*/", sql)
+    return json.loads(match[1]) if match else None
+
+
+def _last_sql(adapter: EngineAdapter) -> str:
+    return adapter.cursor.execute.call_args.args[0]
+
+
+def _run_model(adapter: EngineAdapter, day: int = 1) -> tuple[str, QueryExecutionContext]:
+    with action(
+        "model",
+        "evaluate",
+        snapshot=_snapshot(),
+        interval=_interval(day),
+        batch_index=day - 1,
+    ):
+        context = get_query_execution_context()
+        assert context is not None
+        adapter.execute("SELECT 1")
+        return _last_sql(adapter), context
+
+
+def test_default_console_installs_context_and_derives_nested_phases(
+    make_mocked_engine_adapter: Callable[..., EngineAdapter],
+):
+    adapter = make_mocked_engine_adapter(
+        EngineAdapter,
+        query_execution_context=True,
+        correlation_id=CorrelationId.from_plan_id("plan-1"),
+    )
+    console = SimpleNamespace(get_execution_observer=lambda: None)
+    events = []
+
+    def observer(event, facts, error):
+        events.append((event, dict(facts), get_query_execution_context(), error))
+
+    console.get_execution_observer = lambda: observer
+    with console_observer_scope(console, query_execution_context=True):
+        with action(
+            "model",
+            "evaluate",
+            snapshot=_snapshot(),
+            interval=_interval(1),
+            batch_index=0,
+        ):
+            model_context = get_query_execution_context()
+            assert model_context is not None
+            adapter.execute("SELECT 1")
+            sql = _last_sql(adapter)
+            with action("audit", "audit", audit_name="not_null"):
+                adapter.execute("SELECT 2")
+            audit_context = _payload(_last_sql(adapter))
+            with action("physical", "create", snapshot=_snapshot()):
+                adapter.execute("SELECT 3")
+            materialize_context = _payload(_last_sql(adapter))
+        assert get_query_execution_context() is None
+
+    model_payload = _payload(sql)
+    assert model_payload is not None
+    evaluation = model_context.evaluation
+    assert evaluation is not None
+    assert sql.startswith("/* SQLMESH_PLAN: plan-1 */ /* QUERY_CONTEXT:")
+    assert model_payload["operation"]["id"] == evaluation.evaluation_id
+    assert model_payload["query"] == {"phase": "materialize", "role": "query"}
+    assert audit_context is not None
+    assert audit_context["operation"]["id"] == evaluation.evaluation_id
+    assert audit_context["query"] == {"phase": "audit", "role": "audit"}
+    assert materialize_context is not None
+    assert materialize_context["operation"]["id"] == evaluation.evaluation_id
+    assert materialize_context["query"] == {"phase": "ddl", "role": "query"}
+    model_events = [event for event in events if event[1]["kind"] == "model"]
+    assert [event[0] for event in model_events] == ["start", "finish"]
+    assert all(event[2] == model_context for event in model_events)
+    assert all(event[1]["batch_index"] == 0 for event in model_events)
+
+
+def test_explicit_disabled_scope_suppresses_transport_then_outer_context_resumes(
+    make_mocked_engine_adapter: Callable[..., EngineAdapter],
+):
+    adapter = make_mocked_engine_adapter(EngineAdapter, query_execution_context=True)
+    console = SimpleNamespace(get_execution_observer=lambda: None)
+
+    with console_observer_scope(console, query_execution_context=True):
+        with action(
+            "model",
+            "evaluate",
+            snapshot=_snapshot(),
+            interval=_interval(1),
+            batch_index=0,
+        ):
+            outer = get_query_execution_context()
+            assert outer is not None
+            with observer_scope(None):
+                adapter.execute("SELECT 1")
+                disabled_sql = _last_sql(adapter)
+                assert get_query_execution_context() == outer
+                with console_observer_scope(console, query_execution_context=True):
+                    with action(
+                        "model",
+                        "evaluate",
+                        snapshot=_snapshot("analytics.inner"),
+                        interval=_interval(2),
+                        batch_index=1,
+                    ):
+                        assert get_query_execution_context() == outer
+                        adapter.execute("SELECT nested disabled")
+                        nested_disabled_sql = _last_sql(adapter)
+            adapter.execute("SELECT 2")
+            resumed_sql = _last_sql(adapter)
+
+    assert _payload(disabled_sql) is None
+    assert _payload(nested_disabled_sql) is None
+    resumed = _payload(resumed_sql)
+    assert resumed is not None
+    assert outer.evaluation is not None
+    assert resumed["operation"]["id"] == outer.evaluation.evaluation_id
+
+
+@pytest.mark.parametrize("provider_kind", ["value", "none"])
+def test_custom_provider_remains_authoritative(
+    provider_kind: str,
+    make_mocked_engine_adapter: Callable[..., EngineAdapter],
+):
+    automatic_adapter = make_mocked_engine_adapter(EngineAdapter, query_execution_context=True)
+    if provider_kind == "value":
+        custom = QueryExecutionContext(
+            evaluation=ModelExecutionContext("custom-id", "custom.model", "custom-version")
+        )
+        adapter = make_mocked_engine_adapter(
+            EngineAdapter,
+            query_execution_context=True,
+            query_execution_context_provider=lambda: custom,
+        )
+    else:
+        adapter = make_mocked_engine_adapter(
+            EngineAdapter,
+            query_execution_context=True,
+            query_execution_context_provider=lambda: None,
+        )
+    console = SimpleNamespace(get_execution_observer=lambda: None)
+    assert should_enable_automatic_query_execution_context([automatic_adapter, adapter])
+
+    with (
+        console_observer_scope(console, query_execution_context=True),
+        action(
+            "model",
+            "evaluate",
+            snapshot=_snapshot(),
+            interval=_interval(1),
+            batch_index=0,
+        ),
+    ):
+        adapter.execute("SELECT 1")
+        sql = _last_sql(adapter)
+
+    payload = _payload(sql)
+    if provider_kind == "value":
+        assert payload is not None
+        assert payload["operation"]["id"] == "custom-id"
+    else:
+        assert payload is None
+
+
+def test_auto_producer_requires_enabled_supported_adapter_without_provider(
+    make_mocked_engine_adapter: Callable[..., EngineAdapter],
+):
+    custom = make_mocked_engine_adapter(
+        EngineAdapter,
+        query_execution_context=True,
+        query_execution_context_provider=lambda: None,
+    )
+    disabled = make_mocked_engine_adapter(EngineAdapter, query_execution_context=False)
+    unsupported = make_mocked_engine_adapter(AthenaEngineAdapter, query_execution_context=True)
+
+    assert not should_enable_automatic_query_execution_context([custom, disabled, unsupported])
+
+
+def test_callback_reentry_does_not_emit_automatic_context(
+    make_mocked_engine_adapter: Callable[..., EngineAdapter],
+):
+    adapter = make_mocked_engine_adapter(EngineAdapter, query_execution_context=True)
+    console = SimpleNamespace(get_execution_observer=None)
+    callback_sql = []
+    callback_contexts = []
+
+    def observer(event, facts, error):
+        if facts["kind"] == "model":
+            callback_contexts.append(get_query_execution_context())
+            adapter.execute("SELECT callback")
+            callback_sql.append(_last_sql(adapter))
+
+    console.get_execution_observer = lambda: observer
+    with console_observer_scope(console, query_execution_context=True):
+        regular_sql, model_context = _run_model(adapter)
+
+    assert len(callback_sql) == 2
+    assert all(_payload(sql) is None for sql in callback_sql)
+    assert all(context == model_context for context in callback_contexts)
+    payload = _payload(regular_sql)
+    assert payload is not None, regular_sql
+    assert model_context.evaluation is not None
+    assert payload["operation"]["id"] == model_context.evaluation.evaluation_id
+
+
+def test_worker_frames_are_isolated_for_a_shared_adapter(
+    make_mocked_engine_adapter: Callable[..., EngineAdapter],
+):
+    adapter = make_mocked_engine_adapter(EngineAdapter, query_execution_context=True)
+    console = SimpleNamespace(get_execution_observer=lambda: None)
+    barrier = Barrier(2)
+
+    def evaluate(day: int):
+        with action(
+            "model",
+            "evaluate",
+            snapshot=_snapshot(f"analytics.model_{day}"),
+            interval=_interval(day),
+            batch_index=day - 1,
+        ):
+            context = get_query_execution_context()
+            assert context is not None
+            barrier.wait(timeout=5)
+            sql = adapter._attach_query_execution_context(f"SELECT {day}")
+            barrier.wait(timeout=5)
+            return context, _payload(sql)
+
+    with console_observer_scope(console, query_execution_context=True):
+        contexts = [contextvars.copy_context(), contextvars.copy_context()]
+        with ThreadPoolExecutor(max_workers=2) as executor:
+            results = list(
+                executor.map(lambda pair: pair[0].run(evaluate, pair[1]), zip(contexts, (1, 2)))
+            )
+
+    ids = [context.evaluation.evaluation_id for context, _ in results]
+    assert len(set(ids)) == 2
+    assert [payload["operation"]["id"] for _, payload in results] == ids
+    assert {payload["workload"]["name"] for _, payload in results} == {
+        "analytics.model_1",
+        "analytics.model_2",
+    }
+
+
+@pytest.mark.parametrize("callback_phase", ["start", "finish"])
+def test_base_exception_callback_restores_context(callback_phase: str):
+    interrupt = KeyboardInterrupt()
+    console = SimpleNamespace(get_execution_observer=lambda: None)
+    events = []
+
+    def observer(event, facts, error):
+        events.append((event, get_query_execution_context()))
+        if event == callback_phase:
+            raise interrupt
+
+    console.get_execution_observer = lambda: observer
+    with console_observer_scope(console, query_execution_context=True):
+        with pytest.raises(KeyboardInterrupt) as raised:
+            with action(
+                "model",
+                "evaluate",
+                snapshot=_snapshot(),
+                interval=_interval(1),
+                batch_index=0,
+            ):
+                assert callback_phase == "finish"
+        assert raised.value is interrupt
+        assert get_query_execution_context() is None
+    assert [event for event, _ in events] == (
+        ["start"] if callback_phase == "start" else ["start", "finish"]
+    )
+    assert events[0][1] is not None
+    assert all(context == events[0][1] for _, context in events)
+
+
+def test_observer_failure_and_workload_cancellation_restore_context():
+    error = KeyboardInterrupt()
+    contexts = []
+    console = SimpleNamespace(get_execution_observer=lambda: None)
+
+    def observer(event, facts, failure):
+        contexts.append((event, get_query_execution_context(), failure))
+        raise RuntimeError("observer failure")
+
+    console.get_execution_observer = lambda: observer
+    with console_observer_scope(console, query_execution_context=True):
+        with pytest.raises(KeyboardInterrupt) as raised:
+            with action(
+                "model",
+                "evaluate",
+                snapshot=_snapshot(),
+                interval=_interval(1),
+                batch_index=0,
+            ):
+                raise error
+        assert raised.value is error
+        assert get_query_execution_context() is None
+
+    assert [event for event, _, _ in contexts] == ["start", "finish"]
+    assert contexts[0][1] == contexts[1][1]
+    assert contexts[1][2] is error

From b57bc31fd25833201d8959bf313c0762ff0f16ad Mon Sep 17 00:00:00 2001
From: mday-io 
Date: Sat, 31 Oct 2026 19:23:37 -0400
Subject: [PATCH 4/4] test: cover audit-only context and nested model fallback

Signed-off-by: mday-io 
---
 .../test_query_execution_context_observer.py  | 104 ++++++++++++++++++
 1 file changed, 104 insertions(+)

diff --git a/tests/core/test_query_execution_context_observer.py b/tests/core/test_query_execution_context_observer.py
index 378dc95aff..13c53a4297 100644
--- a/tests/core/test_query_execution_context_observer.py
+++ b/tests/core/test_query_execution_context_observer.py
@@ -120,6 +120,55 @@ def observer(event, facts, error):
     assert all(event[1]["batch_index"] == 0 for event in model_events)
 
 
+@pytest.mark.parametrize(("action_name", "audit_only"), [("audit_only", False), ("evaluate", True)])
+def test_audit_only_context_tags_nested_audits_and_restores(
+    action_name: str,
+    audit_only: bool,
+    make_mocked_engine_adapter: Callable[..., EngineAdapter],
+):
+    adapter = make_mocked_engine_adapter(EngineAdapter, query_execution_context=True)
+    observed = []
+
+    def observer(event, facts, error):
+        if facts["kind"] == "model":
+            observed.append((event, facts["action"], get_query_execution_context()))
+
+    console = SimpleNamespace(get_execution_observer=lambda: observer)
+    with console_observer_scope(console, query_execution_context=True):
+        with action(
+            "model",
+            action_name,
+            snapshot=_snapshot(),
+            interval=_interval(1),
+            batch_index=0,
+            audit_only=audit_only,
+        ):
+            outer = get_query_execution_context()
+            assert outer is not None
+            adapter.execute("SELECT audited model")
+            model_sql = _last_sql(adapter)
+            with action("audit", "audit", audit_name="not_null"):
+                adapter.execute("SELECT audit")
+                audit_sql = _last_sql(adapter)
+            assert get_query_execution_context() == outer
+        assert get_query_execution_context() is None
+
+    evaluation = outer.evaluation
+    assert evaluation is not None
+    assert evaluation.operation_type.value == "audit"
+    model_payload = _payload(model_sql)
+    audit_payload = _payload(audit_sql)
+    assert model_payload is not None
+    assert audit_payload is not None
+    assert model_payload["operation"]["id"] == evaluation.evaluation_id
+    assert audit_payload["operation"]["id"] == evaluation.evaluation_id
+    assert model_payload["operation"]["type"] == "audit"
+    assert model_payload["query"] == {"phase": "audit", "role": "audit"}
+    assert audit_payload["query"] == {"phase": "audit", "role": "audit"}
+    assert [event for event, _, _ in observed] == ["start", "finish"]
+    assert all(context == outer for _, _, context in observed)
+
+
 def test_explicit_disabled_scope_suppresses_transport_then_outer_context_resumes(
     make_mocked_engine_adapter: Callable[..., EngineAdapter],
 ):
@@ -162,6 +211,61 @@ def test_explicit_disabled_scope_suppresses_transport_then_outer_context_resumes
     assert resumed["operation"]["id"] == outer.evaluation.evaluation_id
 
 
+def test_nonmodel_and_malformed_nested_model_clear_then_restore_outer_context(
+    make_mocked_engine_adapter: Callable[..., EngineAdapter],
+):
+    adapter = make_mocked_engine_adapter(EngineAdapter, query_execution_context=True)
+    console = SimpleNamespace(get_execution_observer=lambda: None)
+    with console_observer_scope(console, query_execution_context=True):
+        with action(
+            "model",
+            "evaluate",
+            snapshot=_snapshot(),
+            interval=_interval(1),
+            batch_index=0,
+        ):
+            outer = get_query_execution_context()
+            assert outer is not None
+
+            with action(
+                "model",
+                "evaluate",
+                snapshot=SimpleNamespace(is_model=False),
+                interval=_interval(2),
+                batch_index=1,
+            ):
+                assert get_query_execution_context() is None
+                adapter.execute("SELECT nonmodel")
+                nonmodel_sql = _last_sql(adapter)
+            assert get_query_execution_context() == outer
+
+            with action(
+                "model",
+                "evaluate",
+                snapshot=_snapshot("analytics.malformed"),
+                interval=_interval(3),
+                batch_index=-1,
+            ):
+                assert get_query_execution_context() is None
+                adapter.execute("SELECT malformed model")
+                malformed_sql = _last_sql(adapter)
+            assert get_query_execution_context() == outer
+
+            with action("stage", "unclassified"):
+                assert get_query_execution_context() == outer
+                adapter.execute("SELECT outer context")
+                outer_sql = _last_sql(adapter)
+
+        assert get_query_execution_context() is None
+
+    assert _payload(nonmodel_sql) is None
+    assert _payload(malformed_sql) is None
+    outer_payload = _payload(outer_sql)
+    assert outer_payload is not None
+    assert outer.evaluation is not None
+    assert outer_payload["operation"]["id"] == outer.evaluation.evaluation_id
+
+
 @pytest.mark.parametrize("provider_kind", ["value", "none"])
 def test_custom_provider_remains_authoritative(
     provider_kind: str,