Extend error handling of delta issues in crawlers and hive metastore (#795)

Author: mwojtyczka

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

@@ -8,7 +8,13 @@
 from typing import Any, ClassVar, Generic, Protocol, TypeVar
 
 from databricks.sdk import WorkspaceClient
-from databricks.sdk.errors import BadRequest, NotFound, PermissionDenied, Unknown
+from databricks.sdk.errors import (
+    BadRequest,
+    DataLoss,
+    NotFound,
+    PermissionDenied,
+    Unknown,
+)
 
 from databricks.labs.ucx.mixins.sql import Row, StatementExecutionExt
 
@@ -183,6 +189,10 @@ def _raise_spark_sql_exceptions(error_message: str):
             raise NotFound(error_message) from None
         elif "TABLE_OR_VIEW_NOT_FOUND" in error_message:
             raise NotFound(error_message) from None
+        elif "DELTA_TABLE_NOT_FOUND" in error_message:
+            raise NotFound(error_message) from None
+        elif "DELTA_MISSING_TRANSACTION_LOG" in error_message:
+            raise DataLoss(error_message) from None
         elif "PARSE_SYNTAX_ERROR" in error_message:
             raise BadRequest(error_message) from None
         elif "Operation not allowed" in error_message:

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

@@ -102,12 +102,12 @@ def skip_table(self, schema: str, table: str):
                 f"ALTER TABLE `{schema}`.`{table}` SET TBLPROPERTIES('{self.UCX_SKIP_PROPERTY}' = true)"
             )
         except NotFound as nf:
-            if "[TABLE_OR_VIEW_NOT_FOUND]" in str(nf):
+            if "[TABLE_OR_VIEW_NOT_FOUND]" in str(nf) or "[DELTA_TABLE_NOT_FOUND]" in str(nf):
                 logger.error(f"Failed to apply skip marker for Table {schema}.{table}. Table not found.")
             else:
-                logger.error(nf)
+                logger.error(f"Failed to apply skip marker for Table {schema}.{table}: {nf!s}", exc_info=True)
         except BadRequest as br:
-            logger.error(br)
+            logger.error(f"Failed to apply skip marker for Table {schema}.{table}: {br!s}", exc_info=True)
 
     def skip_schema(self, schema: str):
         # Marks a schema to be skipped in the migration process by applying a table property

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

@@ -279,10 +279,10 @@ def _move_table(
 
             return True
         except NotFound as err:
-            if "[TABLE_OR_VIEW_NOT_FOUND]" in str(err):
+            if "[TABLE_OR_VIEW_NOT_FOUND]" in str(err) or "[DELTA_TABLE_NOT_FOUND]" in str(err):
                 logger.error(f"Could not find table {from_table_name}. Table not found.")
             else:
-                logger.error(err)
+                logger.error(f"Failed to move table {from_table_name}: {err!s}", exc_info=True)
         return False
 
     def _move_view(
@@ -320,5 +320,5 @@ def _move_view(
             if "[TABLE_OR_VIEW_NOT_FOUND]" in str(err):
                 logger.error(f"Could not find view {from_view_name}. View not found.")
             else:
-                logger.error(err)
+                logger.error(f"Failed to move view {from_view_name}: {err!s}", exc_info=True)
         return False

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

@@ -70,10 +70,13 @@ def _safe_get_table_size(self, table_full_name: str) -> int | None:
         try:
             return self._spark._jsparkSession.table(table_full_name).queryExecution().analyzed().stats().sizeInBytes()
         except Exception as e:
-            if "[TABLE_OR_VIEW_NOT_FOUND]" in str(e):
+            if "[TABLE_OR_VIEW_NOT_FOUND]" in str(e) or "[DELTA_TABLE_NOT_FOUND]" in str(e):
                 logger.warning(f"Failed to evaluate {table_full_name} table size. Table not found.")
                 return None
             if "[DELTA_MISSING_TRANSACTION_LOG]" in str(e):
                 logger.warning(f"Delta table {table_full_name} is corrupted: missing transaction log.")
                 return None
-            raise RuntimeError(str(e)) from e
+        except:  # noqa: E722
+            logger.error(f"Failed to evaluate {table_full_name} table size: ", exc_info=True)
+
+        return None

src/databricks/labs/ucx/mixins/sql.py

@@ -7,7 +7,7 @@
 from typing import Any
 
 from databricks.sdk import errors
-from databricks.sdk.errors import NotFound
+from databricks.sdk.errors import DataLoss, NotFound
 from databricks.sdk.service.sql import (
     ColumnInfoTypeName,
     Disposition,
@@ -106,6 +106,10 @@ def _raise_if_needed(status: StatementStatus):
             raise NotFound(error_message)
         if "TABLE_OR_VIEW_NOT_FOUND" in error_message:
             raise NotFound(error_message)
+        if "DELTA_TABLE_NOT_FOUND" in error_message:
+            raise NotFound(error_message)
+        if "DELTA_MISSING_TRANSACTION_LOG" in error_message:
+            raise DataLoss(error_message)
         mapping = {
             ServiceErrorCode.ABORTED: errors.Aborted,
             ServiceErrorCode.ALREADY_EXISTS: errors.AlreadyExists,

tests/unit/framework/test_crawlers.py

@@ -4,7 +4,13 @@
 from unittest import mock
 
 import pytest
-from databricks.sdk.errors import BadRequest, NotFound, PermissionDenied, Unknown
+from databricks.sdk.errors import (
+    BadRequest,
+    DataLoss,
+    NotFound,
+    PermissionDenied,
+    Unknown,
+)
 from databricks.sdk.service import sql
 
 from databricks.labs.ucx.framework.crawlers import (
@@ -260,6 +266,14 @@ def test_raise_spark_sql_exceptions(mocker):
         with pytest.raises(NotFound):
             rb._raise_spark_sql_exceptions(error_message_invalid_table)
 
+        error_message_invalid_table = "DELTA_TABLE_NOT_FOUND foo table does not exist"
+        with pytest.raises(NotFound):
+            rb._raise_spark_sql_exceptions(error_message_invalid_table)
+
+        error_message_invalid_table = "DELTA_MISSING_TRANSACTION_LOG foo table does not exist"
+        with pytest.raises(DataLoss):
+            rb._raise_spark_sql_exceptions(error_message_invalid_table)
+
         error_message_invalid_syntax = "PARSE_SYNTAX_ERROR foo"
         with pytest.raises(BadRequest):
             rb._raise_spark_sql_exceptions(error_message_invalid_syntax)
@@ -293,6 +307,25 @@ def test_execute(mocker):
         with pytest.raises(NotFound):
             rb.execute(sql_query)
 
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "DELTA_TABLE_NOT_FOUND"
+        )
+        with pytest.raises(NotFound):
+            rb.execute(sql_query)
+
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "Unexpected exception"
+        )
+
+        with pytest.raises(BaseException):  # noqa: B017
+            rb.execute(sql_query)
+
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "DELTA_MISSING_TRANSACTION_LOG"
+        )
+        with pytest.raises(DataLoss):
+            rb.execute(sql_query)
+
         pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
             "PARSE_SYNTAX_ERROR"
         )
@@ -332,6 +365,18 @@ def test_fetch(mocker):
         with pytest.raises(NotFound):
             rb.fetch(sql_query)
 
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "DELTA_TABLE_NOT_FOUND"
+        )
+        with pytest.raises(NotFound):
+            rb.fetch(sql_query)
+
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "DELTA_MISSING_TRANSACTION_LOG"
+        )
+        with pytest.raises(DataLoss):
+            rb.fetch(sql_query)
+
         pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
             "PARSE_SYNTAX_ERROR"
         )

tests/unit/hive_metastore/test_table_size.py

@@ -1,7 +1,5 @@
 import sys
 
-import pytest
-
 from databricks.labs.ucx.hive_metastore.table_size import TableSize, TableSizeCrawler
 from tests.unit.framework.mocks import MockBackend
 
@@ -38,7 +36,7 @@ def test_table_size_crawler(mocker):
     assert TableSize("hive_metastore", "db1", "table2", 200) in results
 
 
-def test_table_size_table_not_found_unknown_message(mocker):
+def test_table_size_unknown_error(mocker):
     errors = {}
     rows = {
         "table_size": [],
@@ -53,8 +51,9 @@ def test_table_size_table_not_found_unknown_message(mocker):
     tsc = TableSizeCrawler(backend, "inventory_database")
     tsc._spark._jsparkSession.table().queryExecution().analyzed().stats().sizeInBytes.side_effect = Exception(...)
 
-    with pytest.raises(RuntimeError):
-        tsc.snapshot()
+    results = tsc.snapshot()
+
+    assert len(results) == 0
 
 
 def test_table_size_table_or_view_not_found(mocker):
@@ -81,6 +80,30 @@ def test_table_size_table_or_view_not_found(mocker):
     assert len(results) == 0
 
 
+def test_table_size_delta_table_not_found(mocker):
+    errors = {}
+    rows = {
+        "table_size": [],
+        "hive_metastore.inventory_database.tables": [
+            ("hive_metastore", "db1", "table1", "MANAGED", "DELTA", "dbfs:/location/table", None),
+        ],
+        "SHOW DATABASES": [("db1",)],
+    }
+    backend = MockBackend(fails_on_first=errors, rows=rows)
+    pyspark_sql_session = mocker.Mock()
+    sys.modules["pyspark.sql.session"] = pyspark_sql_session
+    tsc = TableSizeCrawler(backend, "inventory_database")
+
+    # table removed after crawling
+    tsc._spark._jsparkSession.table().queryExecution().analyzed().stats().sizeInBytes.side_effect = Exception(
+        "[DELTA_TABLE_NOT_FOUND]"
+    )
+
+    results = tsc.snapshot()
+
+    assert len(results) == 0
+
+
 def test_table_size_when_table_corrupted(mocker):
     errors = {}
     rows = {