Feat: Support simple model selection expressions (#1338)

Author: izeigerman

sqlmesh/core/context.py

@@ -737,9 +737,9 @@ def plan(
 
         models_override: t.Optional[UniqueKeyDict[str, Model]] = None
         if select_models:
-            models_override = Selector(self.state_reader, self._models, self.path).select_models(
-                select_models, environment, fallback_env_name=create_from or c.PROD
-            )
+            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)
 
         plan = Plan(
             context_diff=self._context_diff(

sqlmesh/core/selector.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import fnmatch
 import typing as t
 from pathlib import Path
 
@@ -18,11 +19,19 @@ def __init__(
         state_reader: StateReader,
         models: UniqueKeyDict[str, Model],
         context_path: Path = Path("."),
+        dag: t.Optional[DAG[str]] = None,
     ):
         self._state_reader = state_reader
         self._models = models
         self._context_path = context_path
 
+        if dag is None:
+            self._dag: DAG[str] = DAG()
+            for model in models.values():
+                self._dag.add(model.name, model.depends_on)
+        else:
+            self._dag = dag
+
     def select_models(
         self,
         model_selections: t.Iterable[str],
@@ -58,8 +67,7 @@ def select_models(
             ).values()
         }
 
-        # TODO: Support selection expressions.
-        all_selected_models = set(model_selections)
+        all_selected_models = self._expand_model_selections(model_selections)
 
         dag: DAG[str] = DAG()
         models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
@@ -82,3 +90,35 @@ def select_models(
         update_model_schemas(dag, models, self._context_path)
 
         return models
+
+    def _expand_model_selections(self, model_selections: t.Iterable[str]) -> t.Set[str]:
+        result: t.Set[str] = set()
+
+        def _add_model(model_name: str, include_upstream: bool, include_downstream: bool) -> None:
+            result.add(model_name)
+            if include_upstream:
+                result.update(self._dag.upstream(model_name))
+            if include_downstream:
+                result.update(self._dag.downstream(model_name))
+
+        for selection in model_selections:
+            if not selection:
+                continue
+
+            include_upstream = False
+            include_downstream = False
+            if selection[0] == "+":
+                selection = selection[1:]
+                include_upstream = True
+            if selection[-1] == "+":
+                selection = selection[:-1]
+                include_downstream = True
+
+            if "*" in selection:
+                for name in self._models:
+                    if fnmatch.fnmatch(name, selection):
+                        _add_model(name, include_upstream, include_downstream)
+            else:
+                _add_model(selection, include_upstream, include_downstream)
+
+        return result

tests/core/test_selector.py

@@ -90,10 +90,34 @@ def test_select_models(mocker: MockerFixture, make_snapshot):
             removed_model.name: removed_model.copy(update={"mapping_schema": added_model_schema}),
         },
     )
+    _assert_models_equal(
+        selector.select_models(["+modified_model"], env_name),
+        {
+            added_model.name: added_model,
+            modified_model_v2.name: modified_model_v2.copy(
+                update={"mapping_schema": added_model_schema}
+            ),
+            removed_model.name: removed_model.copy(update={"mapping_schema": added_model_schema}),
+        },
+    )
+    _assert_models_equal(
+        selector.select_models(["added_model+"], env_name),
+        {
+            added_model.name: added_model,
+            modified_model_v2.name: modified_model_v2.copy(
+                update={"mapping_schema": added_model_schema}
+            ),
+            removed_model.name: removed_model.copy(update={"mapping_schema": added_model_schema}),
+        },
+    )
     _assert_models_equal(
         selector.select_models(["added_model", "modified_model", "removed_model"], env_name),
         local_models,
     )
+    _assert_models_equal(
+        selector.select_models(["*_model", "removed_model"], env_name),
+        local_models,
+    )
 
 
 def test_select_models_missing_env(mocker: MockerFixture, make_snapshot):