From d33ced5304c809a844f6397ae83bcd100945c0fa Mon Sep 17 00:00:00 2001 From: David Dai Date: Mon, 6 Oct 2025 13:08:30 -0700 Subject: [PATCH] feat(experimental): add grants support for DBT custom materializations --- sqlmesh/core/snapshot/evaluator.py | 14 ++++++ tests/dbt/test_custom_materializations.py | 56 +++++++++++++++++++++++ 2 files changed, 70 insertions(+) diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py index 2676709d85..a2c3328d92 100644 --- a/sqlmesh/core/snapshot/evaluator.py +++ b/sqlmesh/core/snapshot/evaluator.py @@ -2908,6 +2908,13 @@ def create( **kwargs, ) + # Apply grants after dbt custom materialization table creation + if not skip_grants: + is_snapshot_deployable = kwargs.get("is_snapshot_deployable", False) + self._apply_grants( + model, table_name, GrantsTargetLayer.PHYSICAL, is_snapshot_deployable + ) + def insert( self, table_name: str, @@ -2926,6 +2933,13 @@ def insert( **kwargs, ) + # Apply grants after custom materialization insert (only on first insert) + if is_first_insert: + is_snapshot_deployable = kwargs.get("is_snapshot_deployable", False) + self._apply_grants( + model, table_name, GrantsTargetLayer.PHYSICAL, is_snapshot_deployable + ) + def append( self, table_name: str, diff --git a/tests/dbt/test_custom_materializations.py b/tests/dbt/test_custom_materializations.py index 9e7a94315c..c1625d0251 100644 --- a/tests/dbt/test_custom_materializations.py +++ b/tests/dbt/test_custom_materializations.py @@ -7,6 +7,7 @@ from sqlmesh import Context from sqlmesh.core.config import ModelDefaultsConfig +from sqlmesh.core.engine_adapter import DuckDBEngineAdapter from sqlmesh.core.model.kind import DbtCustomKind from sqlmesh.dbt.context import DbtContext from sqlmesh.dbt.manifest import ManifestHelper @@ -719,3 +720,58 @@ def test_custom_materialization_lineage_tracking(copy_to_temp_path: t.Callable): # Dev and prod should have the same data as they share physical data assert dev_analytics_result["count"][0] == prod_analytics_result["count"][0] assert dev_analytics_result["unique_waiters"][0] == prod_analytics_result["unique_waiters"][0] + + +@pytest.mark.xdist_group("dbt_manifest") +def test_custom_materialization_grants(copy_to_temp_path: t.Callable, mocker): + path = copy_to_temp_path("tests/fixtures/dbt/sushi_test") + temp_project = path[0] + + models_dir = temp_project / "models" + models_dir.mkdir(parents=True, exist_ok=True) + + grants_model_content = """ +{{ config( + materialized='custom_incremental', + grants={ + 'select': ['user1', 'user2'], + 'insert': ['writer'] + } +) }} + +SELECT + CURRENT_TIMESTAMP as created_at, + 1 as id, + 'grants_test' as test_type +""".strip() + + (models_dir / "test_grants_model.sql").write_text(grants_model_content) + + mocker.patch.object(DuckDBEngineAdapter, "SUPPORTS_GRANTS", True) + mocker.patch.object(DuckDBEngineAdapter, "_get_current_grants_config", return_value={}) + + sync_grants_calls = [] + + def mock_sync_grants(*args, **kwargs): + sync_grants_calls.append((args, kwargs)) + + mocker.patch.object(DuckDBEngineAdapter, "sync_grants_config", side_effect=mock_sync_grants) + + context = Context(paths=path) + + model = context.get_model("sushi.test_grants_model") + assert isinstance(model.kind, DbtCustomKind) + plan = context.plan(select_models=["sushi.test_grants_model"]) + context.apply(plan) + + assert len(sync_grants_calls) == 1 + args = sync_grants_calls[0][0] + assert args + + table = args[0] + grants_config = args[1] + assert table.sql(dialect="duckdb") == "memory.sushi.test_grants_model" + assert grants_config == { + "select": ["user1", "user2"], + "insert": ["writer"], + }