Feat: Make this_env and views available in macros of before after all (#4298)

Author: themisvaltinos

docs/concepts/macros/macro_variables.md

@@ -140,4 +140,5 @@ SQLMesh provides additional predefined variables used to modify model behavior b
 The following variables are also available in [`before_all` and `after_all` statements](../../guides/configuration.md#before_all-and-after_all-statements), as well as in macros invoked within them.
 
 * @this_env - A string value containing the name of the current [environment](../environments.md).
-* @schemas - A list of the schema names of the [virtual layer](../../concepts/glossary.md#virtual-layer) of the current environment.
+* @schemas - A list of the schema names of the [virtual layer](../../concepts/glossary.md#virtual-layer) of the current environment.
+* @views - A list of the view names of the [virtual layer](../../concepts/glossary.md#virtual-layer) of the current environment.

docs/guides/configuration.md

@@ -675,7 +675,7 @@ Configuration for a connection used to run unit tests. An in-memory DuckDB datab
 
 ### Scheduler
 
-Identifies which scheduler backend to use. The scheduler backend is used both for storing metadata and for executing [plans](../concepts/plans.md). By default, the scheduler type is set to `builtin`, which uses the existing SQL engine to store metadata. 
+Identifies which scheduler backend to use. The scheduler backend is used both for storing metadata and for executing [plans](../concepts/plans.md). By default, the scheduler type is set to `builtin`, which uses the existing SQL engine to store metadata.
 
 These options are in the [scheduler](../reference/configuration.md#scheduler) section of the configuration reference page.
 
@@ -1019,20 +1019,18 @@ For example, rather than using an `on_virtual_update` statement in each model to
 
 ```python linenums="1"
 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:
-        mapping = to_view_mapping(
-            evaluator._snapshots.values(), evaluator._environment_naming_info
-        )
+    if evaluator.views:
         return [
-            f"GRANT SELECT ON VIEW {view_name} TO ROLE admin_role;"
-            for view_name in mapping.values()
+            f"GRANT SELECT ON VIEW {view_name} /* sqlglot.meta replace=false */ TO ROLE admin_role;"
+            for view_name in evaluator.views
         ]
 ```
 
+By including the comment `/* sqlglot.meta replace=false */`, you further ensure that the evaluator does not replace the view name with the physical table name during rendering.
+
 ##### Example: Granting Schema Privileges
 
 Similarly, you can define a macro to grant schema usage privileges and, as demonstrated in the configuration above, using `this_env` macro conditionally execute it only in the production environment.
@@ -1042,21 +1040,14 @@ from sqlmesh import macro
 
 @macro()
 def grant_schema_usage(evaluator):
-    if evaluator._environment_naming_info:
-        schemas = {
-            snapshot.qualified_view_name.schema_for_environment(
-                evaluator._environment_naming_info
-            )
-            for snapshot in evaluator._snapshots.values()
-            if snapshot.is_model
-        }
+    if evaluator.this_env == "prod" and evaluator.schemas:
         return [
             f"GRANT USAGE ON SCHEMA {schema} TO admin_role;"
-            for schema in schemas
+            for schema in evaluator.schemas
         ]
 ```
 
-As demonstrated in these examples, the `environment_naming_info` is available within the macro evaluator for macros invoked within the `before_all` and `after_all` statements. Additionally, the macro `this_env` provides access to the current environment name, which can be helpful for more advanced use cases that require fine-grained control over their behaviour.
+As demonstrated in these examples, the `schemas`  and `views` are available within the macro evaluator for macros invoked within the `before_all` and `after_all` statements. Additionally, the macro `this_env` provides access to the current environment name, which can be helpful for more advanced use cases that require fine-grained control over their behaviour.
 
 ### Linting
 

sqlmesh/core/macros.py

@@ -491,6 +491,32 @@ def gateway(self) -> t.Optional[str]:
         """Returns the gateway name."""
         return self.var(c.GATEWAY)
 
+    @property
+    def snapshots(self) -> t.Dict[str, Snapshot]:
+        """Returns the snapshots if available."""
+        return self._snapshots
+
+    @property
+    def this_env(self) -> str:
+        """Returns the name of the current environment in before after all."""
+        if "this_env" not in self.locals:
+            raise SQLMeshError("Environment name is only available in before_all and after_all")
+        return self.locals["this_env"]
+
+    @property
+    def schemas(self) -> t.List[str]:
+        """Returns the schemas of the current environment in before after all macros."""
+        if "schemas" not in self.locals:
+            raise SQLMeshError("Schemas are only available in before_all and after_all")
+        return self.locals["schemas"]
+
+    @property
+    def views(self) -> t.List[str]:
+        """Returns the views of the current environment in before after all macros."""
+        if "views" not in self.locals:
+            raise SQLMeshError("Views are only available in before_all and after_all")
+        return self.locals["views"]
+
     def var(self, var_name: str, default: t.Optional[t.Any] = None) -> t.Optional[t.Any]:
         """Returns the value of the specified variable, or the default value if it doesn't exist."""
         return (self.locals.get(c.SQLMESH_VARS) or {}).get(var_name.lower(), default)

sqlmesh/core/renderer.py

@@ -110,17 +110,23 @@ def _render(
         if environment_naming_info is not None:
             kwargs["this_env"] = getattr(environment_naming_info, "name")
             if snapshots:
-                schemas = set(
-                    [
-                        s.qualified_view_name.schema_for_environment(
-                            environment_naming_info, dialect=self._dialect
+                schemas, views = set(), []
+                for snapshot in snapshots.values():
+                    if snapshot.is_model and not snapshot.is_symbolic:
+                        schemas.add(
+                            snapshot.qualified_view_name.schema_for_environment(
+                                environment_naming_info, dialect=self._dialect
+                            )
+                        )
+                        views.append(
+                            snapshot.display_name(
+                                environment_naming_info, self._default_catalog, self._dialect
+                            )
                         )
-                        for s in snapshots.values()
-                        if s.is_model and not s.is_symbolic
-                    ]
-                )
                 if schemas:
                     kwargs["schemas"] = list(schemas)
+                if views:
+                    kwargs["views"] = views
 
         this_model = kwargs.pop("this_model", None)
 

tests/core/test_context.py

@@ -1417,6 +1417,7 @@ def test_environment_statements(tmp_path: pathlib.Path):
         after_all=[
             "@grant_schema_usage()",
             "@grant_usage_role(@schemas, 'admin')",
+            "@grant_select_privileges()",
         ],
     )
 
@@ -1435,6 +1436,21 @@ def test_environment_statements(tmp_path: pathlib.Path):
         expression,
     )
 
+    create_temp_file(
+        tmp_path,
+        pathlib.Path(macros_dir, "grant_select_privileges.py"),
+        """
+from sqlmesh.core.macros import macro
+@macro()
+def grant_select_privileges(evaluator):
+    if evaluator.this_env and evaluator.views:
+        return [
+            f"GRANT SELECT ON VIEW {view_name} /* sqlglot.meta replace=false */ TO ROLE admin_role;"
+            for view_name in evaluator.views
+        ]
+""",
+    )
+
     create_temp_file(
         tmp_path,
         pathlib.Path(macros_dir, "grant_schema_file.py"),
@@ -1484,6 +1500,7 @@ def grant_usage_role(evaluator, schemas, role):
 
     assert isinstance(python_env["grant_schema_usage"], Executable)
     assert isinstance(python_env["grant_usage_role"], Executable)
+    assert isinstance(python_env["grant_select_privileges"], Executable)
 
     before_all_rendered = render_statements(
         statements=before_all,
@@ -1502,27 +1519,33 @@ def grant_usage_role(evaluator, schemas, role):
         python_env=python_env,
         snapshots=snapshots,
         environment_naming_info=EnvironmentNamingInfo(name="prod"),
-        runtime_stage=RuntimeStage.BEFORE_ALL,
+        runtime_stage=RuntimeStage.AFTER_ALL,
     )
 
-    assert after_all_rendered == [
-        "GRANT USAGE ON SCHEMA db TO user_role",
-        'GRANT USAGE ON SCHEMA "db" TO "admin"',
-    ]
+    assert sorted(after_all_rendered) == sorted(
+        [
+            "GRANT USAGE ON SCHEMA db TO user_role",
+            'GRANT USAGE ON SCHEMA "db" TO "admin"',
+            "GRANT SELECT ON VIEW memory.db.test_after_model /* sqlglot.meta replace=false */ TO ROLE admin_role",
+        ]
+    )
 
     after_all_rendered_dev = render_statements(
         statements=after_all,
         dialect=dialect,
         python_env=python_env,
         snapshots=snapshots,
         environment_naming_info=EnvironmentNamingInfo(name="dev"),
-        runtime_stage=RuntimeStage.BEFORE_ALL,
+        runtime_stage=RuntimeStage.AFTER_ALL,
     )
 
-    assert after_all_rendered_dev == [
-        "GRANT USAGE ON SCHEMA db__dev TO user_role",
-        'GRANT USAGE ON SCHEMA "db__dev" TO "admin"',
-    ]
+    assert sorted(after_all_rendered_dev) == sorted(
+        [
+            "GRANT USAGE ON SCHEMA db__dev TO user_role",
+            'GRANT USAGE ON SCHEMA "db__dev" TO "admin"',
+            "GRANT SELECT ON VIEW memory.db__dev.test_after_model /* sqlglot.meta replace=false */ TO ROLE admin_role",
+        ]
+    )
 
 
 def test_plan_environment_statements(tmp_path: pathlib.Path):