Handle Databricks errors during workspace listings in the table migration status refresher (#3378)

## Changes
Handle Databricks errors during workspace listings in the table
migration status refresher

### Linked issues

Resolves #3262

### Functionality

- [x] modified existing workflow: `assessment`

### Tests

- [x] added unit tests

Author: JCZuurmond

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

@@ -6,11 +6,12 @@
 
 from databricks.labs.lsql.backends import SqlBackend
 from databricks.sdk import WorkspaceClient
-from databricks.sdk.errors import NotFound
+from databricks.sdk.errors import DatabricksError, NotFound
+from databricks.sdk.service.catalog import CatalogInfo
 
 from databricks.labs.ucx.framework.crawlers import CrawlerBase
 from databricks.labs.ucx.framework.utils import escape_sql_identifier
-from databricks.labs.ucx.hive_metastore import TablesCrawler
+from databricks.labs.ucx.hive_metastore.tables import TablesCrawler
 
 logger = logging.getLogger(__name__)
 
@@ -93,9 +94,10 @@ def get_seen_tables(self) -> dict[str, str]:
                 # ws.tables.list returns Iterator[TableInfo], so we need to convert it to a list in order to catch the exception
                 tables = list(self._ws.tables.list(catalog_name=schema.catalog_name, schema_name=schema.name))
             except NotFound:
-                logger.warning(
-                    f"Schema {schema.catalog_name}.{schema.name} no longer exists. Skipping checking its migration status."
-                )
+                logger.warning(f"Schema {schema.full_name} no longer exists. Skipping checking its migration status.")
+                continue
+            except DatabricksError as e:
+                logger.warning(f"Error while listing tables in schema: {schema.full_name}", exc_info=e)
                 continue
             for table in tables:
                 if not table.properties:
@@ -153,10 +155,19 @@ def _try_fetch(self) -> Iterable[TableMigrationStatus]:
         for row in self._fetch(f"SELECT * FROM {escape_sql_identifier(self.full_name)}"):
             yield TableMigrationStatus(*row)
 
+    def _iter_catalogs(self) -> Iterable[CatalogInfo]:
+        try:
+            yield from self._ws.catalogs.list()
+        except DatabricksError as e:
+            logger.error("Cannot list catalogs", exc_info=e)
+
     def _iter_schemas(self):
-        for catalog in self._ws.catalogs.list():
+        for catalog in self._iter_catalogs():
             try:
                 yield from self._ws.schemas.list(catalog_name=catalog.name)
             except NotFound:
                 logger.warning(f"Catalog {catalog.name} no longer exists. Skipping checking its migration status.")
                 continue
+            except DatabricksError as e:
+                logger.warning(f"Error while listing schemas in catalog: {catalog.name}", exc_info=e)
+                continue

tests/unit/hive_metastore/test_table_migrate.py

@@ -1044,7 +1044,10 @@ def test_table_status_seen_tables(caplog):
     client = create_autospec(WorkspaceClient)
     client.catalogs.list.return_value = [CatalogInfo(name="cat1"), CatalogInfo(name="deleted_cat")]
     client.schemas.list.side_effect = [
-        [SchemaInfo(catalog_name="cat1", name="schema1"), SchemaInfo(catalog_name="cat1", name="deleted_schema")],
+        [
+            SchemaInfo(catalog_name="cat1", name="schema1", full_name="cat1.schema1"),
+            SchemaInfo(catalog_name="cat1", name="deleted_schema", full_name="cat1.deleted_schema"),
+        ],
         NotFound(),
     ]
     client.tables.list.side_effect = [

tests/unit/hive_metastore/test_table_migration_status.py

@@ -0,0 +1,66 @@
+from unittest.mock import create_autospec
+
+import pytest
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import BadRequest, DatabricksError, NotFound
+from databricks.sdk.service.catalog import CatalogInfo, SchemaInfo
+
+from databricks.labs.ucx.hive_metastore.tables import TablesCrawler
+from databricks.labs.ucx.hive_metastore.table_migration_status import TableMigrationStatusRefresher
+
+
+def test_table_migration_status_refresher_get_seen_tables_handles_errors_on_catalogs_list(mock_backend) -> None:
+    ws = create_autospec(WorkspaceClient)
+    ws.catalogs.list.side_effect = BadRequest()
+    tables_crawler = create_autospec(TablesCrawler)
+
+    refresher = TableMigrationStatusRefresher(ws, mock_backend, "test", tables_crawler)
+
+    seen_tables = refresher.get_seen_tables()
+
+    assert not seen_tables
+    ws.catalogs.list.assert_called_once()
+    ws.schemas.list.assert_not_called()
+    ws.tables.list.assert_not_called()
+    tables_crawler.snapshot.assert_not_called()
+
+
+@pytest.mark.parametrize("error", [BadRequest(), NotFound()])
+def test_table_migration_status_refresher_get_seen_tables_handles_errors_on_schemas_list(
+    mock_backend, error: DatabricksError
+) -> None:
+    ws = create_autospec(WorkspaceClient)
+    ws.catalogs.list.return_value = [CatalogInfo(name="test")]
+    ws.schemas.list.side_effect = error
+    tables_crawler = create_autospec(TablesCrawler)
+
+    refresher = TableMigrationStatusRefresher(ws, mock_backend, "test", tables_crawler)
+
+    seen_tables = refresher.get_seen_tables()
+
+    assert not seen_tables
+    ws.catalogs.list.assert_called_once()
+    ws.schemas.list.assert_called_once()
+    ws.tables.list.assert_not_called()
+    tables_crawler.snapshot.assert_not_called()
+
+
+@pytest.mark.parametrize("error", [BadRequest(), NotFound()])
+def test_table_migration_status_refresher_get_seen_tables_handles_errors_on_tables_list(
+    mock_backend, error: DatabricksError
+) -> None:
+    ws = create_autospec(WorkspaceClient)
+    ws.catalogs.list.return_value = [CatalogInfo(name="test")]
+    ws.schemas.list.return_value = [SchemaInfo(catalog_name="test", name="test")]
+    ws.tables.list.side_effect = error
+    tables_crawler = create_autospec(TablesCrawler)
+
+    refresher = TableMigrationStatusRefresher(ws, mock_backend, "test", tables_crawler)
+
+    seen_tables = refresher.get_seen_tables()
+
+    assert not seen_tables
+    ws.catalogs.list.assert_called_once()
+    ws.schemas.list.assert_called_once()
+    ws.tables.list.assert_called_once()
+    tables_crawler.snapshot.assert_not_called()