Fix test failure: test_table_migration_job_refreshes_migration_status[regular-migrate-tables]

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

@@ -52,21 +52,16 @@ def __init__(
         self._migrate_grants = migrate_grants
 
     def get_remaining_tables(self) -> list[Table]:
-        self.index_full_refresh()
+        self.index(force_refresh=True)
         table_rows = []
         for crawled_table in self._tc.snapshot():
             if not self._is_migrated(crawled_table.database, crawled_table.name):
                 table_rows.append(crawled_table)
                 logger.warning(f"remained-hive-metastore-table: {crawled_table.key}")
         return table_rows
 
-    def index(self):
-        return self._migration_status_refresher.index()
-
-    def index_full_refresh(self):
-        # when we want the latest up-to-date status, e.g. to determine whether views dependencies have been migrated
-        self._migration_status_refresher.reset()
-        return self._migration_status_refresher.index()
+    def index(self, *, force_refresh: bool = False):
+        return self._migration_status_refresher.index(force_refresh=force_refresh)
 
     def migrate_tables(
         self,
@@ -101,15 +96,17 @@ def _migrate_tables(self, what: What, mounts: list[Mount], hiveserde_in_place_mi
     def _migrate_views(self):
         tables_to_migrate = self._tm.get_tables_to_migrate(self._tc)
         all_tasks = []
-        sequencer = ViewsMigrationSequencer(tables_to_migrate, migration_index=self.index_full_refresh())
+        # Every batch of views to migrate needs an up-to-date table migration index
+        # to determine if the dependencies have been migrated
+        sequencer = ViewsMigrationSequencer(tables_to_migrate, migration_index=self.index(force_refresh=True))
         batches = sequencer.sequence_batches()
         for batch in batches:
             tasks = []
             for view in batch:
                 tasks.append(partial(self._migrate_view, view))
             Threads.strict("migrate views", tasks)
             all_tasks.extend(tasks)
-            self.index_full_refresh()
+            self.index(force_refresh=True)
         return all_tasks
 
     def _migrate_table(self, src_table: TableToMigrate, mounts: list[Mount], hiveserde_in_place_migrate: bool = False):

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

@@ -80,8 +80,8 @@ def __init__(self, ws: WorkspaceClient, sbe: SqlBackend, schema, table_crawler:
         self._ws = ws
         self._table_crawler = table_crawler
 
-    def index(self) -> TableMigrationIndex:
-        return TableMigrationIndex(list(self.snapshot()))
+    def index(self, *, force_refresh: bool = False) -> TableMigrationIndex:
+        return TableMigrationIndex(list(self.snapshot(force_refresh=force_refresh)))
 
     def get_seen_tables(self) -> dict[str, str]:
         seen_tables: dict[str, str] = {}

tests/unit/hive_metastore/test_workflows.py

@@ -1,4 +1,5 @@
 import pytest
+
 from databricks.labs.ucx.hive_metastore.workflows import (
     TableMigration,
     MigrateExternalTablesCTAS,
@@ -63,8 +64,8 @@ def test_migrate_ctas_views(run_workflow):
         MigrateTablesInMounts,
     ],
 )
-def test_update_migration_status(run_workflow, workflow):
+def test_update_migration_status(run_workflow, workflow) -> None:
     """Migration status is refreshed by deleting and showing new tables"""
     ctx = run_workflow(getattr(workflow, "update_migration_status"))
-    assert "TRUNCATE TABLE `hive_metastore`.`ucx`.`migration_status`" in ctx.sql_backend.queries
+    assert ctx.sql_backend.has_rows_written_for("hive_metastore.ucx.migration_status")
     assert "SHOW DATABASES" in ctx.sql_backend.queries