Skip to content

Commit 7020d12

Browse files
Feat: Make this_env and views available in macros of before after all
1 parent fd19aa0 commit 7020d12

5 files changed

Lines changed: 78 additions & 36 deletions

File tree

docs/concepts/macros/macro_variables.md

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,4 +140,5 @@ SQLMesh provides additional predefined variables used to modify model behavior b
140140
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.
141141

142142
* @this_env - A string value containing the name of the current [environment](../environments.md).
143-
* @schemas - A list of the schema names of the [virtual layer](../../concepts/glossary.md#virtual-layer) of the current environment.
143+
* @schemas - A list of the schema names of the [virtual layer](../../concepts/glossary.md#virtual-layer) of the current environment.
144+
* @views - A list of the views names of the [virtual layer](../../concepts/glossary.md#virtual-layer) of the current environment.

docs/guides/configuration.md

Lines changed: 9 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -675,7 +675,7 @@ Configuration for a connection used to run unit tests. An in-memory DuckDB datab
675675

676676
### Scheduler
677677

678-
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.
678+
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.
679679

680680
These options are in the [scheduler](../reference/configuration.md#scheduler) section of the configuration reference page.
681681

@@ -1023,16 +1023,15 @@ from sqlmesh.core.snapshot.definition import to_view_mapping
10231023
10241024
@macro()
10251025
def grant_select_privileges(evaluator):
1026-
if evaluator._environment_naming_info:
1027-
mapping = to_view_mapping(
1028-
evaluator._snapshots.values(), evaluator._environment_naming_info
1029-
)
1026+
if evaluator.views:
10301027
return [
1031-
f"GRANT SELECT ON VIEW {view_name} TO ROLE admin_role;"
1032-
for view_name in mapping.values()
1028+
f"GRANT SELECT ON VIEW {view_name} /* sqlglot.meta replace=false */ TO ROLE admin_role;"
1029+
for view_name in evaluator.views
10331030
]
10341031
```
10351032

1033+
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.
1034+
10361035
##### Example: Granting Schema Privileges
10371036

10381037
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 +1041,14 @@ from sqlmesh import macro
10421041
10431042
@macro()
10441043
def grant_schema_usage(evaluator):
1045-
if evaluator._environment_naming_info:
1046-
schemas = {
1047-
snapshot.qualified_view_name.schema_for_environment(
1048-
evaluator._environment_naming_info
1049-
)
1050-
for snapshot in evaluator._snapshots.values()
1051-
if snapshot.is_model
1052-
}
1044+
if evaluator.this_env == "prod" and evaluator.schemas:
10531045
return [
10541046
f"GRANT USAGE ON SCHEMA {schema} TO admin_role;"
1055-
for schema in schemas
1047+
for schema in evaluator.schemas
10561048
]
10571049
```
10581050

1059-
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.
1051+
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.
10601052

10611053
### Linting
10621054

sqlmesh/core/macros.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -491,6 +491,26 @@ def gateway(self) -> t.Optional[str]:
491491
"""Returns the gateway name."""
492492
return self.var(c.GATEWAY)
493493

494+
@property
495+
def snapshots(self) -> t.Dict[str, Snapshot]:
496+
"""Returns the snapshots if available."""
497+
return self._snapshots
498+
499+
@property
500+
def this_env(self) -> t.Optional[str]:
501+
"""Returns the name of the current environment in before after all."""
502+
return self.locals.get("this_env")
503+
504+
@property
505+
def schemas(self) -> t.Optional[t.List[str]]:
506+
"""Returns the schemas of the current environment in before after all macros."""
507+
return self.locals.get("schemas")
508+
509+
@property
510+
def views(self) -> t.Optional[t.List[str]]:
511+
"""Returns the views of the current environment in before after all macros."""
512+
return self.locals.get("views")
513+
494514
def var(self, var_name: str, default: t.Optional[t.Any] = None) -> t.Optional[t.Any]:
495515
"""Returns the value of the specified variable, or the default value if it doesn't exist."""
496516
return (self.locals.get(c.SQLMESH_VARS) or {}).get(var_name.lower(), default)

sqlmesh/core/renderer.py

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -110,17 +110,23 @@ def _render(
110110
if environment_naming_info is not None:
111111
kwargs["this_env"] = getattr(environment_naming_info, "name")
112112
if snapshots:
113-
schemas = set(
114-
[
115-
s.qualified_view_name.schema_for_environment(
116-
environment_naming_info, dialect=self._dialect
113+
schemas, views = set(), []
114+
for snapshot in snapshots.values():
115+
if snapshot.is_model and not snapshot.is_symbolic:
116+
schemas.add(
117+
snapshot.qualified_view_name.schema_for_environment(
118+
environment_naming_info, dialect=self._dialect
119+
)
120+
)
121+
views.append(
122+
snapshot.display_name(
123+
environment_naming_info, self._default_catalog, self._dialect
124+
)
117125
)
118-
for s in snapshots.values()
119-
if s.is_model and not s.is_symbolic
120-
]
121-
)
122126
if schemas:
123127
kwargs["schemas"] = list(schemas)
128+
if views:
129+
kwargs["views"] = views
124130

125131
this_model = kwargs.pop("this_model", None)
126132

tests/core/test_context.py

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1410,6 +1410,7 @@ def test_environment_statements(tmp_path: pathlib.Path):
14101410
after_all=[
14111411
"@grant_schema_usage()",
14121412
"@grant_usage_role(@schemas, 'admin')",
1413+
"@grant_select_privileges()",
14131414
],
14141415
)
14151416

@@ -1428,6 +1429,21 @@ def test_environment_statements(tmp_path: pathlib.Path):
14281429
expression,
14291430
)
14301431

1432+
create_temp_file(
1433+
tmp_path,
1434+
pathlib.Path(macros_dir, "grant_select_privileges.py"),
1435+
"""
1436+
from sqlmesh.core.macros import macro
1437+
@macro()
1438+
def grant_select_privileges(evaluator):
1439+
if evaluator.this_env and evaluator.views:
1440+
return [
1441+
f"GRANT SELECT ON VIEW {view_name} /* sqlglot.meta replace=false */ TO ROLE admin_role;"
1442+
for view_name in evaluator.views
1443+
]
1444+
""",
1445+
)
1446+
14311447
create_temp_file(
14321448
tmp_path,
14331449
pathlib.Path(macros_dir, "grant_schema_file.py"),
@@ -1477,6 +1493,7 @@ def grant_usage_role(evaluator, schemas, role):
14771493

14781494
assert isinstance(python_env["grant_schema_usage"], Executable)
14791495
assert isinstance(python_env["grant_usage_role"], Executable)
1496+
assert isinstance(python_env["grant_select_privileges"], Executable)
14801497

14811498
before_all_rendered = render_statements(
14821499
statements=before_all,
@@ -1495,27 +1512,33 @@ def grant_usage_role(evaluator, schemas, role):
14951512
python_env=python_env,
14961513
snapshots=snapshots,
14971514
environment_naming_info=EnvironmentNamingInfo(name="prod"),
1498-
runtime_stage=RuntimeStage.BEFORE_ALL,
1515+
runtime_stage=RuntimeStage.AFTER_ALL,
14991516
)
15001517

1501-
assert after_all_rendered == [
1502-
"GRANT USAGE ON SCHEMA db TO user_role",
1503-
'GRANT USAGE ON SCHEMA "db" TO "admin"',
1504-
]
1518+
assert sorted(after_all_rendered) == sorted(
1519+
[
1520+
"GRANT USAGE ON SCHEMA db TO user_role",
1521+
'GRANT USAGE ON SCHEMA "db" TO "admin"',
1522+
"GRANT SELECT ON VIEW memory.db.test_after_model /* sqlglot.meta replace=false */ TO ROLE admin_role",
1523+
]
1524+
)
15051525

15061526
after_all_rendered_dev = render_statements(
15071527
statements=after_all,
15081528
dialect=dialect,
15091529
python_env=python_env,
15101530
snapshots=snapshots,
15111531
environment_naming_info=EnvironmentNamingInfo(name="dev"),
1512-
runtime_stage=RuntimeStage.BEFORE_ALL,
1532+
runtime_stage=RuntimeStage.AFTER_ALL,
15131533
)
15141534

1515-
assert after_all_rendered_dev == [
1516-
"GRANT USAGE ON SCHEMA db__dev TO user_role",
1517-
'GRANT USAGE ON SCHEMA "db__dev" TO "admin"',
1518-
]
1535+
assert sorted(after_all_rendered_dev) == sorted(
1536+
[
1537+
"GRANT USAGE ON SCHEMA db__dev TO user_role",
1538+
'GRANT USAGE ON SCHEMA "db__dev" TO "admin"',
1539+
"GRANT SELECT ON VIEW memory.db__dev.test_after_model /* sqlglot.meta replace=false */ TO ROLE admin_role",
1540+
]
1541+
)
15191542

15201543

15211544
def test_plan_environment_statements(tmp_path: pathlib.Path):

0 commit comments

Comments
 (0)