diff --git a/docs/concepts/query_execution_context.md b/docs/concepts/query_execution_context.md new file mode 100644 index 0000000000..60ff83ff7b --- /dev/null +++ b/docs/concepts/query_execution_context.md @@ -0,0 +1,141 @@ +# 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 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: + local: + connection: + type: duckdb + 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. + +To build a provider from native observer events, see the +[query execution context example](../examples/query_execution_context.md). + +## Configure a provider + +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 +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 also +forwards the same event, facts, and error to the observer returned by the console. + +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. + +## 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 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 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. diff --git a/docs/examples/query_execution_context.md b/docs/examples/query_execution_context.md new file mode 100644 index 0000000000..bff4f0c9ef --- /dev/null +++ b/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. 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..3d20b44189 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 @@ -2593,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() @@ -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/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 new file mode 100644 index 0000000000..a83361f388 --- /dev/null +++ b/tests/core/engine_adapter/test_query_execution_context.py @@ -0,0 +1,559 @@ +# 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 + + +class _QueryCancellation(BaseException): + pass + + +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 + + +@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 cancellation + return _context("eval-after-cancellation") + + adapter = make_mocked_engine_adapter( + EngineAdapter, + query_execution_context=True, + query_execution_context_provider=provider, + ) + + 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 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( + 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
",