Feat: Support expressions when selecting models for restatement (#1486)

Author: izeigerman

sqlmesh/core/context.py

@@ -761,11 +761,16 @@ def plan(
         if include_unmodified is None:
             include_unmodified = self.config.include_unmodified
 
+        model_selector = Selector(self.state_reader, self._models, self.path, dag=self.dag)
+
         models_override: t.Optional[UniqueKeyDict[str, Model]] = None
         if select_models:
-            models_override = Selector(
-                self.state_reader, self._models, self.path, dag=self.dag
-            ).select_models(select_models, environment, fallback_env_name=create_from or c.PROD)
+            models_override = model_selector.select_models(
+                select_models, environment, fallback_env_name=create_from or c.PROD
+            )
+
+        if restate_models is not None:
+            restate_models = model_selector.expand_model_selections(restate_models)
 
         # If no end date is specified, use the max interval end from prod
         # to prevent unintended evaluation of the entire DAG.

sqlmesh/core/selector.py

@@ -67,7 +67,7 @@ def select_models(
             ).values()
         }
 
-        all_selected_models = self._expand_model_selections(model_selections)
+        all_selected_models = self.expand_model_selections(model_selections)
 
         dag: DAG[str] = DAG()
         models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
@@ -91,7 +91,15 @@ def select_models(
 
         return models
 
-    def _expand_model_selections(self, model_selections: t.Iterable[str]) -> t.Set[str]:
+    def expand_model_selections(self, model_selections: t.Iterable[str]) -> t.Set[str]:
+        """Expands a set of model selections into a set of model names.
+
+        Args:
+            model_selections: A set of model selections.
+
+        Returns:
+            A set of model names.
+        """
         result: t.Set[str] = set()
 
         def _add_model(model_name: str, include_upstream: bool, include_downstream: bool) -> None: