feat: allow restating models

Author: ravenac95

dagster_sqlmesh/conftest.py

@@ -106,7 +106,9 @@ def plan_and_run(
         enable_debug_console: bool = False,
         start: t.Optional[TimeLike] = None,
         end: t.Optional[TimeLike] = None,
-        restate_models: t.Optional[t.List[str]] = None,
+        select_models: t.Optional[t.List[str]] = None,
+        restate_selected: bool = False,
+        skip_run: bool = False,
     ):
         """Runs plan and run on SQLMesh with the given configuration and record all of the generated events.
 
@@ -135,19 +137,16 @@ def plan_and_run(
         if execution_time:
             plan_options["execution_time"] = execution_time
             run_options["execution_time"] = execution_time
-        if restate_models:
-            plan_options["restate_models"] = restate_models
-        if start:
-            plan_options["start"] = start
-            run_options["start"] = start
-        if end:
-            plan_options["end"] = end
-            run_options["end"] = end
 
         for event in controller.plan_and_run(
             environment,
+            start=start,
+            end=end,
+            select_models=select_models,
+            restate_selected=restate_selected,
             plan_options=plan_options,
             run_options=run_options,
+            skip_run=skip_run,
         ):
             recorder(event)
 

dagster_sqlmesh/console.py

@@ -7,6 +7,7 @@
 
 from sqlglot.expressions import Alter
 from sqlmesh.core.console import Console
+from sqlmesh.core.model import Model
 from sqlmesh.core.context_diff import ContextDiff
 from sqlmesh.core.environment import EnvironmentNamingInfo
 from sqlmesh.core.linter.rule import RuleViolation
@@ -563,8 +564,7 @@ def show_table_diff_summary(self, table_diff: TableDiff) -> None:
         self.publish(ShowTableDiffSummary(table_diff))
 
     def show_linter_violations(
-        self,
-        violations: list[RuleViolation],
+        self, violations: list[RuleViolation], model: Model, is_error: bool = False
     ) -> None:
         self.publish(LogWarning("Linting violations found", str(violations)))
 

dagster_sqlmesh/controller/base.py

@@ -30,7 +30,7 @@ class PlanOptions(t.TypedDict):
     execution_time: t.NotRequired[TimeLike]
     create_from: t.NotRequired[str]
     skip_tests: t.NotRequired[bool]
-    restate_models: t.NotRequired[t.Iterable[str]]
+    restate_models: t.NotRequired[t.Collection[str]]
     no_gaps: t.NotRequired[bool]
     skip_backfill: t.NotRequired[bool]
     forward_only: t.NotRequired[bool]
@@ -276,19 +276,57 @@ def run_sqlmesh_thread(
 
     def plan_and_run(
         self,
+        *,
+        select_models: list[str] | None = None,
+        restate_selected: bool = False,
+        start: TimeLike | None = None,
+        end: TimeLike | None = None,
         categorizer: t.Optional[SnapshotCategorizer] = None,
         default_catalog: t.Optional[str] = None,
         plan_options: t.Optional[PlanOptions] = None,
         run_options: t.Optional[RunOptions] = None,
+        skip_run: bool = False,
     ):
-        run_options = run_options or {}
-        plan_options = plan_options or {}
+        """Executes a plan and run operation
+
+        This is an opinionated interface for running a plan and run operation in
+        a single thread. It is recommended to use this method for most use cases.
+        """
+        run_options = run_options or RunOptions()
+        plan_options = plan_options or PlanOptions()
+
+        if plan_options.get("select_models") or run_options.get("select_models"):
+            raise ValueError(
+                "select_models should not be set in plan_options or run_options use the `select_models` or `select_models_func` arguments instead"
+            )
+        if plan_options.get("restate_models"):
+            raise ValueError(
+                "restate_models should not be set in plan_options use the `restate_selected` argument with `select_models` or `select_models_func` instead"
+            )
+        select_models = select_models or []
+
+        if start:
+            plan_options["start"] = start
+            run_options["start"] = start
+        if end:
+            plan_options["end"] = end
+            run_options["end"] = end
+
+        if select_models:
+            if restate_selected:
+                plan_options["restate_models"] = select_models
+                plan_options["select_models"] = select_models
+            else:
+                plan_options["select_models"] = select_models
+            run_options["select_models"] = select_models
 
         try:
             self.logger.debug("starting sqlmesh plan")
+            self.logger.debug(f"selected models: {select_models}")
             yield from self.plan(categorizer, default_catalog, **plan_options)
             self.logger.debug("starting sqlmesh run")
-            yield from self.run(**run_options)
+            if not skip_run:
+                yield from self.run(**run_options)
         except Exception as e:
             self.logger.error(f"Error during sqlmesh plan and run: {e}")
             raise e
@@ -442,15 +480,26 @@ def plan(
     def plan_and_run(
         self,
         environment: str,
+        *,
         categorizer: t.Optional[SnapshotCategorizer] = None,
+        select_models: list[str] | None = None,
+        restate_selected: bool = False,
+        start: TimeLike | None = None,
+        end: TimeLike | None = None,
         default_catalog: t.Optional[str] = None,
         plan_options: t.Optional[PlanOptions] = None,
         run_options: t.Optional[RunOptions] = None,
+        skip_run: bool = False,
     ):
         with self.instance(environment, "plan_and_run") as mesh:
             yield from mesh.plan_and_run(
+                start=start,
+                end=end,
+                select_models=select_models,
+                restate_selected=restate_selected,
                 categorizer=categorizer,
                 default_catalog=default_catalog,
                 plan_options=plan_options,
                 run_options=run_options,
+                skip_run=skip_run,
             )

dagster_sqlmesh/resource.py

@@ -8,6 +8,7 @@
 )
 from sqlmesh import Model
 from sqlmesh.utils.dag import DAG
+from sqlmesh.utils.date import TimeLike
 from sqlmesh.core.snapshot import Snapshot
 from sqlmesh.core.context import Context as SQLMeshContext
 
@@ -234,7 +235,11 @@ class SQLMeshResource(ConfigurableResource):
     def run(
         self,
         context: AssetExecutionContext,
+        *,
         environment: str = "dev",
+        start: TimeLike | None = None,
+        end: TimeLike | None = None,
+        restate_selected: bool = False,
         plan_options: t.Optional[PlanOptions] = None,
         run_options: t.Optional[RunOptions] = None,
     ) -> t.Iterable[MaterializeResult]:
@@ -248,9 +253,7 @@ def run(
         with controller.instance(environment) as mesh:
             dag = mesh.models_dag()
 
-            plan_options["select_models"] = []
-            plan_options["backfill_models"] = []
-            run_options["select_models"] = []
+            select_models = []
 
             models = mesh.models()
             models_map = models.copy()
@@ -264,15 +267,17 @@ def run(
                         logger.info(f"selected model: {model.name}")
 
                         models_map[key] = model
-                        plan_options["select_models"].append(model.name)
-                        plan_options["backfill_models"].append(model.name)
-                        run_options["select_models"].append(model.name)
+                        select_models.append(model.name)
 
             event_handler = DagsterSQLMeshEventHandler(
                 context, models_map, dag, "sqlmesh: "
             )
 
             for event in mesh.plan_and_run(
+                start=start,
+                end=end,
+                select_models=select_models,
+                restate_selected=restate_selected,
                 plan_options=plan_options,
                 run_options=run_options,
             ):

dagster_sqlmesh/test_asset.py

@@ -9,4 +9,4 @@ def test_sqlmesh_context_to_asset_outs(sample_sqlmesh_test_context: SQLMeshTestC
     translator = SQLMeshDagsterTranslator()
     outs = controller.to_asset_outs("dev", translator)
     assert len(list(outs.deps)) == 1
-    assert len(outs.outs) == 7
+    assert len(outs.outs) == 9

dagster_sqlmesh/test_sqlmesh_context.py

@@ -122,3 +122,73 @@ def test_sqlmesh_context(sample_sqlmesh_test_context: SQLMeshTestContext):
     """
     )
     assert test_source_model_count[0][0] == 6
+
+
+def test_restating_models(sample_sqlmesh_test_context: SQLMeshTestContext):
+    sample_sqlmesh_test_context.plan_and_run(
+        environment="dev",
+        start="2023-01-01",
+        end="2024-01-01",
+        execution_time="2024-01-02",
+    )
+
+    count_query = sample_sqlmesh_test_context.query(
+        """
+    SELECT COUNT(*) FROM sqlmesh_example__dev.staging_model_4
+    """
+    )
+    assert count_query[0][0] == 366
+
+    feb_sum_query = sample_sqlmesh_test_context.query(
+        """
+    SELECT SUM(value) FROM sqlmesh_example__dev.staging_model_4 WHERE time >= '2023-02-01' AND time < '2023-02-28'
+    """
+    )
+    march_sum_query = sample_sqlmesh_test_context.query(
+        """
+    SELECT SUM(value) FROM sqlmesh_example__dev.staging_model_4 WHERE time >= '2023-03-01' AND time < '2023-03-31'
+    """
+    )
+    intermediate_2_query = sample_sqlmesh_test_context.query(
+        """
+    SELECT * FROM sqlmesh_example__dev.intermediate_model_2
+    """
+    )
+
+    # Restate the model for the month of March
+    sample_sqlmesh_test_context.plan_and_run(
+        environment="dev",
+        start="2023-03-01",
+        end="2023-03-31",
+        execution_time="2024-01-02",
+        select_models=["sqlmesh_example.staging_model_4"],
+        restate_selected=True,
+        skip_run=True,
+    )
+
+    # Check that the sum of values for February and March are the same
+    feb_sum_query_restate = sample_sqlmesh_test_context.query(
+        """
+    SELECT SUM(value) FROM sqlmesh_example__dev.staging_model_4 WHERE time >= '2023-02-01' AND time < '2023-02-28'
+    """
+    )
+    march_sum_query_restate = sample_sqlmesh_test_context.query(
+        """
+    SELECT SUM(value) FROM sqlmesh_example__dev.staging_model_4 WHERE time >= '2023-03-01' AND time < '2023-03-31'
+        """
+    )
+    intermediate_2_query_restate = sample_sqlmesh_test_context.query(
+        """
+    SELECT * FROM sqlmesh_example__dev.intermediate_model_2
+    """
+    )
+
+    assert (
+        feb_sum_query_restate[0][0] == feb_sum_query[0][0]
+    ), "February sum should not change"
+    assert (
+        march_sum_query_restate[0][0] != march_sum_query[0][0]
+    ), "March sum should change"
+    assert (
+        intermediate_2_query_restate[0][0] == intermediate_2_query[0][0]
+    ), "Intermediate model should not change during restate"

pyproject.toml

@@ -10,7 +10,7 @@ readme = "README.md"
 requires-python = ">=3.11,<3.13"
 dependencies = [
     "dagster>=1.7.8",
-    "sqlmesh<1.0",
+    "sqlmesh==0.164.0",
     "pytest>=8.3.2",
     "pyarrow>=18.0.0",
 ]

sample/sqlmesh_project/models/intermediate/intermediate_model_2.sql

@@ -0,0 +1,8 @@
+MODEL (
+  name sqlmesh_example.intermediate_model_2,
+  kind FULL
+);
+
+SELECT
+  SUM(parent.value) as value
+FROM sqlmesh_example.staging_model_4 AS parent

sample/sqlmesh_project/models/staging/staging_model_4.py

@@ -0,0 +1,37 @@
+import typing as t
+from datetime import datetime
+
+import pandas as pd
+from sqlmesh import ExecutionContext, model
+from sqlmesh.core.model import ModelKindName
+import numpy as np
+
+
+@model(
+    name="sqlmesh_example.staging_model_4",
+    is_sql=False,
+    columns={
+        "time": "TIMESTAMP",
+        "value": "DOUBLE",
+    },
+    kind={"name": ModelKindName.INCREMENTAL_BY_TIME_RANGE, "time_column": "time"},
+    start="2023-01-01",
+)
+def staging_model_4(
+    context: ExecutionContext,
+    start: datetime,
+    end: datetime,
+    **kwargs,
+) -> t.Generator[pd.DataFrame, None, None]:
+    # Generates a set of random rows for the model based on the start and end dates
+    date_range = pd.date_range(start=start, end=end, freq="D")
+    num_days = len(date_range)
+
+    data = {
+        "time": date_range,
+        "value": np.random.rand(num_days)
+        * 100,  # Random double values between 0 and 100
+    }
+
+    df = pd.DataFrame(data)
+    yield df