Skip to content

Commit d33ced5

Browse files
committed
feat(experimental): add grants support for DBT custom materializations
1 parent 5e5e2f0 commit d33ced5

File tree

2 files changed

+70
-0
lines changed

2 files changed

+70
-0
lines changed

sqlmesh/core/snapshot/evaluator.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2908,6 +2908,13 @@ def create(
29082908
**kwargs,
29092909
)
29102910

2911+
# Apply grants after dbt custom materialization table creation
2912+
if not skip_grants:
2913+
is_snapshot_deployable = kwargs.get("is_snapshot_deployable", False)
2914+
self._apply_grants(
2915+
model, table_name, GrantsTargetLayer.PHYSICAL, is_snapshot_deployable
2916+
)
2917+
29112918
def insert(
29122919
self,
29132920
table_name: str,
@@ -2926,6 +2933,13 @@ def insert(
29262933
**kwargs,
29272934
)
29282935

2936+
# Apply grants after custom materialization insert (only on first insert)
2937+
if is_first_insert:
2938+
is_snapshot_deployable = kwargs.get("is_snapshot_deployable", False)
2939+
self._apply_grants(
2940+
model, table_name, GrantsTargetLayer.PHYSICAL, is_snapshot_deployable
2941+
)
2942+
29292943
def append(
29302944
self,
29312945
table_name: str,

tests/dbt/test_custom_materializations.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
from sqlmesh import Context
99
from sqlmesh.core.config import ModelDefaultsConfig
10+
from sqlmesh.core.engine_adapter import DuckDBEngineAdapter
1011
from sqlmesh.core.model.kind import DbtCustomKind
1112
from sqlmesh.dbt.context import DbtContext
1213
from sqlmesh.dbt.manifest import ManifestHelper
@@ -719,3 +720,58 @@ def test_custom_materialization_lineage_tracking(copy_to_temp_path: t.Callable):
719720
# Dev and prod should have the same data as they share physical data
720721
assert dev_analytics_result["count"][0] == prod_analytics_result["count"][0]
721722
assert dev_analytics_result["unique_waiters"][0] == prod_analytics_result["unique_waiters"][0]
723+
724+
725+
@pytest.mark.xdist_group("dbt_manifest")
726+
def test_custom_materialization_grants(copy_to_temp_path: t.Callable, mocker):
727+
path = copy_to_temp_path("tests/fixtures/dbt/sushi_test")
728+
temp_project = path[0]
729+
730+
models_dir = temp_project / "models"
731+
models_dir.mkdir(parents=True, exist_ok=True)
732+
733+
grants_model_content = """
734+
{{ config(
735+
materialized='custom_incremental',
736+
grants={
737+
'select': ['user1', 'user2'],
738+
'insert': ['writer']
739+
}
740+
) }}
741+
742+
SELECT
743+
CURRENT_TIMESTAMP as created_at,
744+
1 as id,
745+
'grants_test' as test_type
746+
""".strip()
747+
748+
(models_dir / "test_grants_model.sql").write_text(grants_model_content)
749+
750+
mocker.patch.object(DuckDBEngineAdapter, "SUPPORTS_GRANTS", True)
751+
mocker.patch.object(DuckDBEngineAdapter, "_get_current_grants_config", return_value={})
752+
753+
sync_grants_calls = []
754+
755+
def mock_sync_grants(*args, **kwargs):
756+
sync_grants_calls.append((args, kwargs))
757+
758+
mocker.patch.object(DuckDBEngineAdapter, "sync_grants_config", side_effect=mock_sync_grants)
759+
760+
context = Context(paths=path)
761+
762+
model = context.get_model("sushi.test_grants_model")
763+
assert isinstance(model.kind, DbtCustomKind)
764+
plan = context.plan(select_models=["sushi.test_grants_model"])
765+
context.apply(plan)
766+
767+
assert len(sync_grants_calls) == 1
768+
args = sync_grants_calls[0][0]
769+
assert args
770+
771+
table = args[0]
772+
grants_config = args[1]
773+
assert table.sql(dialect="duckdb") == "memory.sushi.test_grants_model"
774+
assert grants_config == {
775+
"select": ["user1", "user2"],
776+
"insert": ["writer"],
777+
}

0 commit comments

Comments
 (0)