Skip to content

Commit 4345d31

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 a84d809 commit 4345d31

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
@@ -1122,6 +1123,7 @@ def _migrate_target_table(
11221123
rendered_physical_properties=rendered_physical_properties,
11231124
dry_run=False,
11241125
run_pre_post_statements=run_pre_post_statements,
1126+
skip_grants=True, # skip grants for tmp table
11251127
)
11261128
try:
11271129
evaluation_strategy = _evaluation_strategy(snapshot, adapter)
@@ -1387,6 +1389,7 @@ def _execute_create(
13871389
rendered_physical_properties: t.Dict[str, exp.Expression],
13881390
dry_run: bool,
13891391
run_pre_post_statements: bool = True,
1392+
skip_grants: bool = False,
13901393
) -> None:
13911394
adapter = self.get_adapter(snapshot.model.gateway)
13921395
evaluation_strategy = _evaluation_strategy(snapshot, adapter)
@@ -1410,12 +1413,11 @@ def _execute_create(
14101413
is_snapshot_representative=is_snapshot_representative,
14111414
dry_run=dry_run,
14121415
physical_properties=rendered_physical_properties,
1416+
skip_grants=skip_grants,
14131417
)
14141418
if run_pre_post_statements:
14151419
adapter.execute(snapshot.model.render_post_statements(**create_render_kwargs))
14161420

1417-
evaluation_strategy._apply_grants(snapshot.model, table_name, GrantsTargetLayer.PHYSICAL)
1418-
14191421
def _can_clone(self, snapshot: Snapshot, deployability_index: DeployabilityIndex) -> bool:
14201422
adapter = self.get_adapter(snapshot.model.gateway)
14211423
return (
@@ -1683,16 +1685,11 @@ def _apply_grants(
16831685
return
16841686

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

16971694

16981695
class SymbolicStrategy(EvaluationStrategy):
@@ -1857,6 +1854,10 @@ def create(
18571854
column_descriptions=model.column_descriptions if is_table_deployable else None,
18581855
)
18591856

1857+
# Apply grants after table creation (unless explicitly skipped by caller)
1858+
if not kwargs.get("skip_grants", False):
1859+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
1860+
18601861
def migrate(
18611862
self,
18621863
target_table_name: str,
@@ -1882,6 +1883,9 @@ def migrate(
18821883
)
18831884
self.adapter.alter_table(alter_operations)
18841885

1886+
# Apply grants after schema migration
1887+
self._apply_grants(snapshot.model, target_table_name, GrantsTargetLayer.PHYSICAL)
1888+
18851889
def delete(self, name: str, **kwargs: t.Any) -> None:
18861890
_check_table_db_is_physical_schema(name, kwargs["physical_schema"])
18871891
self.adapter.drop_table(name, cascade=kwargs.pop("cascade", False))
@@ -1929,6 +1933,10 @@ def _replace_query_for_model(
19291933
source_columns=source_columns,
19301934
)
19311935

1936+
# Apply grants after table replacement (unless explicitly skipped by caller)
1937+
if not kwargs.get("skip_grants", False):
1938+
self._apply_grants(model, name, GrantsTargetLayer.PHYSICAL)
1939+
19321940
def _get_target_and_source_columns(
19331941
self,
19341942
model: Model,
@@ -2189,12 +2197,18 @@ def create(
21892197
)
21902198
return
21912199

2192-
super().create(table_name, model, is_table_deployable, render_kwargs, **kwargs)
2200+
# Skip grants in parent create call since we'll apply them after data insertion
2201+
kwargs_no_grants = {**kwargs}
2202+
kwargs_no_grants["skip_grants"] = True
2203+
2204+
super().create(table_name, model, is_table_deployable, render_kwargs, **kwargs_no_grants)
21932205
# For seeds we insert data at the time of table creation.
21942206
try:
21952207
for index, df in enumerate(model.render_seed()):
21962208
if index == 0:
2197-
self._replace_query_for_model(model, table_name, df, render_kwargs, **kwargs)
2209+
self._replace_query_for_model(
2210+
model, table_name, df, render_kwargs, **kwargs_no_grants
2211+
)
21982212
else:
21992213
self.adapter.insert_append(
22002214
table_name, df, target_columns_to_types=model.columns_to_types
@@ -2203,6 +2217,9 @@ def create(
22032217
self.adapter.drop_table(table_name)
22042218
raise
22052219

2220+
# Apply grants after seed table creation or data insertion
2221+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
2222+
22062223
def insert(
22072224
self,
22082225
table_name: str,
@@ -2266,6 +2283,9 @@ def create(
22662283
**kwargs,
22672284
)
22682285

2286+
# Apply grants after SCD Type 2 table creation
2287+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
2288+
22692289
def insert(
22702290
self,
22712291
table_name: str,
@@ -2389,7 +2409,7 @@ def insert(
23892409
column_descriptions=model.column_descriptions,
23902410
)
23912411

2392-
# Apply grants after view creation/replacement
2412+
# Apply grants after view creation / replacement
23932413
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
23942414

23952415
def append(
@@ -2414,6 +2434,8 @@ def create(
24142434
# Make sure we don't recreate the view to prevent deletion of downstream views in engines with no late
24152435
# binding support (because of DROP CASCADE).
24162436
logger.info("View '%s' already exists", table_name)
2437+
# Always apply grants when present, even if view exists, to handle grants updates
2438+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
24172439
return
24182440

24192441
logger.info("Creating view '%s'", table_name)
@@ -2437,6 +2459,9 @@ def create(
24372459
column_descriptions=model.column_descriptions if is_table_deployable else None,
24382460
)
24392461

2462+
# Apply grants after view creation
2463+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
2464+
24402465
def migrate(
24412466
self,
24422467
target_table_name: str,
@@ -2617,7 +2642,7 @@ def create(
26172642
is_snapshot_deployable: bool = kwargs["is_snapshot_deployable"]
26182643

26192644
if is_table_deployable and is_snapshot_deployable:
2620-
# We could deploy this to prod; create a proper managed table
2645+
# We cloud deploy this to prod; create a proper managed table
26212646
logger.info("Creating managed table: %s", table_name)
26222647
self.adapter.create_managed_table(
26232648
table_name=table_name,
@@ -2633,17 +2658,23 @@ def create(
26332658

26342659
# Apply grants after managed table creation
26352660
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
2661+
26362662
elif not is_table_deployable:
26372663
# Only create the dev preview table as a normal table.
26382664
# For the main table, if the snapshot is cant be deployed to prod (eg upstream is forward-only) do nothing.
26392665
# Any downstream models that reference it will be updated to point to the dev preview table.
26402666
# If the user eventually tries to deploy it, the logic in insert() will see it doesnt exist and create it
2667+
2668+
# Create preview table but don't apply grants here since the table is not deployable
2669+
# Grants will be applied later when the table becomes deployable
2670+
kwargs_no_grants = {**kwargs}
2671+
kwargs_no_grants["skip_grants"] = True
26412672
super().create(
26422673
table_name=table_name,
26432674
model=model,
26442675
is_table_deployable=is_table_deployable,
26452676
render_kwargs=render_kwargs,
2646-
**kwargs,
2677+
**kwargs_no_grants,
26472678
)
26482679

26492680
def insert(
@@ -2658,7 +2689,6 @@ def insert(
26582689
deployability_index: DeployabilityIndex = kwargs["deployability_index"]
26592690
snapshot: Snapshot = kwargs["snapshot"]
26602691
is_snapshot_deployable = deployability_index.is_deployable(snapshot)
2661-
26622692
if is_first_insert and is_snapshot_deployable and not self.adapter.table_exists(table_name):
26632693
self.adapter.create_managed_table(
26642694
table_name=table_name,
@@ -2671,9 +2701,6 @@ def insert(
26712701
column_descriptions=model.column_descriptions,
26722702
table_format=model.table_format,
26732703
)
2674-
2675-
# Apply grants after managed table creation during first insert
2676-
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
26772704
elif not is_snapshot_deployable:
26782705
# Snapshot isnt deployable; update the preview table instead
26792706
# If the snapshot was deployable, then data would have already been loaded in create() because a managed table would have been created
@@ -2722,6 +2749,10 @@ def migrate(
27222749
f"The schema of the managed model '{target_table_name}' cannot be updated in a forward-only fashion."
27232750
)
27242751

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

0 commit comments

Comments
 (0)