Total Storage Credentials count widget for Assessment Dashboard (#2201)

## Changes
Adding a widget to show the current total storage credentials created in
the workspace (limit 200)

### Linked issues
Introduces #1600

### 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
- [x] manually tested
- [ ] added unit tests
- [ ] added integration tests
- [ ] verified on staging environment (screenshot attached)

---------

Co-authored-by: Serge Smertin <[email protected]>

Author: pritishpai

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

@@ -222,7 +222,8 @@ def get_roles_to_migrate(self) -> list[AWSCredentialCandidate]:
                 roles[role.role_arn].paths.add(external_location.location)
                 if role.privilege == Privilege.WRITE_FILES.value:
                     roles[role.role_arn].privilege = Privilege.WRITE_FILES.value
-
+        if len(roles) > 200:
+            raise RuntimeWarning('Migration will breach UC limits (Storage Credentials > 200).')
         return list(roles.values())
 
     def _get_cluster_policy(self, policy_id: str | None) -> Policy:

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

@@ -317,6 +317,8 @@ def create_access_connectors_for_storage_accounts(self) -> list[tuple[AccessConn
             list[AccessConnector, str] : The access connectors with a storage url to which it has access.
         """
         used_storage_accounts = self._get_storage_accounts()
+        if len(used_storage_accounts) > 200:
+            raise RuntimeWarning('Migration will breach UC limits (Storage Credentials > 200).')
         if len(used_storage_accounts) == 0:
             logger.warning(
                 "There are no external table present with azure storage account. "

src/databricks/labs/ucx/queries/assessment/main/00_7_count_storage_credentials.sql

@@ -0,0 +1,31 @@
+/* --title 'Total Storage Credentials'
+--height 3
+--overrides '{
+               "spec": {
+                 "encodings": {
+                   "value": {
+                     "fieldName": "count_total_storage_credentials",
+                     "rowNumber": 2,
+                     "style": {
+                       "rules": [
+                         {
+                           "condition": {
+                             "operator": ">",
+                             "operand": {
+                               "type": "data-value",
+                               "value": "200"
+                             }
+                           },
+                           "color": "#E92828"
+                         }
+                       ]
+                     },
+                     "displayName": "count_total_storage_credentials"
+                   }
+                 }
+               }
+             }'
+*/
+SELECT
+  COUNT(*) AS count_total_storage_credentials
+FROM inventory.external_locations

src/databricks/labs/ucx/queries/assessment/main/01_0_compatibility.sql

@@ -1,11 +1,10 @@
-/* --title 'UC readiness' --height 4 */
-WITH raw AS (
+/* --title 'UC readiness' --height 3 */
+SELECT
+  COALESCE(CONCAT(ROUND(SUM(ready) / COUNT(*) * 100, 1), '%'), 'N/A') AS readiness
+FROM (
   SELECT
     object_type,
     object_id,
     IF(failures = '[]', 1, 0) AS ready
   FROM inventory.objects
-)
-SELECT
-  COALESCE(CONCAT(ROUND(SUM(ready) / COUNT(*) * 100, 1), '%'), 'N/A') AS readiness
-FROM raw
+)

src/databricks/labs/ucx/queries/assessment/main/01_1_count_jobs.sql

@@ -1,4 +1,4 @@
-/* --title 'Total Jobs' --height 4 */
+/* --title 'Total Jobs' --width 1 --height 3 */
 SELECT
   COUNT(*) AS count_total_jobs
 FROM inventory.jobs

src/databricks/labs/ucx/queries/assessment/main/01_2_count_table_by_storage.sql

@@ -1,4 +1,4 @@
-/* --title 'Table counts by storage' --width 2 --height 4 */
+/* --title 'Table counts by storage' --width 1 --height 3 */
 SELECT
   storage,
   COUNT(*) AS count

src/databricks/labs/ucx/queries/assessment/main/01_3_count_table_by_type_location_format.sql

@@ -1,4 +1,4 @@
-/* --title 'Table counts by type, location and format' --width 2 --height 4 */
+/* --title 'Table counts by type, location and format' --width 2 --height 3 */
 SELECT
   object_type AS source_table_type,
   table_format AS format,

tests/unit/test_cli.py

@@ -2,21 +2,21 @@
 import json
 import time
 from pathlib import Path
-from unittest.mock import create_autospec, patch
+from unittest.mock import create_autospec, patch, Mock
 
 import pytest
 import yaml
-from databricks.sdk.service.provisioning import Workspace
 from databricks.labs.blueprint.tui import MockPrompts
 from databricks.sdk import AccountClient, WorkspaceClient
 from databricks.sdk.errors import NotFound
+from databricks.sdk.errors.platform import BadRequest
 from databricks.sdk.service import iam, jobs, sql
 from databricks.sdk.service.catalog import ExternalLocationInfo
 from databricks.sdk.service.compute import ClusterDetails, ClusterSource
+from databricks.sdk.service.provisioning import Workspace
 from databricks.sdk.service.workspace import ObjectInfo, ObjectType
-from databricks.sdk.errors.platform import BadRequest
 
-from databricks.labs.ucx.assessment.aws import AWSResources
+from databricks.labs.ucx.assessment.aws import AWSResources, AWSRoleAction
 from databricks.labs.ucx.aws.access import AWSResourcePermissions
 from databricks.labs.ucx.azure.access import AzureResourcePermissions
 from databricks.labs.ucx.azure.resources import AzureResources
@@ -31,6 +31,7 @@
     create_uber_principal,
     ensure_assessment_run,
     installations,
+    join_collection,
     logs,
     manual_workspace_info,
     migrate_credentials,
@@ -51,11 +52,11 @@
     validate_external_locations,
     validate_groups_membership,
     workflows,
-    join_collection,
 )
 from databricks.labs.ucx.contexts.account_cli import AccountContext
 from databricks.labs.ucx.contexts.workspace_cli import WorkspaceContext
-from databricks.labs.ucx.hive_metastore import TablesCrawler
+from databricks.labs.ucx.hive_metastore import TablesCrawler, ExternalLocations
+from databricks.labs.ucx.hive_metastore.locations import ExternalLocation
 from databricks.labs.ucx.hive_metastore.tables import Table
 from databricks.labs.ucx.source_code.linters.files import LocalFileMigrator
 
@@ -364,6 +365,79 @@ def test_migrate_credentials_aws(ws):
     ws.storage_credentials.list.assert_called()
 
 
+def test_migrate_credentials_limit_azure(ws):
+    azure_resources = create_autospec(AzureResources)
+    external_locations = create_autospec(ExternalLocations)
+    external_locations_mock = []
+    for i in range(200):
+        external_locations_mock.append(
+            ExternalLocation(
+                location=f"abfss://container{i}@storage{i}.dfs.core.windows.net/folder{i}", table_count=i % 20
+            )
+        )
+    external_locations.snapshot.return_value = external_locations_mock
+    prompts = MockPrompts({'.*': 'yes'})
+    ctx = WorkspaceContext(ws).replace(
+        is_azure=True,
+        azure_cli_authenticated=True,
+        azure_subscription_id='test',
+        azure_resources=azure_resources,
+        external_locations=external_locations,
+    )
+    migrate_credentials(ws, prompts, ctx=ctx)
+    ws.storage_credentials.list.assert_called()
+    azure_resources.storage_accounts.assert_called()
+
+    external_locations_mock.append(
+        ExternalLocation(
+            location=f"abfss://container{201}@storage{201}.dfs.core.windows.net/folder{201}", table_count=25
+        )
+    )
+    with pytest.raises(RuntimeWarning):
+        migrate_credentials(ws, prompts, ctx=ctx)
+
+
+def test_migrate_credentials_limit_aws(ws):
+    aws_resources = create_autospec(AWSResources)
+    external_locations = create_autospec(ExternalLocations)
+
+    external_locations_mock = []
+    aws_role_actions_mock = []
+    for i in range(200):
+        location = f"s3://labsawsbucket/{i}"
+        role = f"arn:aws:iam::123456789012:role/role_name{i}"
+        external_locations_mock.append(ExternalLocation(location=location, table_count=i % 20))
+        aws_role_actions_mock.append(
+            AWSRoleAction(
+                role_arn=role,
+                privilege="READ_FILES",
+                resource_path=location,
+                resource_type="s3",
+            )
+        )
+    external_locations.snapshot.return_value = external_locations_mock
+    aws_resources.validate_connection.return_value = {"Account": "123456789012"}
+
+    prompts = MockPrompts({'.*': 'yes'})
+    AWSResourcePermissions.load_uc_compatible_roles = Mock()
+    AWSResourcePermissions.load_uc_compatible_roles.return_value = aws_role_actions_mock
+    ctx = WorkspaceContext(ws).replace(is_aws=True, aws_resources=aws_resources, external_locations=external_locations)
+    migrate_credentials(ws, prompts, ctx=ctx)
+    ws.storage_credentials.list.assert_called()
+
+    external_locations_mock.append(ExternalLocation(location="s3://labsawsbucket/201", table_count=25))
+    aws_role_actions_mock.append(
+        AWSRoleAction(
+            role_arn="arn:aws:iam::123456789012:role/role_name201",
+            privilege="READ_FILES",
+            resource_path="s3://labsawsbucket/201",
+            resource_type="s3",
+        )
+    )
+    with pytest.raises(RuntimeWarning):
+        migrate_credentials(ws, prompts, ctx=ctx)
+
+
 def test_create_master_principal_not_azure(ws):
     ws.config.is_azure = False
     ws.config.is_aws = False