Added TableMapping functionality to table migrate (#752)

Author: FastLee

src/databricks/labs/ucx/cli.py

@@ -141,7 +141,8 @@ def revert_migrated_tables(schema: str, table: str, *, delete_managed: bool = Fa
     warehouse_id = installation.config.warehouse_id
     sql_backend = StatementExecutionBackend(ws, warehouse_id)
     table_crawler = TablesCrawler(sql_backend, installation.config.inventory_database)
-    tm = TablesMigrate(table_crawler, ws, sql_backend)
+    tmp = TableMapping(ws)
+    tm = TablesMigrate(table_crawler, ws, sql_backend, tmp)
     if tm.print_revert_report(delete_managed=delete_managed) and prompts.confirm(
         "Would you like to continue?", max_attempts=2
     ):

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

@@ -37,6 +37,14 @@ def initial(cls, workspace_name: str, catalog_name: str, table: Table) -> "Rule"
             dst_table=table.name,
         )
 
+    @property
+    def as_uc_table_key(self):
+        return f"{self.catalog_name}.{self.dst_schema}.{self.dst_table}"
+
+    @property
+    def as_hms_table_key(self):
+        return f"hive_metastore.{self.src_schema}.{self.src_table}"
+
 
 class TableMapping:
     UCX_SKIP_PROPERTY = "databricks.labs.ucx.skip"

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

@@ -7,6 +7,7 @@
 from databricks.labs.ucx.framework.crawlers import SqlBackend
 from databricks.labs.ucx.framework.parallel import Threads
 from databricks.labs.ucx.hive_metastore import TablesCrawler
+from databricks.labs.ucx.hive_metastore.mapping import Rule, TableMapping
 from databricks.labs.ucx.hive_metastore.tables import MigrationCount, Table
 
 logger = logging.getLogger(__name__)
@@ -18,62 +19,66 @@ def __init__(
         tc: TablesCrawler,
         ws: WorkspaceClient,
         backend: SqlBackend,
-        default_catalog=None,
-        database_to_catalog_mapping: dict[str, str] | None = None,
+        tm: TableMapping,
     ):
         self._tc = tc
         self._backend = backend
         self._ws = ws
-        self._database_to_catalog_mapping = database_to_catalog_mapping
-        self._default_catalog = self._init_default_catalog(default_catalog)
+        self._tm = tm
         self._seen_tables: dict[str, str] = {}
 
-    @staticmethod
-    def _init_default_catalog(default_catalog):
-        if default_catalog:
-            return default_catalog
-        else:
-            return "ucx_default"  # TODO : Fetch current workspace name and append it to the default catalog.
-
     def migrate_tables(self):
         self._init_seen_tables()
+        mapping_rules = self._get_mapping_rules()
         tasks = []
         for table in self._tc.snapshot():
-            target_catalog = self._default_catalog
-            if self._database_to_catalog_mapping:
-                target_catalog = self._database_to_catalog_mapping[table.database]
-            tasks.append(partial(self._migrate_table, target_catalog, table))
-        _, errors = Threads.gather("migrate tables", tasks)
-        if len(errors) > 0:
-            # TODO: https://github.com/databrickslabs/ucx/issues/406
-            # TODO: pick first X issues in the summary
-            msg = f"Detected {len(errors)} errors: {'. '.join(str(e) for e in errors)}"
-            raise ValueError(msg)
-
-    def _migrate_table(self, target_catalog: str, table: Table):
-        sql = table.uc_create_sql(target_catalog)
-        logger.debug(f"Migrating table {table.key} to using SQL query: {sql}")
-        target = f"{target_catalog}.{table.database}.{table.name}".lower()
-
-        if self._table_already_upgraded(target):
-            logger.info(f"Table {table.key} already upgraded to {self._seen_tables[target]}")
-        elif table.object_type == "MANAGED":
-            self._backend.execute(sql)
-            self._backend.execute(table.sql_alter_to(target_catalog))
-            self._backend.execute(table.sql_alter_from(target_catalog))
-            self._seen_tables[target] = table.key
-        elif table.object_type == "EXTERNAL":
-            result = next(self._backend.fetch(sql))
-            if result.status_code != "SUCCESS":
-                raise ValueError(result.description)
-            self._backend.execute(table.sql_alter_to(target_catalog))
-            self._backend.execute(table.sql_alter_from(target_catalog))
-            self._seen_tables[target] = table.key
-        else:
-            msg = f"Table {table.key} is a {table.object_type} and is not supported for migration yet"
-            raise ValueError(msg)
+            rule = mapping_rules.get(table.key)
+            if not rule:
+                logger.info(f"Skipping table {table.key} table doesn't exist in the mapping table.")
+                continue
+            tasks.append(partial(self._migrate_table, table, rule))
+        Threads.strict("migrate tables", tasks)
+
+    def _migrate_table(self, src_table: Table, rule: Rule):
+        if self._table_already_upgraded(rule.as_uc_table_key):
+            logger.info(f"Table {src_table.key} already upgraded to {rule.as_uc_table_key}")
+            return True
+        if src_table.object_type == "MANAGED":
+            return self._migrate_managed_table(src_table, rule)
+        if src_table.kind == "VIEW":
+            return self._migrate_view(src_table, rule)
+        if src_table.object_type == "EXTERNAL":
+            return self._migrate_external_table(src_table, rule)
+        return True
+
+    def _migrate_external_table(self, src_table: Table, rule: Rule):
+        target_table_key = rule.as_uc_table_key
+        table_migrate_sql = src_table.uc_create_sql(target_table_key)
+        logger.debug(f"Migrating external table {src_table.key} to using SQL query: {table_migrate_sql}")
+        self._backend.execute(table_migrate_sql)
+        return True
+
+    def _migrate_managed_table(self, src_table: Table, rule: Rule):
+        target_table_key = rule.as_uc_table_key
+        table_migrate_sql = src_table.uc_create_sql(target_table_key)
+        logger.debug(f"Migrating managed table {src_table.key} to using SQL query: {table_migrate_sql}")
+        self._backend.execute(table_migrate_sql)
+        self._backend.execute(src_table.sql_alter_to(rule.as_uc_table_key))
+        self._backend.execute(src_table.sql_alter_from(rule.as_uc_table_key))
+        return True
+
+    def _migrate_view(self, src_table: Table, rule: Rule):
+        target_table_key = rule.as_uc_table_key
+        table_migrate_sql = src_table.uc_create_sql(target_table_key)
+        logger.debug(f"Migrating view {src_table.key} to using SQL query: {table_migrate_sql}")
+        self._backend.execute(table_migrate_sql)
+        self._backend.execute(src_table.sql_alter_to(rule.as_uc_table_key))
+        self._backend.execute(src_table.sql_alter_from(rule.as_uc_table_key))
         return True
 
+        msg = f"Table {src_table.key} is a {src_table.object_type} and is not supported for migration yet"
+        logger.info(msg)
+
     def _init_seen_tables(self):
         for catalog in self._ws.catalogs.list():
             for schema in self._ws.schemas.list(catalog_name=catalog.name):
@@ -90,8 +95,6 @@ def _get_tables_to_revert(self, schema: str | None = None, table: str | None = N
         upgraded_tables = []
         if table and not schema:
             logger.error("Cannot accept 'Table' parameter without 'Schema' parameter")
-        if len(self._seen_tables) == 0:
-            self._init_seen_tables()
 
         for cur_table in self._tc.snapshot():
             if schema and cur_table.database != schema:
@@ -105,6 +108,7 @@ def _get_tables_to_revert(self, schema: str | None = None, table: str | None = N
     def revert_migrated_tables(
         self, schema: str | None = None, table: str | None = None, *, delete_managed: bool = False
     ):
+        self._init_seen_tables()
         upgraded_tables = self._get_tables_to_revert(schema=schema, table=table)
         # reverses the _seen_tables dictionary to key by the source table
         reverse_seen = {v: k for (k, v) in self._seen_tables.items()}
@@ -123,10 +127,11 @@ def _revert_migrated_table(self, table: Table, target_table_key: str):
         logger.info(
             f"Reverting {table.object_type} table {table.database}.{table.name} upgraded_to {table.upgraded_to}"
         )
-        self._backend.execute(table.sql_unset_upgraded_to("hive_metastore"))
+        self._backend.execute(table.sql_unset_upgraded_to())
         self._backend.execute(f"DROP {table.kind} IF EXISTS {target_table_key}")
 
     def _get_revert_count(self, schema: str | None = None, table: str | None = None) -> list[MigrationCount]:
+        self._init_seen_tables()
         upgraded_tables = self._get_tables_to_revert(schema=schema, table=table)
 
         table_by_database = defaultdict(list)
@@ -183,3 +188,9 @@ def print_revert_report(self, *, delete_managed: bool) -> bool | None:
             print("Migrated Manged Tables (targets) will be left intact.")
             print("To revert and delete Migrated Tables, add --delete_managed true flag to the command.")
         return True
+
+    def _get_mapping_rules(self) -> dict[str, Rule]:
+        mapping_rules: dict[str, Rule] = {}
+        for rule in self._tm.load():
+            mapping_rules[rule.as_hms_table_key] = rule
+        return mapping_rules

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

@@ -52,43 +52,22 @@ def key(self) -> str:
     def kind(self) -> str:
         return "VIEW" if self.view_text is not None else "TABLE"
 
-    def _sql_external(self, catalog):
-        return f"SYNC TABLE {catalog}.{self.database}.{self.name} FROM {self.key};"
-
-    def _sql_managed(self, catalog):
-        if not self.is_delta:
-            msg = f"{self.key} is not DELTA: {self.table_format}"
-            raise ValueError(msg)
-        return f"CREATE TABLE IF NOT EXISTS {catalog}.{self.database}.{self.name} DEEP CLONE {self.key};"
-
-    def _sql_view(self, catalog):
-        return f"CREATE VIEW IF NOT EXISTS {catalog}.{self.database}.{self.name} AS {self.view_text};"
-
-    def uc_create_sql(self, catalog):
+    def uc_create_sql(self, target_table_key):
         if self.kind == "VIEW":
-            return self._sql_view(catalog)
+            return self._sql_migrate_view(target_table_key)
         elif self.object_type == "EXTERNAL":
-            return self._sql_external(catalog)
+            return self._sql_migrate_external(target_table_key)
         else:
-            return self._sql_managed(catalog)
-
-    def sql_alter_to(self, catalog):
-        return (
-            f"ALTER {self.kind} {self.key} SET"
-            f" TBLPROPERTIES ('upgraded_to' = '{catalog}.{self.database}.{self.name}');"
-        )
-
-    def sql_alter_from(self, catalog):
-        return (
-            f"ALTER {self.kind} {catalog}.{self.database}.{self.name} SET"
-            f" TBLPROPERTIES ('upgraded_from' = '{self.key}');"
-        )
-
-    def sql_unset_upgraded_to(self, catalog):
-        return (
-            f"ALTER {self.kind} `{catalog}`.`{self.database}`.`{self.name}` "
-            f"UNSET TBLPROPERTIES IF EXISTS('upgraded_to');"
-        )
+            return self._sql_migrate_managed(target_table_key)
+
+    def sql_alter_to(self, target_table_key):
+        return f"ALTER {self.kind} {self.key} SET TBLPROPERTIES ('upgraded_to' = '{target_table_key}');"
+
+    def sql_alter_from(self, target_table_key):
+        return f"ALTER {self.kind} {target_table_key} SET TBLPROPERTIES ('upgraded_from' = '{self.key}');"
+
+    def sql_unset_upgraded_to(self):
+        return f"ALTER {self.kind} {self.key} UNSET TBLPROPERTIES IF EXISTS('upgraded_to');"
 
     def is_dbfs_root(self) -> bool:
         if not self.location:
@@ -101,6 +80,18 @@ def is_dbfs_root(self) -> bool:
                 return True
         return False
 
+    def _sql_migrate_external(self, target_table_key):
+        return f"SYNC TABLE {target_table_key} FROM {self.key};"
+
+    def _sql_migrate_managed(self, target_table_key):
+        if not self.is_delta:
+            msg = f"{self.key} is not DELTA: {self.table_format}"
+            raise ValueError(msg)
+        return f"CREATE TABLE IF NOT EXISTS {target_table_key} DEEP CLONE {self.key};"
+
+    def _sql_migrate_view(self, target_table_key):
+        return f"CREATE VIEW IF NOT EXISTS {target_table_key} AS {self.view_text};"
+
 
 @dataclass
 class TableError:

tests/integration/conftest.py

@@ -14,6 +14,7 @@
 
 from databricks.labs.ucx.framework.crawlers import SqlBackend
 from databricks.labs.ucx.hive_metastore import TablesCrawler
+from databricks.labs.ucx.hive_metastore.mapping import Rule
 from databricks.labs.ucx.hive_metastore.tables import Table
 from databricks.labs.ucx.mixins.fixtures import *  # noqa: F403
 from databricks.labs.ucx.workspace_access.groups import MigratedGroup
@@ -134,3 +135,11 @@ def __init__(self, sql_backend: SqlBackend, schema: str, tables: list[TableInfo]
 
     def snapshot(self) -> list[Table]:
         return self._tables
+
+
+class StaticTableMapping:
+    def __init__(self, rules: list[Rule] | None = None):
+        self._rules = rules
+
+    def load(self):
+        return self._rules