Repository: TobikoData/sqlmesh
Pull request #3114 (2024-09-09 23:13:27)

Chore: Use factory method to create builtin scheduler
Introduce a common factory method to create the builtin scheduler instead of directly constructing the `Scheduler` class.


Patch:
diff --git a/sqlmesh/core/config/scheduler.py b/sqlmesh/core/config/scheduler.py
index 5e1110e9..de1da68f 100644
--- a/sqlmesh/core/config/scheduler.py
+++ b/sqlmesh/core/config/scheduler.py
@@ -137,10 +137,9 @@ class BuiltInSchedulerConfig(_EngineAdapterStateSyncSchedulerConfig, BaseConfig)
         return BuiltInPlanEvaluator(
             state_sync=context.state_sync,
             snapshot_evaluator=context.snapshot_evaluator,
+            create_scheduler=context.create_scheduler,
             default_catalog=self.get_default_catalog(context),
-            backfill_concurrent_tasks=context.concurrent_tasks,
             console=context.console,
-            notification_target_manager=context.notification_target_manager,
         )
 
     def get_default_catalog(self, context: GenericContext) -> t.Optional[str]:
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 8d24cb3d..3f4fb43a 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -479,6 +479,17 @@ class GenericContext(BaseContext, t.Generic[C]):
         if not snapshots:
             raise ConfigError("No models were found")
 
+        return self.create_scheduler(snapshots)
+
+    def create_scheduler(self, snapshots: t.Iterable[Snapshot]) -> Scheduler:
+        """Creates the built-in scheduler.
+
+        Args:
+            snapshots: The snapshots to schedule.
+
+        Returns:
+            The built-in scheduler instance.
+        """
         return Scheduler(
             snapshots,
             self.snapshot_evaluator,
diff --git a/sqlmesh/core/plan/evaluator.py b/sqlmesh/core/plan/evaluator.py
index b744989a..5cf57c1b 100644
--- a/sqlmesh/core/plan/evaluator.py
+++ b/sqlmesh/core/plan/evaluator.py
@@ -23,10 +23,9 @@ from sqlmesh.core import constants as c
 from sqlmesh.core.console import Console, get_console
 from sqlmesh.core.notification_target import (
     NotificationTarget,
-    NotificationTargetManager,
 )
 from sqlmesh.core.plan.definition import Plan
-from sqlmesh.core.scheduler import Scheduler, SignalFactory
+from sqlmesh.core.scheduler import Scheduler
 from sqlmesh.core.snapshot import DeployabilityIndex, Snapshot, SnapshotEvaluator, SnapshotIntervals
 from sqlmesh.core.state_sync import StateSync
 from sqlmesh.core.state_sync.base import PromotionResult
@@ -61,19 +60,15 @@ class BuiltInPlanEvaluator(PlanEvaluator):
         self,
         state_sync: StateSync,
         snapshot_evaluator: SnapshotEvaluator,
+        create_scheduler: t.Callable[[t.Iterable[Snapshot]], Scheduler],
         default_catalog: t.Optional[str],
-        backfill_concurrent_tasks: int = 1,
         console: t.Optional[Console] = None,
-        notification_target_manager: t.Optional[NotificationTargetManager] = None,
-        signal_factory: t.Optional[SignalFactory] = None,
     ):
         self.state_sync = state_sync
         self.snapshot_evaluator = snapshot_evaluator
+        self.create_scheduler = create_scheduler
         self.default_catalog = default_catalog
-        self.backfill_concurrent_tasks = backfill_concurrent_tasks
         self.console = console or get_console()
-        self.notification_target_manager = notification_target_manager
-        self.signal_factory = signal_factory
 
     def evaluate(
         self,
@@ -153,16 +148,7 @@ class BuiltInPlanEvaluator(PlanEvaluator):
             return
 
         snapshots = plan.snapshots
-        scheduler = Scheduler(
-            snapshots.values(),
-            self.snapshot_evaluator,
-            self.state_sync,
-            default_catalog=self.default_catalog,
-            max_workers=self.backfill_concurrent_tasks,
-            console=self.console,
-            notification_target_manager=self.notification_target_manager,
-            signal_factory=self.signal_factory,
-        )
+        scheduler = self.create_scheduler(snapshots.values())
         is_run_successful = scheduler.run(
             plan.environment_naming_info,
             plan.start,


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 214a34fc..f7e1c37c 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -229,7 +229,10 @@ def duck_conn() -> duckdb.DuckDBPyConnection:
 
 def push_plan(context: Context, plan: Plan) -> None:
     plan_evaluator = BuiltInPlanEvaluator(
-        context.state_sync, context.snapshot_evaluator, context.default_catalog
+        context.state_sync,
+        context.snapshot_evaluator,
+        context.create_scheduler,
+        context.default_catalog,
     )
     plan_evaluator._push(plan)
     promotion_result = plan_evaluator._promote(plan)
diff --git a/tests/core/test_context.py b/tests/core/test_context.py
index 14e31c37..f305c2fd 100644
--- a/tests/core/test_context.py
+++ b/tests/core/test_context.py
@@ -208,7 +208,10 @@ def test_diff(sushi_context: Context, mocker: MockerFixture):
     success = sushi_context.run(start=yesterday, end=yesterday)
 
     plan_evaluator = BuiltInPlanEvaluator(
-        sushi_context.state_sync, sushi_context.snapshot_evaluator, sushi_context.default_catalog
+        sushi_context.state_sync,
+        sushi_context.snapshot_evaluator,
+        sushi_context.create_scheduler,
+        sushi_context.default_catalog,
     )
 
     plan = PlanBuilder(
diff --git a/tests/core/test_plan_evaluator.py b/tests/core/test_plan_evaluator.py
index 96303588..bec11bba 100644
--- a/tests/core/test_plan_evaluator.py
+++ b/tests/core/test_plan_evaluator.py
@@ -67,6 +67,7 @@ def test_builtin_evaluator_push(sushi_context: Context, make_snapshot):
     evaluator = BuiltInPlanEvaluator(
         sushi_context.state_sync,
         sushi_context.snapshot_evaluator,
+        sushi_context.create_scheduler,
         sushi_context.default_catalog,
         console=sushi_context.console,
     )