Skip to content

Commit ebb1abe

Browse files
committed
fix(db2): implement grants using SYSCAT.TABAUTH instead of INFORMATION_SCHEMA
Db2 does not have INFORMATION_SCHEMA.TABLE_PRIVILEGES (SQL0204N). The correct catalog view is SYSCAT.TABAUTH, which stores per-privilege columns (SELECTAUTH, INSERTAUTH, UPDATEAUTH, DELETEAUTH, ALTERAUTH, INDEXAUTH, CONTROLAUTH) with values 'Y'/'G' (granted) or 'N'. Replace GrantsFromInfoSchemaMixin with native Db2 implementations: - _get_current_grants_config: queries SYSCAT.TABAUTH, unpivots privilege columns into a {privilege: [grantee]} dict, filters by GRANTOR = CURRENT USER - _apply_grants_config_expr: emits GRANT <priv> ON TABLE ... TO <role> - _revoke_grants_config_expr: emits REVOKE <priv> ON TABLE ... FROM <role> Also removes GrantsFromInfoSchemaMixin and CURRENT_USER_OR_ROLE_EXPRESSION from the class since they were only needed by the mixin.
1 parent ae76529 commit ebb1abe

1 file changed

Lines changed: 105 additions & 11 deletions

File tree

  • sqlmesh/core/engine_adapter

‎sqlmesh/core/engine_adapter/db2.py‎

Lines changed: 105 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,7 @@
99
from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
1010

1111
from sqlmesh.core.engine_adapter.base import EngineAdapter, _get_data_object_cache_key
12-
from sqlmesh.core.engine_adapter.mixins import (
13-
GrantsFromInfoSchemaMixin,
14-
PandasNativeFetchDFSupportMixin,
15-
)
12+
from sqlmesh.core.engine_adapter.mixins import PandasNativeFetchDFSupportMixin
1613
from sqlmesh.core.engine_adapter.shared import (
1714
CatalogSupport,
1815
CommentCreationTable,
@@ -46,16 +43,13 @@ def is_db2_error(exception: Exception, error_code: str) -> bool:
4643

4744
@set_catalog()
4845
class Db2EngineAdapter(
49-
GrantsFromInfoSchemaMixin,
5046
PandasNativeFetchDFSupportMixin,
5147
EngineAdapter,
5248
):
5349
DIALECT = "db2"
5450
SUPPORTS_INDEXES = True
5551
SUPPORTS_REPLACE_TABLE = False
5652
SUPPORTS_GRANTS = True
57-
# Db2 uses CURRENT USER special register to identify the grantor.
58-
CURRENT_USER_OR_ROLE_EXPRESSION: exp.Expr = exp.column("CURRENT USER")
5953
COMMENT_CREATION_TABLE = CommentCreationTable.COMMENT_COMMAND_ONLY
6054
COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY
6155
SUPPORTS_QUERY_EXECUTION_TRACKING = True
@@ -330,10 +324,110 @@ def _db2_type_to_sqlglot(self, db2_type: str, length: int, scale: int) -> exp.Da
330324
sqlglot_type = type_mapping.get(db2_type, f"VARCHAR({length})")
331325
return exp.DataType.build(sqlglot_type, dialect="db2")
332326

333-
@staticmethod
334-
def _grant_object_kind(table_type: DataObjectType) -> t.Optional[str]:
335-
"""Db2 GRANT/REVOKE requires TABLE keyword for tables and views."""
336-
return "TABLE"
327+
def _get_current_grants_config(self, table: exp.Table) -> t.Dict[str, t.List[str]]:
328+
"""
329+
Db2 does not have INFORMATION_SCHEMA.TABLE_PRIVILEGES.
330+
Query SYSCAT.TABAUTH which stores per-privilege columns (SELECTAUTH, INSERTAUTH,
331+
etc.) with values 'Y' (granted) or 'G' (granted with grant option).
332+
Filter by GRANTOR = CURRENT USER to return only grants made by the connected user.
333+
"""
334+
schema_name = (table.args.get("db") or self._get_current_schema()).upper() # type: ignore
335+
table_name = table.name.upper()
336+
337+
rows = self.fetchall(
338+
exp.select(
339+
exp.column("GRANTEE"),
340+
exp.column("SELECTAUTH"),
341+
exp.column("INSERTAUTH"),
342+
exp.column("UPDATEAUTH"),
343+
exp.column("DELETEAUTH"),
344+
exp.column("ALTERAUTH"),
345+
exp.column("INDEXAUTH"),
346+
exp.column("CONTROLAUTH"),
347+
)
348+
.from_("SYSCAT.TABAUTH")
349+
.where(
350+
exp.and_(
351+
exp.func("UPPER", exp.column("TABSCHEMA")).eq(
352+
exp.Literal.string(schema_name)
353+
),
354+
exp.func("UPPER", exp.column("TABNAME")).eq(
355+
exp.Literal.string(table_name)
356+
),
357+
exp.column("GRANTOR").eq(
358+
exp.func("UPPER", exp.Anonymous(this="CURRENT USER", expressions=[]))
359+
),
360+
exp.column("GRANTEE").neq(
361+
exp.func("UPPER", exp.Anonymous(this="CURRENT USER", expressions=[]))
362+
),
363+
)
364+
)
365+
)
366+
367+
# SYSCAT column name → SQL privilege name
368+
col_to_priv = {
369+
"SELECTAUTH": "SELECT",
370+
"INSERTAUTH": "INSERT",
371+
"UPDATEAUTH": "UPDATE",
372+
"DELETEAUTH": "DELETE",
373+
"ALTERAUTH": "ALTER",
374+
"INDEXAUTH": "INDEX",
375+
"CONTROLAUTH": "CONTROL",
376+
}
377+
grants: t.Dict[str, t.List[str]] = {}
378+
for row in rows:
379+
grantee = str(row[0]).strip()
380+
for i, (_, priv) in enumerate(col_to_priv.items(), start=1):
381+
val = str(row[i]).strip() if row[i] is not None else "N"
382+
if val in ("Y", "G"):
383+
grants.setdefault(priv, [])
384+
if grantee not in grants[priv]:
385+
grants[priv].append(grantee)
386+
return grants
387+
388+
def _apply_grants_config_expr(
389+
self,
390+
table: exp.Table,
391+
grants_config: t.Dict[str, t.List[str]],
392+
table_type: DataObjectType = DataObjectType.TABLE,
393+
) -> t.List[exp.Expr]:
394+
"""Generate GRANT statements for Db2."""
395+
exprs: t.List[exp.Expr] = []
396+
for privilege, principals in grants_config.items():
397+
for principal in principals:
398+
exprs.append(
399+
exp.Grant(
400+
privileges=[exp.GrantPrivilege(this=exp.Var(this=privilege))],
401+
kind=exp.Var(this="TABLE"),
402+
securable=table.copy(),
403+
principals=[
404+
exp.GrantPrincipal(this=exp.Var(this=principal))
405+
],
406+
)
407+
)
408+
return exprs
409+
410+
def _revoke_grants_config_expr(
411+
self,
412+
table: exp.Table,
413+
grants_config: t.Dict[str, t.List[str]],
414+
table_type: DataObjectType = DataObjectType.TABLE,
415+
) -> t.List[exp.Expr]:
416+
"""Generate REVOKE statements for Db2."""
417+
exprs: t.List[exp.Expr] = []
418+
for privilege, principals in grants_config.items():
419+
for principal in principals:
420+
exprs.append(
421+
exp.Revoke(
422+
privileges=[exp.GrantPrivilege(this=exp.Var(this=privilege))],
423+
kind=exp.Var(this="TABLE"),
424+
securable=table.copy(),
425+
principals=[
426+
exp.GrantPrincipal(this=exp.Var(this=principal))
427+
],
428+
)
429+
)
430+
return exprs
337431

338432
@property
339433
def catalog_support(self) -> CatalogSupport:

0 commit comments

Comments
 (0)