Improved error handling for all queries executed during table migration (#1679)

## Changes
<!-- Summary of your changes that are easy to understand. Add
screenshots when necessary -->
- As we are migrating multiple tables & views in parallel, any error in
SQL execution should not fail the workflow. Added error handling to all
`sql_backend.execute` commands

### Linked issues
<!-- DOC: Link issue with a keyword: close, closes, closed, fix, fixes,
fixed, resolve, resolves, resolved. See
https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword
-->

Resolves #1676

### Tests
<!-- How is this tested? Please see the checklist below and also
describe any other relevant tests -->

- [x] manually tested
- [x] added unit tests
- [x] verified on staging environment (screenshot attached)

Author: nkvuong

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

@@ -29,7 +29,7 @@
     ViewToMigrate,
 )
 from databricks.labs.ucx.workspace_access.groups import GroupManager, MigratedGroup
-from databricks.sdk.errors.platform import BadRequest, NotFound
+from databricks.sdk.errors.platform import DatabricksError
 
 logger = logging.getLogger(__name__)
 
@@ -210,9 +210,15 @@ def _view_can_be_migrated(self, view: ViewToMigrate):
     def _migrate_view_table(self, src_view: ViewToMigrate, grants: list[Grant] | None = None):
         view_migrate_sql = self._sql_migrate_view(src_view)
         logger.debug(f"Migrating view {src_view.src.key} to using SQL query: {view_migrate_sql}")
-        self._backend.execute(view_migrate_sql)
-        self._backend.execute(src_view.src.sql_alter_to(src_view.rule.as_uc_table_key))
-        self._backend.execute(src_view.src.sql_alter_from(src_view.rule.as_uc_table_key, self._ws.get_workspace_id()))
+        try:
+            self._backend.execute(view_migrate_sql)
+            self._backend.execute(src_view.src.sql_alter_to(src_view.rule.as_uc_table_key))
+            self._backend.execute(
+                src_view.src.sql_alter_from(src_view.rule.as_uc_table_key, self._ws.get_workspace_id())
+            )
+        except DatabricksError as e:
+            logger.warning(f"Failed to migrate view {src_view.src.key} to {src_view.rule.as_uc_table_key}: {e}")
+            return False
         return self._migrate_acl(src_view.src, src_view.rule, grants)
 
     def _sql_migrate_view(self, src_view: ViewToMigrate) -> str:
@@ -231,7 +237,8 @@ def _migrate_external_table(self, src_table: Table, rule: Rule, grants: list[Gra
         sync_result = next(iter(self._backend.fetch(table_migrate_sql)))
         if sync_result.status_code != "SUCCESS":
             logger.warning(
-                f"SYNC command failed to migrate {src_table.key} to {target_table_key}. Status code: {sync_result.status_code}. Description: {sync_result.description}"
+                f"SYNC command failed to migrate table {src_table.key} to {target_table_key}. "
+                f"Status code: {sync_result.status_code}. Description: {sync_result.description}"
             )
             return False
         self._backend.execute(src_table.sql_alter_from(rule.as_uc_table_key, self._ws.get_workspace_id()))
@@ -271,9 +278,13 @@ def _migrate_external_table_hiveserde_in_place(
         logger.debug(
             f"Migrating external table {src_table.key} to {rule.as_uc_table_key} using SQL query: {table_migrate_sql}"
         )
-        self._backend.execute(table_migrate_sql)
-        self._backend.execute(src_table.sql_alter_to(rule.as_uc_table_key))
-        self._backend.execute(src_table.sql_alter_from(rule.as_uc_table_key, self._ws.get_workspace_id()))
+        try:
+            self._backend.execute(table_migrate_sql)
+            self._backend.execute(src_table.sql_alter_to(rule.as_uc_table_key))
+            self._backend.execute(src_table.sql_alter_from(rule.as_uc_table_key, self._ws.get_workspace_id()))
+        except DatabricksError as e:
+            logger.warning(f"Failed to migrate table {src_table.key} to {rule.as_uc_table_key}: {e}")
+            return False
         return self._migrate_acl(src_table, rule, grants)
 
     def _migrate_dbfs_root_table(self, src_table: Table, rule: Rule, grants: list[Grant] | None = None):
@@ -282,9 +293,13 @@ def _migrate_dbfs_root_table(self, src_table: Table, rule: Rule, grants: list[Gr
         logger.debug(
             f"Migrating managed table {src_table.key} to {rule.as_uc_table_key} using SQL query: {table_migrate_sql}"
         )
-        self._backend.execute(table_migrate_sql)
-        self._backend.execute(src_table.sql_alter_to(rule.as_uc_table_key))
-        self._backend.execute(src_table.sql_alter_from(rule.as_uc_table_key, self._ws.get_workspace_id()))
+        try:
+            self._backend.execute(table_migrate_sql)
+            self._backend.execute(src_table.sql_alter_to(rule.as_uc_table_key))
+            self._backend.execute(src_table.sql_alter_from(rule.as_uc_table_key, self._ws.get_workspace_id()))
+        except DatabricksError as e:
+            logger.warning(f"Failed to migrate table {src_table.key} to {rule.as_uc_table_key}: {e}")
+            return False
         return self._migrate_acl(src_table, rule, grants)
 
     def _migrate_table_create_ctas(self, src_table: Table, rule: Rule, grants: list[Grant], mounts: list[Mount]):
@@ -299,9 +314,13 @@ def _migrate_table_create_ctas(self, src_table: Table, rule: Rule, grants: list[
                 dst_table_location = ExternalLocations.resolve_mount(src_table.location, mounts) + "_ctas_migrated"
             table_migrate_sql = src_table.sql_migrate_ctas_external(rule.as_uc_table_key, dst_table_location)
         logger.debug(f"Migrating table {src_table.key} to {rule.as_uc_table_key} using SQL query: {table_migrate_sql}")
-        self._backend.execute(table_migrate_sql)
-        self._backend.execute(src_table.sql_alter_to(rule.as_uc_table_key))
-        self._backend.execute(src_table.sql_alter_from(rule.as_uc_table_key, self._ws.get_workspace_id()))
+        try:
+            self._backend.execute(table_migrate_sql)
+            self._backend.execute(src_table.sql_alter_to(rule.as_uc_table_key))
+            self._backend.execute(src_table.sql_alter_from(rule.as_uc_table_key, self._ws.get_workspace_id()))
+        except DatabricksError as e:
+            logger.warning(f"Failed to migrate table {src_table.key} to {rule.as_uc_table_key}: {e}")
+            return False
         return self._migrate_acl(src_table, rule, grants)
 
     def _migrate_acl(self, src: Table, rule: Rule, grants: list[Grant] | None):
@@ -315,7 +334,7 @@ def _migrate_acl(self, src: Table, rule: Rule, grants: list[Grant] | None):
             logger.debug(f"Migrating acls on {rule.as_uc_table_key} using SQL query: {acl_migrate_sql}")
             try:
                 self._backend.execute(acl_migrate_sql)
-            except (BadRequest, NotFound) as e:
+            except DatabricksError as e:
                 logger.warning(f"Failed to migrate ACL for {src.key} to {rule.as_uc_table_key}: {e}")
         return True
 
@@ -360,8 +379,11 @@ def _revert_migrated_table(self, table: Table, target_table_key: str):
         logger.info(
             f"Reverting {table.object_type} table {table.database}.{table.name} upgraded_to {table.upgraded_to}"
         )
-        self._backend.execute(table.sql_unset_upgraded_to())
-        self._backend.execute(f"DROP {table.kind} IF EXISTS {target_table_key}")
+        try:
+            self._backend.execute(table.sql_unset_upgraded_to())
+            self._backend.execute(f"DROP {table.kind} IF EXISTS {target_table_key}")
+        except DatabricksError as e:
+            logger.warning(f"Failed to revert table {table.key}: {e}")
 
     def _get_revert_count(self, schema: str | None = None, table: str | None = None) -> list[MigrationCount]:
         self._init_seen_tables()

tests/unit/hive_metastore/test_table_migrate.py

@@ -40,7 +40,7 @@
 
 
 @pytest.fixture
-def ws() -> WorkspaceClient:
+def ws():
     client = create_autospec(WorkspaceClient)
     client.get_workspace_id.return_value = "12345"
     return client
@@ -230,7 +230,7 @@ def test_migrate_external_table_failed_sync(ws, caplog):
 
 
 @pytest.mark.parametrize(
-    'hiveserde_in_place_migrate, describe, ddl, migrated, expected_value',
+    'hiveserde_in_place_migrate, describe, ddl, errors, migrated, expected_value',
     [
         # test migrate parquet hiveserde table in place
         (
@@ -245,6 +245,7 @@ def test_migrate_external_table_failed_sync(ws, caplog):
                     "CREATE TABLE hive_metastore.schema.test_parquet (id INT) USING PARQUET LOCATION 'dbfs:/mnt/test/table1'"
                 ),
             ],
+            {},
             True,
             "CREATE TABLE ucx_default.db1_dst.external_dst (id INT) USING PARQUET LOCATION 's3://test/folder/table1'",
         ),
@@ -253,6 +254,7 @@ def test_migrate_external_table_failed_sync(ws, caplog):
             True,
             MockBackend.rows("col_name", "data_type", "comment")[("dummy", "dummy", None)],
             MockBackend.rows("createtab_stmt")[("dummy"),],
+            {},
             False,
             "hive_metastore.db1_src.external_src table can only be migrated using CTAS.",
         ),
@@ -265,6 +267,7 @@ def test_migrate_external_table_failed_sync(ws, caplog):
                 ("OutputFormat", "org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat", None),
             ],
             MockBackend.rows("createtab_stmt")[("dummy"),],
+            {},
             False,
             "Failed to generate in-place migration DDL for hive_metastore.db1_src.external_src, skip the in-place migration. It can be migrated in CTAS workflow",
         ),
@@ -273,20 +276,39 @@ def test_migrate_external_table_failed_sync(ws, caplog):
             False,
             None,
             None,
+            {},
             False,
             "",
         ),
+        # test failed migration
+        (
+            True,
+            MockBackend.rows("col_name", "data_type", "comment")[
+                ("Serde Library", "org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe", None),
+                ("InputFormat", "org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat", None),
+                ("OutputFormat", "org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat", None),
+            ],
+            MockBackend.rows("createtab_stmt")[
+                (
+                    "CREATE TABLE hive_metastore.schema.test_parquet (id INT) USING PARQUET LOCATION 'dbfs:/mnt/test/table1'"
+                ),
+            ],
+            {"CREATE TABLE ucx": "error"},
+            False,
+            "Failed to migrate table hive_metastore.db1_src.external_src to ucx_default.db1_dst.external_dst: error",
+        ),
     ],
 )
 def test_migrate_external_hiveserde_table_in_place(
-    ws, caplog, hiveserde_in_place_migrate, describe, ddl, migrated, expected_value
+    ws, caplog, hiveserde_in_place_migrate, describe, ddl, errors, migrated, expected_value
 ):
     caplog.set_level(logging.INFO)
     backend = MockBackend(
         rows={
             "DESCRIBE TABLE EXTENDED *": describe,
             "SHOW CREATE TABLE *": ddl,
-        }
+        },
+        fails_on_first=errors,
     )
     table_crawler = TablesCrawler(backend, "inventory_database")
     udf_crawler = UdfsCrawler(backend, "inventory_database")
@@ -632,7 +654,7 @@ def get_table_migrator(backend: SqlBackend) -> TablesMigrator:
     return table_migrate
 
 
-def test_revert_migrated_tables_skip_managed(ws):
+def test_revert_migrated_tables_skip_managed():
     errors = {}
     rows = {}
     backend = MockBackend(fails_on_first=errors, rows=rows)
@@ -651,7 +673,7 @@ def test_revert_migrated_tables_skip_managed(ws):
     assert "DROP VIEW IF EXISTS cat1.schema1.dest_view1" in revert_queries
 
 
-def test_revert_migrated_tables_including_managed(ws):
+def test_revert_migrated_tables_including_managed():
     errors = {}
     rows = {}
     backend = MockBackend(fails_on_first=errors, rows=rows)
@@ -682,7 +704,6 @@ def test_no_migrated_tables(ws):
     backend = MockBackend(fails_on_first=errors, rows=rows)
     table_crawler = create_autospec(TablesCrawler)
     grant_crawler = create_autospec(GrantsCrawler)
-    ws = create_autospec(WorkspaceClient)
     table_mapping = create_autospec(TableMapping)
     table_mapping.load.return_value = [
         Rule("workspace", "catalog_1", "db1", "db1", "managed", "managed"),
@@ -1144,3 +1165,82 @@ def test_migrate_views_should_be_properly_sequenced(ws):
     assert table_keys.index("hive_metastore.db1_src.v1_src") > table_keys.index("hive_metastore.db1_src.v3_src")
     assert table_keys.index("hive_metastore.db1_src.v3_src") > table_keys.index("hive_metastore.db1_src.v2_src")
     assert next((key for key in table_keys if key == "hive_metastore.db1_src.t1_src"), None) is None
+
+
+def test_migrate_view_failed(ws, caplog):
+    errors = {"CREATE OR REPLACE VIEW": "error"}
+    create = "CREATE OR REPLACE VIEW hive_metastore.db1_src.view_src (a,b) AS SELECT * FROM db1_src.managed_dbfs"
+    rows = {"SHOW CREATE TABLE": [{"createtab_stmt": create}]}
+    backend = MockBackend(fails_on_first=errors, rows=rows)
+    table_crawler = TablesCrawler(backend, "inventory_database")
+    udf_crawler = UdfsCrawler(backend, "inventory_database")
+    grant_crawler = GrantsCrawler(table_crawler, udf_crawler)
+    table_mapping = table_mapping_mock(["managed_dbfs", "view"])
+    group_manager = GroupManager(backend, ws, "inventory_database")
+    migration_status_refresher = create_autospec(MigrationStatusRefresher)
+    migration_status_refresher.get_seen_tables.return_value = {
+        "ucx_default.db1_dst.managed_dbfs": "hive_metastore.db1_src.managed_dbfs"
+    }
+    migration_index = MigrationIndex(
+        [
+            MigrationStatus("db1_src", "managed_dbfs", "ucx_default", "db1_dst", "new_managed_dbfs"),
+            MigrationStatus("db1_src", "view_src", "ucx_default", "db1_dst", "view_dst"),
+        ]
+    )
+    migration_status_refresher.index.return_value = migration_index
+    principal_grants = create_autospec(PrincipalACL)
+    table_migrate = TablesMigrator(
+        table_crawler,
+        grant_crawler,
+        ws,
+        backend,
+        table_mapping,
+        group_manager,
+        migration_status_refresher,
+        principal_grants,
+    )
+    table_migrate.migrate_tables(what=What.VIEW)
+
+    assert (
+        "Failed to migrate view hive_metastore.db1_src.view_src to ucx_default.db1_dst.view_dst: error" in caplog.text
+    )
+    principal_grants.get_interactive_cluster_grants.assert_not_called()
+
+
+def test_migrate_dbfs_root_tables_failed(ws, caplog):
+    errors = {"CREATE TABLE IF NOT EXISTS": "error"}
+    backend = MockBackend(fails_on_first=errors, rows={})
+    table_crawler = TablesCrawler(backend, "inventory_database")
+    udf_crawler = UdfsCrawler(backend, "inventory_database")
+    grant_crawler = GrantsCrawler(table_crawler, udf_crawler)
+    table_mapping = table_mapping_mock(["managed_dbfs"])
+    group_manager = GroupManager(backend, ws, "inventory_database")
+    migration_status_refresher = MigrationStatusRefresher(ws, backend, "inventory_database", table_crawler)
+    principal_grants = create_autospec(PrincipalACL)
+    table_migrate = TablesMigrator(
+        table_crawler,
+        grant_crawler,
+        ws,
+        backend,
+        table_mapping,
+        group_manager,
+        migration_status_refresher,
+        principal_grants,
+    )
+    table_migrate.migrate_tables(what=What.DBFS_ROOT_DELTA)
+
+    principal_grants.get_interactive_cluster_grants.assert_not_called()
+
+    assert (
+        "Failed to migrate table hive_metastore.db1_src.managed_dbfs to ucx_default.db1_dst.managed_dbfs: error"
+        in caplog.text
+    )
+
+
+def test_revert_migrated_tables_failed(caplog):
+    errors = {"ALTER TABLE": "error"}
+    rows = {}
+    backend = MockBackend(fails_on_first=errors, rows=rows)
+    table_migrate = get_table_migrator(backend)
+    table_migrate.revert_migrated_tables(schema="test_schema1")
+    assert "Failed to revert table hive_metastore.test_schema1.test_table1: error" in caplog.text