feat: add automatic model query context

docs/concepts/query_execution_context.md

@@ -0,0 +1,190 @@
+# 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.
+
+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:
+  local:
+    connection:
+      type: duckdb
+      query_execution_context: true
+```
+
+## Automatic model context
+
+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.
+
+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.
+
+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
+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")
+```
+
+`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
+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 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, and the
+payload includes `batch_index` when available. 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"}
+}
+```
+
+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
+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
+opaque and does not change that identity.
+
+## SQL comment behavior
+
+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 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.

docs/examples/query_execution_context.md

@@ -0,0 +1,25 @@
+# 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 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
+payload, provider contract, and transport limits.

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

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

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

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

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