Add clean up activities when migrate-credentials cmd fails intermittently (#2479)

<!-- REMOVE IRRELEVANT COMMENTS BEFORE CREATING A PULL REQUEST -->
## Changes
This PR adds functionality to delete access connectors and storage
credentials created during the migrate-credentials cmd for azure and
when it fails intermittently.

<!-- Summary of your changes that are easy to understand. Add
screenshots when necessary -->

### 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 #2362 

### Functionality

- [ ] added relevant user documentation
- [ ] added new CLI command
- [ ] modified existing command: `databricks labs ucx ...`
- [ ] added a new workflow
- [ ] modified existing workflow: `...`
- [ ] added a new table
- [ ] modified existing table: `...`

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

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

---------

Co-authored-by: Cor <[email protected]>

Author: HariGS-DB

src/databricks/labs/ucx/azure/access.py

@@ -512,9 +512,23 @@ def create_access_connectors_for_storage_accounts(self) -> list[tuple[AccessConn
         thread_name = "Creating access connectors for storage accounts"
         results, errors = Threads.gather(thread_name, tasks)
         if len(errors) > 0:
+            logger.error(
+                "Error creating access connectors. Please review the error message and resolve the issue before trying again. "
+                "Removing successfully created access connectors."
+            )
+            delete_access_connectors = [access_connector for access_connector, _ in results]
+            self.delete_access_connectors(*delete_access_connectors)
             raise ManyError(errors)
         return list(results)
 
+    def delete_storage_credential(self, *storage_credentials: Any):
+        for storage_credential in storage_credentials:
+            self._ws.storage_credentials.delete(storage_credential.name)
+
+    def delete_access_connectors(self, *access_connectors: Any):
+        for access_connector in access_connectors:
+            self._azurerm.delete_access_connector(str(access_connector.id))
+
     def _apply_storage_permission(
         self,
         principal_id: str,

src/databricks/labs/ucx/azure/credentials.py

@@ -5,7 +5,7 @@
 from databricks.labs.blueprint.tui import Prompts
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import BadRequest
-from databricks.sdk.errors.platform import InvalidParameterValue
+from databricks.sdk.errors.platform import InvalidParameterValue, PermissionDenied, NotFound
 from databricks.sdk.service.catalog import (
     AzureManagedIdentityRequest,
     AzureServicePrincipal,
@@ -264,19 +264,29 @@ def _migrate_service_principals(
                     f"Service principal '{spn.permission_mapping.principal}' accesses storage account "
                     f"'{spn.permission_mapping.prefix}' with non-Allow network configuration"
                 )
+            try:
 
-            storage_credential_info = self._storage_credential_manager.create_with_client_secret(spn)
-            validation_results = self._storage_credential_manager.validate(
-                storage_credential_info,
-                spn.permission_mapping.prefix,
-            )
-            execution_result.append(validation_results)
+                storage_credential_info = self._storage_credential_manager.create_with_client_secret(spn)
+                validation_results = self._storage_credential_manager.validate(
+                    storage_credential_info,
+                    spn.permission_mapping.prefix,
+                )
+                execution_result.append(validation_results)
+            except (PermissionDenied, NotFound, BadRequest) as e:
+                logger.error(
+                    "Error creating storage credentials. Please check error message and try again."
+                    "Rolling back storage credential that are created."
+                )
+                self._resource_permissions.delete_storage_credential(execution_result)
+                raise e
         return execution_result
 
-    def _create_storage_credentials_for_storage_accounts(self) -> list[StorageCredentialValidationResult]:
+    def _create_storage_credentials_for_storage_accounts(
+        self,
+    ) -> list[StorageCredentialValidationResult]:
         access_connectors = self._resource_permissions.create_access_connectors_for_storage_accounts()
 
-        execution_results = []
+        execution_results: list[StorageCredentialValidationResult] = []
         for access_connector, url in access_connectors:
             storage_credential_info = self._ws.storage_credentials.create(
                 access_connector.name,
@@ -288,6 +298,15 @@ def _create_storage_credentials_for_storage_accounts(self) -> list[StorageCreden
                 validation_results = self._storage_credential_manager.validate(storage_credential_info, url)
             except BadRequest:
                 logger.warning(f"Could not validate storage credential {storage_credential_info.name} for url {url}")
+            except (PermissionDenied, NotFound):
+                logger.error(
+                    "Error creating access connectors. Please review the error message and resolve the issue before "
+                    "trying again. "
+                    "Removing successfully created storage credentials."
+                )
+                self._resource_permissions.delete_storage_credential(*execution_results)
+                delete_access_connectors = [access_connector for access_connector, _ in access_connectors]
+                self._resource_permissions.delete_access_connectors(*delete_access_connectors)
             else:
                 execution_results.append(validation_results)
 
@@ -298,7 +317,7 @@ def run(self, prompts: Prompts, include_names: set[str] | None = None) -> list[S
             "[RECOMMENDED] Please confirm to create an access connector with a managed identity for each storage "
             "account."
         )
-        ac_results = []
+        ac_results: list[StorageCredentialValidationResult] = []
         if plan_confirmed:
             ac_results = self._create_storage_credentials_for_storage_accounts()
 

src/databricks/labs/ucx/azure/resources.py

@@ -761,11 +761,10 @@ def create_or_update_access_connector(
 
         return access_connector
 
-    def delete_access_connector(self, subscription_id: str, resource_group_name: str, name: str) -> None:
+    def delete_access_connector(self, url: str) -> None:
         """Delete an access connector.
 
         Docs:
             https://learn.microsoft.com/en-us/rest/api/databricks/access-connectors/delete?view=rest-databricks-2023-05-01&tabs=HTTP
         """
-        url = f"/subscriptions/{subscription_id}/resourceGroups/{resource_group_name}/providers/Microsoft.Databricks/accessConnectors/{name}"
         self._mgmt.delete(url, api_version="2023-05-01")

tests/unit/azure/test_access.py

@@ -15,6 +15,7 @@
     GetWorkspaceWarehouseConfigResponseSecurityPolicy,
     SetWorkspaceWarehouseConfigRequestSecurityPolicy,
 )
+from databricks.labs.blueprint.parallel import ManyError
 from databricks.sdk.service.workspace import GetSecretResponse
 
 from databricks.labs.ucx.azure.access import AzureResourcePermissions
@@ -1064,3 +1065,62 @@ def test_create_access_connectors_for_storage_accounts_log_permission_applied(ca
     w.cluster_policies.get.assert_not_called()
     w.secrets.get_secret.assert_not_called()
     w.secrets.create_scope.assert_not_called()
+
+
+def test_create_access_connectors_for_storage_accounts_rollback(caplog):
+    """Log that the permissions for the access connector are applied."""
+    w = create_autospec(WorkspaceClient)
+
+    rows = {
+        "SELECT \\* FROM `hive_metastore`.`ucx`.`external_locations`": [
+            ["abfss://[email protected]/folder1", "1"],
+            ["abfss://[email protected]/folder2", "1"],
+        ]
+    }
+    backend = MockBackend(rows=rows)
+
+    location = ExternalLocations(w, backend, "ucx")
+    installation = MockInstallation()
+
+    azure_resources = create_autospec(AzureResources)
+    azure_resources.storage_accounts.return_value = [
+        StorageAccount(
+            id=AzureResource('/subscriptions/abc/providers/Microsoft.Storage/storageAccounts/storage1'),
+            name="storage1",
+            location="westeu",
+            default_network_action="Allow",
+        ),
+        StorageAccount(
+            id=AzureResource('/subscriptions/abc/providers/Microsoft.Storage/storageAccounts/storage2'),
+            name="storage2",
+            location="westeu",
+            default_network_action="Allow",
+        ),
+    ]
+
+    access_connector_id = AzureResource(
+        "/subscriptions/test/resourceGroups/rg-test/providers/Microsoft.Databricks/accessConnectors/ac-test"
+    )
+    azure_resources.create_or_update_access_connector.side_effect = [
+        AccessConnector(
+            id=access_connector_id,
+            name="ac-test",
+            location="westeu",
+            provisioning_state="Succeeded",
+            identity_type="SystemAssigned",
+            principal_id="test",
+            tenant_id="test",
+        ),
+        PermissionDenied(),
+    ]
+
+    azure_resource_permission = AzureResourcePermissions(installation, w, azure_resources, location)
+
+    with pytest.raises(ManyError):
+        azure_resource_permission.create_access_connectors_for_storage_accounts()
+        assert any("STORAGE_BLOB_DATA_CONTRIBUTOR" in message for message in caplog.messages)
+
+    azure_resources.delete_access_connector.assert_called_with(
+        '/subscriptions/test/resourceGroups/rg-test/providers/Microsoft.Databricks/accessConnectors/ac-test'
+    )
+    w.secrets.create_scope.assert_not_called()