|
9 | 9 | from sqlglot.optimizer.normalize_identifiers import normalize_identifiers |
10 | 10 |
|
11 | 11 | 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 |
16 | 13 | from sqlmesh.core.engine_adapter.shared import ( |
17 | 14 | CatalogSupport, |
18 | 15 | CommentCreationTable, |
@@ -46,16 +43,13 @@ def is_db2_error(exception: Exception, error_code: str) -> bool: |
46 | 43 |
|
47 | 44 | @set_catalog() |
48 | 45 | class Db2EngineAdapter( |
49 | | - GrantsFromInfoSchemaMixin, |
50 | 46 | PandasNativeFetchDFSupportMixin, |
51 | 47 | EngineAdapter, |
52 | 48 | ): |
53 | 49 | DIALECT = "db2" |
54 | 50 | SUPPORTS_INDEXES = True |
55 | 51 | SUPPORTS_REPLACE_TABLE = False |
56 | 52 | 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") |
59 | 53 | COMMENT_CREATION_TABLE = CommentCreationTable.COMMENT_COMMAND_ONLY |
60 | 54 | COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY |
61 | 55 | SUPPORTS_QUERY_EXECUTION_TRACKING = True |
@@ -330,10 +324,110 @@ def _db2_type_to_sqlglot(self, db2_type: str, length: int, scale: int) -> exp.Da |
330 | 324 | sqlglot_type = type_mapping.get(db2_type, f"VARCHAR({length})") |
331 | 325 | return exp.DataType.build(sqlglot_type, dialect="db2") |
332 | 326 |
|
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 |
337 | 431 |
|
338 | 432 | @property |
339 | 433 | def catalog_support(self) -> CatalogSupport: |
|
0 commit comments