Back out changes relating to the way the migration-status information is passed into the table encoder.

This has been backed out to reduce clutter on the PR; they have been moved to PR #3270.

Author: asnare

src/databricks/labs/ucx/contexts/workflow_task.py

@@ -22,10 +22,10 @@
 from databricks.labs.ucx.contexts.application import GlobalContext
 from databricks.labs.ucx.hive_metastore import TablesInMounts, TablesCrawler
 from databricks.labs.ucx.hive_metastore.table_size import TableSizeCrawler
-from databricks.labs.ucx.hive_metastore.tables import FasterTableScanCrawler, Table
+from databricks.labs.ucx.hive_metastore.tables import FasterTableScanCrawler
 from databricks.labs.ucx.hive_metastore.udfs import Udf
 from databricks.labs.ucx.installer.logs import TaskRunWarningRecorder
-from databricks.labs.ucx.progress.grants import GrantProgressEncoder, Grant
+from databricks.labs.ucx.progress.grants import GrantProgressEncoder
 from databricks.labs.ucx.progress.history import ProgressEncoder
 from databricks.labs.ucx.progress.jobs import JobsProgressEncoder
 from databricks.labs.ucx.progress.tables import TableProgressEncoder
@@ -189,7 +189,7 @@ def policies_progress(self) -> ProgressEncoder[PolicyInfo]:
         )
 
     @cached_property
-    def grants_progress(self) -> ProgressEncoder[Grant]:
+    def grants_progress(self) -> GrantProgressEncoder:
         return GrantProgressEncoder(
             self.sql_backend,
             self.grant_ownership,
@@ -221,11 +221,11 @@ def pipelines_progress(self) -> ProgressEncoder[PipelineInfo]:
         )
 
     @cached_property
-    def tables_progress(self) -> ProgressEncoder[Table]:
+    def tables_progress(self) -> TableProgressEncoder:
         return TableProgressEncoder(
             self.sql_backend,
             self.table_ownership,
-            self.migration_status_refresher,
+            self.migration_status_refresher.index(force_refresh=False),
             self.parent_run_id,
             self.workspace_id,
             self.config.ucx_catalog,

src/databricks/labs/ucx/progress/history.py

@@ -5,7 +5,7 @@
 import logging
 from enum import Enum, EnumMeta
 from collections.abc import Iterable, Sequence
-from typing import Any, ClassVar, Generic, Protocol, TypeVar, get_type_hints
+from typing import Any, ClassVar, Generic, Protocol, TypeVar, get_type_hints, final
 
 from databricks.labs.lsql.backends import SqlBackend
 
@@ -279,6 +279,7 @@ def __init__(
     def full_name(self) -> str:
         return f"{self._catalog}.{self._schema}.{self._table}"
 
+    @final
     def append_inventory_snapshot(self, snapshot: Iterable[Record]) -> None:
         history_records = [self._encode_record_as_historical(record) for record in snapshot]
         logger.debug(f"Appending {len(history_records)} {self._klass} record(s) to history.")

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

@@ -1,19 +1,13 @@
-import logging
-from collections.abc import Iterable
 from dataclasses import replace
 
 from databricks.labs.lsql.backends import SqlBackend
-from databricks.labs.ucx.framework.crawlers import CrawlerBase
-from databricks.labs.ucx.framework.utils import escape_sql_identifier
 
 from databricks.labs.ucx.hive_metastore.tables import Table
-from databricks.labs.ucx.hive_metastore.table_migration_status import TableMigrationIndex, TableMigrationStatus
+from databricks.labs.ucx.hive_metastore.table_migration_status import TableMigrationIndex
 from databricks.labs.ucx.hive_metastore.ownership import TableOwnership
 from databricks.labs.ucx.progress.history import ProgressEncoder
 from databricks.labs.ucx.progress.install import Historical
 
-logger = logging.getLogger(__name__)
-
 
 class TableProgressEncoder(ProgressEncoder[Table]):
     """Encoder class:Table to class:History.
@@ -27,7 +21,7 @@ def __init__(
         self,
         sql_backend: SqlBackend,
         ownership: TableOwnership,
-        migration_status_refresher: CrawlerBase[TableMigrationStatus],
+        table_migration_index: TableMigrationIndex,
         run_id: int,
         workspace_id: int,
         catalog: str,
@@ -44,24 +38,17 @@ def __init__(
             schema,
             table,
         )
-        self._migration_status_refresher = migration_status_refresher
-
-    def append_inventory_snapshot(self, snapshot: Iterable[Table]) -> None:
-        migration_index = TableMigrationIndex(self._migration_status_refresher.snapshot())
-        history_records = [self._encode_table_as_historical(record, migration_index) for record in snapshot]
-        logger.debug(f"Appending {len(history_records)} {self._klass} table record(s) to history.")
-        # This is the only writer, and the mode is 'append'. This is documented as conflict-free.
-        self._sql_backend.save_table(escape_sql_identifier(self.full_name), history_records, Historical, mode="append")
+        self._table_migration_index = table_migration_index
 
-    def _encode_table_as_historical(self, record: Table, migration_index: TableMigrationIndex) -> Historical:
-        """Encode a table record, enriching with the migration status.
+    def _encode_record_as_historical(self, record: Table) -> Historical:
+        """Encode record as historical.
 
         A table failure means that the table is pending migration. Grants are purposefully left out, because a grant
         might not be mappable to UC, like `READ_METADATA`, thus possibly resulting in false "pending migration" failure
         for tables that are migrated to UC with their relevant grants also being migrated.
         """
         historical = super()._encode_record_as_historical(record)
         failures = []
-        if not migration_index.is_migrated(record.database, record.name):
+        if not self._table_migration_index.is_migrated(record.database, record.name):
             failures.append("Pending migration")
         return replace(historical, failures=historical.failures + failures)

tests/unit/progress/test_tables.py

@@ -4,11 +4,9 @@
 
 from databricks.labs.ucx.framework.owners import Ownership
 from databricks.labs.ucx.framework.utils import escape_sql_identifier
-from databricks.labs.ucx.hive_metastore.table_migration_status import (
-    TableMigrationStatusRefresher,
-    TableMigrationStatus,
-)
+from databricks.labs.ucx.hive_metastore.table_migration_status import TableMigrationIndex
 from databricks.labs.ucx.hive_metastore.tables import Table
+from databricks.labs.ucx.progress.grants import GrantProgressEncoder
 from databricks.labs.ucx.progress.tables import TableProgressEncoder
 
 
@@ -21,21 +19,21 @@
 def test_table_progress_encoder_no_failures(mock_backend, table: Table) -> None:
     ownership = create_autospec(Ownership)
     ownership.owner_of.return_value = "user"
-    migration_status_crawler = create_autospec(TableMigrationStatusRefresher)
-    migration_status_crawler.snapshot.return_value = (
-        TableMigrationStatus(table.database, table.name, "main", "default", table.name, update_ts=None),
-    )
+    table_migration_index = create_autospec(TableMigrationIndex)
+    table_migration_index.is_migrated.return_value = True
+    grant_progress_encoder = create_autospec(GrantProgressEncoder)
     encoder = TableProgressEncoder(
-        mock_backend, ownership, migration_status_crawler, run_id=1, workspace_id=123456789, catalog="test"
+        mock_backend, ownership, table_migration_index, run_id=1, workspace_id=123456789, catalog="test"
     )
 
     encoder.append_inventory_snapshot([table])
 
     rows = mock_backend.rows_written_for(escape_sql_identifier(encoder.full_name), "append")
-    assert rows, f"No rows written for: {encoder.full_name}"
+    assert len(rows) > 0, f"No rows written for: {encoder.full_name}"
     assert len(rows[0].failures) == 0
     ownership.owner_of.assert_called_once()
-    migration_status_crawler.snapshot.assert_called_once()
+    table_migration_index.is_migrated.assert_called_with(table.database, table.name)
+    grant_progress_encoder.assert_not_called()
 
 
 @pytest.mark.parametrize(
@@ -47,12 +45,11 @@ def test_table_progress_encoder_no_failures(mock_backend, table: Table) -> None:
 def test_table_progress_encoder_pending_migration_failure(mock_backend, table: Table) -> None:
     ownership = create_autospec(Ownership)
     ownership.owner_of.return_value = "user"
-    migration_status_crawler = create_autospec(TableMigrationStatusRefresher)
-    migration_status_crawler.snapshot.return_value = (
-        TableMigrationStatus(table.database, table.name),  # No destination: therefore not yet migrated.
-    )
+    table_migration_index = create_autospec(TableMigrationIndex)
+    table_migration_index.is_migrated.return_value = False
+    grant_progress_encoder = create_autospec(GrantProgressEncoder)
     encoder = TableProgressEncoder(
-        mock_backend, ownership, migration_status_crawler, run_id=1, workspace_id=123456789, catalog="test"
+        mock_backend, ownership, table_migration_index, run_id=1, workspace_id=123456789, catalog="test"
     )
 
     encoder.append_inventory_snapshot([table])
@@ -61,4 +58,5 @@ def test_table_progress_encoder_pending_migration_failure(mock_backend, table: T
     assert len(rows) > 0, f"No rows written for: {encoder.full_name}"
     assert rows[0].failures == ["Pending migration"]
     ownership.owner_of.assert_called_once()
-    migration_status_crawler.snapshot.assert_called_once()
+    table_migration_index.is_migrated.assert_called_with(table.database, table.name)
+    grant_progress_encoder.assert_not_called()