Added databricks labs ucx revert-migrated-table command (#729)

FastLee · web-flow · commit 473f13741d59 · 2023-12-28T10:02:49.000+01:00
diff --git a/labs.yml b/labs.yml
@@ -49,3 +49,13 @@ commands:
 
   - name: validate-external-locations
     description: validates and provides mapping to external table to external location and shared generation tf scripts
+
+  - name: revert-migrated-tables
+    description: remove notation on a migrated table for re-migration
+    flags:
+      - name: schema
+        description: Schema to revert (if left blank all schemas in the workspace will be reverted)
+      - name: table
+        description: Table to revert (if left blank all tables in the schema will be reverted). Requires schema parameter to be specified.)
+      - name: delete_managed
+        description: Revert and delete managed tables
diff --git a/src/databricks/labs/ucx/cli.py b/src/databricks/labs/ucx/cli.py
@@ -11,6 +11,7 @@
 from databricks.labs.ucx.framework.tui import Prompts
 from databricks.labs.ucx.hive_metastore import ExternalLocations, TablesCrawler
 from databricks.labs.ucx.hive_metastore.mapping import TableMapping
+from databricks.labs.ucx.hive_metastore.tables import TablesMigrate
 from databricks.labs.ucx.install import WorkspaceInstaller
 from databricks.labs.ucx.installer import InstallationManager
 
@@ -113,6 +114,31 @@ def ensure_assessment_run():
         workspace_installer.validate_and_run("assessment")
 
 
+def revert_migrated_tables(schema: str, table: str, *, delete_managed: bool = False):
+    ws = WorkspaceClient()
+    prompts = Prompts()
+    installation_manager = InstallationManager(ws)
+    installation = installation_manager.for_user(ws.current_user.me())
+    if not schema and not table:
+        if not prompts.confirm(
+            "You haven't specified a schema or a table. All migrated tables will be reverted."
+            " Would you like to continue?",
+            max_attempts=2,
+        ):
+            return None
+    if not installation:
+        logger.error(CANT_FIND_UCX_MSG)
+        return None
+    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)
+    if tm.print_revert_report(delete_managed=delete_managed) and prompts.confirm(
+        "Would you like to continue?", max_attempts=2
+    ):
+        tm.revert_migrated_tables(schema, table, delete_managed=delete_managed)
+
+
 MAPPING = {
     "open-remote-config": open_remote_config,
     "installations": list_installations,
@@ -123,6 +149,7 @@ def ensure_assessment_run():
     "validate-external-locations": validate_external_locations,
     "ensure-assessment-run": ensure_assessment_run,
     "skip": skip,
+    "revert-migrated-tables": revert_migrated_tables,
 }
 
 
diff --git a/src/databricks/labs/ucx/hive_metastore/tables.py b/src/databricks/labs/ucx/hive_metastore/tables.py
@@ -1,5 +1,6 @@
 import logging
 import re
+from collections import defaultdict
 from collections.abc import Iterable, Iterator
 from dataclasses import dataclass
 from functools import partial
@@ -73,6 +74,12 @@ def sql_alter_from(self, catalog):
             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');"
+        )
+
 
 @dataclass
 class TableError:
@@ -82,6 +89,14 @@ class TableError:
     error: str | None = None
 
 
+@dataclass
+class MigrationCount:
+    database: str
+    managed_tables: int = 0
+    external_tables: int = 0
+    views: int = 0
+
+
 class TablesCrawler(CrawlerBase):
     def __init__(self, backend: SqlBackend, schema):
         """
@@ -252,3 +267,103 @@ def _init_seen_tables(self):
 
     def _table_already_upgraded(self, target) -> bool:
         return target in self._seen_tables
+
+    def _get_tables_to_revert(self, schema: str | None = None, table: str | None = None) -> list[Table]:
+        schema = schema.lower() if schema else None
+        table = table.lower() if table else None
+        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:
+                continue
+            if table and cur_table.name != table:
+                continue
+            if cur_table.key in self._seen_tables.values():
+                upgraded_tables.append(cur_table)
+        return upgraded_tables
+
+    def revert_migrated_tables(
+        self, schema: str | None = None, table: str | None = None, *, delete_managed: bool = False
+    ):
+        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()}
+        tasks = []
+        for upgraded_table in upgraded_tables:
+            if upgraded_table.kind == "VIEW" or upgraded_table.object_type == "EXTERNAL" or delete_managed:
+                tasks.append(partial(self._revert_migrated_table, upgraded_table, reverse_seen[upgraded_table.key]))
+                continue
+            logger.info(
+                f"Skipping {upgraded_table.object_type} Table {upgraded_table.database}.{upgraded_table.name} "
+                f"upgraded_to {upgraded_table.upgraded_to}"
+            )
+        Threads.strict("revert migrated tables", tasks)
+
+    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(f"DROP {table.kind} IF EXISTS {target_table_key}")
+
+    def _get_revert_count(self, schema: str | None = None, table: str | None = None) -> list[MigrationCount]:
+        upgraded_tables = self._get_tables_to_revert(schema=schema, table=table)
+
+        table_by_database = defaultdict(list)
+        for cur_table in upgraded_tables:
+            table_by_database[cur_table.database].append(cur_table)
+
+        migration_list = []
+        for cur_database in table_by_database.keys():
+            external_tables = 0
+            managed_tables = 0
+            views = 0
+            for current_table in table_by_database[cur_database]:
+                if current_table.upgraded_to is not None:
+                    if current_table.kind == "VIEW":
+                        views += 1
+                        continue
+                    if current_table.object_type == "EXTERNAL":
+                        external_tables += 1
+                        continue
+                    if current_table.object_type == "MANAGED":
+                        managed_tables += 1
+                        continue
+            migration_list.append(
+                MigrationCount(
+                    database=cur_database, managed_tables=managed_tables, external_tables=external_tables, views=views
+                )
+            )
+        return migration_list
+
+    def is_upgraded(self, schema: str, table: str) -> bool:
+        result = self._backend.fetch(f"SHOW TBLPROPERTIES `{schema}`.`{table}`")
+        for value in result:
+            if value["key"] == "upgraded_to":
+                logger.info(f"{schema}.{table} is set as upgraded")
+                return True
+        logger.info(f"{schema}.{table} is set as not upgraded")
+        return False
+
+    def print_revert_report(self, *, delete_managed: bool) -> bool | None:
+        migrated_count = self._get_revert_count()
+        if not migrated_count:
+            logger.info("No migrated tables were found.")
+            return False
+        print("The following is the count of migrated tables and views found in scope:")
+        print("Database                      | External Tables  | Managed Table    | Views            |")
+        print("=" * 88)
+        for count in migrated_count:
+            print(f"{count.database:<30}| {count.external_tables:16} | {count.managed_tables:16} | {count.views:16} |")
+        print("=" * 88)
+        print("Migrated External Tables and Views (targets) will be deleted")
+        if delete_managed:
+            print("Migrated Manged Tables (targets) will be deleted")
+        else:
+            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
diff --git a/tests/integration/hive_metastore/test_migrate.py b/tests/integration/hive_metastore/test_migrate.py
@@ -101,3 +101,39 @@ def test_migrate_external_table(ws, sql_backend, inventory_schema, make_catalog,
 
     target_tables = list(sql_backend.fetch(f"SHOW TABLES IN {dst_schema.full_name}"))
     assert len(target_tables) == 1
+
+
+@retried(on=[NotFound], timeout=timedelta(minutes=5))
+def test_revert_migrated_table(ws, sql_backend, inventory_schema, make_schema, make_table, make_catalog):
+    src_schema1 = make_schema(catalog_name="hive_metastore")
+    src_schema2 = make_schema(catalog_name="hive_metastore")
+    table_to_revert = make_table(schema_name=src_schema1.name)
+    table_not_migrated = make_table(schema_name=src_schema1.name)
+    table_to_not_revert = make_table(schema_name=src_schema2.name)
+    all_tables = [table_to_revert, table_not_migrated, table_to_not_revert]
+
+    dst_catalog = make_catalog()
+    dst_schema1 = make_schema(catalog_name=dst_catalog.name, name=src_schema1.name)
+    dst_schema2 = make_schema(catalog_name=dst_catalog.name, name=src_schema2.name)
+
+    static_crawler = StaticTablesCrawler(sql_backend, inventory_schema, all_tables)
+    tm = TablesMigrate(static_crawler, ws, sql_backend, dst_catalog.name)
+    tm.migrate_tables()
+
+    tm.revert_migrated_tables(src_schema1.name, delete_managed=True)
+
+    # Checking that two of the tables were reverted and one was left intact.
+    # The first two table belongs to schema 1 and should have not "upgraded_to" property
+    assert not tm.is_upgraded(table_to_revert.schema_name, table_to_revert.name)
+    # The second table didn't have the "upgraded_to" property set and should remain that way.
+    assert not tm.is_upgraded(table_not_migrated.schema_name, table_not_migrated.name)
+    # The third table belongs to schema2 and had the "upgraded_to" property set and should remain that way.
+    assert tm.is_upgraded(table_to_not_revert.schema_name, table_to_not_revert.name)
+
+    target_tables_schema1 = list(sql_backend.fetch(f"SHOW TABLES IN {dst_schema1.full_name}"))
+    assert len(target_tables_schema1) == 0
+
+    target_tables_schema2 = list(sql_backend.fetch(f"SHOW TABLES IN {dst_schema2.full_name}"))
+    assert len(target_tables_schema2) == 1
+    assert target_tables_schema2[0]["database"] == dst_schema2.name
+    assert target_tables_schema2[0]["tableName"] == table_to_not_revert.name
diff --git a/tests/unit/hive_metastore/test_migrate.py b/tests/unit/hive_metastore/test_migrate.py
