Added migration_status table to capture a snapshot of migrated tables (#1041)

Author: FastLee

src/databricks/labs/ucx/framework/crawlers.py

@@ -47,6 +47,13 @@ def full_name(self) -> str:
         """
         return f"{self._catalog}.{self._schema}.{self._table}"
 
+    def reset(self):
+        """
+        Delete the content of the inventory table.
+        The next call to `snapshot` will re-populate the table.
+        """
+        self._exec(f"DELETE FROM {self.full_name}")
+
     @staticmethod
     def _valid(name: str) -> str:
         """

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

@@ -1,5 +1,8 @@
+import datetime
 import logging
 from collections import defaultdict
+from collections.abc import Iterable
+from dataclasses import dataclass
 from functools import partial
 
 from databricks.labs.blueprint.installation import Installation
@@ -15,25 +18,38 @@
 )
 
 from databricks.labs.ucx.config import WorkspaceConfig
+from databricks.labs.ucx.framework.crawlers import CrawlerBase
 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, What
 
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class MigrationStatus:
+    src_schema: str
+    src_table: str
+    dst_catalog: str | None = None
+    dst_schema: str | None = None
+    dst_table: str | None = None
+    update_ts: str | None = None
+
+
 class TablesMigrate:
     def __init__(
         self,
         tables_crawler: TablesCrawler,
         ws: WorkspaceClient,
         backend: SqlBackend,
         table_mapping: TableMapping,
+        migration_status_refresher,
     ):
         self._tc = tables_crawler
         self._backend = backend
         self._ws = ws
         self._tm = table_mapping
+        self._migration_status_refresher = migration_status_refresher
         self._seen_tables: dict[str, str] = {}
 
     @classmethod
@@ -43,7 +59,8 @@ def for_cli(cls, ws: WorkspaceClient, product='ucx'):
         sql_backend = StatementExecutionBackend(ws, config.warehouse_id)
         table_crawler = TablesCrawler(sql_backend, config.inventory_database)
         table_mapping = TableMapping(installation, ws, sql_backend)
-        return cls(table_crawler, ws, sql_backend, table_mapping)
+        migration_status_refresher = MigrationStatusRefresher(ws, sql_backend, config.inventory_database, table_crawler)
+        return cls(table_crawler, ws, sql_backend, table_mapping, migration_status_refresher)
 
     def migrate_tables(self, *, what: What | None = None):
         self._init_seen_tables()
@@ -93,19 +110,6 @@ def _migrate_view(self, src_table: Table, rule: Rule):
         self._backend.execute(src_table.sql_alter_from(rule.as_uc_table_key, self._ws.get_workspace_id()))
         return True
 
-    def _iter_schemas(self):
-        for catalog in self._ws.catalogs.list():
-            yield from self._ws.schemas.list(catalog_name=catalog.name)
-
-    def _init_seen_tables(self):
-        for schema in self._iter_schemas():
-            for table in self._ws.tables.list(catalog_name=schema.catalog_name, schema_name=schema.name):
-                if table.properties is None:
-                    continue
-                if "upgraded_from" not in table.properties:
-                    continue
-                self._seen_tables[table.full_name.lower()] = table.properties["upgraded_from"].lower()
-
     def _table_already_upgraded(self, target) -> bool:
         return target in self._seen_tables
 
@@ -169,13 +173,7 @@ def _get_revert_count(self, schema: str | None = None, table: str | None = None)
         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
+        return self._migration_status_refresher.is_upgraded(schema, table)
 
     def print_revert_report(self, *, delete_managed: bool) -> bool | None:
         migrated_count = self._get_revert_count()
@@ -215,6 +213,9 @@ def print_revert_report(self, *, delete_managed: bool) -> bool | None:
             print("To revert and delete Migrated Tables, add --delete_managed true flag to the command")
         return True
 
+    def _init_seen_tables(self):
+        self._seen_tables = self._migration_status_refresher.get_seen_tables()
+
 
 class TableMove:
     def __init__(self, ws: WorkspaceClient, backend: SqlBackend):
@@ -458,3 +459,62 @@ def _recreate_view(self, to_view_name, view_text):
         create_sql = f"CREATE VIEW {to_view_name} AS {view_text}"
         logger.info(f"Creating view {to_view_name}")
         self._backend.execute(create_sql)
+
+
+class MigrationStatusRefresher(CrawlerBase[MigrationStatus]):
+    def __init__(self, ws: WorkspaceClient, sbe: SqlBackend, schema, table_crawler: TablesCrawler):
+        super().__init__(sbe, "hive_metastore", schema, "migration_status", MigrationStatus)
+        self._ws = ws
+        self._table_crawler = table_crawler
+
+    def snapshot(self) -> Iterable[MigrationStatus]:
+        return self._snapshot(self._try_fetch, self._crawl)
+
+    def get_seen_tables(self) -> dict[str, str]:
+        seen_tables: dict[str, str] = {}
+        for schema in self._iter_schemas():
+            for table in self._ws.tables.list(catalog_name=schema.catalog_name, schema_name=schema.name):
+                if not table.properties:
+                    continue
+                if "upgraded_from" not in table.properties:
+                    continue
+                if not table.full_name:
+                    logger.warning(f"The table {table.name} in {schema.name} has no full name")
+                    continue
+                seen_tables[table.full_name.lower()] = table.properties["upgraded_from"].lower()
+        return seen_tables
+
+    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 _crawl(self) -> Iterable[MigrationStatus]:
+        all_tables = self._table_crawler.snapshot()
+        reverse_seen = {v: k for k, v in self.get_seen_tables().items()}
+        timestamp = datetime.datetime.now(datetime.timezone.utc).timestamp()
+        for table in all_tables:
+            table_migration_status = MigrationStatus(
+                src_schema=table.database,
+                src_table=table.name,
+                update_ts=str(timestamp),
+            )
+            if table.key in reverse_seen and self.is_upgraded(table.database, table.name):
+                target_table = reverse_seen[table.key]
+                if len(target_table.split(".")) == 3:
+                    table_migration_status.dst_catalog = target_table.split(".")[0]
+                    table_migration_status.dst_schema = target_table.split(".")[1]
+                    table_migration_status.dst_table = target_table.split(".")[2]
+            yield table_migration_status
+
+    def _try_fetch(self) -> Iterable[MigrationStatus]:
+        for row in self._fetch(f"SELECT * FROM {self._schema}.{self._table}"):
+            yield MigrationStatus(*row)
+
+    def _iter_schemas(self):
+        for catalog in self._ws.catalogs.list():
+            yield from self._ws.schemas.list(catalog_name=catalog.name)

src/databricks/labs/ucx/install.py

@@ -66,6 +66,7 @@
 from databricks.labs.ucx.framework.tasks import _TASKS, Task
 from databricks.labs.ucx.hive_metastore.grants import Grant
 from databricks.labs.ucx.hive_metastore.locations import ExternalLocation, Mount
+from databricks.labs.ucx.hive_metastore.table_migrate import MigrationStatus
 from databricks.labs.ucx.hive_metastore.table_size import TableSize
 from databricks.labs.ucx.hive_metastore.tables import Table, TableError
 from databricks.labs.ucx.installer.hms_lineage import HiveMetastoreLineageEnabler
@@ -159,6 +160,7 @@ def deploy_schema(sql_backend: SqlBackend, inventory_schema: str):
             functools.partial(table, "workspace_objects", WorkspaceObjectInfo),
             functools.partial(table, "permissions", Permissions),
             functools.partial(table, "submit_runs", SubmitRunInfo),
+            functools.partial(table, "migration_status", MigrationStatus),
         ],
     )
     deployer.deploy_view("objects", "queries/views/objects.sql")

tests/integration/hive_metastore/test_migrate.py

@@ -6,7 +6,10 @@
 from databricks.sdk.retries import retried
 
 from databricks.labs.ucx.hive_metastore.mapping import Rule
-from databricks.labs.ucx.hive_metastore.table_migrate import TablesMigrate
+from databricks.labs.ucx.hive_metastore.table_migrate import (
+    MigrationStatusRefresher,
+    TablesMigrate,
+)
 from databricks.labs.ucx.hive_metastore.tables import Table
 
 from ..conftest import StaticTableMapping, StaticTablesCrawler
@@ -38,7 +41,8 @@ def test_migrate_managed_tables(ws, sql_backend, inventory_schema, make_catalog,
         ),
     ]
     table_mapping = StaticTableMapping(ws, sql_backend, rules=rules)
-    table_migrate = TablesMigrate(table_crawler, ws, sql_backend, table_mapping)
+    migration_status_refresher = MigrationStatusRefresher(ws, sql_backend, inventory_schema, table_crawler)
+    table_migrate = TablesMigrate(table_crawler, ws, sql_backend, table_mapping, migration_status_refresher)
 
     table_migrate.migrate_tables()
 
@@ -94,7 +98,8 @@ def test_migrate_tables_with_cache_should_not_create_table(
         ),
     ]
     table_mapping = StaticTableMapping(ws, sql_backend, rules=rules)
-    table_migrate = TablesMigrate(table_crawler, ws, sql_backend, table_mapping)
+    migration_status_refresher = MigrationStatusRefresher(ws, sql_backend, inventory_schema, table_crawler)
+    table_migrate = TablesMigrate(table_crawler, ws, sql_backend, table_mapping, migration_status_refresher)
 
     # FIXME: flaky: databricks.sdk.errors.platform.NotFound: Catalog 'ucx_cjazg' does not exist.
     table_migrate.migrate_tables()
@@ -110,16 +115,11 @@ def test_migrate_external_table(ws, sql_backend, inventory_schema, make_catalog,
     if not ws.config.is_azure:
         pytest.skip("temporary: only works in azure test env")
     src_schema = make_schema(catalog_name="hive_metastore")
-
     mounted_location = f'dbfs:/mnt/{env_or_skip("TEST_MOUNT_NAME")}/a/b/c'
     src_external_table = make_table(schema_name=src_schema.name, external_csv=mounted_location)
-
     dst_catalog = make_catalog()
     dst_schema = make_schema(catalog_name=dst_catalog.name, name=src_schema.name)
-
     logger.info(f"dst_catalog={dst_catalog.name}, external_table={src_external_table.full_name}")
-
-    # crawler = TablesCrawler(sql_backend, inventory_schema)
     table_crawler = StaticTablesCrawler(sql_backend, inventory_schema, [src_external_table])
     rules = [
         Rule(
@@ -131,8 +131,10 @@ def test_migrate_external_table(ws, sql_backend, inventory_schema, make_catalog,
             src_external_table.name,
         ),
     ]
-    table_mapping = StaticTableMapping(ws, sql_backend, rules=rules)
-    table_migrate = TablesMigrate(table_crawler, ws, sql_backend, table_mapping)
+    migration_status_refresher = MigrationStatusRefresher(ws, sql_backend, inventory_schema, table_crawler)
+    table_migrate = TablesMigrate(
+        table_crawler, ws, sql_backend, StaticTableMapping(ws, sql_backend, rules=rules), migration_status_refresher
+    )
 
     table_migrate.migrate_tables()
 
@@ -142,6 +144,14 @@ def test_migrate_external_table(ws, sql_backend, inventory_schema, make_catalog,
     assert target_table_properties["upgraded_from"] == src_external_table.full_name
     assert target_table_properties[Table.UPGRADED_FROM_WS_PARAM] == str(ws.get_workspace_id())
 
+    migration_status = MigrationStatusRefresher(ws, sql_backend, inventory_schema, table_crawler).snapshot()
+    assert len(migration_status) == 1
+    assert migration_status[0].src_schema == src_external_table.schema_name
+    assert migration_status[0].src_table == src_external_table.name
+    assert migration_status[0].dst_catalog == dst_catalog.name
+    assert migration_status[0].dst_schema == dst_schema.name
+    assert migration_status[0].dst_table == src_external_table.name
+
 
 @retried(on=[NotFound], timeout=timedelta(minutes=5))
 def test_revert_migrated_table(
@@ -178,7 +188,8 @@ def test_revert_migrated_table(
         ),
     ]
     table_mapping = StaticTableMapping(ws, sql_backend, rules=rules)
-    table_migrate = TablesMigrate(table_crawler, ws, sql_backend, table_mapping)
+    migration_status_refresher = MigrationStatusRefresher(ws, sql_backend, inventory_schema, table_crawler)
+    table_migrate = TablesMigrate(table_crawler, ws, sql_backend, table_mapping, migration_status_refresher)
     table_migrate.migrate_tables()
 
     table_migrate.revert_migrated_tables(src_schema1.name, delete_managed=True)
@@ -284,7 +295,8 @@ def test_mapping_reverts_table(
         ),
     ]
     table_mapping = StaticTableMapping(ws, sql_backend, rules=rules)
-    table_migrate = TablesMigrate(table_crawler, ws, sql_backend, table_mapping)
+    migration_status_refresher = MigrationStatusRefresher(ws, sql_backend, inventory_schema, table_crawler)
+    table_migrate = TablesMigrate(table_crawler, ws, sql_backend, table_mapping, migration_status_refresher)
     table_migrate.migrate_tables()
 
     target_table_properties = ws.tables.get(f"{dst_schema.full_name}.{table_to_skip.name}").properties