fix: adding in run_id injection

sqlmesh/core/context.py

@@ -118,7 +118,7 @@
     filter_tests_by_patterns,
 )
 from sqlmesh.core.user import User
-from sqlmesh.utils import UniqueKeyDict, Verbosity
+from sqlmesh.utils import CorrelationId, UniqueKeyDict, Verbosity
 from sqlmesh.utils.concurrency import concurrent_apply_to_values
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.date import (
@@ -811,6 +811,9 @@ def run(
             engine_type=self.snapshot_evaluator.adapter.dialect,
             state_sync_type=self.state_sync.state_type(),
         )
+        snapshot_evaluator = self.snapshot_evaluator.set_correlation_id(
+            CorrelationId.from_run_id(analytics_run_id)
+        )
         self._load_materializations()
 
         env_check_attempts_num = max(
@@ -863,6 +866,7 @@ def _has_environment_changed() -> bool:
                     select_models=select_models,
                     circuit_breaker=_has_environment_changed,
                     no_auto_upstream=no_auto_upstream,
+                    snapshot_evaluator=snapshot_evaluator,
                 )
                 done = True
             except CircuitBreakerError:
@@ -2586,8 +2590,9 @@ def _run(
         select_models: t.Optional[t.Collection[str]],
         circuit_breaker: t.Optional[t.Callable[[], bool]],
         no_auto_upstream: bool,
+        snapshot_evaluator: t.Optional[SnapshotEvaluator] = None,
     ) -> CompletionStatus:
-        scheduler = self.scheduler(environment=environment)
+        scheduler = self.scheduler(environment=environment, snapshot_evaluator=snapshot_evaluator)
         snapshots = scheduler.snapshots
 
         if select_models is not None:

sqlmesh/utils/__init__.py

@@ -409,6 +409,10 @@ def __str__(self) -> str:
     def from_plan_id(cls, plan_id: str) -> CorrelationId:
         return CorrelationId(JobType.PLAN, plan_id)
 
+    @classmethod
+    def from_run_id(cls, run_id: str) -> CorrelationId:
+        return CorrelationId(JobType.RUN, run_id)
+
 
 def get_source_columns_to_types(
     columns_to_types: t.Dict[str, exp.DataType],

tests/core/integration/test_model_kinds.py

@@ -2332,6 +2332,48 @@ def _correlation_id_in_sqls(correlation_id: CorrelationId, mock_logger):
         assert _correlation_id_in_sqls(correlation_id, mock_logger)
 
 
+@time_machine.travel("2023-01-08 15:00:00 UTC")
+def test_run_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)
+
+    run_id = "test_run_id"
+    ctx = Context(paths=[tmp_path], config=Config())
+
+    create_temp_file(
+        tmp_path,
+        Path("models", "test.sql"),
+        """
+MODEL (
+    name test.a,
+    kind INCREMENTAL_BY_TIME_RANGE (
+        time_column event_ts
+    ),
+    cron '@daily',
+    start '2023-01-07'
+);
+SELECT 1 AS col, '2023-01-07' AS event_ts
+""",
+    )
+
+    ctx.load()
+    ctx.plan(auto_apply=True, no_prompts=True)
+
+    with time_machine.travel("2023-01-09 00:00:00 UTC"):
+        with mock.patch(
+            "sqlmesh.core.context.analytics.collector.on_run_start", return_value=run_id
+        ):
+            with mock.patch(
+                "sqlmesh.core.engine_adapter.base.EngineAdapter._log_sql"
+            ) as mock_logger:
+                ctx.run()
+
+    correlation_id = CorrelationId.from_run_id(run_id)
+    assert str(correlation_id) == f"SQLMESH_RUN: {run_id}"
+    assert _correlation_id_in_sqls(correlation_id, mock_logger)
+
+
 @time_machine.travel("2023-01-08 15:00:00 UTC")
 def test_scd_type_2_regular_run_with_offset(init_and_plan_context: t.Callable):
     context, plan = init_and_plan_context("examples/sushi")