Fixed handling for OWN table permissions (#571)

Closes #558

src/databricks/labs/ucx/hive_metastore/grants.py

@@ -76,7 +76,10 @@ def this_type_and_key(self):
     def hive_grant_sql(self) -> str:
         object_type, object_key = self.this_type_and_key()
         # See https://docs.databricks.com/en/sql/language-manual/security-grant.html
-        return f"GRANT {self.action_type} ON {object_type} {object_key} TO `{self.principal}`"
+        if self.action_type.upper() == "OWN":
+            return f"ALTER {object_key} OWNER TO `{self.principal}`"
+        else:
+            return f"GRANT {self.action_type} ON {object_type} {object_key} TO `{self.principal}`"
 
     def hive_revoke_sql(self) -> str:
         object_type, object_key = self.this_type_and_key()

src/databricks/labs/ucx/workspace_access/tacl.py

@@ -37,11 +37,16 @@ def get_crawler_tasks(self) -> Iterator[Callable[..., Permissions | None]]:
         # * GRANT MODIFY ON TABLE hive_metastore.db_a.table_a TO group_a
         # will be folded and executed in one statement/transaction:
         # * GRANT SELECT, MODIFY ON TABLE hive_metastore.db_a.table_a TO group_a
+        # The exception is OWN permission which are set with ALTER table
 
         folded_actions = collections.defaultdict(set)
+        own_permissions = set()
         for grant in self._grants_crawler.snapshot():
             key = (grant.principal, grant.this_type_and_key())
-            folded_actions[key].add(grant.action_type)
+            if grant.action_type.upper() == "OWN":
+                own_permissions.add(key)
+            else:
+                folded_actions[key].add(grant.action_type)
 
         def inner(object_type: str, object_id: str, grant: Grant) -> Permissions:
             return Permissions(object_type=object_type, object_id=object_id, raw=json.dumps(dataclasses.asdict(grant)))
@@ -50,6 +55,10 @@ def inner(object_type: str, object_id: str, grant: Grant) -> Permissions:
             grant = self._from_reduced(object_type, object_id, principal, ", ".join(sorted(actions)))
             yield functools.partial(inner, object_type=object_type, object_id=object_id, grant=grant)
 
+        for principal, (object_type, object_id) in own_permissions:
+            grant = self._from_reduced(object_type, object_id, principal, "OWN")
+            yield functools.partial(inner, object_type=object_type, object_id=object_id, grant=grant)
+
     def _from_reduced(self, object_type: str, object_id: str, principal: str, action_type: str):
         match object_type:
             case "TABLE":

tests/integration/workspace_access/test_groups.py

@@ -276,3 +276,107 @@ def test_migration_state_should_be_saved_without_missing_anything(sql_backend, m
     new_state = state.fetch_migration_state(sql_backend, inventory_database.name)
 
     assert len(new_state.groups) == len(state.groups)
+
+
+def test_set_owner_permission(
+    ws,
+    sql_backend,
+    inventory_schema,
+    make_ucx_group,
+    make_group,
+    make_acc_group,
+    make_cluster_policy,
+    make_table,
+):
+    ws_group, _ = make_ucx_group()
+
+    logger.info("Testing setting ownership on table.")
+    dummy_table = make_table()
+    logger.info(f"Table name {dummy_table.full_name} group name {ws_group.display_name}")
+    sql_backend.execute(f"GRANT SELECT, MODIFY ON TABLE {dummy_table.full_name} TO `{ws_group.display_name}`")
+    sql_backend.execute(f"ALTER {dummy_table.full_name} OWNER TO `{ws_group.display_name}`")
+
+    group_manager = GroupManager(ws, GroupsConfig(auto=True))
+    group_manager.prepare_groups_in_environment()
+
+    group_info = group_manager.migration_state.get_by_workspace_group_name(ws_group.display_name)
+
+    tables = TablesCrawler(sql_backend, inventory_schema)
+    grants = GrantsCrawler(tables)
+    tacl = TableAclSupport(grants, sql_backend)
+    permission_manager = PermissionManager(sql_backend, inventory_schema, [tacl])
+
+    permission_manager.inventorize_permissions()
+
+    dummy_grants = list(permission_manager.load_all_for("TABLE", dummy_table.full_name, Grant))
+    assert 2 == len(dummy_grants)
+
+    table_permissions = grants.for_table_info(dummy_table)
+    assert ws_group.display_name in table_permissions
+    assert "MODIFY" in table_permissions[ws_group.display_name]
+    assert "SELECT" in table_permissions[ws_group.display_name]
+    assert "OWN" in table_permissions[ws_group.display_name]
+
+    permission_manager.apply_group_permissions(group_manager.migration_state, destination="backup")
+
+    @retried(on=[AssertionError], timeout=timedelta(seconds=30))
+    def check_permissions_for_backup_group():
+        logger.info("check_permissions_for_backup_group()")
+
+        table_permissions = grants.for_table_info(dummy_table)
+        assert group_info.workspace.display_name in table_permissions
+        assert group_info.backup.display_name in table_permissions
+        assert "MODIFY" in table_permissions[group_info.workspace.display_name]
+        assert "SELECT" in table_permissions[group_info.workspace.display_name]
+        assert "MODIFY" in table_permissions[group_info.backup.display_name]
+        assert "SELECT" in table_permissions[group_info.backup.display_name]
+        assert "OWN" in table_permissions[group_info.backup.display_name]
+
+    check_permissions_for_backup_group()
+
+    group_manager.replace_workspace_groups_with_account_groups(group_manager.migration_state)
+
+    @retried(on=[AssertionError], timeout=timedelta(minutes=1))
+    def check_permissions_after_replace():
+        logger.info("check_permissions_after_replace()")
+
+        table_permissions = grants.for_table_info(dummy_table)
+        assert group_info.account.display_name in table_permissions
+        assert group_info.backup.display_name in table_permissions
+        assert "MODIFY" in table_permissions[group_info.backup.display_name]
+        assert "SELECT" in table_permissions[group_info.backup.display_name]
+        assert "OWN" in table_permissions[group_info.backup.display_name]
+
+    check_permissions_after_replace()
+
+    permission_manager.apply_group_permissions(group_manager.migration_state, destination="account")
+
+    @retried(on=[AssertionError], timeout=timedelta(seconds=30))
+    def check_permissions_for_account_group():
+        logger.info("check_permissions_for_account_group()")
+
+        table_permissions = grants.for_table_info(dummy_table)
+        assert group_info.account.display_name in table_permissions
+        assert group_info.backup.display_name in table_permissions
+        assert "MODIFY" in table_permissions[group_info.backup.display_name]
+        assert "SELECT" in table_permissions[group_info.backup.display_name]
+        assert "MODIFY" in table_permissions[group_info.account.display_name]
+        assert "SELECT" in table_permissions[group_info.account.display_name]
+        assert "OWN" in table_permissions[group_info.account.display_name]
+
+    check_permissions_for_account_group()
+
+    for _info in group_manager.migration_state.groups:
+        ws.groups.delete(_info.backup.id)
+
+    @retried(on=[AssertionError], timeout=timedelta(minutes=1))
+    def check_table_permissions_after_backup_delete():
+        logger.info("check_table_permissions_after_backup_delete()")
+
+        table_permissions = grants.for_table_info(dummy_table)
+        assert group_info.account.display_name in table_permissions
+        assert "MODIFY" in table_permissions[group_info.account.display_name]
+        assert "SELECT" in table_permissions[group_info.account.display_name]
+        assert "OWN" in table_permissions[group_info.account.display_name]
+
+    check_table_permissions_after_backup_delete()

tests/unit/hive_metastore/test_grants.py

@@ -77,6 +77,11 @@ def test_hive_sql():
     assert grant.hive_revoke_sql() == "REVOKE SELECT ON TABLE hive_metastore.mydb.mytable FROM `user`"
 
 
+def test_hive_own_sql():
+    grant = Grant(principal="user", action_type="OWN", catalog="hive_metastore", database="mydb", table="mytable")
+    assert grant.hive_grant_sql() == "ALTER hive_metastore.mydb.mytable OWNER TO `user`"
+
+
 def test_hive_revoke_sql():
     grant = Grant(principal="user", action_type="SELECT", catalog="hive_metastore", database="mydb", table="mytable")
     assert grant.hive_revoke_sql() == "REVOKE SELECT ON TABLE hive_metastore.mydb.mytable FROM `user`"

tests/unit/workspace_access/test_tacl.py

@@ -69,7 +69,7 @@ def test_tacl_crawler_multiple_permissions():
     assert "catalog_a.database_b.table_c" == permissions.object_id
     assert Grant(
         principal="foo@example.com",
-        action_type="MODIFY, OWN, SELECT",
+        action_type="MODIFY, SELECT",
         catalog="catalog_a",
         database="database_b",
         table="table_c",
@@ -183,6 +183,21 @@ def test_tacl_crawler_multiple_permissions():
         anonymous_function=True,
     ) == Grant(**json.loads(permissions.raw))
 
+    permissions = next(crawler_tasks)()
+
+    assert "TABLE" == permissions.object_type
+    assert "catalog_a.database_b.table_c" == permissions.object_id
+    assert Grant(
+        principal="foo@example.com",
+        action_type="OWN",
+        catalog="catalog_a",
+        database="database_b",
+        table="table_c",
+        view=None,
+        any_file=False,
+        anonymous_function=False,
+    ) == Grant(**json.loads(permissions.raw))
+
 
 def test_tacl_applier(mocker):
     sql_backend = MockBackend()