Skip to content

Commit 6198359

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 9cdabed commit 6198359

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
@@ -1056,6 +1056,7 @@ def _clone_snapshot_in_dev(
10561056
allow_additive_snapshots=allow_additive_snapshots,
10571057
run_pre_post_statements=run_pre_post_statements,
10581058
)
1059+
10591060
except Exception:
10601061
adapter.drop_table(target_table_name)
10611062
raise
@@ -1156,6 +1157,7 @@ def _migrate_target_table(
11561157
rendered_physical_properties=rendered_physical_properties,
11571158
dry_run=False,
11581159
run_pre_post_statements=run_pre_post_statements,
1160+
skip_grants=True, # skip grants for tmp table
11591161
)
11601162
try:
11611163
evaluation_strategy = _evaluation_strategy(snapshot, adapter)
@@ -1421,6 +1423,7 @@ def _execute_create(
14211423
rendered_physical_properties: t.Dict[str, exp.Expression],
14221424
dry_run: bool,
14231425
run_pre_post_statements: bool = True,
1426+
skip_grants: bool = False,
14241427
) -> None:
14251428
adapter = self.get_adapter(snapshot.model.gateway)
14261429
evaluation_strategy = _evaluation_strategy(snapshot, adapter)
@@ -1444,12 +1447,11 @@ def _execute_create(
14441447
is_snapshot_representative=is_snapshot_representative,
14451448
dry_run=dry_run,
14461449
physical_properties=rendered_physical_properties,
1450+
skip_grants=skip_grants,
14471451
)
14481452
if run_pre_post_statements:
14491453
adapter.execute(snapshot.model.render_post_statements(**create_render_kwargs))
14501454

1451-
evaluation_strategy._apply_grants(snapshot.model, table_name, GrantsTargetLayer.PHYSICAL)
1452-
14531455
def _can_clone(self, snapshot: Snapshot, deployability_index: DeployabilityIndex) -> bool:
14541456
adapter = self.get_adapter(snapshot.model.gateway)
14551457
return (
@@ -1718,16 +1720,11 @@ def _apply_grants(
17181720
return
17191721

17201722
logger.info(f"Applying grants for model {model.name} to table {table_name}")
1721-
1722-
try:
1723-
self.adapter.sync_grants_config(
1724-
exp.to_table(table_name, dialect=model.dialect),
1725-
grants_config,
1726-
model.grants_table_type,
1727-
)
1728-
except Exception:
1729-
# Log error but don't fail evaluation if grants fail
1730-
logger.error(f"Failed to apply grants for model {model.name}", exc_info=True)
1723+
self.adapter.sync_grants_config(
1724+
exp.to_table(table_name, dialect=self.adapter.dialect),
1725+
grants_config,
1726+
model.grants_table_type,
1727+
)
17311728

17321729

17331730
class SymbolicStrategy(EvaluationStrategy):
@@ -1892,6 +1889,10 @@ def create(
18921889
column_descriptions=model.column_descriptions if is_table_deployable else None,
18931890
)
18941891

1892+
# Apply grants after table creation (unless explicitly skipped by caller)
1893+
if not kwargs.get("skip_grants", False):
1894+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
1895+
18951896
def migrate(
18961897
self,
18971898
target_table_name: str,
@@ -1917,6 +1918,9 @@ def migrate(
19171918
)
19181919
self.adapter.alter_table(alter_operations)
19191920

1921+
# Apply grants after schema migration
1922+
self._apply_grants(snapshot.model, target_table_name, GrantsTargetLayer.PHYSICAL)
1923+
19201924
def delete(self, name: str, **kwargs: t.Any) -> None:
19211925
_check_table_db_is_physical_schema(name, kwargs["physical_schema"])
19221926
self.adapter.drop_table(name, cascade=kwargs.pop("cascade", False))
@@ -1964,6 +1968,10 @@ def _replace_query_for_model(
19641968
source_columns=source_columns,
19651969
)
19661970

1971+
# Apply grants after table replacement (unless explicitly skipped by caller)
1972+
if not kwargs.get("skip_grants", False):
1973+
self._apply_grants(model, name, GrantsTargetLayer.PHYSICAL)
1974+
19671975
def _get_target_and_source_columns(
19681976
self,
19691977
model: Model,
@@ -2224,12 +2232,18 @@ def create(
22242232
)
22252233
return
22262234

2227-
super().create(table_name, model, is_table_deployable, render_kwargs, **kwargs)
2235+
# Skip grants in parent create call since we'll apply them after data insertion
2236+
kwargs_no_grants = {**kwargs}
2237+
kwargs_no_grants["skip_grants"] = True
2238+
2239+
super().create(table_name, model, is_table_deployable, render_kwargs, **kwargs_no_grants)
22282240
# For seeds we insert data at the time of table creation.
22292241
try:
22302242
for index, df in enumerate(model.render_seed()):
22312243
if index == 0:
2232-
self._replace_query_for_model(model, table_name, df, render_kwargs, **kwargs)
2244+
self._replace_query_for_model(
2245+
model, table_name, df, render_kwargs, **kwargs_no_grants
2246+
)
22332247
else:
22342248
self.adapter.insert_append(
22352249
table_name, df, target_columns_to_types=model.columns_to_types
@@ -2238,6 +2252,9 @@ def create(
22382252
self.adapter.drop_table(table_name)
22392253
raise
22402254

2255+
# Apply grants after seed table creation or data insertion
2256+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
2257+
22412258
def insert(
22422259
self,
22432260
table_name: str,
@@ -2301,6 +2318,9 @@ def create(
23012318
**kwargs,
23022319
)
23032320

2321+
# Apply grants after SCD Type 2 table creation
2322+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
2323+
23042324
def insert(
23052325
self,
23062326
table_name: str,
@@ -2424,7 +2444,7 @@ def insert(
24242444
column_descriptions=model.column_descriptions,
24252445
)
24262446

2427-
# Apply grants after view creation/replacement
2447+
# Apply grants after view creation / replacement
24282448
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
24292449

24302450
def append(
@@ -2449,6 +2469,8 @@ def create(
24492469
# Make sure we don't recreate the view to prevent deletion of downstream views in engines with no late
24502470
# binding support (because of DROP CASCADE).
24512471
logger.info("View '%s' already exists", table_name)
2472+
# Always apply grants when present, even if view exists, to handle grants updates
2473+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
24522474
return
24532475

24542476
logger.info("Creating view '%s'", table_name)
@@ -2472,6 +2494,9 @@ def create(
24722494
column_descriptions=model.column_descriptions if is_table_deployable else None,
24732495
)
24742496

2497+
# Apply grants after view creation
2498+
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
2499+
24752500
def migrate(
24762501
self,
24772502
target_table_name: str,
@@ -2652,7 +2677,7 @@ def create(
26522677
is_snapshot_deployable: bool = kwargs["is_snapshot_deployable"]
26532678

26542679
if is_table_deployable and is_snapshot_deployable:
2655-
# We could deploy this to prod; create a proper managed table
2680+
# We cloud deploy this to prod; create a proper managed table
26562681
logger.info("Creating managed table: %s", table_name)
26572682
self.adapter.create_managed_table(
26582683
table_name=table_name,
@@ -2668,17 +2693,23 @@ def create(
26682693

26692694
# Apply grants after managed table creation
26702695
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
2696+
26712697
elif not is_table_deployable:
26722698
# Only create the dev preview table as a normal table.
26732699
# For the main table, if the snapshot is cant be deployed to prod (eg upstream is forward-only) do nothing.
26742700
# Any downstream models that reference it will be updated to point to the dev preview table.
26752701
# If the user eventually tries to deploy it, the logic in insert() will see it doesnt exist and create it
2702+
2703+
# Create preview table but don't apply grants here since the table is not deployable
2704+
# Grants will be applied later when the table becomes deployable
2705+
kwargs_no_grants = {**kwargs}
2706+
kwargs_no_grants["skip_grants"] = True
26762707
super().create(
26772708
table_name=table_name,
26782709
model=model,
26792710
is_table_deployable=is_table_deployable,
26802711
render_kwargs=render_kwargs,
2681-
**kwargs,
2712+
**kwargs_no_grants,
26822713
)
26832714

26842715
def insert(
@@ -2693,7 +2724,6 @@ def insert(
26932724
deployability_index: DeployabilityIndex = kwargs["deployability_index"]
26942725
snapshot: Snapshot = kwargs["snapshot"]
26952726
is_snapshot_deployable = deployability_index.is_deployable(snapshot)
2696-
26972727
if is_first_insert and is_snapshot_deployable and not self.adapter.table_exists(table_name):
26982728
self.adapter.create_managed_table(
26992729
table_name=table_name,
@@ -2706,9 +2736,6 @@ def insert(
27062736
column_descriptions=model.column_descriptions,
27072737
table_format=model.table_format,
27082738
)
2709-
2710-
# Apply grants after managed table creation during first insert
2711-
self._apply_grants(model, table_name, GrantsTargetLayer.PHYSICAL)
27122739
elif not is_snapshot_deployable:
27132740
# Snapshot isnt deployable; update the preview table instead
27142741
# If the snapshot was deployable, then data would have already been loaded in create() because a managed table would have been created
@@ -2757,6 +2784,10 @@ def migrate(
27572784
f"The schema of the managed model '{target_table_name}' cannot be updated in a forward-only fashion."
27582785
)
27592786

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

0 commit comments

Comments
 (0)