fix: adding in run_id injection (#5893)

Signed-off-by: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com>

Author: cmgoffena13

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:
@@ -2605,8 +2609,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/engine_adapter/integration/test_integration_bigquery.py

@@ -1,5 +1,8 @@
 import typing as t
+from unittest import mock
+
 import pytest
+import time_machine
 from pathlib import Path
 from sqlglot import exp
 from sqlglot.optimizer.qualify_columns import quote_identifiers
@@ -441,7 +444,7 @@ def test_table_diff_table_name_matches_column_name(ctx: TestContext):
     assert row_diff.full_match_count == 1
 
 
-def test_correlation_id_in_job_labels(ctx: TestContext):
+def test_plan_correlation_id_in_job_labels(ctx: TestContext):
     model_name = ctx.table("test")
 
     sqlmesh = ctx.create_context()
@@ -469,3 +472,55 @@ def test_correlation_id_in_job_labels(ctx: TestContext):
     labels = adapter._job_params.get("labels")
     correlation_id = CorrelationId.from_plan_id(plan.plan_id)
     assert labels == {correlation_id.job_type.value.lower(): correlation_id.job_id}
+
+
+@time_machine.travel("2023-01-08 15:00:00 UTC")
+def test_run_correlation_id_in_job_labels(ctx: TestContext):
+    run_id = "test_run_id"
+    model_name = ctx.table("run_test")
+
+    sqlmesh = ctx.create_context()
+    sqlmesh.upsert_model(
+        load_sql_based_model(
+            d.parse(
+                f"""
+MODEL (
+    name {model_name},
+    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
+"""
+            )
+        )
+    )
+    sqlmesh.plan(auto_apply=True, no_prompts=True)
+
+    captured_evaluators: t.List = []
+    original_scheduler = sqlmesh.scheduler
+
+    def scheduler_wrapper(
+        environment: t.Optional[str] = None,
+        snapshot_evaluator: t.Optional[t.Any] = None,
+    ):
+        if snapshot_evaluator is not None:
+            captured_evaluators.append(snapshot_evaluator)
+        return original_scheduler(environment=environment, snapshot_evaluator=snapshot_evaluator)
+
+    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.object(sqlmesh, "scheduler", scheduler_wrapper):
+                sqlmesh.run()
+
+    assert captured_evaluators
+    adapter = t.cast(BigQueryEngineAdapter, captured_evaluators[-1].adapter)
+
+    assert adapter.correlation_id is not None
+    labels = adapter._job_params.get("labels")
+    correlation_id = CorrelationId.from_run_id(run_id)
+    assert labels == {correlation_id.job_type.value.lower(): correlation_id.job_id}

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")