Added assign-owner-group command (#3111)

replaces #3075 
close #2890

---------

Co-authored-by: Serge Smertin <serge.smertin@databricks.com>

Author: FastLee

README.md

@@ -118,6 +118,7 @@ See [contributing instructions](CONTRIBUTING.md) to help improve this project.
   * [`skip` command](#skip-command)
   * [`unskip` command](#unskip-command)
   * [`create-catalogs-schemas` command](#create-catalogs-schemas-command)
+  * [`assign-owner-group` command](#assign-owner-group-command)
   * [`migrate-tables` command](#migrate-tables-command)
   * [`revert-migrated-tables` command](#revert-migrated-tables-command)
   * [`move` command](#move-command)
@@ -1312,6 +1313,7 @@ access the configuration file from the command line. Here's the description of c
   * `num_threads`: An optional integer representing the number of threads to use for migration.
   * `database_to_catalog_mapping`: An optional dictionary mapping source database names to target catalog names.
   * `default_catalog`: An optional string representing the default catalog name.
+  * `default_owner_group`: Assigns this group to all migrated objects (catalogs, databases, tables, views, etc.). The group has to be an account group and the user running the migration has to be a member of this group.
   * `log_level`: An optional string representing the log level.
   * `workspace_start_path`: A string representing the starting path for notebooks and directories crawler in the workspace.
   * `instance_profile`: An optional string representing the name of the instance profile.
@@ -1708,6 +1710,20 @@ to the bucket. It then maps the bucket to the tables which has external location
 the schema and catalog if at least one such table is migrated to it.
 [[back to top](#databricks-labs-ucx)]
 
+## `assign-owner-group` command
+
+```text
+databricks labs ucx assign-owner-group
+```
+
+This command assigns the owner group to the UCX catalog and all migrated tables.
+The owner group is an account group that will be designated as an owner to all migrated objects (catalogs, schemas, tables, views).
+The principal running the command and later, the migration workflows, is required to be a member of this group.
+The command will list all the groups the principal is a member of and allow the selection of the owner group.
+It sets the default_owner_group property in the config.yml file.
+
+[[back to top](#databricks-labs-ucx)]
+
 ## `migrate-tables` command
 
 ```text

labs.yml

@@ -358,4 +358,5 @@ commands:
   - name: enable-hms-federation
     description: (EXPERIMENTAL) Enable HMS federation based migration flow. When this is enabled, UCX will create a federated HMS catalog which syncs from the workspace HMS.
 
-
+  - name: assign-owner-group
+    description: Assign owner group to the workspace. This group will be assigned as an owner to all migrated tables and views.

src/databricks/labs/ucx/cli.py

@@ -1,3 +1,4 @@
+from dataclasses import replace
 from io import BytesIO
 import json
 import webbrowser
@@ -634,6 +635,32 @@ def create_ucx_catalog(w: WorkspaceClient, prompts: Prompts, ctx: WorkspaceConte
     workspace_context.verify_progress_tracking.verify()
 
 
+@ucx.command
+def assign_owner_group(
+    w: WorkspaceClient,
+    prompts: Prompts,
+    *,
+    ctx: WorkspaceContext | None = None,
+    run_as_collection: bool = False,
+    a: AccountClient | None = None,
+) -> None:
+    """
+    Pick owner group. This group will be assigned as owner to all the migrated tables.
+    """
+    if ctx:
+        workspace_contexts = [ctx]
+    else:
+        workspace_contexts = _get_workspace_contexts(w, a, run_as_collection)
+
+    owner_group = workspace_contexts[0].group_manager.pick_owner_group(prompts)
+    if not owner_group:
+        return
+    for workspace_context in workspace_contexts:
+        config = workspace_context.installation.load(WorkspaceConfig)
+        config = replace(config, default_owner_group=owner_group)
+        workspace_context.installation.save(config)
+
+
 @ucx.command
 def migrate_tables(
     w: WorkspaceClient,

src/databricks/labs/ucx/config.py

@@ -83,6 +83,9 @@ class WorkspaceConfig:  # pylint: disable=too-many-instance-attributes
     # [INTERNAL ONLY] If specified, the large-scale scanners will only list the specified number of objects.
     debug_listing_upper_limit: int | None = None
 
+    # Default table owner, assigned to a group
+    default_owner_group: str | None = None
+
     def replace_inventory_variable(self, text: str) -> str:
         return text.replace("$inventory", f"hive_metastore.{self.inventory_database}")
 

src/databricks/labs/ucx/contexts/application.py

@@ -44,7 +44,12 @@
 )
 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,
+    DefaultSecurableOwnership,
+)
 from databricks.labs.ucx.hive_metastore.table_migrate import (
     TableMigrationStatusRefresher,
     TablesMigrator,
@@ -272,6 +277,18 @@ def table_ownership(self) -> TableOwnership:
             self.workspace_path_ownership,
         )
 
+    @cached_property
+    def default_securable_ownership(self) -> DefaultSecurableOwnership:
+        # validate that the default_owner_group is set and is a valid group (the current user is a member)
+
+        return DefaultSecurableOwnership(
+            self.administrator_locator,
+            self.tables_crawler,
+            self.group_manager,
+            self.config.default_owner_group,
+            lambda: self.workspace_client.current_user.me().user_name,
+        )
+
     @cached_property
     def workspace_path_ownership(self) -> WorkspacePathOwnership:
         return WorkspacePathOwnership(self.administrator_locator, self.workspace_client)
@@ -316,9 +333,15 @@ def acl_migrator(self):
             self.config.inventory_database,
         )
 
+    @cached_property
+    def table_ownership_grant_loader(self) -> TableOwnershipGrantLoader:
+        return TableOwnershipGrantLoader(self.tables_crawler, self.default_securable_ownership)
+
     @cached_property
     def migrate_grants(self) -> MigrateGrants:
+        # owner grants have to come first
         grant_loaders: list[Callable[[], Iterable[Grant]]] = [
+            self.default_securable_ownership.load,
             self.grants_crawler.snapshot,
             self.principal_acl.get_interactive_cluster_grants,
         ]

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

@@ -119,12 +119,12 @@ def object_key(self) -> str:
 
     @property
     def order(self) -> int:
-        """Order of the grants to be upheld when applying."""
-        match self.action_type:
-            case "OWN":  # Apply ownership as last to avoid losing permissions for applying grants
-                return 1
-            case _:
-                return 0
+        """Apply Ownership First, then the rest.
+        Consider Revising When we properly apply manage permission"""
+
+        if self.action_type == "OWN":  # Apply ownership as last to avoid losing permissions for applying grants
+            return 0
+        return 1
 
     def this_type_and_key(self):
         return self.type_and_key(
@@ -780,10 +780,12 @@ def __init__(
     ):
         self._sql_backend = sql_backend
         self._group_manager = group_manager
+        # List of grant loaders, the assumption that the first one is a loader for ownership grants
         self._grant_loaders = grant_loaders
 
     def apply(self, src: SecurableObject, dst: SecurableObject) -> bool:
-        for grant in self._match_grants(src):
+        grants = self._match_grants(src)
+        for grant in grants:
             acl_migrate_sql = grant.uc_grant_sql(dst.kind, dst.full_name)
             if acl_migrate_sql is None:
                 logger.warning(
@@ -801,8 +803,9 @@ def apply(self, src: SecurableObject, dst: SecurableObject) -> bool:
         return True
 
     def retrieve(self, src: SecurableObject, dst: SecurableObject) -> list[Grant]:
-        grants = []
-        for grant in self._match_grants(src):
+        grants = self._match_grants(src)
+        discover_grants = []
+        for grant in grants:
             acl_migrate_sql = grant.uc_grant_sql(dst.kind, dst.full_name)
             if acl_migrate_sql is None:
                 logger.warning(
@@ -811,18 +814,23 @@ def retrieve(self, src: SecurableObject, dst: SecurableObject) -> list[Grant]:
                 )
                 continue
             logger.debug(f"Retrieving acls on {dst.full_name} using SQL query: {acl_migrate_sql}")
-            grants.append(grant)
-        return grants
+            discover_grants.append(grant)
+        return discover_grants
 
     @cached_property
     def _workspace_to_account_group_names(self) -> dict[str, str]:
         return {g.name_in_workspace: g.name_in_account for g in self._group_manager.snapshot()}
 
     @cached_property
     def _grants(self) -> list[Grant]:
+        # Accumulate grants from all loaders skips ownership grants
         grants = []
-        for loader in self._grant_loaders:
+        for index, loader in enumerate(self._grant_loaders):
             for grant in loader():
+                # Skip ownership grants for all loaders other than the first one
+                # The assumption is that the first one will be designated as the ownership loader
+                if index != 0 and grant.action_type == "OWN":
+                    continue
                 grants.append(grant)
         return grants
 
@@ -841,6 +849,26 @@ def _replace_account_group(self, grant: Grant) -> Grant:
             return grant
         return replace(grant, principal=target_principal)
 
+    def _apply_ownership(self, dst: SecurableObject, owner: str):
+        owner = escape_sql_identifier(owner)
+        destination = escape_sql_identifier(dst.full_name)
+        if dst.kind == "TABLE":
+            sql = f"ALTER TABLE {destination} " f"SET OWNER TO {owner}"
+        elif dst.kind == "VIEW":
+            sql = f"ALTER VIEW {destination} " f"SET OWNER TO {owner}"
+        elif dst.kind in {"SCHEMA", "DATABASE"}:
+            sql = f"ALTER SCHEMA {destination} " f"SET OWNER TO {owner}"
+        elif dst.kind == "CATALOG":
+            sql = f"ALTER CATALOG {destination} " f"SET OWNER TO {owner}"
+        else:
+            logger.warning(f"Unknown object type {dst.kind} for ownership migration")
+            return
+        logger.debug(f"Applying ownership on {dst.full_name} using SQL query: {sql}")
+        try:
+            self._sql_backend.execute(sql)
+        except DatabricksError as e:
+            logger.warning(f"Failed to apply ownership on {dst.full_name}", exc_info=e)
+
 
 class ACLMigrator(CrawlerBase[Grant]):
     def __init__(

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

@@ -1,4 +1,5 @@
 import logging
+from collections.abc import Iterable, Callable
 from functools import cached_property
 
 from databricks.labs.ucx.framework.owners import (
@@ -8,11 +9,12 @@
     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
 from databricks.labs.ucx.source_code.used_table import UsedTablesCrawler
+from databricks.labs.ucx.workspace_access.groups import GroupManager
 
 logger = logging.getLogger(__name__)
 
@@ -85,6 +87,116 @@ def _grants_snapshot(self):
         return self._grants_crawler.snapshot()
 
 
+class DefaultSecurableOwnership(Ownership[Table]):
+    """Determine ownership of tables in the inventory based on the following rules:
+    -- If a global owner group is defined, then all tables are owned by that group.
+    -- If the user running the application can be identified, then all tables are owned by that user.
+    -- Otherwise, the table is owned by the administrator.
+    """
+
+    def __init__(
+        self,
+        administrator_locator: AdministratorLocator,
+        table_crawler: TablesCrawler,
+        group_manager: GroupManager,
+        default_owner_group: str | None,
+        app_principal_resolver: Callable[[], str | None],
+    ) -> None:
+        super().__init__(administrator_locator)
+        self._tables_crawler = table_crawler
+        self._group_manager = group_manager
+        self._default_owner_group = default_owner_group
+        self._app_principal_resolver = app_principal_resolver
+
+    @cached_property
+    def _application_principal(self) -> str | None:
+        return self._app_principal_resolver()
+
+    @cached_property
+    def _static_owner(self) -> str | None:
+        # If the default owner group is not valid, fall back to the application principal
+        if self._default_owner_group and self._group_manager.validate_owner_group(self._default_owner_group):
+            logger.warning("Default owner group is not valid, falling back to administrator ownership.")
+            return self._default_owner_group
+        return self._application_principal
+
+    def load(self) -> Iterable[Grant]:
+        databases = set()
+        catalogs = set()
+        owner = self._static_owner
+        if not owner:
+            logger.warning("No owner found for tables and databases")
+            return
+        for table in self._tables_crawler.snapshot():
+            table_name, view_name = self._names(table)
+
+            if table.database not in databases:
+                databases.add(table.database)
+            if table.catalog not in catalogs:
+                catalogs.add(table.catalog)
+            yield Grant(
+                principal=owner,
+                action_type='OWN',
+                catalog=table.catalog,
+                database=table.database,
+                table=table_name,
+                view=view_name,
+            )
+        for database in databases:
+            yield Grant(
+                principal=owner,
+                action_type='OWN',
+                catalog="hive_metastore",
+                database=database,
+                table=None,
+                view=None,
+            )
+
+        for catalog in catalogs:
+            yield Grant(
+                principal=owner,
+                action_type='OWN',
+                catalog=catalog,
+                database=None,
+                table=None,
+                view=None,
+            )
+
+    @staticmethod
+    def _names(table: Table) -> tuple[str | None, str | None]:
+        if table.view_text:
+            return None, table.name
+        return table.name, None
+
+    def _maybe_direct_owner(self, record: Table) -> str | None:
+        return self._static_owner
+
+
+class TableOwnershipGrantLoader:
+    def __init__(self, tables_crawler: TablesCrawler, table_ownership: Ownership[Table]) -> 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.