Added scope for UC crawling.

Author: FastLee

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

@@ -123,7 +123,7 @@ def _migrate_tables(
     ):
         tables_to_migrate = self._table_mapping.get_tables_to_migrate(self._tables_crawler, check_uc_table)
         tables_in_scope = filter(lambda t: t.src.what == what, tables_to_migrate)
-        self._init_seen_tables(scope={t.rule.as_uc_table_key for t in tables_to_migrate})
+        self._init_seen_tables(scope={t.rule.as_uc_table for t in tables_to_migrate})
         tasks = []
         for table in tables_in_scope:
             tasks.append(
@@ -597,7 +597,7 @@ def print_revert_report(
             print("To revert and delete Migrated Tables, add --delete_managed true flag to the command")
         return True
 
-    def _init_seen_tables(self, *, scope: set[str] | None = None):
+    def _init_seen_tables(self, *, scope: set[Table] | None = None):
         self._seen_tables = self._migration_status_refresher.get_seen_tables(scope=scope)
 
     def _sql_alter_to(self, table: Table, target_table_key: str):

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

@@ -13,10 +13,11 @@
 
 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 TablesCrawler
+from databricks.labs.ucx.hive_metastore.tables import TablesCrawler, Table
 
 logger = logging.getLogger(__name__)
 
+
 @dataclass
 class TableMigrationStatus:
     src_schema: str
@@ -95,11 +96,17 @@ def __init__(self, ws: WorkspaceClient, sql_backend: SqlBackend, schema, tables_
     def index(self, *, force_refresh: bool = False) -> TableMigrationIndex:
         return TableMigrationIndex(self.snapshot(force_refresh=force_refresh))
 
-    def get_seen_tables(self, *, scope: set[str] | None = None) -> dict[str, str]:
+    def get_seen_tables(self, *, scope: set[Table] | None = None) -> dict[str, str]:
         seen_tables: dict[str, str] = {}
         tasks = []
+        schema_scope = {(table.catalog.lower(), table.database.lower()) for table in scope} if scope else None
+        table_scope = {table.full_name.lower() for table in scope} if scope else None
         for schema in self._iter_schemas():
-            if schema.catalog_name is None or schema.name is None:
+            if (
+                schema.catalog_name is None
+                or schema.name is None
+                or (schema_scope and (schema.catalog_name.lower(), schema.name.lower()) not in schema_scope)
+            ):
                 continue
             tasks.append(
                 partial(
@@ -121,7 +128,7 @@ def get_seen_tables(self, *, scope: set[str] | None = None) -> dict[str, str]:
                 if not table.full_name:
                     logger.warning(f"The table {table.name} in {schema.name} has no full name")
                     continue
-                if scope and table.full_name not in scope:
+                if table_scope and table.full_name.lower() not in table_scope:
                     continue
                 if not table.properties or "upgraded_from" not in table.properties:
                     continue
@@ -183,7 +190,10 @@ def _iter_catalogs(self) -> Iterable[CatalogInfo]:
 
     def _iter_schemas(self) -> Iterable[SchemaInfo]:
         for catalog in self._iter_catalogs():
-            if catalog.name is None or catalog.catalog_type in (CatalogType.DELTASHARING_CATALOG, CatalogType.SYSTEM_CATALOG):
+            if catalog.name is None or catalog.catalog_type in (
+                CatalogType.DELTASHARING_CATALOG,
+                CatalogType.SYSTEM_CATALOG,
+            ):
                 continue
             try:
                 yield from self._ws.schemas.list(catalog_name=catalog.name)
@@ -199,7 +209,9 @@ def _iter_tables(self, catalog_name: str, schema_name: str) -> list[TableInfo]:
             # ws.tables.list returns Iterator[TableInfo], so we need to convert it to a list in order to catch the exception
             return list(self._ws.tables.list(catalog_name=catalog_name, schema_name=schema_name))
         except NotFound:
-            logger.warning(f"Schema {catalog_name}.{schema_name} no longer exists. Skipping checking its migration status.")
+            logger.warning(
+                f"Schema {catalog_name}.{schema_name} no longer exists. Skipping checking its migration status."
+            )
             return []
         except DatabricksError as e:
             logger.warning(f"Error while listing tables in schema: {schema_name}", exc_info=e)

tests/unit/hive_metastore/test_table_migrate.py

@@ -1043,24 +1043,24 @@ def test_table_status_seen_tables(caplog):
     table_crawler = create_autospec(TablesCrawler)
     client = create_autospec(WorkspaceClient)
     client.catalogs.list.return_value = [CatalogInfo(name="cat1"), CatalogInfo(name="deleted_cat")]
-    schemas = { "cat1":
-        [
+    schemas = {
+        "cat1": [
             SchemaInfo(catalog_name="cat1", name="schema1", full_name="cat1.schema1"),
             SchemaInfo(catalog_name="cat1", name="deleted_schema", full_name="cat1.deleted_schema"),
         ],
-        "deleted_cat":
-        None,
-}
+        "deleted_cat": None,
+    }
 
     def schema_list(catalog_name):
         schema = schemas[catalog_name]
         if not schema:
             raise NotFound()
         return schema
+
     client.schemas.list = schema_list
 
-    tables = { ("cat1", "schema1"):
-        [
+    tables = {
+        ("cat1", "schema1"): [
             TableInfo(
                 catalog_name="cat1",
                 schema_name="schema1",
@@ -1097,13 +1097,13 @@ def schema_list(catalog_name):
         ],
         ("cat1", "deleted_schema"): None,
     }
+
     def table_list(catalog_name, schema_name):
         table = tables[(catalog_name, schema_name)]
         if not table:
             raise NotFound()
         return table
 
-
     client.tables.list = table_list
     table_status_crawler = TableMigrationStatusRefresher(client, backend, "ucx", table_crawler)
     seen_tables = table_status_crawler.get_seen_tables()

tests/unit/hive_metastore/test_table_migration_status.py

@@ -6,7 +6,7 @@
 from databricks.sdk.errors import BadRequest, DatabricksError, NotFound
 from databricks.sdk.service.catalog import CatalogInfoSecurableKind, CatalogInfo, SchemaInfo, TableInfo
 
-from databricks.labs.ucx.hive_metastore.tables import TablesCrawler
+from databricks.labs.ucx.hive_metastore.tables import TablesCrawler, Table
 from databricks.labs.ucx.hive_metastore.table_migration_status import TableMigrationStatusRefresher
 
 
@@ -163,8 +163,9 @@ def tables_list(catalog_name: str, schema_name: str) -> Iterable[TableInfo]:
     tables_crawler = create_autospec(TablesCrawler)
     refresher = TableMigrationStatusRefresher(ws, mock_backend, "test", tables_crawler)
 
+    scope_table = Table("test1", "test1", "test1", "Table", "Delta")
     # Test with scope
-    assert refresher.get_seen_tables(scope={"test1.test1.test1"}) == {"test1.test1.test1": "test1"}
+    assert refresher.get_seen_tables(scope={scope_table}) == {"test1.test1.test1": "test1"}
     # Test without scope
     assert refresher.get_seen_tables() == {"test1.test1.test1": "test1", "test2.test2.test2": "test2"}
     ws.tables.list.assert_called()