Fix: Set the plan DAG start date at request time (#1244)

Author: izeigerman

sqlmesh/schedulers/airflow/common.py

@@ -78,6 +78,7 @@ class PlanDagSpec(PydanticModel):
     is_dev: bool
     forward_only: t.Optional[bool]
     environment_expiration_ts: t.Optional[int]
+    dag_start_ts: t.Optional[int]
 
 
 class EnvironmentsResponse(PydanticModel):

sqlmesh/schedulers/airflow/dag_generator.py

@@ -25,7 +25,7 @@
 from sqlmesh.schedulers.airflow.operators.notification import (
     BaseNotificationOperatorProvider,
 )
-from sqlmesh.utils.date import TimeLike, to_datetime
+from sqlmesh.utils.date import TimeLike, to_datetime, yesterday_timestamp
 from sqlmesh.utils.errors import SQLMeshError
 
 logger = logging.getLogger(__name__)
@@ -136,7 +136,9 @@ def _create_plan_application_dag(self, plan_dag_spec: common.PlanDagSpec) -> DAG
         with DAG(
             dag_id=dag_id,
             schedule_interval="@once",
-            start_date=pendulum.now(tz="UTC"),
+            start_date=pendulum.instance(
+                to_datetime(plan_dag_spec.dag_start_ts or yesterday_timestamp())
+            ),
             max_active_tasks=plan_dag_spec.backfill_concurrent_tasks,
             catchup=False,
             is_paused_upon_creation=False,

sqlmesh/schedulers/airflow/plan.py

@@ -7,7 +7,7 @@
 from sqlmesh.core.snapshot import SnapshotTableInfo
 from sqlmesh.core.state_sync import StateSync
 from sqlmesh.schedulers.airflow import common
-from sqlmesh.utils.date import now
+from sqlmesh.utils.date import now, now_timestamp
 from sqlmesh.utils.errors import SQLMeshError
 
 
@@ -97,6 +97,7 @@ def create_plan_dag_spec(
         is_dev=request.is_dev,
         forward_only=request.forward_only,
         environment_expiration_ts=request.environment.expiration_ts,
+        dag_start_ts=now_timestamp(),
     )
 
 

tests/schedulers/airflow/test_plan.py

@@ -125,7 +125,11 @@ def test_create_plan_dag_spec(
     state_sync_mock.get_environment.return_value = old_environment
     state_sync_mock.get_snapshot_intervals.return_value = []
 
-    plan_spec = create_plan_dag_spec(plan_request, state_sync_mock)
+    with mock.patch(
+        "sqlmesh.schedulers.airflow.plan.now_timestamp",
+        side_effect=lambda: to_timestamp("2023-01-01"),
+    ):
+        plan_spec = create_plan_dag_spec(plan_request, state_sync_mock)
     assert plan_spec == common.PlanDagSpec(
         request_id="test_request_id",
         environment_naming_info=EnvironmentNamingInfo(
@@ -152,6 +156,7 @@ def test_create_plan_dag_spec(
         users=[],
         is_dev=False,
         forward_only=True,
+        dag_start_ts=to_timestamp("2023-01-01"),
     )
 
     state_sync_mock.get_snapshots.assert_called_once()
@@ -240,9 +245,11 @@ def test_restatement(
             dev_intervals=[],
         )
     ]
+    now_value = "2022-01-09T23:59:59+00:00"
     with mock.patch(
-        "sqlmesh.schedulers.airflow.plan.now",
-        side_effect=lambda: to_datetime("2022-01-09T23:59:59+00:00"),
+        "sqlmesh.schedulers.airflow.plan.now", side_effect=lambda: to_datetime(now_value)
+    ), mock.patch(
+        "sqlmesh.schedulers.airflow.plan.now_timestamp", side_effect=lambda: to_timestamp(now_value)
     ):
         plan_spec = create_plan_dag_spec(plan_request, state_sync_mock)
     assert plan_spec == common.PlanDagSpec(
@@ -271,6 +278,7 @@ def test_restatement(
         users=[],
         is_dev=False,
         forward_only=True,
+        dag_start_ts=to_timestamp(now_value),
     )
 
     state_sync_mock.get_snapshots.assert_called_once()