Perf: add opt-in project-index loading for render

Signed-off-by: Andreas Fredhøi <andreas.fredhoi@fresio.no>

Author: fresioAS

docs/reference/cli.md

@@ -447,6 +447,9 @@ Options:
                               only they will be expanded as raw queries.
   --dialect TEXT              The SQL dialect to render the query as.
   --no-format                 Disable fancy formatting of the query.
+  --use-project-index         Use the persistent project index to load and
+                              render only the target model and its upstream
+                              dependencies.
   --max-text-width INTEGER    The max number of characters in a segment before
                               creating new lines in pretty mode.
   --leading-comma             Determines whether or not the comma is leading

sqlmesh/cli/main.py

@@ -141,8 +141,8 @@ def cli(
         if ctx.invoked_subcommand in SKIP_LOAD_COMMANDS:
             load = False
 
-    # Unlike the other commands above, lint can scope its own load for multi-project contexts.
-    if ctx.invoked_subcommand == "lint":
+    # These commands can scope their own load for multi-project contexts.
+    if ctx.invoked_subcommand in ("lint", "render"):
         load = False
 
     configs = load_configs(config, Context.CONFIG_TYPE, paths, dotenv_path=dotenv)
@@ -284,6 +284,11 @@ def init(
     help="The SQL dialect to render the query as.",
 )
 @click.option("--no-format", is_flag=True, help="Disable fancy formatting of the query.")
+@click.option(
+    "--use-project-index",
+    is_flag=True,
+    help="Use the persistent project index to load and render only the target model and its upstream dependencies.",
+)
 @opt.format_options
 @click.pass_context
 @error_handler
@@ -297,19 +302,20 @@ def render(
     expand: t.Optional[t.Union[bool, t.Iterable[str]]] = None,
     dialect: t.Optional[str] = None,
     no_format: bool = False,
+    use_project_index: bool = False,
     **format_kwargs: t.Any,
 ) -> None:
     """Render a model's query, optionally expanding referenced models."""
-    model = ctx.obj.get_model(model, raise_if_missing=True)
-
     rendered = ctx.obj.render(
         model,
         start=start,
         end=end,
         execution_time=execution_time,
         expand=expand,
+        use_project_index=use_project_index,
     )
 
+    model = ctx.obj.get_model(model, raise_if_missing=True)
     format_config = ctx.obj.config_for_node(model).format
     format_kwargs = {
         **format_config.generator_options,

sqlmesh/core/context.py

@@ -1186,6 +1186,7 @@ def render(
         end: t.Optional[TimeLike] = None,
         execution_time: t.Optional[TimeLike] = None,
         expand: t.Union[bool, t.Iterable[str]] = False,
+        use_project_index: bool = False,
         **kwargs: t.Any,
     ) -> exp.Expr:
         """Renders a model's query, expanding macros with provided kwargs, and optionally expanding referenced models.
@@ -1198,12 +1199,20 @@ def render(
             expand: Whether or not to use expand materialized models, defaults to False.
                 If True, all referenced models are expanded as raw queries.
                 If a list, only referenced models are expanded as raw queries.
+            use_project_index: Whether to use the persistent project index to load and
+                render only the target model and its transitive upstream dependencies.
 
         Returns:
             The rendered expression.
         """
         execution_time = execution_time or now()
 
+        if not self._loaded:
+            target_fqns = (
+                {self._node_or_snapshot_to_fqn(model_or_snapshot)} if use_project_index else None
+            )
+            self.load(model_fqns=target_fqns, use_project_index=use_project_index)
+
         model = self.get_model(model_or_snapshot, raise_if_missing=True)
 
         if expand and not isinstance(expand, bool):
@@ -1232,7 +1241,19 @@ def render(
             )
             return next(pandas_to_sql(t.cast(pd.DataFrame, df), model.columns_to_types))
 
-        snapshots = self.snapshots
+        if use_project_index:
+            # Only the target model and its transitive upstream dependencies can be referenced
+            # by the rendered query, so there is no need to create snapshots for the rest.
+            upstream_fqns = {model.fqn, *self.dag.upstream(model.fqn)}
+            upstream_models: UniqueKeyDict[str, Model] = UniqueKeyDict(
+                "models", {fqn: m for fqn, m in self._models.items() if fqn in upstream_fqns}
+            )
+            snapshots = self._snapshots(
+                upstream_models,
+                include_standalone_audits=False,
+            )
+        else:
+            snapshots = self.snapshots
         deployability_index = DeployabilityIndex.create(snapshots.values(), start=start)
 
         return model.render_query_or_raise(
@@ -2986,9 +3007,13 @@ def _get_engine_adapter(self, gateway: t.Optional[str] = None) -> EngineAdapter:
         return self.engine_adapter
 
     def _snapshots(
-        self, models_override: t.Optional[UniqueKeyDict[str, Model]] = None
+        self,
+        models_override: t.Optional[UniqueKeyDict[str, Model]] = None,
+        include_standalone_audits: bool = True,
     ) -> t.Dict[str, Snapshot]:
-        nodes = {**(models_override or self._models), **self._standalone_audits}
+        nodes: t.Dict[str, Node] = dict(models_override or self._models)
+        if include_standalone_audits:
+            nodes.update(self._standalone_audits)
         snapshots = self._nodes_to_snapshots(nodes)
         stored_snapshots = self.state_reader.get_snapshots(snapshots.values())
 

tests/cli/test_cli.py

@@ -2073,6 +2073,19 @@ def test_render(runner: CliRunner, tmp_path: Path):
 
     assert expected in cleaned_output
 
+    indexed_result = runner.invoke(
+        cli,
+        [
+            "--paths",
+            str(tmp_path),
+            "render",
+            "sqlmesh_example.full_model",
+            "--use-project-index",
+            "--no-format",
+        ],
+    )
+    assert indexed_result.exit_code == 0
+
 
 @time_machine.travel(FREEZE_TIME)
 def test_signals(runner: CliRunner, tmp_path: Path):

tests/core/test_context.py

@@ -274,6 +274,82 @@ def test_render_seed_model(sushi_context, assert_exp_eq):
     )
 
 
+@pytest.mark.slow
+def test_render_only_creates_snapshots_for_upstream_models(sushi_context: Context):
+    model = sushi_context.get_model("sushi.top_waiters", raise_if_missing=True)
+    upstream_fqns = {model.fqn, *sushi_context.dag.upstream(model.fqn)}
+
+    # Sanity check that the project contains models outside of the target model's subgraph.
+    assert set(sushi_context.models) - upstream_fqns
+
+    with patch.object(
+        sushi_context.state_reader,
+        "get_snapshots",
+        wraps=sushi_context.state_reader.get_snapshots,
+    ) as default_get_snapshots_mock:
+        sushi_context.render("sushi.top_waiters")
+
+    default_requested_names = {
+        snapshot.name
+        for call_args in default_get_snapshots_mock.call_args_list
+        for snapshot in call_args.args[0]
+    }
+    assert set(sushi_context.models) <= default_requested_names
+
+    with patch.object(
+        sushi_context.state_reader,
+        "get_snapshots",
+        wraps=sushi_context.state_reader.get_snapshots,
+    ) as get_snapshots_mock:
+        sushi_context.render("sushi.top_waiters", use_project_index=True)
+
+    requested_names = {
+        snapshot.name
+        for call_args in get_snapshots_mock.call_args_list
+        for snapshot in call_args.args[0]
+    }
+    assert model.fqn in requested_names
+    assert requested_names == upstream_fqns
+
+
+def test_render_only_loads_upstream_model_files(tmp_path: pathlib.Path) -> None:
+    create_temp_file(
+        tmp_path,
+        pathlib.Path("models", "a.sql"),
+        "MODEL(name a, kind FULL); SELECT 1 AS col;",
+    )
+    create_temp_file(
+        tmp_path,
+        pathlib.Path("models", "b.sql"),
+        "MODEL(name b, kind FULL); SELECT col FROM a;",
+    )
+    create_temp_file(
+        tmp_path,
+        pathlib.Path("models", "c.sql"),
+        "MODEL(name c, kind FULL); SELECT col FROM b;",
+    )
+    config = Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))
+
+    # Populate the persistent model path/dependency index.
+    Context(config=config, paths=tmp_path, load=False).load(use_project_index=True)
+
+    ctx = Context(config=config, paths=tmp_path, load=False)
+    loader = t.cast(SqlMeshLoader, ctx._loaders[0])
+    with patch.object(
+        loader,
+        "_load_sql_models",
+        wraps=loader._load_sql_models,
+    ) as load_sql_models_mock:
+        ctx.render("b", use_project_index=True)
+
+    selected_paths = load_sql_models_mock.call_args.kwargs["selected_paths"]
+    assert {path.name for path in selected_paths} == {"a.sql", "b.sql"}
+    assert set(ctx.models) == {
+        ctx.get_model("a", raise_if_missing=True).fqn,
+        ctx.get_model("b", raise_if_missing=True).fqn,
+    }
+
+
 @pytest.mark.slow
 def test_diff(sushi_context: Context, mocker: MockerFixture):
     mock_console = mocker.Mock()