Skip scanning objects that were removed on platform side since the last scan time, so that integration tests are less flaky (#922)

## Changes
Adds handling of objects that may not be present and the API is throwing
an exception to indicate this. While it is not expected to happen,
during testing this is triggered at times and causes failures of the
integration tests. Log instead and continue.

### Linked issues

Resolves #897

Author: larsgeorge-db

src/databricks/labs/ucx/assessment/init_scripts.py

@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import ResourceDoesNotExist
 
 from databricks.labs.ucx.assessment.crawlers import (
     _AZURE_SP_CONF_FAILURE_MSG,
@@ -62,7 +63,11 @@ def _assess_global_init_scripts(self, all_global_init_scripts):
                 failures="[]",
             )
             failures = []
-            script = self._ws.global_init_scripts.get(gis.script_id)
+            try:
+                script = self._ws.global_init_scripts.get(gis.script_id)
+            except ResourceDoesNotExist:
+                logger.warning(f"removed on the backend {gis.script_id}")
+                continue
             global_init_script = base64.b64decode(script.script).decode("utf-8")
             if not global_init_script:
                 continue

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

@@ -379,7 +379,11 @@ def _alias_table(
             return False
 
     def _reapply_grants(self, from_table_name, to_table_name, *, target_view: bool = False):
-        grants = self._ws.grants.get(SecurableType.TABLE, from_table_name)
+        try:
+            grants = self._ws.grants.get(SecurableType.TABLE, from_table_name)
+        except NotFound:
+            logger.warning(f"removed on the backend {from_table_name}")
+            return
         if grants.privilege_assignments is not None:
             logger.info(f"Applying grants on table {to_table_name}")
             grants_changes = []

src/databricks/labs/ucx/install.py

@@ -782,7 +782,11 @@ def _cluster_node_type(self, spec: compute.ClusterSpec) -> compute.ClusterSpec:
         return replace(spec, gcp_attributes=compute.GcpAttributes(availability=compute.GcpAvailability.ON_DEMAND_GCP))
 
     def _check_policy_has_instance_pool(self, policy_id):
-        policy = self._ws.cluster_policies.get(policy_id=policy_id)
+        try:
+            policy = self._ws.cluster_policies.get(policy_id=policy_id)
+        except NotFound:
+            logger.warning(f"removed on the backend {policy_id}")
+            return False
         def_json = json.loads(policy.definition)
         instance_pool = def_json.get("instance_pool_id")
         if instance_pool is not None:

src/databricks/labs/ucx/installer/hms_lineage.py

@@ -4,6 +4,7 @@
 
 from databricks.labs.blueprint.tui import Prompts
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import InvalidParameterValue
 from databricks.sdk.service.compute import GlobalInitScriptDetailsWithContent
 
 global_init_script = """if [[ $DB_IS_DRIVER = "TRUE" ]]; then
@@ -48,7 +49,11 @@ def _check_lineage_spark_config_exists(self) -> GlobalInitScriptDetailsWithConte
         for script in self._ws.global_init_scripts.list():
             if not script.script_id:
                 continue
-            script_content = self._ws.global_init_scripts.get(script_id=script.script_id)
+            try:
+                script_content = self._ws.global_init_scripts.get(script_id=script.script_id)
+            except InvalidParameterValue as err:
+                logger.warning(f"Failed to get init script {script.script_id}: {err}")
+                continue
             if not script_content:
                 continue
             content = script_content.script

src/databricks/labs/ucx/workspace_access/listing.py

@@ -5,7 +5,7 @@
 from itertools import groupby
 
 from databricks.sdk import WorkspaceClient
-from databricks.sdk.errors import InternalError, NotFound
+from databricks.sdk.errors import InternalError, NotFound, ResourceDoesNotExist
 from databricks.sdk.retries import retried
 from databricks.sdk.service.workspace import ObjectInfo, ObjectType
 
@@ -76,7 +76,11 @@ def is_dir(x: ObjectInfo) -> bool:
     def walk(self, start_path="/"):
         self.start_time = dt.datetime.now()
         logger.info(f"Recursive WorkspaceFS listing started at {self.start_time}")
-        root_object = self._ws.workspace.get_status(start_path)
+        try:
+            root_object = self._ws.workspace.get_status(start_path)
+        except ResourceDoesNotExist:
+            logger.warning(f"removed on the backend {start_path}")
+            return self.results
         self.results.append(root_object)
 
         with ThreadPoolExecutor(self._num_threads) as executor:

src/databricks/labs/ucx/workspace_access/verification.py

@@ -2,7 +2,7 @@
 from typing import Literal
 
 from databricks.sdk import WorkspaceClient
-from databricks.sdk.errors import PermissionDenied
+from databricks.sdk.errors import NotFound, PermissionDenied, ResourceDoesNotExist
 
 from databricks.labs.ucx.workspace_access.groups import MigrationState
 from databricks.labs.ucx.workspace_access.secrets import SecretScopesSupport
@@ -33,7 +33,11 @@ def verify_applied_permissions(
         target: Literal["backup", "account"],
     ):
         base_attr = "temporary_name" if target == "backup" else "name_in_account"
-        op = self._ws.permissions.get(object_type, object_id)
+        try:
+            op = self._ws.permissions.get(object_type, object_id)
+        except ResourceDoesNotExist:
+            logger.warning(f"removed on backend: {object_type}.{object_id}")
+            return
         for info in migration_state.groups:
             if not op.access_control_list:
                 continue
@@ -78,8 +82,16 @@ def verify_roles_and_entitlements(self, migration_state: MigrationState, target:
             comparison_base = getattr(el, "id_in_workspace" if target == "backup" else "id_in_workspace")
             comparison_target = getattr(el, target_attr)
 
-            base_group_info = self._ws.groups.get(comparison_base)
-            target_group_info = self._ws.groups.get(comparison_target)
+            try:
+                base_group_info = self._ws.groups.get(comparison_base)
+            except NotFound:
+                logger.warning(f"removed on backend: {comparison_base}")
+                continue
+            try:
+                target_group_info = self._ws.groups.get(comparison_target)
+            except NotFound:
+                logger.warning(f"removed on backend: {comparison_target}")
+                continue
 
             assert base_group_info.roles == target_group_info.roles
             assert base_group_info.entitlements == target_group_info.entitlements

tests/unit/assessment/test_init_scripts.py

@@ -1,5 +1,6 @@
 import base64
 
+from databricks.sdk.errors import ResourceDoesNotExist
 from databricks.sdk.service.compute import GlobalInitScriptDetails
 
 from databricks.labs.ucx.assessment.init_scripts import (
@@ -102,3 +103,27 @@ def test_init_script_without_config_should_have_empty_creator_name(mocker):
             script_id="222", script_name="newscript", enabled=False, created_by=None, success=1, failures="[]"
         ),
     ]
+
+
+def test_missing_global_init_script(mocker, caplog):
+    mock_ws = mocker.Mock()
+    mock_ws.global_init_scripts.list.return_value = [
+        GlobalInitScriptDetails(
+            created_at=111,
+            created_by=None,
+            enabled=False,
+            name="newscript",
+            position=4,
+            script_id="222",
+            updated_at=111,
+            updated_by="[email protected]",
+        )
+    ]
+    mock_ws.global_init_scripts.get.side_effect = ResourceDoesNotExist("RESOURCE_DOES_NOT_EXIST")
+    mockbackend = MockBackend()
+    crawler = GlobalInitScriptCrawler(mock_ws, mockbackend, schema="ucx")
+    result = crawler.snapshot()
+    result = mockbackend.rows_written_for("hive_metastore.ucx.global_init_scripts", "append")
+
+    assert len(result) == 0
+    assert "removed on the backend 222" in caplog.messages

tests/unit/hive_metastore/test_table_move.py

@@ -204,6 +204,36 @@ def test_alias_tables_not_found_view_unknown_error(caplog):
     assert len([rec.message for rec in caplog.records if "unknown error" in rec.message]) == 1
 
 
+def test_move_tables_get_grants_fails_because_table_removed(caplog):
+    client = create_autospec(WorkspaceClient)
+
+    client.tables.list.return_value = [
+        TableInfo(
+            catalog_name="SrcC",
+            schema_name="SrcS",
+            name="table1",
+            full_name="SrcC.SrcS.table1",
+            table_type=TableType.EXTERNAL,
+        ),
+    ]
+
+    rows = {
+        "SHOW CREATE TABLE SrcC.SrcS.table1": [
+            ["CREATE TABLE SrcC.SrcS.table1 (name string)"],
+        ]
+    }
+
+    client.grants.get.side_effect = NotFound('TABLE_DOES_NOT_EXIST')
+    client.schemas.get.side_effect = [SchemaInfo(), SchemaInfo()]
+    client.tables.get.side_effect = [NotFound(), NotFound(), NotFound(), NotFound()]
+    client.grants.update = MagicMock()
+    backend = MockBackend(rows=rows)
+    tm = TableMove(client, backend)
+    tm.move_tables("SrcC", "SrcS", "table1", "TgtC", "TgtS", False)
+
+    assert "removed on the backend SrcC.SrcS.table1" in caplog.messages
+
+
 def test_move_all_tables_and_drop_source():
     client = create_autospec(WorkspaceClient)
 

tests/unit/installer/test_hms_lineage.py

@@ -2,6 +2,7 @@
 
 from databricks.labs.blueprint.tui import MockPrompts
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import InvalidParameterValue
 from databricks.sdk.service.compute import (
     GlobalInitScriptDetails,
     GlobalInitScriptDetailsWithContent,
@@ -127,3 +128,26 @@ def test_disabled_and_then_enabled(caplog):
     hmle.apply(MockPrompts({r'HMS lineage collection init script is disabled.*': 'yes'}))
 
     ws.global_init_scripts.update.assert_called_once()
+
+
+def test_get_script_fails_missing_script(caplog):
+    ws = create_autospec(WorkspaceClient)
+    ginit_scripts = [
+        GlobalInitScriptDetails(
+            created_at=1695045723722,
+            created_by="[email protected]",
+            enabled=True,
+            name="test123",
+            position=0,
+            script_id="12345",
+            updated_at=1695046359612,
+            updated_by="[email protected]",
+        )
+    ]
+    ws.global_init_scripts.list.return_value = ginit_scripts
+    ws.global_init_scripts.get.side_effect = InvalidParameterValue("INVALID_PARAMETER_VALUE")
+
+    with caplog.at_level('WARN'):
+        hmle = HiveMetastoreLineageEnabler(ws)
+        hmle.apply(MockPrompts({r'No HMS lineage collection init script exists.*': 'yes'}))
+        assert "Failed to get init script 12345: INVALID_PARAMETER_VALUE" in caplog.messages

tests/unit/test_install.py

@@ -834,3 +834,19 @@ def test_repair_run_result_state(ws, caplog, mock_installation_with_jobs, any_pr
 
     workspace_installation.repair_run("assessment")
     assert "Please try after sometime" in caplog.text
+
+
+def test_check_policy_has_instance_pool(ws, mock_installation, any_prompt, caplog):
+    ws.cluster_policies.get = MagicMock()
+    ws.cluster_policies.get.side_effect = NotFound("NO_POLICY")
+
+    sql_backend = MockBackend()
+    wheels = create_autospec(WheelsV2)
+    config = WorkspaceConfig(inventory_database='ucx')
+    timeout = timedelta(seconds=1)
+    workspace_installation = WorkspaceInstallation(
+        config, mock_installation_with_jobs, sql_backend, wheels, ws, any_prompt, timeout
+    )
+    res = workspace_installation._check_policy_has_instance_pool(1)
+    assert not res
+    assert "removed on the backend 1" in caplog.messages