Refactor create_scheduler

Author: VaggelisD

sqlmesh/core/context.py

@@ -540,10 +540,10 @@ def scheduler(self, environment: t.Optional[str] = None) -> Scheduler:
         if not snapshots:
             raise ConfigError("No models were found")
 
-        return self.create_scheduler(snapshots)
+        return self.create_scheduler(snapshots, self.snapshot_evaluator())
 
     def create_scheduler(
-        self, snapshots: t.Iterable[Snapshot], correlation_id: t.Optional[CorrelationId] = None
+        self, snapshots: t.Iterable[Snapshot], snapshot_evaluator: SnapshotEvaluator
     ) -> Scheduler:
         """Creates the built-in scheduler.
 
@@ -555,7 +555,7 @@ def create_scheduler(
         """
         return Scheduler(
             snapshots,
-            self.snapshot_evaluator(correlation_id),
+            snapshot_evaluator,
             self.state_sync,
             default_catalog=self.default_catalog,
             max_workers=self.concurrent_tasks,

sqlmesh/core/plan/evaluator.py

@@ -41,7 +41,6 @@
 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__)
@@ -73,7 +72,7 @@ class BuiltInPlanEvaluator(PlanEvaluator):
     def __init__(
         self,
         state_sync: StateSync,
-        create_scheduler: t.Callable[[t.Iterable[Snapshot], t.Optional[CorrelationId]], Scheduler],
+        create_scheduler: t.Callable[[t.Iterable[Snapshot], SnapshotEvaluator], Scheduler],
         default_catalog: t.Optional[str],
         console: t.Optional[Console] = None,
     ):
@@ -231,9 +230,7 @@ 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(), CorrelationId.from_plan_id(plan.plan_id)
-        )
+        scheduler = self.create_scheduler(stage.all_snapshots.values(), self.snapshot_evaluator)
         errors, _ = scheduler.run_merged_intervals(
             merged_intervals=stage.snapshot_to_intervals,
             deployability_index=stage.deployability_index,
@@ -254,7 +251,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, CorrelationId.from_plan_id(plan.plan_id))
+        scheduler = self.create_scheduler(audit_snapshots, self.snapshot_evaluator)
         completion_status = scheduler.audit(
             plan.environment,
             plan.start,