Fix: Expand the parent's model query during query rendering if the parent is not categorized (#4254)

Author: izeigerman

sqlmesh/core/context.py

@@ -901,13 +901,7 @@ def get_snapshot(
         """
         if isinstance(node_or_snapshot, Snapshot):
             return node_or_snapshot
-        if isinstance(node_or_snapshot, str) and not self.standalone_audits.get(node_or_snapshot):
-            node_or_snapshot = normalize_model_name(
-                node_or_snapshot,
-                dialect=self.default_dialect,
-                default_catalog=self.default_catalog,
-            )
-        fqn = node_or_snapshot if isinstance(node_or_snapshot, str) else node_or_snapshot.fqn
+        fqn = self._node_or_snapshot_to_fqn(node_or_snapshot)
         snapshot = self.snapshots.get(fqn)
 
         if raise_if_missing and not snapshot:
@@ -1052,14 +1046,30 @@ def evaluate(
             execution_time: The date/time time reference to use for execution time.
             limit: A limit applied to the model.
         """
-        snapshot = self.get_snapshot(model_or_snapshot, raise_if_missing=True)
+        snapshots = self.snapshots
+        fqn = self._node_or_snapshot_to_fqn(model_or_snapshot)
+        if fqn not in snapshots:
+            raise SQLMeshError(f"Cannot find snapshot for '{fqn}'")
+        snapshot = snapshots[fqn]
+
+        # Expand all uncategorized parents since physical tables don't exist for them yet
+        expand = [
+            parent
+            for parent in self.dag.upstream(snapshot.model.fqn)
+            if (parent_snapshot := snapshots.get(parent))
+            and parent_snapshot.is_model
+            and parent_snapshot.model.is_sql
+            and not parent_snapshot.categorized
+        ]
+
         df = self.snapshot_evaluator.evaluate_and_fetch(
             snapshot,
             start=start,
             end=end,
             execution_time=execution_time,
             snapshots=self.snapshots,
             limit=limit or c.DEFAULT_MAX_LIMIT,
+            expand=expand,
         )
 
         if df is None:
@@ -2448,6 +2458,19 @@ def _nodes_to_snapshots(self, nodes: t.Dict[str, Node]) -> t.Dict[str, Snapshot]
             snapshots[snapshot.name] = snapshot
         return snapshots
 
+    def _node_or_snapshot_to_fqn(self, node_or_snapshot: NodeOrSnapshot) -> str:
+        if isinstance(node_or_snapshot, Snapshot):
+            return node_or_snapshot.name
+        if isinstance(node_or_snapshot, str) and not self.standalone_audits.get(node_or_snapshot):
+            return normalize_model_name(
+                node_or_snapshot,
+                dialect=self.default_dialect,
+                default_catalog=self.default_catalog,
+            )
+        if not isinstance(node_or_snapshot, str):
+            return node_or_snapshot.fqn
+        return node_or_snapshot
+
     @property
     def _plan_preview_enabled(self) -> bool:
         if self.config.plan.enable_preview is not None:

sqlmesh/core/engine_adapter/base.py

@@ -152,6 +152,7 @@ def with_log_level(self, level: int) -> EngineAdapter:
             register_comments=self._register_comments,
             null_connection=True,
             multithreaded=self._multithreaded,
+            pretty_sql=self._pretty_sql,
             **self._extra_config,
         )
 

sqlmesh/core/snapshot/definition.py

@@ -1014,6 +1014,11 @@ def categorize_as(self, category: SnapshotChangeCategory) -> None:
 
         self.change_category = category
 
+    @property
+    def categorized(self) -> bool:
+        """Whether the snapshot has been categorized."""
+        return self.change_category is not None and self.version is not None
+
     def table_name(self, is_deployable: bool = True) -> str:
         """Full table name pointing to the materialized location of the snapshot.
 

tests/core/test_context.py

@@ -1409,7 +1409,6 @@ def test_environment_statements(tmp_path: pathlib.Path):
         ],
         after_all=[
             "@grant_schema_usage()",
-            "@grant_select_privileges()",
             "@grant_usage_role(@schemas, 'admin')",
         ],
     )
@@ -1429,29 +1428,6 @@ def test_environment_statements(tmp_path: pathlib.Path):
         expression,
     )
 
-    create_temp_file(
-        tmp_path,
-        pathlib.Path(macros_dir, "grant_select_file.py"),
-        """
-from sqlmesh.core.macros import macro
-from sqlmesh.core.snapshot.definition import to_view_mapping
-
-@macro()
-def grant_select_privileges(evaluator):
-    if evaluator._environment_naming_info and evaluator.runtime_stage == 'before_all':
-        mapping = to_view_mapping(
-            evaluator._snapshots.values(), evaluator._environment_naming_info
-        )
-        return [
-            stmt
-            for stmt in [
-                f"GRANT SELECT ON VIEW {view_name} TO ROLE admin_role;"
-                for view_name in mapping.values()
-            ]
-        ]
-""",
-    )
-
     create_temp_file(
         tmp_path,
         pathlib.Path(macros_dir, "grant_schema_file.py"),
@@ -1499,8 +1475,8 @@ def grant_usage_role(evaluator, schemas, role):
     after_all = environment_statements.after_all
     python_env = environment_statements.python_env
 
-    assert isinstance(python_env["to_view_mapping"], Executable)
-    assert isinstance(python_env["grant_select_privileges"], Executable)
+    assert isinstance(python_env["grant_schema_usage"], Executable)
+    assert isinstance(python_env["grant_usage_role"], Executable)
 
     before_all_rendered = render_statements(
         statements=before_all,
@@ -1524,7 +1500,6 @@ def grant_usage_role(evaluator, schemas, role):
 
     assert after_all_rendered == [
         "GRANT USAGE ON SCHEMA db TO user_role",
-        "GRANT SELECT ON VIEW memory.db.test_after_model TO ROLE admin_role",
         'GRANT USAGE ON SCHEMA "db" TO "admin"',
     ]
 
@@ -1539,7 +1514,6 @@ def grant_usage_role(evaluator, schemas, role):
 
     assert after_all_rendered_dev == [
         "GRANT USAGE ON SCHEMA db__dev TO user_role",
-        "GRANT SELECT ON VIEW memory.db__dev.test_after_model TO ROLE admin_role",
         'GRANT USAGE ON SCHEMA "db__dev" TO "admin"',
     ]
 

tests/core/test_integration.py

@@ -4075,6 +4075,25 @@ def test_no_backfill_for_model_downstream_of_metadata_change(init_and_plan_conte
     )
 
 
+@time_machine.travel("2023-01-08 15:00:00 UTC")
+def test_evaluate_uncategorized_snapshot(init_and_plan_context: t.Callable):
+    context, plan = init_and_plan_context("examples/sushi")
+    context.apply(plan)
+
+    # Add a new projection
+    model = context.get_model("sushi.waiter_revenue_by_day")
+    context.upsert_model(add_projection_to_model(t.cast(SqlModel, model)))
+
+    # Downstream model references the new projection
+    downstream_model = context.get_model("sushi.top_waiters")
+    context.upsert_model(add_projection_to_model(t.cast(SqlModel, downstream_model), literal=False))
+
+    df = context.evaluate(
+        "sushi.top_waiters", start="2023-01-05", end="2023-01-06", execution_time=now()
+    )
+    assert set(df["one"].tolist()) == {1}
+
+
 @time_machine.travel("2023-01-08 15:00:00 UTC")
 def test_dbt_requirements(sushi_dbt_context: Context):
     assert set(sushi_dbt_context.requirements) == {"dbt-core", "dbt-duckdb"}