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