Forward correlation_id to create_scheduler calls

Author: VaggelisD

sqlmesh/core/engine_adapter/bigquery.py

@@ -208,6 +208,11 @@ def _begin_session(self, properties: SessionProperties) -> None:
                 "Invalid value for `session_properties.query_label`. Must be an array or tuple."
             )
 
+        if self.correlation_id:
+            parsed_query_label.append(
+                (self.correlation_id.job_type.value.lower(), self.correlation_id.job_id)
+            )
+
         if parsed_query_label:
             query_label_str = ",".join([":".join(label) for label in parsed_query_label])
             query = f'SET @@query_label = "{query_label_str}";SELECT 1;'

sqlmesh/core/plan/evaluator.py

@@ -41,6 +41,7 @@
 from sqlmesh.utils.concurrency import NodeExecutionFailedError
 from sqlmesh.utils.errors import PlanError, SQLMeshError
 from sqlmesh.utils.dag import DAG
+from sqlmesh.utils import CorrelationId
 from sqlmesh.utils.date import now
 
 logger = logging.getLogger(__name__)
@@ -72,7 +73,7 @@ class BuiltInPlanEvaluator(PlanEvaluator):
     def __init__(
         self,
         state_sync: StateSync,
-        create_scheduler: t.Callable[[t.Iterable[Snapshot]], Scheduler],
+        create_scheduler: t.Callable[[t.Iterable[Snapshot], t.Optional[CorrelationId]], Scheduler],
         default_catalog: t.Optional[str],
         console: t.Optional[Console] = None,
     ):
@@ -230,7 +231,9 @@ def visit_backfill_stage(self, stage: stages.BackfillStage, plan: EvaluatablePla
             self.console.log_success("SKIP: No model batches to execute")
             return
 
-        scheduler = self.create_scheduler(stage.all_snapshots.values())
+        scheduler = self.create_scheduler(
+            stage.all_snapshots.values(), CorrelationId.from_plan_id(plan.plan_id)
+        )
         errors, _ = scheduler.run_merged_intervals(
             merged_intervals=stage.snapshot_to_intervals,
             deployability_index=stage.deployability_index,
@@ -251,7 +254,7 @@ def visit_audit_only_run_stage(
             return
 
         # If there are any snapshots to be audited, we'll reuse the scheduler's internals to audit them
-        scheduler = self.create_scheduler(audit_snapshots)
+        scheduler = self.create_scheduler(audit_snapshots, CorrelationId.from_plan_id(plan.plan_id))
         completion_status = scheduler.audit(
             plan.environment,
             plan.start,