Feat: Support model selection when constructing a plan (#1335)

Author: izeigerman

sqlmesh/cli/main.py

@@ -275,13 +275,20 @@ def diff(ctx: click.Context, environment: t.Optional[str] = None) -> None:
     help="Include unmodified models in the target environment.",
     default=None,
 )
+@click.option(
+    "--select-model",
+    type=str,
+    multiple=True,
+    help="Select specific model changes that should be included in the plan.",
+)
 @click.pass_context
 @error_handler
 def plan(ctx: click.Context, environment: t.Optional[str] = None, **kwargs: t.Any) -> None:
     """Plan a migration of the current context's models with the given environment."""
     context = ctx.obj
     restate_models = kwargs.pop("restate_model", None)
-    context.plan(environment, restate_models=restate_models, **kwargs)
+    select_models = kwargs.pop("select_model", None)
+    context.plan(environment, restate_models=restate_models, select_models=select_models, **kwargs)
 
 
 @cli.command("run")

sqlmesh/core/context.py

@@ -71,6 +71,7 @@
 from sqlmesh.core.plan import Plan
 from sqlmesh.core.scheduler import Scheduler
 from sqlmesh.core.schema_loader import create_schema_file
+from sqlmesh.core.selector import Selector
 from sqlmesh.core.snapshot import (
     Snapshot,
     SnapshotEvaluator,
@@ -545,58 +546,7 @@ def snapshots(self) -> t.Dict[str, Snapshot]:
         If one of the snapshots has been previosly stored in the persisted state, the stored
         instance will be returned.
         """
-        prod = self.state_reader.get_environment(c.PROD)
-        remote_snapshots = (
-            {
-                snapshot.name: snapshot
-                for snapshot in self.state_reader.get_snapshots(prod.snapshots).values()
-            }
-            if prod
-            else {}
-        )
-
-        fingerprint_cache: t.Dict[str, SnapshotFingerprint] = {}
-        models = self._models.copy()
-        audits = self._audits.copy()
-        projects = {config.project for config in self.configs.values()}
-
-        for name, snapshot in remote_snapshots.items():
-            if name not in models and snapshot.project not in projects:
-                models[name] = snapshot.model
-
-                for audit in snapshot.audits:
-                    if name not in audits:
-                        audits[name] = audit
-
-        snapshots = {}
-
-        for model in models.values():
-            if model.name not in self._models and model.name in remote_snapshots:
-                snapshot = remote_snapshots[model.name]
-                ttl = snapshot.ttl
-                project = snapshot.project
-            else:
-                config = self.config_for_model(model)
-                ttl = config.snapshot_ttl
-                project = config.project
-
-            snapshot = Snapshot.from_model(
-                model,
-                nodes=models,
-                audits=audits,
-                cache=fingerprint_cache,
-                ttl=ttl,
-                project=project,
-            )
-            snapshots[model.name] = snapshot
-
-        stored_snapshots = self.state_reader.get_snapshots(snapshots.values())
-
-        for snapshot in stored_snapshots.values():
-            # Keep the original model instance to preserve the query cache.
-            snapshot.node = snapshots[snapshot.name].node
-
-        return {name: stored_snapshots.get(s.snapshot_id, s) for name, s in snapshots.items()}
+        return self._snapshots()
 
     def render(
         self,
@@ -728,6 +678,7 @@ def plan(
         no_auto_categorization: t.Optional[bool] = None,
         effective_from: t.Optional[TimeLike] = None,
         include_unmodified: t.Optional[bool] = None,
+        select_models: t.Optional[t.Collection[str]] = None,
     ) -> Plan:
         """Interactively create a migration plan.
 
@@ -762,6 +713,7 @@ def plan(
                 option determines the behavior.
             effective_from: The effective date from which to apply forward-only changes on production.
             include_unmodified: Indicates whether to include unmodified models in the target development environment.
+            model_selections: A list of model selection strings to filter the models that should be included into this plan.
 
         Returns:
             The populated Plan object.
@@ -783,8 +735,18 @@ def plan(
         if include_unmodified is None:
             include_unmodified = self.config.include_unmodified
 
+        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
+            )
+
         plan = Plan(
-            context_diff=self._context_diff(environment or c.PROD, create_from=create_from),
+            context_diff=self._context_diff(
+                environment or c.PROD,
+                snapshots=self._snapshots(models_override),
+                create_from=create_from,
+            ),
             start=start,
             end=end,
             execution_time=execution_time,
@@ -1178,6 +1140,62 @@ def _model_tables(self) -> t.Dict[str, str]:
             for name, snapshot in self.snapshots.items()
         }
 
+    def _snapshots(
+        self, models_override: t.Optional[UniqueKeyDict[str, Model]] = None
+    ) -> t.Dict[str, Snapshot]:
+        prod = self.state_reader.get_environment(c.PROD)
+        remote_snapshots = (
+            {
+                snapshot.name: snapshot
+                for snapshot in self.state_reader.get_snapshots(prod.snapshots).values()
+            }
+            if prod
+            else {}
+        )
+
+        fingerprint_cache: t.Dict[str, SnapshotFingerprint] = {}
+        models = (models_override or self._models).copy()
+        audits = self._audits.copy()
+        projects = {config.project for config in self.configs.values()}
+
+        for name, snapshot in remote_snapshots.items():
+            if name not in models and snapshot.project not in projects:
+                models[name] = snapshot.model
+
+                for audit in snapshot.audits:
+                    if name not in audits:
+                        audits[name] = audit
+
+        snapshots = {}
+
+        for model in models.values():
+            if model.name not in self._models and model.name in remote_snapshots:
+                snapshot = remote_snapshots[model.name]
+                ttl = snapshot.ttl
+                project = snapshot.project
+            else:
+                config = self.config_for_model(model)
+                ttl = config.snapshot_ttl
+                project = config.project
+
+            snapshot = Snapshot.from_model(
+                model,
+                nodes=models,
+                audits=audits,
+                cache=fingerprint_cache,
+                ttl=ttl,
+                project=project,
+            )
+            snapshots[model.name] = snapshot
+
+        stored_snapshots = self.state_reader.get_snapshots(snapshots.values())
+
+        for snapshot in stored_snapshots.values():
+            # Keep the original model instance to preserve the query cache.
+            snapshot.node = snapshots[snapshot.name].node
+
+        return {name: stored_snapshots.get(s.snapshot_id, s) for name, s in snapshots.items()}
+
     def _context_diff(
         self,
         environment: str,

sqlmesh/core/selector.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import typing as t
+from pathlib import Path
+
+from sqlmesh.core.environment import Environment
+from sqlmesh.core.loader import update_model_schemas
+from sqlmesh.core.model import Model
+from sqlmesh.core.state_sync import StateReader
+from sqlmesh.utils import UniqueKeyDict
+from sqlmesh.utils.dag import DAG
+from sqlmesh.utils.errors import SQLMeshError
+
+
+class Selector:
+    def __init__(
+        self,
+        state_reader: StateReader,
+        models: UniqueKeyDict[str, Model],
+        context_path: Path = Path("."),
+    ):
+        self._state_reader = state_reader
+        self._models = models
+        self._context_path = context_path
+
+    def select_models(
+        self,
+        model_selections: t.Iterable[str],
+        target_env_name: str,
+        fallback_env_name: t.Optional[str] = None,
+    ) -> UniqueKeyDict[str, Model]:
+        """Given a set of selections returns models from the current state with names matching the
+        selection while sourcing the remaining models from the target environment.
+
+        Args:
+            model_selections: A set of selections.
+            target_env_name: The name of the target environment.
+            fallback_env_name: The name of the fallback environment that will be used if the target
+                environment doesn't exist.
+
+        Returns:
+            A dictionary of models.
+        """
+        target_env = self._state_reader.get_environment(Environment.normalize_name(target_env_name))
+        if not target_env and fallback_env_name:
+            target_env = self._state_reader.get_environment(
+                Environment.normalize_name(fallback_env_name)
+            )
+        if not target_env:
+            raise SQLMeshError(
+                f"Either the '{target_env_name}' or the '{fallback_env_name}' environment must exist in order to apply model selection."
+            )
+
+        env_models = {
+            s.name: s.model
+            for s in self._state_reader.get_snapshots(
+                target_env.snapshots, hydrate_seeds=True
+            ).values()
+        }
+
+        # TODO: Support selection expressions.
+        all_selected_models = set(model_selections)
+
+        dag: DAG[str] = DAG()
+        models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
+        all_model_names = set(self._models) | set(env_models)
+        for name in all_model_names:
+            model: t.Optional[Model] = None
+            if name not in all_selected_models and name in env_models:
+                # Unselected modified or added model.
+                model = env_models[name]
+            elif name in all_selected_models and name in self._models:
+                # Selected modified or removed model.
+                model = self._models[name]
+
+            if model:
+                # model.copy() can't be used here due to a cached state that can be a part of a model instance.
+                model = type(model).parse_obj(model.dict(exclude={"mapping_schema"}))
+                models[name] = model
+                dag.add(model.name, model.depends_on)
+
+        update_model_schemas(dag, models, self._context_path)
+
+        return models

sqlmesh/magics.py

@@ -306,6 +306,12 @@ def test(self, line: str, test_def_raw: t.Optional[str] = None) -> None:
         help="Include unmodified models in the target environment.",
         default=None,
     )
+    @argument(
+        "--select-model",
+        type=str,
+        nargs="*",
+        help="Select specific model changes that should be included in the plan.",
+    )
     @line_magic
     def plan(self, line: str) -> None:
         """Goes through a set of prompts to both establish a plan and apply it"""
@@ -332,6 +338,7 @@ def plan(self, line: str) -> None:
             no_auto_categorization=args.no_auto_categorization,
             effective_from=args.effective_from,
             include_unmodified=args.include_unmodified,
+            select_models=args.select_model,
         )
         self._context.console = console