add airflow depends_on_past (#940)

Author: eakmanrq

setup.cfg

@@ -72,6 +72,9 @@ ignore_missing_imports = True
 [mypy-langchain.*]
 ignore_missing_imports = True
 
+[mypy-pytest_lazyfixture.*]
+ignore_missing_imports = True
+
 [autoflake]
 in-place = True
 expand-star-imports = True

setup.py

@@ -78,6 +78,7 @@
             "PyGithub",
             "pytest",
             "pytest-asyncio",
+            "pytest-lazy-fixture",
             "pytest-mock",
             "pyspark>=3.4.0",
             "pytz",

sqlmesh/schedulers/airflow/dag_generator.py

@@ -349,7 +349,7 @@ def _create_backfill_tasks(
             snapshot_end_task = EmptyOperator(
                 task_id=f"snapshot_backfill__{snapshot.name}__{snapshot.identifier}__end"
             )
-            if snapshot.is_incremental_by_unique_key:
+            if snapshot.depends_on_past:
                 baseoperator.chain(snapshot_start_task, *tasks, snapshot_end_task)
             else:
                 snapshot_start_task >> tasks >> snapshot_end_task

sqlmesh/schedulers/airflow/plan.py

@@ -62,7 +62,7 @@ def create_plan_dag_spec(
             intervals=state_sync.get_snapshot_intervals(all_snapshots.values()),
             start=request.environment.start_at,
             end=end,
-            latest=end,
+            latest=now(),
             is_dev=request.is_dev,
             restatements=request.restatements,
         )

tests/conftest.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import datetime
 import typing as t
 from pathlib import Path
 from shutil import rmtree
@@ -13,15 +14,39 @@
 from sqlglot.helper import ensure_list
 
 from sqlmesh.core.context import Context
+from sqlmesh.core.engine_adapter.base import EngineAdapter
 from sqlmesh.core.model import Model
 from sqlmesh.core.plan import BuiltInPlanEvaluator, Plan
 from sqlmesh.core.snapshot import Snapshot
 from sqlmesh.utils import random_id
-from sqlmesh.utils.date import TimeLike
+from sqlmesh.utils.date import TimeLike, to_date, to_ds
 
 pytest_plugins = ["tests.common_fixtures"]
 
 
+class SushiDataValidator:
+    def __init__(self, engine_adapter: EngineAdapter):
+        self.engine_adapter = engine_adapter
+
+    @classmethod
+    def from_context(cls, context: Context):
+        return cls(engine_adapter=context.engine_adapter)
+
+    def validate(self, model_name: str, start: TimeLike, end: TimeLike) -> t.Dict[t.Any, t.Any]:
+        if model_name == "sushi.customer_revenue_lifetime":
+            query = "SELECT ds, count(*) AS the_count FROM sushi.customer_revenue_lifetime group by 1 order by 2 desc, 1 desc"
+            results = self.engine_adapter.fetchdf(query).to_dict()
+            start_date, end_date = to_date(start), to_date(end)
+            num_days_diff = (end_date - start_date).days + 1
+            assert len(results["ds"]) == num_days_diff
+            assert list(results["ds"].values()) == [
+                to_ds(end_date - datetime.timedelta(days=x)) for x in range(num_days_diff)
+            ]
+            return results
+        else:
+            raise NotImplementedError(f"Unknown model_name: {model_name}")
+
+
 # Ignore all local config files
 @pytest.fixture(scope="session", autouse=True)
 def ignore_local_config_files():
@@ -142,6 +167,16 @@ def random_name() -> t.Callable:
     return lambda: f"generated_{random_id()}"
 
 
+@pytest.fixture
+def sushi_data_validator(sushi_context: Context) -> SushiDataValidator:
+    return SushiDataValidator.from_context(sushi_context)
+
+
+@pytest.fixture
+def sushi_fixed_date_data_validator(sushi_context_fixed_date: Context) -> SushiDataValidator:
+    return SushiDataValidator.from_context(sushi_context_fixed_date)
+
+
 def delete_cache(project_paths: str | t.List[str]) -> None:
     for path in ensure_list(project_paths):
         try:

tests/core/test_integration.py

@@ -1,4 +1,3 @@
-import datetime
 import typing as t
 from collections import Counter
 
@@ -26,7 +25,8 @@
     SnapshotInfoLike,
     SnapshotTableInfo,
 )
-from sqlmesh.utils.date import TimeLike, to_date, to_ds, to_timestamp, yesterday
+from sqlmesh.utils.date import TimeLike, to_date, to_ds, to_timestamp
+from tests.conftest import SushiDataValidator
 
 
 @pytest.fixture(autouse=True)
@@ -655,18 +655,15 @@ def test_multi(mocker):
 
 @pytest.mark.integration
 @pytest.mark.core_integration
-def test_incremental_time_self_reference(mocker: MockerFixture, sushi_context: Context):
+def test_incremental_time_self_reference(
+    mocker: MockerFixture, sushi_context: Context, sushi_data_validator: SushiDataValidator
+):
+    start_date, end_date = to_date("1 week ago"), to_date("yesterday")
     df = sushi_context.engine_adapter.fetchdf("SELECT MIN(ds) FROM sushi.customer_revenue_lifetime")
-    assert df.iloc[0, 0] == to_ds("1 week ago")
+    assert df.iloc[0, 0] == to_ds(start_date)
     df = sushi_context.engine_adapter.fetchdf("SELECT MAX(ds) FROM sushi.customer_revenue_lifetime")
-    assert df.iloc[0, 0] == to_ds("yesterday")
-    query_get_date_and_count = "SELECT ds, count(*) AS the_count FROM sushi.customer_revenue_lifetime group by 1 order by 2 desc, 1 desc"
-    results = sushi_context.engine_adapter.fetchdf(query_get_date_and_count).to_dict()
-    # Validate that both rows increase over time and all days are present
-    assert len(results["ds"]) == 7
-    assert list(results["ds"].values()) == [
-        to_ds(yesterday() - datetime.timedelta(days=x)) for x in range(7)
-    ]
+    assert df.iloc[0, 0] == to_ds(end_date)
+    results = sushi_data_validator.validate("sushi.customer_revenue_lifetime", start_date, end_date)
     plan = sushi_context.plan(
         restate_models=["sushi.customer_revenue_lifetime", "sushi.customer_revenue_by_day"],
         no_prompts=True,
@@ -709,7 +706,9 @@ def test_incremental_time_self_reference(mocker: MockerFixture, sushi_context: C
         "sushi.customer_revenue_by_day": 1,
     }
     # Validate that the results are the same as before the restate
-    assert results == sushi_context.engine_adapter.fetchdf(query_get_date_and_count).to_dict()
+    assert results == sushi_data_validator.validate(
+        "sushi.customer_revenue_lifetime", start_date, end_date
+    )
 
 
 def initial_add(context: Context, environment: str):

tests/schedulers/airflow/test_end_to_end.py

@@ -4,7 +4,8 @@
 
 from sqlmesh.core.context import Context
 from sqlmesh.schedulers.airflow.client import AirflowClient
-from sqlmesh.utils.date import now, yesterday_ds
+from sqlmesh.utils.date import now
+from tests.conftest import SushiDataValidator
 
 
 @pytest.fixture(autouse=True)
@@ -19,12 +20,13 @@ def get_receiver_dag() -> None:
 @pytest.mark.integration
 @pytest.mark.airflow_integration
 def test_sushi(mocker: MockerFixture, is_docker: bool):
-    start = yesterday_ds()
+    start = "1 week ago"
     end = now()
     latest = end
 
     airflow_config = "airflow_config_docker" if is_docker else "airflow_config"
     context = Context(paths="./examples/sushi", config=airflow_config)
+    data_validator = SushiDataValidator.from_context(context)
 
     context.plan(
         environment="test_dev",
@@ -36,6 +38,8 @@ def test_sushi(mocker: MockerFixture, is_docker: bool):
         auto_apply=True,
     )
 
+    data_validator.validate("sushi.customer_revenue_lifetime", start, end)
+
     # Ensure that the plan has been applied successfully.
     no_change_plan = context.plan(
         environment="test_dev_two",