Refactor APIs

Author: VaggelisD

sqlmesh/core/config/scheduler.py

@@ -28,14 +28,11 @@ class SchedulerConfig(abc.ABC):
     """Abstract base class for Scheduler configurations."""
 
     @abc.abstractmethod
-    def create_plan_evaluator(
-        self, context: GenericContext, job_id: t.Optional[str] = None
-    ) -> PlanEvaluator:
+    def create_plan_evaluator(self, context: GenericContext) -> PlanEvaluator:
         """Creates a Plan Evaluator instance.
 
         Args:
             context: The SQLMesh Context.
-            job_id: The plan ID.
         """
 
     @abc.abstractmethod
@@ -130,12 +127,9 @@ class BuiltInSchedulerConfig(_EngineAdapterStateSyncSchedulerConfig, BaseConfig)
 
     type_: t.Literal["builtin"] = Field(alias="type", default="builtin")
 
-    def create_plan_evaluator(
-        self, context: GenericContext, job_id: t.Optional[str] = None
-    ) -> PlanEvaluator:
+    def create_plan_evaluator(self, context: GenericContext) -> PlanEvaluator:
         return BuiltInPlanEvaluator(
             state_sync=context.state_sync,
-            snapshot_evaluator=context.snapshot_evaluator(job_id),
             create_scheduler=context.create_scheduler,
             default_catalog=context.default_catalog,
             console=context.console,

sqlmesh/core/context.py

@@ -453,7 +453,7 @@ def snapshot_evaluator(self, job_id: t.Optional[str] = None) -> SnapshotEvaluato
         ):
             self._snapshot_evaluator = SnapshotEvaluator(
                 {
-                    gateway: adapter.with_log_level(logging.INFO, job_id)
+                    gateway: adapter.with_settings(level=logging.INFO, job_id=job_id)
                     for gateway, adapter in self.engine_adapters.items()
                 },
                 ddl_concurrent_tasks=self.concurrent_tasks,
@@ -1592,7 +1592,10 @@ def apply(
                 default_catalog=self.default_catalog,
                 console=self.console,
             )
-            explainer.evaluate(plan.to_evaluatable())
+            explainer.evaluate(
+                plan.to_evaluatable(),
+                snapshot_evaluator=self.snapshot_evaluator(job_id=plan.plan_id),
+            )
             return
 
         self.notification_target_manager.notify(
@@ -1905,7 +1908,7 @@ def _table_diff(
             )
 
         return TableDiff(
-            adapter=adapter.with_log_level(logger.getEffectiveLevel()),
+            adapter=adapter.with_settings(logger.getEffectiveLevel()),
             source=source,
             target=target,
             on=on,
@@ -2392,8 +2395,10 @@ def _run(
         return completion_status
 
     def _apply(self, plan: Plan, circuit_breaker: t.Optional[t.Callable[[], bool]]) -> None:
-        self._scheduler.create_plan_evaluator(self, job_id=plan.plan_id).evaluate(
-            plan.to_evaluatable(), circuit_breaker=circuit_breaker
+        self._scheduler.create_plan_evaluator(self).evaluate(
+            plan.to_evaluatable(),
+            snapshot_evaluator=self.snapshot_evaluator(job_id=plan.plan_id),
+            circuit_breaker=circuit_breaker,
         )
 
     @python_api_analytics

sqlmesh/core/engine_adapter/athena.py

@@ -46,7 +46,7 @@ def __init__(
         self, *args: t.Any, s3_warehouse_location: t.Optional[str] = None, **kwargs: t.Any
     ):
         # Need to pass s3_warehouse_location to the superclass so that it goes into _extra_config
-        # which means that EngineAdapter.with_log_level() keeps this property when it makes a clone
+        # which means that EngineAdapter.with_settings() keeps this property when it makes a clone
         super().__init__(*args, s3_warehouse_location=s3_warehouse_location, **kwargs)
         self.s3_warehouse_location = s3_warehouse_location
 

sqlmesh/core/engine_adapter/base.py

@@ -147,7 +147,7 @@ def __init__(
         self._multithreaded = multithreaded
         self._job_id = job_id
 
-    def with_log_level(self, level: int, job_id: t.Optional[str] = None) -> EngineAdapter:
+    def with_settings(self, level: int, **kwargs: t.Any) -> EngineAdapter:
         adapter = self.__class__(
             self._connection_pool,
             dialect=self.dialect,
@@ -158,8 +158,8 @@ def with_log_level(self, level: int, job_id: t.Optional[str] = None) -> EngineAd
             null_connection=True,
             multithreaded=self._multithreaded,
             pretty_sql=self._pretty_sql,
-            job_id=job_id,
             **self._extra_config,
+            **kwargs,
         )
 
         return adapter

sqlmesh/core/plan/evaluator.py

@@ -49,7 +49,10 @@
 class PlanEvaluator(abc.ABC):
     @abc.abstractmethod
     def evaluate(
-        self, plan: EvaluatablePlan, circuit_breaker: t.Optional[t.Callable[[], bool]] = None
+        self,
+        plan: EvaluatablePlan,
+        snapshot_evaluator: SnapshotEvaluator,
+        circuit_breaker: t.Optional[t.Callable[[], bool]] = None,
     ) -> None:
         """Evaluates a plan by pushing snapshots and backfilling data.
 
@@ -60,20 +63,20 @@ def evaluate(
 
         Args:
             plan: The plan to evaluate.
+            snapshot_evaluator: The snapshot evaluator to use.
+            circuit_breaker: The circuit breaker to use.
         """
 
 
 class BuiltInPlanEvaluator(PlanEvaluator):
     def __init__(
         self,
         state_sync: StateSync,
-        snapshot_evaluator: SnapshotEvaluator,
         create_scheduler: t.Callable[[t.Iterable[Snapshot]], Scheduler],
         default_catalog: t.Optional[str],
         console: t.Optional[Console] = None,
     ):
         self.state_sync = state_sync
-        self.snapshot_evaluator = snapshot_evaluator
         self.create_scheduler = create_scheduler
         self.default_catalog = default_catalog
         self.console = console or get_console()
@@ -82,9 +85,12 @@ def __init__(
     def evaluate(
         self,
         plan: EvaluatablePlan,
+        snapshot_evaluator: SnapshotEvaluator,
         circuit_breaker: t.Optional[t.Callable[[], bool]] = None,
     ) -> None:
         self._circuit_breaker = circuit_breaker
+        self.snapshot_evaluator = snapshot_evaluator
+
         self.console.start_plan_evaluation(plan)
         analytics.collector.on_plan_apply_start(
             plan=plan,

sqlmesh/core/plan/explainer.py

@@ -20,6 +20,7 @@
 from sqlmesh.utils import Verbosity, rich as srich, to_snake_case
 from sqlmesh.utils.date import to_ts
 from sqlmesh.utils.errors import SQLMeshError
+from sqlmesh.core.snapshot.evaluator import SnapshotEvaluator
 
 
 logger = logging.getLogger(__name__)
@@ -37,7 +38,10 @@ def __init__(
         self.console = console or get_console()
 
     def evaluate(
-        self, plan: EvaluatablePlan, circuit_breaker: t.Optional[t.Callable[[], bool]] = None
+        self,
+        plan: EvaluatablePlan,
+        snapshot_evaluator: SnapshotEvaluator,
+        circuit_breaker: t.Optional[t.Callable[[], bool]] = None,
     ) -> None:
         plan_stages = stages.build_plan_stages(plan, self.state_reader, self.default_catalog)
         explainer_console = _get_explainer_console(

tests/conftest.py

@@ -266,10 +266,10 @@ def duck_conn() -> duckdb.DuckDBPyConnection:
 def push_plan(context: Context, plan: Plan) -> None:
     plan_evaluator = BuiltInPlanEvaluator(
         context.state_sync,
-        context.snapshot_evaluator(),
         context.create_scheduler,
         context.default_catalog,
     )
+    plan_evaluator.snapshot_evaluator = context.snapshot_evaluator(job_id=plan.plan_id)
     deployability_index = DeployabilityIndex.create(context.snapshots.values())
     evaluatable_plan = plan.to_evaluatable()
     stages = plan_stages.build_plan_stages(

tests/core/test_plan_evaluator.py

@@ -59,11 +59,12 @@ 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,
     )
+    evaluator.snapshot_evaluator = sushi_context.snapshot_evaluator(job_id=plan.plan_id)
+
     evaluatable_plan = plan.to_evaluatable()
     stages = plan_stages.build_plan_stages(
         evaluatable_plan, sushi_context.state_sync, sushi_context.default_catalog

tests/core/test_table_diff.py

@@ -335,11 +335,11 @@ def test_generated_sql(sushi_context_fixed_date: Context, mocker: MockerFixture)
     sample_query_sql = 'WITH "source_only" AS (SELECT \'source_only\' AS "__sqlmesh_sample_type", "s__key", "s__value", "s____sqlmesh_join_key", "t__key", "t__value", "t____sqlmesh_join_key" FROM "memory"."sqlmesh_temp_test"."__temp_diff_abcdefgh" WHERE "s_exists" = 1 AND "row_joined" = 0 ORDER BY "s__key" NULLS FIRST LIMIT 20), "target_only" AS (SELECT \'target_only\' AS "__sqlmesh_sample_type", "s__key", "s__value", "s____sqlmesh_join_key", "t__key", "t__value", "t____sqlmesh_join_key" FROM "memory"."sqlmesh_temp_test"."__temp_diff_abcdefgh" WHERE "t_exists" = 1 AND "row_joined" = 0 ORDER BY "t__key" NULLS FIRST LIMIT 20), "common_rows" AS (SELECT \'common_rows\' AS "__sqlmesh_sample_type", "s__key", "s__value", "s____sqlmesh_join_key", "t__key", "t__value", "t____sqlmesh_join_key" FROM "memory"."sqlmesh_temp_test"."__temp_diff_abcdefgh" WHERE "row_joined" = 1 AND "row_full_match" = 0 ORDER BY "s__key" NULLS FIRST, "t__key" NULLS FIRST LIMIT 20) SELECT "__sqlmesh_sample_type", "s__key", "s__value", "s____sqlmesh_join_key", "t__key", "t__value", "t____sqlmesh_join_key" FROM "source_only" UNION ALL SELECT "__sqlmesh_sample_type", "s__key", "s__value", "s____sqlmesh_join_key", "t__key", "t__value", "t____sqlmesh_join_key" FROM "target_only" UNION ALL SELECT "__sqlmesh_sample_type", "s__key", "s__value", "s____sqlmesh_join_key", "t__key", "t__value", "t____sqlmesh_join_key" FROM "common_rows"'
     drop_sql = 'DROP TABLE IF EXISTS "memory"."sqlmesh_temp_test"."__temp_diff_abcdefgh"'
 
-    # make with_log_level() return the current instance of engine_adapter so we can still spy on _execute
+    # make with_settings() return the current instance of engine_adapter so we can still spy on _execute
     mocker.patch.object(
-        engine_adapter, "with_log_level", new_callable=lambda: lambda _: engine_adapter
+        engine_adapter, "with_settings", new_callable=lambda: lambda _: engine_adapter
     )
-    assert engine_adapter.with_log_level(1) == engine_adapter
+    assert engine_adapter.with_settings(1) == engine_adapter
 
     spy_execute = mocker.spy(engine_adapter, "_execute")
     mocker.patch("sqlmesh.core.engine_adapter.base.random_id", return_value="abcdefgh")