fix: execution_time was not getting passed in from plan (#2573)

Author: tobymao

sqlmesh/core/plan/evaluator.py

@@ -151,6 +151,7 @@ def _backfill(
             plan.environment_naming_info,
             plan.start,
             plan.end,
+            execution_time=plan.execution_time,
             restatements=plan.restatements,
             selected_snapshots=selected_snapshots,
             deployability_index=deployability_index,
@@ -278,7 +279,6 @@ def _restate(self, plan: Plan) -> None:
                 (plan.context_diff.snapshots[s_id], interval)
                 for s_id, interval in plan.restatements.items()
             ],
-            plan.execution_time,
             remove_shared_versions=not plan.is_dev,
         )
 
@@ -373,6 +373,7 @@ def _apply_plan(self, plan: Plan, plan_request_id: str) -> None:
                 for change_source, snapshots in plan.indirectly_modified.items()
             },
             removed_snapshots=list(plan.context_diff.removed_snapshots),
+            execution_time=plan.execution_time,
         )
         plan_dag_spec = create_plan_dag_spec(plan_application_request, self.state_sync)
         PlanDagState.from_state_sync(self.state_sync).add_dag_spec(plan_dag_spec)
@@ -449,6 +450,7 @@ def _apply_plan(self, plan: Plan, plan_request_id: str) -> None:
                 for change_source, snapshots in plan.indirectly_modified.items()
             },
             removed_snapshots=list(plan.context_diff.removed_snapshots),
+            execution_time=plan.execution_time,
         )
 
 

sqlmesh/core/snapshot/definition.py

@@ -656,14 +656,11 @@ def add_interval(self, start: TimeLike, end: TimeLike, is_dev: bool = False) ->
         else:
             self.intervals = merged_intervals
 
-    def remove_interval(
-        self, interval: Interval, execution_time: t.Optional[TimeLike] = None
-    ) -> None:
+    def remove_interval(self, interval: Interval) -> None:
         """Remove an interval from the snapshot.
 
         Args:
             interval: The interval to remove.
-            execution_time: The date/time time reference to use for execution time. Defaults to now.
         """
         self.intervals = remove_interval(self.intervals, *interval)
         self.dev_intervals = remove_interval(self.dev_intervals, *interval)
@@ -1504,7 +1501,7 @@ def missing_intervals(
             snapshot_start_date, snapshot_end_date = (to_datetime(i) for i in interval)
             snapshot = snapshot.copy()
             snapshot.intervals = snapshot.intervals.copy()
-            snapshot.remove_interval(interval, execution_time)
+            snapshot.remove_interval(interval)
 
         missing_interval_end_date = snapshot_end_date
         node_end_date = snapshot.node.end

sqlmesh/core/state_sync/base.py

@@ -307,7 +307,6 @@ def invalidate_environment(self, name: str) -> None:
     def remove_interval(
         self,
         snapshot_intervals: t.Sequence[t.Tuple[SnapshotInfoLike, Interval]],
-        execution_time: t.Optional[TimeLike] = None,
         remove_shared_versions: bool = False,
     ) -> None:
         """Remove an interval from a list of snapshots and sync it to the store.

sqlmesh/core/state_sync/cache.py

@@ -136,12 +136,11 @@ def _add_snapshot_intervals(self, snapshot_intervals: SnapshotIntervals) -> None
     def remove_interval(
         self,
         snapshot_intervals: t.Sequence[t.Tuple[SnapshotInfoLike, Interval]],
-        execution_time: t.Optional[TimeLike] = None,
         remove_shared_versions: bool = False,
     ) -> None:
         for s, _ in snapshot_intervals:
             self.snapshot_cache.pop(s.snapshot_id, None)
-        self.state_sync.remove_interval(snapshot_intervals, execution_time, remove_shared_versions)
+        self.state_sync.remove_interval(snapshot_intervals, remove_shared_versions)
 
     def unpause_snapshots(
         self, snapshots: t.Collection[SnapshotInfoLike], unpaused_dt: TimeLike

sqlmesh/core/state_sync/engine_adapter.py

@@ -673,7 +673,6 @@ def remove_partial_intervals(
     def remove_interval(
         self,
         snapshot_intervals: t.Sequence[t.Tuple[SnapshotInfoLike, Interval]],
-        execution_time: t.Optional[TimeLike] = None,
         remove_shared_versions: bool = False,
     ) -> None:
         intervals_to_remove: t.Sequence[

sqlmesh/schedulers/airflow/client.py

@@ -17,6 +17,7 @@
 from sqlmesh.core.user import User
 from sqlmesh.schedulers.airflow import common
 from sqlmesh.utils import unique
+from sqlmesh.utils.date import TimeLike
 from sqlmesh.utils.errors import (
     ApiClientError,
     ApiServerError,
@@ -201,6 +202,7 @@ def apply_plan(
         directly_modified_snapshots: t.Optional[t.List[SnapshotId]] = None,
         indirectly_modified_snapshots: t.Optional[t.Dict[str, t.List[SnapshotId]]] = None,
         removed_snapshots: t.Optional[t.List[SnapshotId]] = None,
+        execution_time: t.Optional[TimeLike] = None,
     ) -> None:
         request = common.PlanApplicationRequest(
             new_snapshots=list(new_snapshots),
@@ -221,6 +223,7 @@ def apply_plan(
             directly_modified_snapshots=directly_modified_snapshots or [],
             indirectly_modified_snapshots=indirectly_modified_snapshots or {},
             removed_snapshots=removed_snapshots or [],
+            execution_time=execution_time,
         )
 
         response = self._session.post(

sqlmesh/schedulers/airflow/common.py

@@ -55,6 +55,7 @@ class PlanApplicationRequest(PydanticModel):
     directly_modified_snapshots: t.List[SnapshotId]
     indirectly_modified_snapshots: t.Dict[str, t.List[SnapshotId]]
     removed_snapshots: t.List[SnapshotId]
+    execution_time: t.Optional[TimeLike] = None
 
     def is_selected_for_backfill(self, model_fqn: str) -> bool:
         return self.models_to_backfill is None or model_fqn in self.models_to_backfill
@@ -89,6 +90,7 @@ class PlanDagSpec(PydanticModel):
     directly_modified_snapshots: t.Optional[t.List[SnapshotId]] = None
     indirectly_modified_snapshots: t.Optional[t.Dict[str, t.List[SnapshotId]]] = None
     removed_snapshots: t.Optional[t.List[SnapshotId]] = None
+    execution_time: t.Optional[TimeLike] = None
 
 
 class EnvironmentsResponse(PydanticModel):

sqlmesh/schedulers/airflow/dag_generator.py

@@ -207,6 +207,7 @@ def _create_plan_application_dag(self, plan_dag_spec: common.PlanDagSpec) -> DAG
                 plan_dag_spec.deployability_index,
                 plan_dag_spec.environment.plan_id,
                 "before_promote",
+                plan_dag_spec.execution_time,
             )
 
             (
@@ -218,6 +219,7 @@ def _create_plan_application_dag(self, plan_dag_spec: common.PlanDagSpec) -> DAG
                 plan_dag_spec.deployability_index,
                 plan_dag_spec.environment.plan_id,
                 "after_promote",
+                plan_dag_spec.execution_time,
             )
 
             (
@@ -425,6 +427,7 @@ def _create_backfill_tasks(
         deployability_index: DeployabilityIndex,
         plan_id: str,
         task_id_suffix: str,
+        execution_time: t.Optional[TimeLike],
     ) -> t.Tuple[BaseOperator, BaseOperator]:
         snapshot_to_tasks = {}
         for intervals_per_snapshot in backfill_intervals:
@@ -456,6 +459,7 @@ def _create_backfill_tasks(
                     end=end,
                     deployability_index=deployability_index,
                     plan_id=plan_id,
+                    execution_time=execution_time,
                 )
                 external_sensor_task = self._create_hwm_external_sensor(
                     snapshot, start=start, end=end
@@ -593,6 +597,7 @@ def _create_snapshot_evaluation_operator(
         task_id: str,
         start: t.Optional[TimeLike] = None,
         end: t.Optional[TimeLike] = None,
+        execution_time: t.Optional[TimeLike] = None,
         deployability_index: t.Optional[DeployabilityIndex] = None,
         plan_id: t.Optional[str] = None,
     ) -> BaseOperator:
@@ -607,6 +612,7 @@ def _create_snapshot_evaluation_operator(
                 end=end,
                 deployability_index=deployability_index or DeployabilityIndex.all_deployable(),
                 plan_id=plan_id,
+                execution_time=execution_time,
             ),
             task_id=task_id,
         )

sqlmesh/schedulers/airflow/plan.py

@@ -130,7 +130,7 @@ def create_plan_dag_spec(
             [s for s in all_snapshots.values() if request.is_selected_for_backfill(s.name)],
             start=request.environment.start_at,
             end=end,
-            execution_time=now(),
+            execution_time=request.execution_time or now(),
             deployability_index=deployability_index_for_evaluation,
             restatements=restatements,
             end_bounded=request.end_bounded,
@@ -182,6 +182,7 @@ def create_plan_dag_spec(
         directly_modified_snapshots=request.directly_modified_snapshots,
         indirectly_modified_snapshots=request.indirectly_modified_snapshots,
         removed_snapshots=request.removed_snapshots,
+        execution_time=request.execution_time,
     )
 
 

sqlmesh/schedulers/airflow/state_sync.py

@@ -225,7 +225,6 @@ def _add_snapshot_intervals(self, snapshot_intervals: SnapshotIntervals) -> None
     def remove_interval(
         self,
         snapshot_intervals: t.Sequence[t.Tuple[SnapshotInfoLike, Interval]],
-        execution_time: t.Optional[TimeLike] = None,
         remove_shared_versions: bool = False,
     ) -> None:
         """Remove an interval from a list of snapshots and sync it to the store.