Added databricks labs ucx alias command to create a view of tables from one schema/catalog in another schema/catalog (#837)

Author: FastLee

labs.yml

@@ -80,6 +80,23 @@ commands:
       - name: to-schema
         description: target schema to migrate tables to
 
+  - name: alias
+    description: |
+      alias tables across schema/catalog withing a UC metastore 
+      create a view pointing to the "from" table
+      if a view is aliased, recreates the same view in the target schema/catalog
+    flags:
+      - name: from-catalog
+        description: from catalog name
+      - name: from-schema
+        description: from schema
+      - name: from-table
+        description: table names to alias. enter * to migrate all tables
+      - name: to-catalog
+        description: target catalog to migrate schema to
+      - name: to-schema
+        description: target schema to migrate tables to
+
   - name: save-azure-storage-accounts
     description: Identifies all storage account used by tables, identify spn and its permission on each storage accounts
     flags:

src/databricks/labs/ucx/cli.py

@@ -237,6 +237,36 @@ def move(
     tables.move_tables(from_catalog, from_schema, from_table, to_catalog, to_schema, del_table)
 
 
+@ucx.command
+def alias(
+    w: WorkspaceClient,
+    from_catalog: str,
+    from_schema: str,
+    from_table: str,
+    to_catalog: str,
+    to_schema: str,
+):
+    """move a uc table/tables from one schema to another schema in same or different catalog"""
+    installation_manager = InstallationManager(w)
+    installation = installation_manager.for_user(w.current_user.me())
+    if not installation:
+        logger.error(CANT_FIND_UCX_MSG)
+        return
+    sql_backend = StatementExecutionBackend(w, installation.config.warehouse_id)
+    tables = TableMove(w, sql_backend)
+    if from_catalog == "" or to_catalog == "":
+        logger.error("Please enter from_catalog and to_catalog details")
+        return
+    if from_schema == "" or to_schema == "" or from_table == "":
+        logger.error("Please enter from_schema, to_schema and from_table(enter * for migrating all tables) details.")
+        return
+    if from_catalog == to_catalog and from_schema == to_schema:
+        logger.error("please select a different schema or catalog to migrate to")
+        return
+    logger.info(f"aliasing table {from_table} from {from_catalog}.{from_schema} to {to_catalog}.{to_schema}")
+    tables.alias_tables(from_catalog, from_schema, from_table, to_catalog, to_schema)
+
+
 @ucx.command
 def save_azure_storage_accounts(w: WorkspaceClient, subscription_id: str):
     """identifies all azure storage account used by external tables

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

@@ -5,7 +5,12 @@
 from databricks.labs.blueprint.parallel import Threads
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import NotFound
-from databricks.sdk.service.catalog import PermissionsChange, SecurableType, TableType
+from databricks.sdk.service.catalog import (
+    PermissionsChange,
+    Privilege,
+    SecurableType,
+    TableType,
+)
 
 from databricks.labs.ucx.framework.crawlers import SqlBackend
 from databricks.labs.ucx.hive_metastore import TablesCrawler
@@ -229,7 +234,8 @@ def move_tables(
                             self._move_table, from_catalog, from_schema, table.name, to_catalog, to_schema, del_table
                         )
                     )
-                else:
+                    continue
+                if table.view_definition:
                     view_tasks.append(
                         partial(
                             self._move_view,
@@ -242,11 +248,70 @@ def move_tables(
                             table.view_definition,
                         )
                     )
+                    continue
+                logger.warning(
+                    f"table {from_table} was not identified as a valid table or view. skipping this table..."
+                )
         Threads.strict("Creating tables", table_tasks)
         logger.info(f"Moved {len(list(table_tasks))} tables to the new schema {to_schema}.")
         Threads.strict("Creating views", view_tasks)
         logger.info(f"Moved {len(list(view_tasks))} views to the new schema {to_schema}.")
 
+    def alias_tables(
+        self,
+        from_catalog: str,
+        from_schema: str,
+        from_table: str,
+        to_catalog: str,
+        to_schema: str,
+    ):
+        try:
+            self._ws.schemas.get(f"{from_catalog}.{from_schema}")
+        except NotFound:
+            logger.error(f"schema {from_schema} not found in catalog {from_catalog}, enter correct schema details.")
+            return
+        try:
+            self._ws.schemas.get(f"{to_catalog}.{to_schema}")
+        except NotFound:
+            logger.warning(f"schema {to_schema} not found in {to_catalog}, creating...")
+            self._ws.schemas.create(to_schema, to_catalog)
+
+        tables = self._ws.tables.list(from_catalog, from_schema)
+        alias_tasks = []
+        filtered_tables = [table for table in tables if from_table in [table.name, "*"]]
+        for table in filtered_tables:
+            try:
+                self._ws.tables.get(f"{to_catalog}.{to_schema}.{table.name}")
+                logger.warning(
+                    f"table {from_table} already present in {from_catalog}.{from_schema}. skipping this table..."
+                )
+                continue
+            except NotFound:
+                if table.table_type and table.table_type in (TableType.EXTERNAL, TableType.MANAGED):
+                    alias_tasks.append(
+                        partial(self._alias_table, from_catalog, from_schema, table.name, to_catalog, to_schema)
+                    )
+                    continue
+                if table.view_definition:
+                    alias_tasks.append(
+                        partial(
+                            self._move_view,
+                            from_catalog,
+                            from_schema,
+                            table.name,
+                            to_catalog,
+                            to_schema,
+                            False,
+                            table.view_definition,
+                        )
+                    )
+                    continue
+                logger.warning(
+                    f"table {from_table} was not identified as a valid table or view. skipping this table..."
+                )
+        Threads.strict("Creating aliases", alias_tasks)
+        logger.info(f"Created {len(list(alias_tasks))} table and view aliases in the new schema {to_schema}.")
+
     def _move_table(
         self,
         from_catalog: str,
@@ -273,13 +338,45 @@ def _move_table(
                 logger.error(f"Failed to move table {from_table_name}: {err!s}", exc_info=True)
         return False
 
-    def _reapply_grants(self, from_table_name, to_table_name):
+    def _alias_table(
+        self,
+        from_catalog: str,
+        from_schema: str,
+        from_table: str,
+        to_catalog: str,
+        to_schema: str,
+    ) -> bool:
+        from_table_name = f"{from_catalog}.{from_schema}.{from_table}"
+        to_table_name = f"{to_catalog}.{to_schema}.{from_table}"
+        try:
+            self._create_alias_view(from_table_name, to_table_name)
+            self._reapply_grants(from_table_name, to_table_name, target_view=True)
+            return True
+        except NotFound as err:
+            if "[TABLE_OR_VIEW_NOT_FOUND]" in str(err) or "[DELTA_TABLE_NOT_FOUND]" in str(err):
+                logger.error(f"Could not find table {from_table_name}. Table not found.")
+            else:
+                logger.error(f"Failed to alias table {from_table_name}: {err!s}", exc_info=True)
+            return False
+
+    def _reapply_grants(self, from_table_name, to_table_name, *, target_view: bool = False):
         grants = self._ws.grants.get(SecurableType.TABLE, from_table_name)
         if grants.privilege_assignments is not None:
             logger.info(f"Applying grants on table {to_table_name}")
-            grants_changes = [
-                PermissionsChange(pair.privileges, pair.principal) for pair in grants.privilege_assignments
-            ]
+            grants_changes = []
+            for permission in grants.privilege_assignments:
+                if not permission.privileges:
+                    continue
+                if not target_view:
+                    grants_changes.append(PermissionsChange(list(permission.privileges), permission.principal))
+                    continue
+                privileges = set()
+                for privilege in permission.privileges:
+                    if privilege != Privilege.MODIFY:
+                        privileges.add(privilege)
+                if privileges:
+                    grants_changes.append(PermissionsChange(list(privileges), permission.principal))
+
             self._ws.grants.update(SecurableType.TABLE, to_table_name, changes=grants_changes)
 
     def _recreate_table(self, from_table_name, to_table_name):
@@ -288,6 +385,11 @@ def _recreate_table(self, from_table_name, to_table_name):
         logger.info(f"Creating table {to_table_name}")
         self._backend.execute(create_table_sql)
 
+    def _create_alias_view(self, from_table_name, to_table_name):
+        create_view_sql = f"CREATE VIEW {to_table_name} AS SELECT * FROM {from_table_name}"
+        logger.info(f"Creating view {to_table_name} on {from_table_name}")
+        self._backend.execute(create_view_sql)
+
     def _move_view(
         self,
         from_catalog: str,

tests/integration/hive_metastore/test_table_move.py

@@ -136,3 +136,57 @@ def test_move_views(ws, sql_backend, make_catalog, make_schema, make_table, make
 
     for t in from_views:
         assert t.name in [from_view_not_to_migrate.name, from_table.name]
+
+
+@retried(on=[NotFound], timeout=timedelta(minutes=2))
+def test_alias_tables(ws, sql_backend, make_catalog, make_schema, make_table, make_acc_group):
+    tm = TableMove(ws, sql_backend)
+    group_a = make_acc_group()
+    group_b = make_acc_group()
+    from_catalog = make_catalog()
+    from_schema = make_schema(catalog_name=from_catalog.name)
+    from_table_1 = make_table(catalog_name=from_catalog.name, schema_name=from_schema.name)
+    from_table_2 = make_table(catalog_name=from_catalog.name, schema_name=from_schema.name)
+    from_table_3 = make_table(catalog_name=from_catalog.name, schema_name=from_schema.name)
+    from_view_1 = make_table(
+        catalog_name=from_catalog.name,
+        schema_name=from_schema.name,
+        view=True,
+        ctas=f"select * from {from_table_2.full_name} where 1=2",
+    )
+    to_catalog = make_catalog()
+    to_schema = make_schema(catalog_name=to_catalog.name)
+    # creating a table in target schema to test skipping
+    to_table_3 = make_table(catalog_name=to_catalog.name, schema_name=to_schema.name, name=from_table_3.name)
+    sql_backend.execute(f"GRANT SELECT ON TABLE {from_table_1.full_name} TO `{group_a.display_name}`")
+    sql_backend.execute(f"GRANT SELECT,MODIFY ON TABLE {from_table_2.full_name} TO `{group_b.display_name}`")
+    sql_backend.execute(f"GRANT SELECT ON VIEW {from_view_1.full_name} TO `{group_b.display_name}`")
+    sql_backend.execute(f"GRANT SELECT ON TABLE {to_table_3.full_name} TO `{group_a.display_name}`")
+
+    tm.alias_tables(from_catalog.name, from_schema.name, "*", to_catalog.name, to_schema.name)
+
+    to_tables = ws.tables.list(catalog_name=to_catalog.name, schema_name=to_schema.name)
+    table_1_grant = ws.grants.get(
+        securable_type=SecurableType.TABLE, full_name=f"{to_catalog.name}.{to_schema.name}.{from_table_1.name}"
+    )
+    table_2_grant = ws.grants.get(
+        securable_type=SecurableType.TABLE, full_name=f"{to_catalog.name}.{to_schema.name}.{from_table_2.name}"
+    )
+    table_3_grant = ws.grants.get(
+        securable_type=SecurableType.TABLE, full_name=f"{to_catalog.name}.{to_schema.name}.{from_table_3.name}"
+    )
+    view_1_grant = ws.grants.get(
+        securable_type=SecurableType.TABLE, full_name=f"{to_catalog.name}.{to_schema.name}.{from_view_1.name}"
+    )
+    for t in to_tables:
+        assert t.name in [from_table_1.name, from_table_2.name, from_table_3.name, from_view_1.name]
+
+    expected_table_1_grant = [PrivilegeAssignment(group_a.display_name, [Privilege.SELECT])]
+    expected_table_2_grant = [PrivilegeAssignment(group_b.display_name, [Privilege.SELECT])]
+    expected_table_3_grant = [PrivilegeAssignment(group_a.display_name, [Privilege.SELECT])]
+    expected_view_1_grant = [PrivilegeAssignment(group_b.display_name, [Privilege.SELECT])]
+
+    assert table_1_grant.privilege_assignments == expected_table_1_grant
+    assert table_2_grant.privilege_assignments == expected_table_2_grant
+    assert table_3_grant.privilege_assignments == expected_table_3_grant
+    assert view_1_grant.privilege_assignments == expected_view_1_grant