From bacf1481a68f047c73dfc119f58497fb80a3914f Mon Sep 17 00:00:00 2001 From: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com> Date: Wed, 8 Jul 2026 21:35:22 -0600 Subject: [PATCH 1/3] fix: adding in run_id injection Signed-off-by: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com> --- sqlmesh/core/context.py | 9 +++++++-- sqlmesh/utils/__init__.py | 4 ++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py index 5902977331..f3f37d92bb 100644 --- a/sqlmesh/core/context.py +++ b/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: diff --git a/sqlmesh/utils/__init__.py b/sqlmesh/utils/__init__.py index 59605893ba..da8d4f85a4 100644 --- a/sqlmesh/utils/__init__.py +++ b/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], From 85d9f9a1a23cfef535b1fd4800f120a4e038655b Mon Sep 17 00:00:00 2001 From: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com> Date: Wed, 8 Jul 2026 21:45:52 -0600 Subject: [PATCH 2/3] adding in test for run_in injection Signed-off-by: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com> --- tests/core/integration/test_model_kinds.py | 42 ++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tests/core/integration/test_model_kinds.py b/tests/core/integration/test_model_kinds.py index 1cc1bf7aeb..108fd1cb01 100644 --- a/tests/core/integration/test_model_kinds.py +++ b/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") From d1e9ff7b036fd0f5f6947bcfa531a22bd338ac62 Mon Sep 17 00:00:00 2001 From: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com> Date: Thu, 9 Jul 2026 16:22:54 -0600 Subject: [PATCH 3/3] adding in bigquery label test for run_id Signed-off-by: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com> --- .../integration/test_integration_bigquery.py | 57 ++++++++++++++++++- 1 file changed, 56 insertions(+), 1 deletion(-) diff --git a/tests/core/engine_adapter/integration/test_integration_bigquery.py b/tests/core/engine_adapter/integration/test_integration_bigquery.py index 0a6dd6b2a4..ce31c255ad 100644 --- a/tests/core/engine_adapter/integration/test_integration_bigquery.py +++ b/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}