Skip to content

Commit

Permalink
Use table ownership heuristics for migrating ACLs
Browse files Browse the repository at this point in the history
This PR prepends `OWN` grants when migrating ACLs for tables. For detailed logic on ownership heuristics, see #3066
  • Loading branch information
nfx authored and FastLee committed Nov 12, 2024
1 parent c24523e commit ccf20c5
Show file tree
Hide file tree
Showing 3 changed files with 87 additions and 2 deletions.
12 changes: 11 additions & 1 deletion src/databricks/labs/ucx/contexts/application.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,11 @@
)
from databricks.labs.ucx.hive_metastore.mapping import TableMapping
from databricks.labs.ucx.hive_metastore.table_migration_status import TableMigrationIndex
from databricks.labs.ucx.hive_metastore.ownership import TableMigrationOwnership, TableOwnership
from databricks.labs.ucx.hive_metastore.ownership import (
TableMigrationOwnership,
TableOwnership,
TableOwnershipGrantLoader,
)
from databricks.labs.ucx.hive_metastore.table_migrate import (
TableMigrationStatusRefresher,
TablesMigrator,
Expand Down Expand Up @@ -316,9 +320,15 @@ def acl_migrator(self):
self.config.inventory_database,
)

@cached_property
def table_ownership_grant_loader(self) -> TableOwnershipGrantLoader:
return TableOwnershipGrantLoader(self.tables_crawler, self.table_ownership)

@cached_property
def migrate_grants(self) -> MigrateGrants:
# owner grants have to come first
grant_loaders: list[Callable[[], Iterable[Grant]]] = [
self.table_ownership_grant_loader.load,
self.grants_crawler.snapshot,
self.principal_acl.get_interactive_cluster_grants,
]
Expand Down
28 changes: 27 additions & 1 deletion src/databricks/labs/ucx/hive_metastore/ownership.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import logging
from collections.abc import Iterable
from functools import cached_property

from databricks.labs.ucx.framework.owners import (
Expand All @@ -8,7 +9,7 @@
WorkspacePathOwnership,
)
from databricks.labs.ucx.hive_metastore import TablesCrawler
from databricks.labs.ucx.hive_metastore.grants import GrantsCrawler
from databricks.labs.ucx.hive_metastore.grants import GrantsCrawler, Grant
from databricks.labs.ucx.hive_metastore.table_migration_status import TableMigrationStatus
from databricks.labs.ucx.hive_metastore.tables import Table
from databricks.labs.ucx.source_code.base import UsedTable
Expand Down Expand Up @@ -85,6 +86,31 @@ def _grants_snapshot(self):
return self._grants_crawler.snapshot()


class TableOwnershipGrantLoader:
def __init__(self, tables_crawler: TablesCrawler, table_ownership: TableOwnership) -> None:
self._tables_crawler = tables_crawler
self._table_ownership = table_ownership

def load(self) -> Iterable[Grant]:
for table in self._tables_crawler.snapshot():
owner = self._table_ownership.owner_of(table)
table_name, view_name = self._names(table)
yield Grant(
principal=owner,
action_type='OWN',
catalog=table.catalog,
database=table.database,
table=table_name,
view=view_name,
)

@staticmethod
def _names(table: Table) -> tuple[str | None, str | None]:
if table.view_text:
return None, table.name
return table.name, None


class TableMigrationOwnership(Ownership[TableMigrationStatus]):
"""Determine ownership of table migration records in the inventory.
Expand Down
49 changes: 49 additions & 0 deletions tests/unit/progress/test_ownership.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
from unittest.mock import create_autospec

from databricks.labs.ucx.hive_metastore.grants import Grant
from databricks.labs.ucx.hive_metastore.ownership import TableOwnership, TableOwnershipGrantLoader
from databricks.labs.ucx.hive_metastore.tables import Table, TablesCrawler


def test_table_ownership_grants():
tables_crawler = create_autospec(TablesCrawler)
tables_crawler.snapshot.return_value = [
Table(
catalog='a',
database='b',
name='c',
object_type='EXTERNAL',
table_format='PARQUET',
),
Table(
catalog='b',
database='c',
name='d',
object_type='VIEW',
table_format='...',
view_text='...',
),
]
table_ownership = create_autospec(TableOwnership)
table_ownership.owner_of.return_value = 'somebody'

loader = TableOwnershipGrantLoader(tables_crawler, table_ownership)

grants = list(loader.load())

assert grants == [
Grant(
principal='somebody',
action_type='OWN',
catalog='a',
database='b',
table='c',
),
Grant(
principal='somebody',
action_type='OWN',
catalog='b',
database='c',
view='d',
),
]

0 comments on commit ccf20c5

Please sign in to comment.