Feat: Tag all queries under a plan

Author: VaggelisD

sqlmesh/core/config/scheduler.py

@@ -28,11 +28,14 @@ class SchedulerConfig(abc.ABC):
     """Abstract base class for Scheduler configurations."""
 
     @abc.abstractmethod
-    def create_plan_evaluator(self, context: GenericContext) -> PlanEvaluator:
+    def create_plan_evaluator(
+        self, context: GenericContext, job_id: t.Optional[str] = None
+    ) -> PlanEvaluator:
         """Creates a Plan Evaluator instance.
 
         Args:
             context: The SQLMesh Context.
+            job_id: The plan ID.
         """
 
     @abc.abstractmethod
@@ -127,10 +130,12 @@ class BuiltInSchedulerConfig(_EngineAdapterStateSyncSchedulerConfig, BaseConfig)
 
     type_: t.Literal["builtin"] = Field(alias="type", default="builtin")
 
-    def create_plan_evaluator(self, context: GenericContext) -> PlanEvaluator:
+    def create_plan_evaluator(
+        self, context: GenericContext, job_id: t.Optional[str] = None
+    ) -> PlanEvaluator:
         return BuiltInPlanEvaluator(
             state_sync=context.state_sync,
-            snapshot_evaluator=context.snapshot_evaluator,
+            snapshot_evaluator=context.snapshot_evaluator(job_id),
             create_scheduler=context.create_scheduler,
             default_catalog=context.default_catalog,
             console=context.console,

sqlmesh/core/context.py

@@ -446,12 +446,14 @@ def engine_adapter(self) -> EngineAdapter:
             self._engine_adapter = self.connection_config.create_engine_adapter()
         return self._engine_adapter
 
-    @property
-    def snapshot_evaluator(self) -> SnapshotEvaluator:
-        if not self._snapshot_evaluator:
+    def snapshot_evaluator(self, job_id: t.Optional[str] = None) -> SnapshotEvaluator:
+        # Cache snapshot evaluators by job_id to avoid old job_ids being attached to future Context operations
+        if not self._snapshot_evaluator or any(
+            adapter._job_id != job_id for adapter in self._snapshot_evaluator.adapters.values()
+        ):
             self._snapshot_evaluator = SnapshotEvaluator(
                 {
-                    gateway: adapter.with_log_level(logging.INFO)
+                    gateway: adapter.with_log_level(logging.INFO, job_id)
                     for gateway, adapter in self.engine_adapters.items()
                 },
                 ddl_concurrent_tasks=self.concurrent_tasks,
@@ -538,7 +540,9 @@ def scheduler(self, environment: t.Optional[str] = None) -> Scheduler:
 
         return self.create_scheduler(snapshots)
 
-    def create_scheduler(self, snapshots: t.Iterable[Snapshot]) -> Scheduler:
+    def create_scheduler(
+        self, snapshots: t.Iterable[Snapshot], job_id: t.Optional[str] = None
+    ) -> Scheduler:
         """Creates the built-in scheduler.
 
         Args:
@@ -549,7 +553,7 @@ def create_scheduler(self, snapshots: t.Iterable[Snapshot]) -> Scheduler:
         """
         return Scheduler(
             snapshots,
-            self.snapshot_evaluator,
+            self.snapshot_evaluator(job_id),
             self.state_sync,
             default_catalog=self.default_catalog,
             max_workers=self.concurrent_tasks,
@@ -714,7 +718,7 @@ def run(
             NotificationEvent.RUN_START, environment=environment
         )
         analytics_run_id = analytics.collector.on_run_start(
-            engine_type=self.snapshot_evaluator.adapter.dialect,
+            engine_type=self.snapshot_evaluator().adapter.dialect,
             state_sync_type=self.state_sync.state_type(),
         )
         self._load_materializations()
@@ -1076,7 +1080,7 @@ def evaluate(
             and not parent_snapshot.categorized
         ]
 
-        df = self.snapshot_evaluator.evaluate_and_fetch(
+        df = self.snapshot_evaluator().evaluate_and_fetch(
             snapshot,
             start=start,
             end=end,
@@ -2110,7 +2114,7 @@ def audit(
         errors = []
         skipped_count = 0
         for snapshot in snapshots:
-            for audit_result in self.snapshot_evaluator.audit(
+            for audit_result in self.snapshot_evaluator().audit(
                 snapshot=snapshot,
                 start=start,
                 end=end,
@@ -2142,7 +2146,7 @@ def audit(
             self.console.log_status_update(f"Got {error.count} results, expected 0.")
             if error.query:
                 self.console.show_sql(
-                    f"{error.query.sql(dialect=self.snapshot_evaluator.adapter.dialect)}"
+                    f"{error.query.sql(dialect=self.snapshot_evaluator().adapter.dialect)}"
                 )
 
         self.console.log_status_update("Done.")
@@ -2388,7 +2392,7 @@ def _run(
         return completion_status
 
     def _apply(self, plan: Plan, circuit_breaker: t.Optional[t.Callable[[], bool]]) -> None:
-        self._scheduler.create_plan_evaluator(self).evaluate(
+        self._scheduler.create_plan_evaluator(self, job_id=plan.plan_id).evaluate(
             plan.to_evaluatable(), circuit_breaker=circuit_breaker
         )
 
@@ -2682,7 +2686,7 @@ def _run_janitor(self, ignore_ttl: bool = False) -> None:
         )
 
         # Remove the expired snapshots tables
-        self.snapshot_evaluator.cleanup(
+        self.snapshot_evaluator().cleanup(
             target_snapshots=cleanup_targets,
             on_complete=self.console.update_cleanup_progress,
         )

sqlmesh/core/engine_adapter/base.py

@@ -123,6 +123,7 @@ def __init__(
         pre_ping: bool = False,
         pretty_sql: bool = False,
         shared_connection: bool = False,
+        job_id: t.Optional[str] = None,
         **kwargs: t.Any,
     ):
         self.dialect = dialect.lower() or self.DIALECT
@@ -144,8 +145,9 @@ def __init__(
         self._pre_ping = pre_ping
         self._pretty_sql = pretty_sql
         self._multithreaded = multithreaded
+        self._job_id = job_id
 
-    def with_log_level(self, level: int) -> EngineAdapter:
+    def with_log_level(self, level: int, job_id: t.Optional[str] = None) -> EngineAdapter:
         adapter = self.__class__(
             self._connection_pool,
             dialect=self.dialect,
@@ -156,6 +158,7 @@ def with_log_level(self, level: int) -> EngineAdapter:
             null_connection=True,
             multithreaded=self._multithreaded,
             pretty_sql=self._pretty_sql,
+            job_id=job_id,
             **self._extra_config,
         )
 
@@ -2175,6 +2178,9 @@ def execute(
                 else:
                     sql = t.cast(str, e)
 
+                if self._job_id:
+                    sql = f"/* sqlmesh_ref: {self._job_id} */ {sql}"
+
                 self._log_sql(
                     sql,
                     expression=e if isinstance(e, exp.Expression) else None,

tests/conftest.py

@@ -266,7 +266,7 @@ def duck_conn() -> duckdb.DuckDBPyConnection:
 def push_plan(context: Context, plan: Plan) -> None:
     plan_evaluator = BuiltInPlanEvaluator(
         context.state_sync,
-        context.snapshot_evaluator,
+        context.snapshot_evaluator(),
         context.create_scheduler,
         context.default_catalog,
     )

tests/core/test_integration.py

@@ -1137,7 +1137,7 @@ def test_non_breaking_change_after_forward_only_in_dev(
     init_and_plan_context: t.Callable, has_view_binding: bool
 ):
     context, plan = init_and_plan_context("examples/sushi")
-    context.snapshot_evaluator.adapter.HAS_VIEW_BINDING = has_view_binding
+    context.snapshot_evaluator().adapter.HAS_VIEW_BINDING = has_view_binding
     context.apply(plan)
 
     model = context.get_model("sushi.waiter_revenue_by_day")
@@ -6467,3 +6467,25 @@ def plan_with_output(ctx: Context, environment: str):
     for environment in ["dev", "prod"]:
         context_diff = ctx._context_diff(environment)
         assert context_diff.environment == environment
+
+
+def test_plan_evaluator_job_id(tmp_path: Path, mocker: MockerFixture):
+    def _to_sqls(mock_logger):
+        return [call[0][0] for call in mock_logger.call_args_list]
+
+    create_temp_file(
+        tmp_path, Path("models") / "test.sql", "MODEL (name test.a, kind FULL); SELECT 1 AS col"
+    )
+
+    # Case 1: Ensure that the job id (plan_id) is included in the SQL
+    with mock.patch("sqlmesh.core.engine_adapter.base.EngineAdapter._log_sql") as mock_logger:
+        ctx = Context(paths=[tmp_path], config=Config())
+        plan = ctx.plan(auto_apply=True, no_prompts=True)
+
+    assert any(f"/* sqlmesh_ref: {plan.plan_id} */" in sql for sql in _to_sqls(mock_logger))
+
+    # Case 2: Ensure that the previous job id is not included in the SQL for other operations
+    with mock.patch("sqlmesh.core.engine_adapter.base.EngineAdapter._log_sql") as mock_logger:
+        ctx.snapshot_evaluator().adapter.execute("SELECT 1")
+
+    assert not any(f"/* sqlmesh_ref: {plan.plan_id} */" in sql for sql in _to_sqls(mock_logger))

tests/core/test_plan_evaluator.py

@@ -59,7 +59,7 @@ def test_builtin_evaluator_push(sushi_context: Context, make_snapshot):
 
     evaluator = BuiltInPlanEvaluator(
         sushi_context.state_sync,
-        sushi_context.snapshot_evaluator,
+        sushi_context.snapshot_evaluator(),
         sushi_context.create_scheduler,
         sushi_context.default_catalog,
         console=sushi_context.console,