Repository: TobikoData/sqlmesh
Pull request #2573 (2024-05-07 03:04:54)

fix: execution_time was not getting passed in from plan

izeigerman: How about this spot:https://github.com/TobikoData/sqlmesh/blob/8e237207b546f30d06366495a39ffa0b5ecced36/sqlmesh/schedulers/airflow/plan.py#L133
izeigerman: Also I think we need to pass it here as well: https://github.com/TobikoData/sqlmesh/blob/8e237207b546f30d06366495a39ffa0b5ecced36/sqlmesh/schedulers/airflow/plan.py#L116

Patch:
diff --git a/sqlmesh/core/plan/evaluator.py b/sqlmesh/core/plan/evaluator.py
index a4cdf66b..6fe25221 100644
--- a/sqlmesh/core/plan/evaluator.py
+++ b/sqlmesh/core/plan/evaluator.py
@@ -151,6 +151,7 @@ class BuiltInPlanEvaluator(PlanEvaluator):
             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 @@ class BuiltInPlanEvaluator(PlanEvaluator):
                 (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 @@ class StateBasedAirflowPlanEvaluator(BaseAirflowPlanEvaluator):
                 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 @@ class AirflowPlanEvaluator(StateBasedAirflowPlanEvaluator):
                 for change_source, snapshots in plan.indirectly_modified.items()
             },
             removed_snapshots=list(plan.context_diff.removed_snapshots),
+            execution_time=plan.execution_time,
         )
 
 
diff --git a/sqlmesh/core/snapshot/definition.py b/sqlmesh/core/snapshot/definition.py
index 20912651..3f9df01c 100644
--- a/sqlmesh/core/snapshot/definition.py
+++ b/sqlmesh/core/snapshot/definition.py
@@ -656,14 +656,11 @@ class Snapshot(PydanticModel, SnapshotInfoMixin):
         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
diff --git a/sqlmesh/core/state_sync/base.py b/sqlmesh/core/state_sync/base.py
index f603da36..b0e47f7f 100644
--- a/sqlmesh/core/state_sync/base.py
+++ b/sqlmesh/core/state_sync/base.py
@@ -307,7 +307,6 @@ class StateSync(StateReader, abc.ABC):
     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.
diff --git a/sqlmesh/core/state_sync/cache.py b/sqlmesh/core/state_sync/cache.py
index 209b67e5..e022120a 100644
--- a/sqlmesh/core/state_sync/cache.py
+++ b/sqlmesh/core/state_sync/cache.py
@@ -136,12 +136,11 @@ class CachingStateSync(DelegatingStateSync):
     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
diff --git a/sqlmesh/core/state_sync/engine_adapter.py b/sqlmesh/core/state_sync/engine_adapter.py
index 797aeaaa..e151bbb5 100644
--- a/sqlmesh/core/state_sync/engine_adapter.py
+++ b/sqlmesh/core/state_sync/engine_adapter.py
@@ -673,7 +673,6 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
     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[
diff --git a/sqlmesh/schedulers/airflow/client.py b/sqlmesh/schedulers/airflow/client.py
index 8234de17..6cecadb8 100644
--- a/sqlmesh/schedulers/airflow/client.py
+++ b/sqlmesh/schedulers/airflow/client.py
@@ -17,6 +17,7 @@ from sqlmesh.core.state_sync import Versions
 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 @@ class AirflowClient(BaseAirflowClient):
         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 @@ class AirflowClient(BaseAirflowClient):
             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(
diff --git a/sqlmesh/schedulers/airflow/common.py b/sqlmesh/schedulers/airflow/common.py
index d4685b31..5e31c3e2 100644
--- a/sqlmesh/schedulers/airflow/common.py
+++ b/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):
diff --git a/sqlmesh/schedulers/airflow/dag_generator.py b/sqlmesh/schedulers/airflow/dag_generator.py
index 9fed8796..8d377f93 100644
--- a/sqlmesh/schedulers/airflow/dag_generator.py
+++ b/sqlmesh/schedulers/airflow/dag_generator.py
@@ -207,6 +207,7 @@ class SnapshotDagGenerator:
                 plan_dag_spec.deployability_index,
                 plan_dag_spec.environment.plan_id,
                 "before_promote",
+                plan_dag_spec.execution_time,
             )
 
             (
@@ -218,6 +219,7 @@ class SnapshotDagGenerator:
                 plan_dag_spec.deployability_index,
                 plan_dag_spec.environment.plan_id,
                 "after_promote",
+                plan_dag_spec.execution_time,
             )
 
             (
@@ -425,6 +427,7 @@ class SnapshotDagGenerator:
         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 @@ class SnapshotDagGenerator:
                     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 @@ class SnapshotDagGenerator:
         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 @@ class SnapshotDagGenerator:
                 end=end,
                 deployability_index=deployability_index or DeployabilityIndex.all_deployable(),
                 plan_id=plan_id,
+                execution_time=execution_time,
             ),
             task_id=task_id,
         )
diff --git a/sqlmesh/schedulers/airflow/plan.py b/sqlmesh/schedulers/airflow/plan.py
index 4cf23162..c0efe345 100644
--- a/sqlmesh/schedulers/airflow/plan.py
+++ b/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,
     )
 
 
diff --git a/sqlmesh/schedulers/airflow/state_sync.py b/sqlmesh/schedulers/airflow/state_sync.py
index 33d58ea1..a8416c6d 100644
--- a/sqlmesh/schedulers/airflow/state_sync.py
+++ b/sqlmesh/schedulers/airflow/state_sync.py
@@ -225,7 +225,6 @@ class HttpStateSync(StateSync):
     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.


Test patch:
diff --git a/tests/core/test_context.py b/tests/core/test_context.py
index a8c1add3..4aaee522 100644
--- a/tests/core/test_context.py
+++ b/tests/core/test_context.py
@@ -254,6 +254,36 @@ def test_evaluate_limit():
     assert context.evaluate("without_limit", "2020-01-01", "2020-01-02", "2020-01-02", 2).size == 2
 
 
+def test_plan_execution_time():
+    context = Context(config=Config())
+    context.upsert_model(
+        load_sql_based_model(
+            parse(
+                """
+                MODEL(
+                    name db.x, 
+                    start '2024-01-01',
+                    kind FULL
+                );
+
+                SELECT @execution_date AS execution_date
+                """
+            )
+        )
+    )
+
+    context.plan(
+        "dev",
+        execution_time="2024-01-02",
+        auto_apply=True,
+        no_prompts=True,
+    )
+    assert (
+        str(list(context.fetchdf("select * from db__dev.x")["execution_date"])[0])
+        == "2024-01-02 00:00:00"
+    )
+
+
 def test_clear_caches(tmp_path: pathlib.Path):
     models_dir = tmp_path / "models"
 
diff --git a/tests/core/test_plan_evaluator.py b/tests/core/test_plan_evaluator.py
index 597ca075..9fc3f570 100644
--- a/tests/core/test_plan_evaluator.py
+++ b/tests/core/test_plan_evaluator.py
@@ -104,6 +104,7 @@ def test_airflow_evaluator(sushi_plan: Plan, mocker: MockerFixture):
         directly_modified_snapshots=[],
         indirectly_modified_snapshots={},
         removed_snapshots=[],
+        execution_time=None,
     )
 
     airflow_client_mock.wait_for_dag_run_completion.assert_called_once()