feat: allow restating models and any downstream deps (#40)

* feat: allows restating models and affected downstream models

* chore: update uv lock

* chore: uv sync before a release

Author: ravenac95

.release-it.json

@@ -30,6 +30,7 @@
   },
   "hooks": {
     "before:init": ["rm -rf dist"],
+    "before:release": "uv sync",
     "after:release": "uv build && uv publish"
   }
 }

dagster_sqlmesh/controller/base.py

@@ -290,6 +290,7 @@ def plan_and_run(
         self,
         *,
         select_models: list[str] | None = None,
+        restate_models: list[str] | None = None,
         restate_selected: bool = False,
         start: TimeLike | None = None,
         end: TimeLike | None = None,
@@ -316,6 +317,7 @@ def plan_and_run(
                 "restate_models should not be set in plan_options use the `restate_selected` argument with `select_models` instead"
             )
         select_models = select_models or []
+        restate_models = restate_models or []
 
         if start:
             plan_options["start"] = start
@@ -324,14 +326,14 @@ def plan_and_run(
             plan_options["end"] = end
             run_options["end"] = end
 
+        if restate_models:
+            plan_options["restate_models"] = restate_models
         if select_models:
+            plan_options["select_models"] = select_models
+            run_options["select_models"] = 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}")

dagster_sqlmesh/resource.py

@@ -266,6 +266,8 @@ def run(
         environment: str = "dev",
         start: TimeLike | None = None,
         end: TimeLike | None = None,
+        restate_models: list[str] | None = None,
+        select_models: list[str] | None = None,
         restate_selected: bool = False,
         skip_run: bool = False,
         plan_options: PlanOptions | None = None,
@@ -317,6 +319,7 @@ def run(
                 start=start,
                 end=end,
                 select_models=select_models,
+                restate_models=restate_models,
                 restate_selected=restate_selected,
                 skip_run=skip_run,
                 plan_options=plan_options,