Skip to content

Commit db4c90b

Browse files
committed
Push apply grants down to evaluation strategies
and always apply grants on migrate(). This way, grants will always be applied even when grants are the only metadata that changed.
1 parent 648b11d commit db4c90b

4 files changed

Lines changed: 927 additions & 86 deletions

File tree

sqlmesh/core/snapshot/evaluator.py

Lines changed: 52 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1046,6 +1046,7 @@ def _clone_snapshot_in_dev(
10461046
allow_destructive_snapshots=allow_destructive_snapshots,
10471047
allow_additive_snapshots=allow_additive_snapshots,
10481048
)
1049+
10491050
except Exception:
10501051
adapter.drop_table(target_table_name)
10511052
raise
@@ -1125,6 +1126,7 @@ def _migrate_target_table(
11251126
rendered_physical_properties=rendered_physical_properties,
11261127
dry_run=False,
11271128
run_pre_post_statements=run_pre_post_statements,
1129+
skip_grants=True, # skip grants for tmp table
11281130
)
11291131
try:
11301132
evaluation_strategy = _evaluation_strategy(snapshot, adapter)
@@ -1390,6 +1392,7 @@ def _execute_create(
13901392
rendered_physical_properties: t.Dict[str, exp.Expression],
13911393
dry_run: bool,
13921394
run_pre_post_statements: bool = True,
1395+
skip_grants: bool = False,
13931396
) -> None:
13941397
adapter = self.get_adapter(snapshot.model.gateway)
13951398
evaluation_strategy = _evaluation_strategy(snapshot, adapter)
@@ -1413,12 +1416,11 @@ def _execute_create(
14131416
is_snapshot_representative=is_snapshot_representative,
14141417
dry_run=dry_run,
14151418
physical_properties=rendered_physical_properties,
1419+
skip_grants=skip_grants,
14161420
)
14171421
if run_pre_post_statements:
14181422
adapter.execute(snapshot.model.render_post_statements(**create_render_kwargs))
14191423

1420-
evaluation_strategy._apply_grants(snapshot.model, table_name, GrantsTargetLayer.PHYSICAL)
1421-
14221424
def _can_clone(self, snapshot: Snapshot, deployability_index: DeployabilityIndex) -> bool:
14231425
adapter = self.get_adapter(snapshot.model.gateway)
14241426
return (
@@ -1686,16 +1688,11 @@ def _apply_grants(
16861688
return
16871689

16881690
logger.info(f"Applying grants for model {model.name} to table {table_name}")
1689-
1690-
try:
1691-
self.adapter.sync_grants_config(
1692-
exp.to_table(table_name, dialect=model.dialect),
1693-
grants_config,
1694-
model.grants_table_type,
1695-
)
1696-
except Exception:
1697-
# Log error but don't fail evaluation if grants fail
1698-
logger.error(f"Failed to apply grants for model {model.name}", exc_info=True)
1691+
self.adapter.sync_grants_config(
1692+
exp.to_table(table_name, dialect=self.adapter.dialect),
1693+
grants_config,
1694+
model.grants_table_type,
1695+
)
16991696

17001697

17011698
class SymbolicStrategy(EvaluationStrategy):
@@ -1860,6 +1857,10 @@ def create(
18601857
column_descriptions=model.column_descriptions if is_table_deployable else None,
18611858
)
18621859

1860+
# Apply grants after table creation (unless explicitly skipped by caller)
1861+
if not kwargs.get("skip_grants", False):
1862+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
1863+
18631864
def migrate(
18641865
self,
18651866
target_table_name: str,
@@ -1885,6 +1886,9 @@ def migrate(
18851886
)
18861887
self.adapter.alter_table(alter_operations)
18871888

1889+
# Apply grants after schema migration
1890+
self._apply_grants(snapshot.model, target_table_name, GrantsTargetLayer.PHYSICAL)
1891+
18881892
def delete(self, name: str, **kwargs: t.Any) -> None:
18891893
_check_table_db_is_physical_schema(name, kwargs["physical_schema"])
18901894
self.adapter.drop_table(name, cascade=kwargs.pop("cascade", False))
@@ -1932,6 +1936,10 @@ def _replace_query_for_model(
19321936
source_columns=source_columns,
19331937
)
19341938

1939+
# Apply grants after table replacement (unless explicitly skipped by caller)
1940+
if not kwargs.get("skip_grants", False):
1941+
self._apply_grants(model, name, GrantsTargetLayer.PHYSICAL)
1942+
19351943
def _get_target_and_source_columns(
19361944
self,
19371945
model: Model,
@@ -2192,12 +2200,18 @@ def create(
21922200
)
21932201
return
21942202

2195-
super().create(table_name, model, is_table_deployable, render_kwargs, **kwargs)
2203+
# Skip grants in parent create call since we'll apply them after data insertion
2204+
kwargs_no_grants = {**kwargs}
2205+
kwargs_no_grants["skip_grants"] = True
2206+
2207+
super().create(table_name, model, is_table_deployable, render_kwargs, **kwargs_no_grants)
21962208
# For seeds we insert data at the time of table creation.
21972209
try:
21982210
for index, df in enumerate(model.render_seed()):
21992211
if index == 0:
2200-
self._replace_query_for_model(model, table_name, df, render_kwargs, **kwargs)
2212+
self._replace_query_for_model(
2213+
model, table_name, df, render_kwargs, **kwargs_no_grants
2214+
)
22012215
else:
22022216
self.adapter.insert_append(
22032217
table_name, df, target_columns_to_types=model.columns_to_types
@@ -2206,6 +2220,9 @@ def create(
22062220
self.adapter.drop_table(table_name)
22072221
raise
22082222

2223+
# Apply grants after seed table creation or data insertion
2224+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
2225+
22092226
def insert(
22102227
self,
22112228
table_name: str,
@@ -2269,6 +2286,9 @@ def create(
22692286
**kwargs,
22702287
)
22712288

2289+
# Apply grants after SCD Type 2 table creation
2290+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
2291+
22722292
def insert(
22732293
self,
22742294
table_name: str,
@@ -2392,7 +2412,7 @@ def insert(
23922412
column_descriptions=model.column_descriptions,
23932413
)
23942414

2395-
# Apply grants after view creation/replacement
2415+
# Apply grants after view creation / replacement
23962416
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
23972417

23982418
def append(
@@ -2417,6 +2437,8 @@ def create(
24172437
# Make sure we don't recreate the view to prevent deletion of downstream views in engines with no late
24182438
# binding support (because of DROP CASCADE).
24192439
logger.info("View '%s' already exists", table_name)
2440+
# Always apply grants when present, even if view exists, to handle grants updates
2441+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
24202442
return
24212443

24222444
logger.info("Creating view '%s'", table_name)
@@ -2440,6 +2462,9 @@ def create(
24402462
column_descriptions=model.column_descriptions if is_table_deployable else None,
24412463
)
24422464

2465+
# Apply grants after view creation
2466+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
2467+
24432468
def migrate(
24442469
self,
24452470
target_table_name: str,
@@ -2620,7 +2645,7 @@ def create(
26202645
is_snapshot_deployable: bool = kwargs["is_snapshot_deployable"]
26212646

26222647
if is_table_deployable and is_snapshot_deployable:
2623-
# We could deploy this to prod; create a proper managed table
2648+
# We cloud deploy this to prod; create a proper managed table
26242649
logger.info("Creating managed table: %s", table_name)
26252650
self.adapter.create_managed_table(
26262651
table_name=table_name,
@@ -2636,17 +2661,23 @@ def create(
26362661

26372662
# Apply grants after managed table creation
26382663
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
2664+
26392665
elif not is_table_deployable:
26402666
# Only create the dev preview table as a normal table.
26412667
# For the main table, if the snapshot is cant be deployed to prod (eg upstream is forward-only) do nothing.
26422668
# Any downstream models that reference it will be updated to point to the dev preview table.
26432669
# If the user eventually tries to deploy it, the logic in insert() will see it doesnt exist and create it
2670+
2671+
# Create preview table but don't apply grants here since the table is not deployable
2672+
# Grants will be applied later when the table becomes deployable
2673+
kwargs_no_grants = {**kwargs}
2674+
kwargs_no_grants["skip_grants"] = True
26442675
super().create(
26452676
table_name=table_name,
26462677
model=model,
26472678
is_table_deployable=is_table_deployable,
26482679
render_kwargs=render_kwargs,
2649-
**kwargs,
2680+
**kwargs_no_grants,
26502681
)
26512682

26522683
def insert(
@@ -2661,7 +2692,6 @@ def insert(
26612692
deployability_index: DeployabilityIndex = kwargs["deployability_index"]
26622693
snapshot: Snapshot = kwargs["snapshot"]
26632694
is_snapshot_deployable = deployability_index.is_deployable(snapshot)
2664-
26652695
if is_first_insert and is_snapshot_deployable and not self.adapter.table_exists(table_name):
26662696
self.adapter.create_managed_table(
26672697
table_name=table_name,
@@ -2674,9 +2704,6 @@ def insert(
26742704
column_descriptions=model.column_descriptions,
26752705
table_format=model.table_format,
26762706
)
2677-
2678-
# Apply grants after managed table creation during first insert
2679-
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
26802707
elif not is_snapshot_deployable:
26812708
# Snapshot isnt deployable; update the preview table instead
26822709
# If the snapshot was deployable, then data would have already been loaded in create() because a managed table would have been created
@@ -2725,6 +2752,10 @@ def migrate(
27252752
f"The schema of the managed model '{target_table_name}' cannot be updated in a forward-only fashion."
27262753
)
27272754

2755+
# Apply grants after verifying no schema changes
2756+
# This ensures metadata-only changes (grants) are applied
2757+
self._apply_grants(snapshot.model, target_table_name, GrantsTargetLayer.PHYSICAL)
2758+
27282759
def delete(self, name: str, **kwargs: t.Any) -> None:
27292760
# a dev preview table is created as a normal table, so it needs to be dropped as a normal table
27302761
_check_table_db_is_physical_schema(name, kwargs["physical_schema"])

0 commit comments

Comments
 (0)