Skip to content

Commit 151bbb5

Browse files
committed
feat(experimental): add official support for grants
1 parent dbc7de6 commit 151bbb5

18 files changed

Lines changed: 1664 additions & 4 deletions

.gitignore

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,12 @@ dmypy.json
138138
*~
139139
*#
140140

141+
# Vim
142+
*.swp
143+
*.swo
144+
.null-ls*
145+
146+
141147
*.duckdb
142148
*.duckdb.wal
143149

@@ -158,3 +164,4 @@ spark-warehouse/
158164

159165
# claude
160166
.claude/
167+

sqlmesh/core/_typing.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
SessionProperties = t.Dict[str, t.Union[exp.Expression, str, int, float, bool]]
1212
CustomMaterializationProperties = t.Dict[str, t.Union[exp.Expression, str, int, float, bool]]
1313

14+
1415
if sys.version_info >= (3, 11):
1516
from typing import Self as Self
1617
else:

sqlmesh/core/engine_adapter/_typing.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,3 +30,4 @@
3030
]
3131

3232
QueryOrDF = t.Union[Query, DF]
33+
GrantsConfig = t.Dict[str, t.List[str]]

sqlmesh/core/engine_adapter/base.py

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@
6363
from sqlmesh.core.engine_adapter._typing import (
6464
DF,
6565
BigframeSession,
66+
GrantsConfig,
6667
PySparkDataFrame,
6768
PySparkSession,
6869
Query,
@@ -79,6 +80,9 @@
7980
KEY_FOR_CREATABLE_TYPE = "CREATABLE_TYPE"
8081

8182

83+
# Use existing DataObjectType from shared module for grants
84+
85+
8286
@set_catalog()
8387
class EngineAdapter:
8488
"""Base class wrapping a Database API compliant connection.
@@ -114,6 +118,7 @@ class EngineAdapter:
114118
SUPPORTS_TUPLE_IN = True
115119
HAS_VIEW_BINDING = False
116120
SUPPORTS_REPLACE_TABLE = True
121+
SUPPORTS_GRANTS = False
117122
DEFAULT_CATALOG_TYPE = DIALECT
118123
QUOTE_IDENTIFIERS_IN_VIEWS = True
119124
MAX_IDENTIFIER_LENGTH: t.Optional[int] = None
@@ -2391,6 +2396,193 @@ def wap_publish(self, table_name: TableName, wap_id: str) -> None:
23912396
"""
23922397
raise NotImplementedError(f"Engine does not support WAP: {type(self)}")
23932398

2399+
def _get_current_grants_config(self, table: exp.Table) -> GrantsConfig:
2400+
"""Returns current grants for a table as a dictionary.
2401+
2402+
This method queries the database and returns the current grants/permissions
2403+
for the given table, parsed into a dictionary format. The it handles
2404+
case-insensitive comparison between these current grants and the desired
2405+
grants from model configuration.
2406+
2407+
Args:
2408+
table: The table/view to query grants for.
2409+
2410+
Returns:
2411+
Dictionary mapping permissions to lists of grantees. Permission names
2412+
should be returned as the database provides them (typically uppercase
2413+
for standard SQL permissions, but engine-specific roles may vary).
2414+
2415+
Raises:
2416+
NotImplementedError: If the engine does not support grants.
2417+
"""
2418+
if not self.SUPPORTS_GRANTS:
2419+
raise NotImplementedError(f"Engine does not support grants: {type(self)}")
2420+
raise NotImplementedError("Subclass must implement get_current_grants")
2421+
2422+
def _sync_grants_config(
2423+
self,
2424+
table: exp.Table,
2425+
grants_config: GrantsConfig,
2426+
table_type: DataObjectType = DataObjectType.TABLE,
2427+
) -> None:
2428+
"""Applies the grants_config to a table authoritatively.
2429+
It first compares the specified grants against the current grants, and then
2430+
applies the diffs to the table by revoking and granting privileges as needed.
2431+
2432+
Args:
2433+
table: The table/view to apply grants to.
2434+
grants_config: Dictionary mapping privileges to lists of grantees.
2435+
table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW).
2436+
"""
2437+
if not self.SUPPORTS_GRANTS:
2438+
raise NotImplementedError(f"Engine does not support grants: {type(self)}")
2439+
2440+
current_grants = self._get_current_grants_config(table)
2441+
new_grants, revoked_grants = self._diff_grants_configs(grants_config, current_grants)
2442+
revoke_exprs = self._revoke_grants_config_expr(table, revoked_grants, table_type)
2443+
grant_exprs = self._apply_grants_config_expr(table, new_grants, table_type)
2444+
dcl_exprs = revoke_exprs + grant_exprs
2445+
2446+
if dcl_exprs:
2447+
self.execute(dcl_exprs)
2448+
2449+
def _apply_grants_config(
2450+
self,
2451+
table: exp.Table,
2452+
grants_config: GrantsConfig,
2453+
table_type: DataObjectType = DataObjectType.TABLE,
2454+
) -> None:
2455+
"""Applies grants to a table.
2456+
2457+
Args:
2458+
table: The table/view to grant permissions on.
2459+
grants_config: Dictionary mapping privileges to lists of grantees.
2460+
table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW).
2461+
2462+
Raises:
2463+
NotImplementedError: If the engine does not support grants.
2464+
"""
2465+
2466+
if grants := self._apply_grants_config_expr(table, grants_config, table_type):
2467+
self.execute(grants)
2468+
2469+
def _revoke_grants_config(
2470+
self,
2471+
table: exp.Table,
2472+
grants_config: GrantsConfig,
2473+
table_type: DataObjectType = DataObjectType.TABLE,
2474+
) -> None:
2475+
"""Revokes grants from a table.
2476+
2477+
Args:
2478+
table: The table/view to revoke privileges from.
2479+
grants_config: Dictionary mapping privileges to lists of grantees.
2480+
table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW).
2481+
2482+
Raises:
2483+
NotImplementedError: If the engine does not support grants.
2484+
"""
2485+
if revokes := self._revoke_grants_config_expr(table, grants_config, table_type):
2486+
self.execute(revokes)
2487+
2488+
def _apply_grants_config_expr(
2489+
self,
2490+
table: exp.Table,
2491+
grant_config: GrantsConfig,
2492+
table_type: DataObjectType = DataObjectType.TABLE,
2493+
) -> t.List[exp.Grant]:
2494+
"""Returns SQLGlot Grant expressions to apply grants to a table.
2495+
2496+
Args:
2497+
table: The table/view to grant permissions on.
2498+
grant_config: Dictionary mapping permissions to lists of grantees.
2499+
table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW).
2500+
2501+
Returns:
2502+
List of SQLGlot Grant expressions.
2503+
2504+
Raises:
2505+
NotImplementedError: If the engine does not support grants.
2506+
"""
2507+
if not self.SUPPORTS_GRANTS:
2508+
raise NotImplementedError(f"Engine does not support grants: {type(self)}")
2509+
raise NotImplementedError("Subclass must implement _apply_grants_config_expr")
2510+
2511+
def _revoke_grants_config_expr(
2512+
self,
2513+
table: exp.Table,
2514+
grant_config: GrantsConfig,
2515+
table_type: DataObjectType = DataObjectType.TABLE,
2516+
) -> t.List[exp.Expression]:
2517+
"""Returns SQLGlot expressions to revoke grants from a table.
2518+
2519+
Note: SQLGlot doesn't yet have a Revoke expression type, so implementations
2520+
may return other expression types or handle revokes as strings.
2521+
2522+
Args:
2523+
table: The table/view to revoke permissions from.
2524+
grant_config: Dictionary mapping permissions to lists of grantees.
2525+
table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW).
2526+
2527+
Returns:
2528+
List of SQLGlot expressions for revoke operations.
2529+
2530+
Raises:
2531+
NotImplementedError: If the engine does not support grants.
2532+
"""
2533+
if not self.SUPPORTS_GRANTS:
2534+
raise NotImplementedError(f"Engine does not support grants: {type(self)}")
2535+
raise NotImplementedError("Subclass must implement _revoke_grants_config_expr")
2536+
2537+
@classmethod
2538+
def _diff_grants_configs(
2539+
cls, new_config: GrantsConfig, old_config: GrantsConfig
2540+
) -> t.Tuple[GrantsConfig, GrantsConfig]:
2541+
"""Compute additions and removals between two grants configurations.
2542+
2543+
This method compares new (desired) and old (current) GrantsConfigs case-insensitively
2544+
for both privilege keys and grantees, while preserving original casing
2545+
in the output GrantsConfigs.
2546+
2547+
Args:
2548+
new_config: Desired grants configuration (specified by the user).
2549+
old_config: Current grants configuration (returned by the database).
2550+
2551+
Returns:
2552+
A tuple of (additions, removals) GrantsConfig where:
2553+
- additions contains privileges/grantees present in new_config but not in old_config
2554+
- additions uses keys and grantee strings from new_config (user-specified casing)
2555+
- removals contains privileges/grantees present in old_config but not in new_config
2556+
- removals uses keys and grantee strings from old_config (database-returned casing)
2557+
2558+
Notes:
2559+
- Comparison is case-insensitive using casefold(); original casing is preserved in results.
2560+
- Overlapping grantees (case-insensitive) are excluded from the results.
2561+
"""
2562+
2563+
def _diffs(config1: GrantsConfig, config2: GrantsConfig) -> GrantsConfig:
2564+
diffs: GrantsConfig = {}
2565+
cf_config2 = {k.casefold(): {g.casefold() for g in v} for k, v in config2.items()}
2566+
for key, grantees in config1.items():
2567+
cf_key = key.casefold()
2568+
2569+
# Missing key (add all grantees)
2570+
if cf_key not in cf_config2:
2571+
diffs[key] = grantees.copy()
2572+
continue
2573+
2574+
# Include only grantees not in config2
2575+
cf_grantees2 = cf_config2[cf_key]
2576+
diff_grantees = []
2577+
for grantee in grantees:
2578+
if grantee.casefold() not in cf_grantees2:
2579+
diff_grantees.append(grantee)
2580+
if diff_grantees:
2581+
diffs[key] = diff_grantees
2582+
return diffs
2583+
2584+
return _diffs(new_config, old_config), _diffs(old_config, new_config)
2585+
23942586
@contextlib.contextmanager
23952587
def transaction(
23962588
self,

sqlmesh/core/engine_adapter/base_postgres.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ class BasePostgresEngineAdapter(EngineAdapter):
2626
COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY
2727
SUPPORTS_QUERY_EXECUTION_TRACKING = True
2828
SUPPORTED_DROP_CASCADE_OBJECT_KINDS = ["SCHEMA", "TABLE", "VIEW"]
29+
CURRENT_SCHEMA_EXPRESSION = exp.func("current_schema")
2930

3031
def columns(
3132
self, table_name: TableName, include_pseudo_columns: bool = False
@@ -58,6 +59,7 @@ def columns(
5859
raise SQLMeshError(
5960
f"Could not get columns for table '{table.sql(dialect=self.dialect)}'. Table not found."
6061
)
62+
6163
return {
6264
column_name: exp.DataType.build(data_type, dialect=self.dialect, udt=True)
6365
for column_name, data_type in resp
@@ -188,3 +190,10 @@ def _get_data_objects(
188190
)
189191
for row in df.itertuples()
190192
]
193+
194+
def get_current_schema(self) -> str:
195+
"""Returns the current default schema for the connection."""
196+
result = self.fetchone(exp.select(self.CURRENT_SCHEMA_EXPRESSION))
197+
if result and result[0]:
198+
return result[0]
199+
return "public"

sqlmesh/core/engine_adapter/postgres.py

Lines changed: 82 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from functools import cached_property, partial
77
from sqlglot import exp
88

9+
from sqlmesh.core.engine_adapter.shared import DataObjectType
910
from sqlmesh.core.engine_adapter.base_postgres import BasePostgresEngineAdapter
1011
from sqlmesh.core.engine_adapter.mixins import (
1112
GetCurrentCatalogFromFunctionMixin,
@@ -17,7 +18,9 @@
1718

1819
if t.TYPE_CHECKING:
1920
from sqlmesh.core._typing import TableName
20-
from sqlmesh.core.engine_adapter._typing import DF, QueryOrDF
21+
from sqlmesh.core.engine_adapter._typing import DF, GrantsConfig, QueryOrDF
22+
23+
DCL = t.TypeVar("DCL", exp.Grant, exp.Revoke)
2124

2225
logger = logging.getLogger(__name__)
2326

@@ -30,6 +33,7 @@ class PostgresEngineAdapter(
3033
RowDiffMixin,
3134
):
3235
DIALECT = "postgres"
36+
SUPPORTS_GRANTS = True
3337
SUPPORTS_INDEXES = True
3438
HAS_VIEW_BINDING = True
3539
CURRENT_CATALOG_EXPRESSION = exp.column("current_catalog")
@@ -135,3 +139,80 @@ def server_version(self) -> t.Tuple[int, int]:
135139
if match:
136140
return int(match.group(1)), int(match.group(2))
137141
return 0, 0
142+
143+
def _dcl_grants_config_expr(
144+
self,
145+
dcl_cmd: t.Type[DCL],
146+
relation: exp.Expression,
147+
grant_config: GrantsConfig,
148+
) -> t.Union[t.List[exp.Grant], t.List[exp.Revoke]]:
149+
expressions = []
150+
for privilege, principals in grant_config.items():
151+
if not principals:
152+
continue
153+
154+
grant = dcl_cmd(
155+
privileges=[exp.GrantPrivilege(this=exp.Var(this=privilege))],
156+
securable=relation,
157+
principals=principals, # use original strings so user can to choose quote or not
158+
)
159+
expressions.append(grant)
160+
161+
return expressions
162+
163+
def _apply_grants_config_expr(
164+
self,
165+
table: exp.Table,
166+
grant_config: GrantsConfig,
167+
table_type: DataObjectType = DataObjectType.TABLE,
168+
) -> t.List[exp.Grant]:
169+
# https://www.postgresql.org/docs/current/sql-grant.html
170+
return t.cast(
171+
t.List[exp.Grant],
172+
self._dcl_grants_config_expr(exp.Grant, table, grant_config),
173+
)
174+
175+
def _revoke_grants_config_expr(
176+
self,
177+
table: exp.Table,
178+
grant_config: GrantsConfig,
179+
table_type: DataObjectType = DataObjectType.TABLE,
180+
) -> t.List[exp.Expression]:
181+
# https://www.postgresql.org/docs/current/sql-revoke.html
182+
return t.cast(
183+
t.List[exp.Expression],
184+
self._dcl_grants_config_expr(exp.Revoke, table, grant_config),
185+
)
186+
187+
def _get_current_grants_config(self, table: exp.Table) -> GrantsConfig:
188+
"""Returns current grants for a Postgres table as a dictionary."""
189+
table_schema = table.db or self.get_current_schema()
190+
table_name = table.name
191+
192+
# https://www.postgresql.org/docs/current/infoschema-role-table-grants.html
193+
grant_expr = (
194+
exp.select("privilege_type", "grantee")
195+
.from_(exp.table_("role_table_grants", db="information_schema"))
196+
.where(
197+
exp.and_(
198+
exp.column("table_schema").eq(exp.Literal.string(table_schema)),
199+
exp.column("table_name").eq(exp.Literal.string(table_name)),
200+
exp.column("grantor").eq(exp.column("current_role")),
201+
exp.column("grantee").neq(exp.column("current_role")),
202+
)
203+
)
204+
)
205+
results = self.fetchall(grant_expr)
206+
207+
grants_dict: t.Dict[str, t.List[str]] = {}
208+
for row in results:
209+
privilege = str(row[0])
210+
grantee = str(row[1])
211+
212+
if privilege not in grants_dict:
213+
grants_dict[privilege] = []
214+
215+
if grantee not in grants_dict[privilege]:
216+
grants_dict[privilege].append(grantee)
217+
218+
return grants_dict

sqlmesh/core/model/common.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -641,6 +641,7 @@ def parse_strings_with_macro_refs(value: t.Any, dialect: DialectType) -> t.Any:
641641
"physical_properties_",
642642
"virtual_properties_",
643643
"materialization_properties_",
644+
"grants_",
644645
mode="before",
645646
check_fields=False,
646647
)(parse_properties)

0 commit comments

Comments
 (0)