Skip to content

Commit ca1b1e6

Browse files
committed
feat: add automatic model query context
Signed-off-by: mday-io <mdaytn@gmail.com>
1 parent 41f5d27 commit ca1b1e6

10 files changed

Lines changed: 714 additions & 46 deletions

File tree

‎docs/concepts/query_execution_context.md‎

Lines changed: 73 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,11 @@ SQLMesh can add structured metadata to warehouse statements so a query-history
44
tool can connect a statement to the plan or run, model, and evaluation that issued
55
it. The metadata is a versioned JSON object in a `QUERY_CONTEXT` SQL comment.
66

7-
Enable `query_execution_context` on a connection and set a provider on its
8-
engine adapter. The provider returns a `QueryExecutionContext`. The flag defaults
9-
to `false`. For YAML configuration, set the option under the gateway's
10-
`connection`:
7+
Set `query_execution_context` to `true` on a connection to enable bounded query
8+
context comments on supported adapters. The flag defaults to `false`. With this
9+
flag enabled, `Context.apply()` and built-in scheduler runs add model context to
10+
SQL on supported adapters without a custom provider. For YAML configuration,
11+
set the option under the gateway's `connection`:
1112

1213
```yaml
1314
gateways:
@@ -17,20 +18,59 @@ gateways:
1718
query_execution_context: true
1819
```
1920
20-
The flag enables context comments but does not infer model context. Without a
21-
provider, SQLMesh adds parent metadata only when a plan or run correlation ID is
22-
available. SQLMesh calls the provider after rendering each statement and before
23-
executing it. The call runs in the current worker's context, so concurrent model
24-
evaluations can return different IDs for a shared adapter.
21+
## Automatic model context
2522
26-
To build a provider from native observer events, see the
27-
[query execution context example](../examples/query_execution_context.md).
23+
For each model batch, SQLMesh creates a context with a new operation ID, the
24+
model's canonical name, an opaque snapshot ID, and the batch interval and index
25+
when available. Intervals use UTC and an exclusive end.
2826
29-
## Configure a provider
27+
The context stays active for SQL inside the batch, including nested audits.
28+
Automatic context uses the broad phases `materialize`, `audit`, and `ddl`. The
29+
`role` field does not identify the engine operation for each SQL statement.
3030

31-
Set the provider on warehouse adapters before loading a Python `Context`. When
32-
`Context` creates its snapshot evaluator, it clones those adapters and preserves
33-
the provider and flag.
31+
Outside a model batch, SQLMesh can add parent metadata when a plan or run
32+
correlation ID is available. Use a custom provider when you need caller-defined
33+
context for other statements.
34+
35+
## Read the context from an observer
36+
37+
SQLMesh establishes the model context before it calls an observer's `start`
38+
callback. It restores the context after the `finish` callback. Both callbacks can
39+
read the active immutable value with `get_query_execution_context()`:
40+
41+
```python
42+
from sqlmesh.core.query_execution import get_query_execution_context
43+
44+
45+
def observe(event, facts, error):
46+
context = get_query_execution_context()
47+
if context is not None and context.evaluation is not None:
48+
print(context.evaluation.evaluation_id)
49+
```
50+
51+
For console observer registration and callback behavior, see
52+
[Execution observations](../reference/execution_observations.md).
53+
54+
SQLMesh forwards each observer event to the configured console observer once. If
55+
the console has no observer, model context still works. Nested observer scopes
56+
restore the outer observer and context on exit. Use `observer_scope(None)` to
57+
disable the built-in observer for its block.
58+
59+
An observer callback can read the active context with
60+
`get_query_execution_context()`. SQL issued inside the callback does not use that
61+
automatic context. A custom adapter provider remains authoritative. Without one,
62+
an active plan or run correlation ID can still produce parent-only metadata.
63+
64+
## Configure a custom provider
65+
66+
Use a custom provider when you need caller-defined operation IDs or context
67+
outside a model batch. Set it on warehouse adapters before loading a Python
68+
`Context`. When `Context` creates its snapshot evaluator, it clones those
69+
adapters and preserves the provider and flag.
70+
71+
SQLMesh calls the provider after rendering each statement and before executing
72+
it. The provider runs in the current worker's context, so concurrent workers can
73+
return different values for a shared adapter.
3474

3575
```python
3676
from sqlmesh.core.context import Context
@@ -53,23 +93,32 @@ with observer_scope(observer):
5393

5494
`QueryContextObserver` returns the immutable context active in its worker, or
5595
`None` when no context applies. It assigns an evaluation ID when SQLMesh reports
56-
a model batch and restores the previous context when that event finishes. It also
96+
a model batch and restores the previous context when that event finishes. It
5797
forwards the same event, facts, and error to the observer returned by the console.
5898

99+
A custom provider replaces the built-in provider for its adapter. If it returns
100+
`None`, SQLMesh does not use automatic model context for that statement. SQLMesh
101+
can still add parent-only metadata when the adapter has a correlation ID.
102+
59103
Adapters created with `with_settings` inherit the flag and provider. Pass
60104
`query_execution_context_provider=None` to clear the inherited provider. Setting
61105
the flag to `false` disables provider calls and comment output.
62106

63107
When `state_connection` is omitted, the built-in scheduler uses the warehouse
64-
connection for its state adapter and disables query-context comments on that
65-
adapter. To add context to state SQL, configure `state_connection` explicitly
66-
and set a provider on its adapter.
108+
connection for its state adapter and disables query context on that adapter. An
109+
explicitly configured `state_connection` uses its own `query_execution_context`
110+
setting. The automatic observer activates only if a supported warehouse adapter
111+
has `query_execution_context` enabled and no custom provider. An enabled state
112+
adapter without a custom provider can use that context during a model batch.
113+
Setting the flag only on `state_connection` does not activate the observer. Use a
114+
custom provider for state SQL outside a model batch.
67115

68116
## Read the payload
69117

70118
The version-one payload uses the SQLMesh model's canonical name and opaque
71-
snapshot identifier. Operation intervals use UTC half-open boundaries. Fields
72-
that are unavailable are omitted.
119+
snapshot identifier. Operation intervals use UTC half-open boundaries, and the
120+
payload includes `batch_index` when available. Fields that are unavailable are
121+
omitted.
73122

74123
```json
75124
{
@@ -93,9 +142,9 @@ that are unavailable are omitted.
93142
}
94143
```
95144

96-
The provider assigns an operation ID to the statements for one model evaluation,
97-
including nested audits. A later evaluation gets a new ID. The adapter preserves
98-
the provider's ID. Consumers can compare it with the execution observer's record.
145+
Each model batch gets one operation ID, which nested audits reuse. The built-in
146+
provider creates this ID. A custom provider supplies its own. The adapter preserves
147+
the provider's ID when it adds the context comment.
99148

100149
The adapter's correlation ID sets `parent` and replaces any parent supplied by
101150
the provider. If the adapter has no correlation ID, SQLMesh clears the provider's

‎docs/reference/configuration.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,7 @@ Most parameters are specific to the connection engine `type` - see [below](#engi
212212
| `register_comments` | Whether SQLMesh should register model comments with the SQL engine (if the engine supports it). (Default: `true`.) | bool | N |
213213
| `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 |
214214
| `pretty_sql` | If SQL should be formatted before being executed, not recommended in a production setting. (Default: `false`.) | bool | N |
215-
| `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 |
215+
| `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 |
216216

217217
#### Engine-specific
218218

‎sqlmesh/core/engine_adapter/base.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2706,7 +2706,14 @@ def _attach_query_execution_context(self, sql: str) -> str:
27062706

27072707
provider = self.query_execution_context_provider
27082708
try:
2709-
context = provider() if provider is not None else None
2709+
if provider is None:
2710+
from sqlmesh.core.query_execution_context_observer import (
2711+
get_automatic_query_execution_context,
2712+
)
2713+
2714+
context = get_automatic_query_execution_context()
2715+
else:
2716+
context = provider()
27102717
except Exception:
27112718
self._warn_query_execution_context_once(
27122719
"provider",

‎sqlmesh/core/execution_observation.py‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,14 +119,33 @@ def execution_context_factory() -> t.Optional[t.Callable[[], contextvars.Context
119119

120120

121121
@contextlib.contextmanager
122-
def console_observer_scope(console: t.Any) -> t.Iterator[None]:
122+
def console_observer_scope(
123+
console: t.Any, *, query_execution_context: bool = False
124+
) -> t.Iterator[None]:
123125
"""Resolve once at the outer boundary; an explicit None scope disables hooks."""
124126
if _scoped.get():
127+
observer = _observer.get()
128+
if query_execution_context and observer is not None:
129+
from sqlmesh.core.query_execution_context_observer import (
130+
with_automatic_query_execution_observer,
131+
)
132+
133+
automatic_observer = with_automatic_query_execution_observer(observer)
134+
if automatic_observer is not observer:
135+
with observer_scope(automatic_observer):
136+
yield
137+
return
125138
yield
126139
return
127140
try:
128141
observer = console.get_execution_observer()
129142
except Exception:
130143
observer = None
144+
if query_execution_context:
145+
from sqlmesh.core.query_execution_context_observer import (
146+
with_automatic_query_execution_observer,
147+
)
148+
149+
observer = with_automatic_query_execution_observer(observer)
131150
with observer_scope(observer):
132151
yield

‎sqlmesh/core/plan/evaluator.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@
1919
import typing as t
2020
from sqlmesh.core import analytics
2121
from sqlmesh.core.execution_observation import action, console_observer_scope
22+
from sqlmesh.core.query_execution_context_observer import (
23+
should_enable_automatic_query_execution_context,
24+
)
2225
from sqlmesh.core import constants as c
2326
from sqlmesh.core.console import Console, get_console
2427
from sqlmesh.core.environment import EnvironmentNamingInfo, execute_environment_statements
@@ -103,7 +106,12 @@ def evaluate(
103106

104107
try:
105108
plan_stages = stages.build_plan_stages(plan, self.state_sync, self.default_catalog)
106-
with console_observer_scope(self.console):
109+
with console_observer_scope(
110+
self.console,
111+
query_execution_context=should_enable_automatic_query_execution_context(
112+
self.snapshot_evaluator.adapters.values()
113+
),
114+
):
107115
self._evaluate_stages(plan_stages, plan)
108116
except Exception as e:
109117
analytics.collector.on_plan_apply_end(plan_id=plan.plan_id, error=e)

‎sqlmesh/core/query_execution.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -170,10 +170,10 @@ def _query_execution_context_transport_scope() -> t.Iterator[None]:
170170

171171

172172
@contextlib.contextmanager
173-
def query_execution_scope(context: QueryExecutionContext) -> t.Iterator[None]:
174-
"""Set a query execution context for this scope and restore it on exit."""
175-
if not isinstance(context, QueryExecutionContext):
176-
raise TypeError("context must be a QueryExecutionContext")
173+
def query_execution_scope(context: t.Optional[QueryExecutionContext]) -> t.Iterator[None]:
174+
"""Set or clear the query execution context for this scope."""
175+
if context is not None and not isinstance(context, QueryExecutionContext):
176+
raise TypeError("context must be a QueryExecutionContext or None")
177177
token = _current_query_execution_context.set(context)
178178
try:
179179
yield

0 commit comments

Comments
 (0)