Chore: Reintroduce tagging queries with correlation ID

Author: VaggelisD

sqlmesh/core/context.py

@@ -448,14 +448,8 @@ def engine_adapter(self) -> EngineAdapter:
     @property
     def snapshot_evaluator(self) -> SnapshotEvaluator:
         if not self._snapshot_evaluator:
-            self._snapshot_evaluator = SnapshotEvaluator(
-                {
-                    gateway: adapter.with_settings(log_level=logging.INFO)
-                    for gateway, adapter in self.engine_adapters.items()
-                },
-                ddl_concurrent_tasks=self.concurrent_tasks,
-                selected_gateway=self.selected_gateway,
-            )
+            self._snapshot_evaluator = self._create_snapshot_evaluator(log_level=logging.INFO)
+
         return self._snapshot_evaluator
 
     def execution_context(
@@ -517,7 +511,11 @@ def upsert_model(self, model: t.Union[str, Model], **kwargs: t.Any) -> Model:
 
         return model
 
-    def scheduler(self, environment: t.Optional[str] = None) -> Scheduler:
+    def scheduler(
+        self,
+        environment: t.Optional[str] = None,
+        snapshot_evaluator: t.Optional[SnapshotEvaluator] = None,
+    ) -> Scheduler:
         """Returns the built-in scheduler.
 
         Args:
@@ -539,9 +537,11 @@ def scheduler(self, environment: t.Optional[str] = None) -> Scheduler:
         if not snapshots:
             raise ConfigError("No models were found")
 
-        return self.create_scheduler(snapshots)
+        return self.create_scheduler(snapshots, snapshot_evaluator or self.snapshot_evaluator)
 
-    def create_scheduler(self, snapshots: t.Iterable[Snapshot]) -> Scheduler:
+    def create_scheduler(
+        self, snapshots: t.Iterable[Snapshot], snapshot_evaluator: SnapshotEvaluator
+    ) -> Scheduler:
         """Creates the built-in scheduler.
 
         Args:
@@ -552,7 +552,7 @@ def create_scheduler(self, snapshots: t.Iterable[Snapshot]) -> Scheduler:
         """
         return Scheduler(
             snapshots,
-            self.snapshot_evaluator,
+            snapshot_evaluator,
             self.state_sync,
             default_catalog=self.default_catalog,
             max_workers=self.concurrent_tasks,
@@ -2960,6 +2960,16 @@ def load_model_tests(
 
         return model_tests
 
+    def _create_snapshot_evaluator(self, **kwargs: t.Any) -> SnapshotEvaluator:
+        return SnapshotEvaluator(
+            {
+                gateway: adapter.with_settings(**kwargs)
+                for gateway, adapter in self.engine_adapters.items()
+            },
+            ddl_concurrent_tasks=self.concurrent_tasks,
+            selected_gateway=self.selected_gateway,
+        )
+
 
 class Context(GenericContext[Config]):
     CONFIG_TYPE = Config

sqlmesh/core/engine_adapter/base.py

@@ -147,15 +147,15 @@ def __init__(
         self._multithreaded = multithreaded
         self.correlation_id = correlation_id
 
-    def with_settings(self, log_level: int, **kwargs: t.Any) -> EngineAdapter:
+    def with_settings(self, log_level: int = logging.DEBUG, **kwargs: t.Any) -> EngineAdapter:
         adapter = self.__class__(
             self._connection_pool,
             dialect=self.dialect,
             sql_gen_kwargs=self._sql_gen_kwargs,
             default_catalog=self._default_catalog,
             execute_log_level=log_level,
             register_comments=self._register_comments,
-            null_connection=True,
+            null_connection=self._extra_config.pop("null_connection", True),
             multithreaded=self._multithreaded,
             pretty_sql=self._pretty_sql,
             **self._extra_config,

sqlmesh/core/plan/evaluator.py

@@ -38,6 +38,7 @@
 )
 from sqlmesh.utils import to_snake_case
 from sqlmesh.core.state_sync import StateSync
+from sqlmesh.utils import CorrelationId
 from sqlmesh.utils.concurrency import NodeExecutionFailedError
 from sqlmesh.utils.errors import PlanError, SQLMeshError
 from sqlmesh.utils.dag import DAG
@@ -71,7 +72,7 @@ def __init__(
         self,
         state_sync: StateSync,
         snapshot_evaluator: SnapshotEvaluator,
-        create_scheduler: t.Callable[[t.Iterable[Snapshot]], Scheduler],
+        create_scheduler: t.Callable[[t.Iterable[Snapshot], SnapshotEvaluator], Scheduler],
         default_catalog: t.Optional[str],
         console: t.Optional[Console] = None,
     ):
@@ -89,6 +90,7 @@ def evaluate(
     ) -> None:
         self._circuit_breaker = circuit_breaker
 
+        self.set_correlation_id(CorrelationId.from_plan_id(plan.plan_id))
         self.console.start_plan_evaluation(plan)
         analytics.collector.on_plan_apply_start(
             plan=plan,
@@ -228,7 +230,7 @@ def visit_backfill_stage(self, stage: stages.BackfillStage, plan: EvaluatablePla
             self.console.log_success("SKIP: No model batches to execute")
             return
 
-        scheduler = self.create_scheduler(stage.all_snapshots.values())
+        scheduler = self.create_scheduler(stage.all_snapshots.values(), self.snapshot_evaluator)
         errors, _ = scheduler.run_merged_intervals(
             merged_intervals=stage.snapshot_to_intervals,
             deployability_index=stage.deployability_index,
@@ -249,7 +251,7 @@ def visit_audit_only_run_stage(
             return
 
         # If there are any snapshots to be audited, we'll reuse the scheduler's internals to audit them
-        scheduler = self.create_scheduler(audit_snapshots)
+        scheduler = self.create_scheduler(audit_snapshots, self.snapshot_evaluator)
         completion_status = scheduler.audit(
             plan.environment,
             plan.start,
@@ -348,6 +350,13 @@ def visit_finalize_environment_stage(
     ) -> None:
         self.state_sync.finalize(plan.environment)
 
+    def set_correlation_id(self, correlation_id: CorrelationId) -> None:
+        for key, adapter in self.snapshot_evaluator.adapters.items():
+            if correlation_id != adapter.correlation_id:
+                self.snapshot_evaluator.adapters[key] = adapter.with_settings(
+                    correlation_id=correlation_id
+                )
+
     def _promote_snapshots(
         self,
         plan: EvaluatablePlan,

sqlmesh/core/snapshot/evaluator.py

@@ -122,11 +122,7 @@ def __init__(
         self.adapters = (
             adapters if isinstance(adapters, t.Dict) else {selected_gateway or "": adapters}
         )
-        self.adapter = (
-            next(iter(self.adapters.values()))
-            if not selected_gateway
-            else self.adapters[selected_gateway]
-        )
+        self.selected_gateway = selected_gateway
         self.ddl_concurrent_tasks = ddl_concurrent_tasks
 
     def evaluate(
@@ -603,6 +599,14 @@ def close(self) -> None:
         except Exception:
             logger.exception("Failed to close Snapshot Evaluator")
 
+    @property
+    def adapter(self) -> EngineAdapter:
+        return (
+            next(iter(self.adapters.values()))
+            if not self.selected_gateway
+            else self.adapters[self.selected_gateway]
+        )
+
     def _evaluate_snapshot(
         self,
         snapshot: Snapshot,

tests/core/test_integration.py

@@ -71,6 +71,7 @@
 from sqlmesh.utils.errors import NoChangesPlanError, SQLMeshError, PlanError, ConfigError
 from sqlmesh.utils.pydantic import validate_string
 from tests.conftest import DuckDBMetadata, SushiDataValidator
+from sqlmesh.utils import CorrelationId
 from tests.utils.test_helpers import use_terminal_console
 from tests.utils.test_filesystem import create_temp_file
 
@@ -6815,3 +6816,28 @@ def test_scd_type_2_full_restatement_no_start_date(init_and_plan_context: t.Call
             # valid_from should be the epoch, valid_to should be NaT
             assert str(row["valid_from"]) == "1970-01-01 00:00:00"
             assert pd.isna(row["valid_to"])
+
+
+def test_plan_evaluator_correlation_id(tmp_path: Path):
+    def _correlation_id_in_sqls(correlation_id: CorrelationId, mock_logger):
+        sqls = [call[0][0] for call in mock_logger.call_args_list]
+        return any(f"/* {correlation_id} */" in sql for sql in sqls)
+
+    ctx = Context(paths=[tmp_path], config=Config())
+
+    # Case: Ensure that the correlation id (plan_id) is included in the SQL for each plan
+    for i in range(2):
+        create_temp_file(
+            tmp_path,
+            Path("models", "test.sql"),
+            f"MODEL (name test.a, kind FULL); SELECT {i} AS col",
+        )
+
+        with mock.patch("sqlmesh.core.engine_adapter.base.EngineAdapter._log_sql") as mock_logger:
+            ctx.load()
+            plan = ctx.plan(auto_apply=True, no_prompts=True)
+
+        correlation_id = CorrelationId.from_plan_id(plan.plan_id)
+        assert str(correlation_id) == f"SQLMESH_PLAN: {plan.plan_id}"
+
+        assert _correlation_id_in_sqls(correlation_id, mock_logger)