feat(experimental): add official support for grants

Author: newtonapple

.gitignore

@@ -138,6 +138,12 @@ dmypy.json
 *~
 *#
 
+# Vim
+*.swp
+*.swo
+.null-ls*
+
+
 *.duckdb
 *.duckdb.wal
 
@@ -158,3 +164,4 @@ spark-warehouse/
 
 # claude
 .claude/
+

pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
     "requests",
     "rich[jupyter]",
     "ruamel.yaml",
-    "sqlglot[rs]~=27.9.0",
+    "sqlglot[rs]~=27.10.0",
     "tenacity",
     "time-machine",
     "json-stream"

sqlmesh/core/_typing.py

@@ -11,6 +11,7 @@
     SessionProperties = t.Dict[str, t.Union[exp.Expression, str, int, float, bool]]
     CustomMaterializationProperties = t.Dict[str, t.Union[exp.Expression, str, int, float, bool]]
 
+
 if sys.version_info >= (3, 11):
     from typing import Self as Self
 else:

sqlmesh/core/engine_adapter/_typing.py

@@ -30,3 +30,4 @@
     ]
 
     QueryOrDF = t.Union[Query, DF]
+    GrantsConfig = t.Dict[str, t.List[str]]

sqlmesh/core/engine_adapter/base.py

@@ -63,6 +63,7 @@
     from sqlmesh.core.engine_adapter._typing import (
         DF,
         BigframeSession,
+        GrantsConfig,
         PySparkDataFrame,
         PySparkSession,
         Query,
@@ -79,6 +80,9 @@
 KEY_FOR_CREATABLE_TYPE = "CREATABLE_TYPE"
 
 
+# Use existing DataObjectType from shared module for grants
+
+
 @set_catalog()
 class EngineAdapter:
     """Base class wrapping a Database API compliant connection.
@@ -114,6 +118,7 @@ class EngineAdapter:
     SUPPORTS_TUPLE_IN = True
     HAS_VIEW_BINDING = False
     SUPPORTS_REPLACE_TABLE = True
+    SUPPORTS_GRANTS = False
     DEFAULT_CATALOG_TYPE = DIALECT
     QUOTE_IDENTIFIERS_IN_VIEWS = True
     MAX_IDENTIFIER_LENGTH: t.Optional[int] = None
@@ -2345,6 +2350,193 @@ def wap_publish(self, table_name: TableName, wap_id: str) -> None:
         """
         raise NotImplementedError(f"Engine does not support WAP: {type(self)}")
 
+    def _get_current_grants_config(self, table: exp.Table) -> GrantsConfig:
+        """Returns current grants for a table as a dictionary.
+
+        This method queries the database and returns the current grants/permissions
+        for the given table, parsed into a dictionary format. The it handles
+        case-insensitive comparison between these current grants and the desired
+        grants from model configuration.
+
+        Args:
+            table: The table/view to query grants for.
+
+        Returns:
+            Dictionary mapping permissions to lists of grantees. Permission names
+            should be returned as the database provides them (typically uppercase
+            for standard SQL permissions, but engine-specific roles may vary).
+
+        Raises:
+            NotImplementedError: If the engine does not support grants.
+        """
+        if not self.SUPPORTS_GRANTS:
+            raise NotImplementedError(f"Engine does not support grants: {type(self)}")
+        raise NotImplementedError("Subclass must implement get_current_grants")
+
+    def _sync_grants_config(
+        self,
+        table: exp.Table,
+        grants_config: GrantsConfig,
+        table_type: DataObjectType = DataObjectType.TABLE,
+    ) -> None:
+        """Applies the grants_config to a table authoritatively.
+        It first compares the specified grants against the current grants, and then
+        applies the diffs to the table by revoking and granting privileges as needed.
+
+        Args:
+            table: The table/view to apply grants to.
+            grants_config: Dictionary mapping privileges to lists of grantees.
+            table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW).
+        """
+        if not self.SUPPORTS_GRANTS:
+            raise NotImplementedError(f"Engine does not support grants: {type(self)}")
+
+        current_grants = self._get_current_grants_config(table)
+        new_grants, revoked_grants = self._diff_grants_configs(grants_config, current_grants)
+        revoke_exprs = self._revoke_grants_config_expr(table, revoked_grants, table_type)
+        grant_exprs = self._apply_grants_config_expr(table, new_grants, table_type)
+        dcl_exprs = revoke_exprs + grant_exprs
+
+        if dcl_exprs:
+            self.execute(dcl_exprs)
+
+    def _apply_grants_config(
+        self,
+        table: exp.Table,
+        grants_config: GrantsConfig,
+        table_type: DataObjectType = DataObjectType.TABLE,
+    ) -> None:
+        """Applies grants to a table.
+
+        Args:
+            table: The table/view to grant permissions on.
+            grants_config: Dictionary mapping privileges to lists of grantees.
+            table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW).
+
+        Raises:
+            NotImplementedError: If the engine does not support grants.
+        """
+
+        if grants := self._apply_grants_config_expr(table, grants_config, table_type):
+            self.execute(grants)
+
+    def _revoke_grants_config(
+        self,
+        table: exp.Table,
+        grants_config: GrantsConfig,
+        table_type: DataObjectType = DataObjectType.TABLE,
+    ) -> None:
+        """Revokes grants from a table.
+
+        Args:
+            table: The table/view to revoke privileges from.
+            grants_config: Dictionary mapping privileges to lists of grantees.
+            table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW).
+
+        Raises:
+            NotImplementedError: If the engine does not support grants.
+        """
+        if revokes := self._revoke_grants_config_expr(table, grants_config, table_type):
+            self.execute(revokes)
+
+    def _apply_grants_config_expr(
+        self,
+        table: exp.Table,
+        grant_config: GrantsConfig,
+        table_type: DataObjectType = DataObjectType.TABLE,
+    ) -> t.List[exp.Grant]:
+        """Returns SQLGlot Grant expressions to apply grants to a table.
+
+        Args:
+            table: The table/view to grant permissions on.
+            grant_config: Dictionary mapping permissions to lists of grantees.
+            table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW).
+
+        Returns:
+            List of SQLGlot Grant expressions.
+
+        Raises:
+            NotImplementedError: If the engine does not support grants.
+        """
+        if not self.SUPPORTS_GRANTS:
+            raise NotImplementedError(f"Engine does not support grants: {type(self)}")
+        raise NotImplementedError("Subclass must implement _apply_grants_config_expr")
+
+    def _revoke_grants_config_expr(
+        self,
+        table: exp.Table,
+        grant_config: GrantsConfig,
+        table_type: DataObjectType = DataObjectType.TABLE,
+    ) -> t.List[exp.Expression]:
+        """Returns SQLGlot expressions to revoke grants from a table.
+
+        Note: SQLGlot doesn't yet have a Revoke expression type, so implementations
+        may return other expression types or handle revokes as strings.
+
+        Args:
+            table: The table/view to revoke permissions from.
+            grant_config: Dictionary mapping permissions to lists of grantees.
+            table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW).
+
+        Returns:
+            List of SQLGlot expressions for revoke operations.
+
+        Raises:
+            NotImplementedError: If the engine does not support grants.
+        """
+        if not self.SUPPORTS_GRANTS:
+            raise NotImplementedError(f"Engine does not support grants: {type(self)}")
+        raise NotImplementedError("Subclass must implement _revoke_grants_config_expr")
+
+    @classmethod
+    def _diff_grants_configs(
+        cls, new_config: GrantsConfig, old_config: GrantsConfig
+    ) -> t.Tuple[GrantsConfig, GrantsConfig]:
+        """Compute additions and removals between two grants configurations.
+
+        This method compares new (desired) and old (current) GrantsConfigs case-insensitively
+        for both privilege keys and grantees, while preserving original casing
+        in the output GrantsConfigs.
+
+        Args:
+            new_config: Desired grants configuration (specified by the user).
+            old_config: Current grants configuration (returned by the database).
+
+        Returns:
+            A tuple of (additions, removals) GrantsConfig where:
+            - additions contains privileges/grantees present in new_config but not in old_config
+            - additions uses keys and grantee strings from new_config (user-specified casing)
+            - removals contains privileges/grantees present in old_config but not in new_config
+            - removals uses keys and grantee strings from old_config (database-returned casing)
+
+        Notes:
+            - Comparison is case-insensitive using casefold(); original casing is preserved in results.
+            - Overlapping grantees (case-insensitive) are excluded from the results.
+        """
+
+        def _diffs(config1: GrantsConfig, config2: GrantsConfig) -> GrantsConfig:
+            diffs: GrantsConfig = {}
+            cf_config2 = {k.casefold(): {g.casefold() for g in v} for k, v in config2.items()}
+            for key, grantees in config1.items():
+                cf_key = key.casefold()
+
+                # Missing key (add all grantees)
+                if cf_key not in cf_config2:
+                    diffs[key] = grantees.copy()
+                    continue
+
+                # Include only grantees not in config2
+                cf_grantees2 = cf_config2[cf_key]
+                diff_grantees = []
+                for grantee in grantees:
+                    if grantee.casefold() not in cf_grantees2:
+                        diff_grantees.append(grantee)
+                if diff_grantees:
+                    diffs[key] = diff_grantees
+            return diffs
+
+        return _diffs(new_config, old_config), _diffs(old_config, new_config)
+
     @contextlib.contextmanager
     def transaction(
         self,

sqlmesh/core/engine_adapter/base_postgres.py

@@ -26,6 +26,7 @@ class BasePostgresEngineAdapter(EngineAdapter):
     COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY
     SUPPORTS_QUERY_EXECUTION_TRACKING = True
     SUPPORTED_DROP_CASCADE_OBJECT_KINDS = ["SCHEMA", "TABLE", "VIEW"]
+    CURRENT_SCHEMA_EXPRESSION = exp.func("current_schema")
 
     def columns(
         self, table_name: TableName, include_pseudo_columns: bool = False
@@ -58,6 +59,7 @@ def columns(
             raise SQLMeshError(
                 f"Could not get columns for table '{table.sql(dialect=self.dialect)}'. Table not found."
             )
+
         return {
             column_name: exp.DataType.build(data_type, dialect=self.dialect, udt=True)
             for column_name, data_type in resp
@@ -188,3 +190,10 @@ def _get_data_objects(
             )
             for row in df.itertuples()
         ]
+
+    def get_current_schema(self) -> str:
+        """Returns the current default schema for the connection."""
+        result = self.fetchone(exp.select(self.CURRENT_SCHEMA_EXPRESSION))
+        if result and result[0]:
+            return result[0]
+        return "public"

sqlmesh/core/engine_adapter/postgres.py

@@ -6,6 +6,7 @@
 from functools import cached_property, partial
 from sqlglot import exp
 
+from sqlmesh.core.engine_adapter.shared import DataObjectType
 from sqlmesh.core.engine_adapter.base_postgres import BasePostgresEngineAdapter
 from sqlmesh.core.engine_adapter.mixins import (
     GetCurrentCatalogFromFunctionMixin,
@@ -17,7 +18,9 @@
 
 if t.TYPE_CHECKING:
     from sqlmesh.core._typing import TableName
-    from sqlmesh.core.engine_adapter._typing import DF, QueryOrDF
+    from sqlmesh.core.engine_adapter._typing import DF, GrantsConfig, QueryOrDF
+
+    DCL = t.TypeVar("DCL", exp.Grant, exp.Revoke)
 
 logger = logging.getLogger(__name__)
 
@@ -30,6 +33,7 @@ class PostgresEngineAdapter(
     RowDiffMixin,
 ):
     DIALECT = "postgres"
+    SUPPORTS_GRANTS = True
     SUPPORTS_INDEXES = True
     HAS_VIEW_BINDING = True
     CURRENT_CATALOG_EXPRESSION = exp.column("current_catalog")
@@ -135,3 +139,80 @@ def server_version(self) -> t.Tuple[int, int]:
             if match:
                 return int(match.group(1)), int(match.group(2))
         return 0, 0
+
+    def _dcl_grants_config_expr(
+        self,
+        dcl_cmd: t.Type[DCL],
+        relation: exp.Expression,
+        grant_config: GrantsConfig,
+    ) -> t.Union[t.List[exp.Grant], t.List[exp.Revoke]]:
+        expressions = []
+        for privilege, principals in grant_config.items():
+            if not principals:
+                continue
+
+            grant = dcl_cmd(
+                privileges=[exp.GrantPrivilege(this=exp.Var(this=privilege))],
+                securable=relation,
+                principals=principals,  # use original strings so user can to choose quote or not
+            )
+            expressions.append(grant)
+
+        return expressions
+
+    def _apply_grants_config_expr(
+        self,
+        table: exp.Table,
+        grant_config: GrantsConfig,
+        table_type: DataObjectType = DataObjectType.TABLE,
+    ) -> t.List[exp.Grant]:
+        # https://www.postgresql.org/docs/current/sql-grant.html
+        return t.cast(
+            t.List[exp.Grant],
+            self._dcl_grants_config_expr(exp.Grant, table, grant_config),
+        )
+
+    def _revoke_grants_config_expr(
+        self,
+        table: exp.Table,
+        grant_config: GrantsConfig,
+        table_type: DataObjectType = DataObjectType.TABLE,
+    ) -> t.List[exp.Expression]:
+        # https://www.postgresql.org/docs/current/sql-revoke.html
+        return t.cast(
+            t.List[exp.Expression],
+            self._dcl_grants_config_expr(exp.Revoke, table, grant_config),
+        )
+
+    def _get_current_grants_config(self, table: exp.Table) -> GrantsConfig:
+        """Returns current grants for a Postgres table as a dictionary."""
+        table_schema = table.db or self.get_current_schema()
+        table_name = table.name
+
+        # https://www.postgresql.org/docs/current/infoschema-role-table-grants.html
+        grant_expr = (
+            exp.select("privilege_type", "grantee")
+            .from_(exp.table_("role_table_grants", db="information_schema"))
+            .where(
+                exp.and_(
+                    exp.column("table_schema").eq(exp.Literal.string(table_schema)),
+                    exp.column("table_name").eq(exp.Literal.string(table_name)),
+                    exp.column("grantor").eq(exp.column("current_role")),
+                    exp.column("grantee").neq(exp.column("current_role")),
+                )
+            )
+        )
+        results = self.fetchall(grant_expr)
+
+        grants_dict: t.Dict[str, t.List[str]] = {}
+        for row in results:
+            privilege = str(row[0])
+            grantee = str(row[1])
+
+            if privilege not in grants_dict:
+                grants_dict[privilege] = []
+
+            if grantee not in grants_dict[privilege]:
+                grants_dict[privilege].append(grantee)
+
+        return grants_dict