Crawl dashboards, queries, and alerts

.gitignore

@@ -145,4 +145,6 @@ cython_debug/
 # dev files and scratches
 dev/cleanup.py
 
-Support
+Support
+
+.python-version

src/databricks/labs/ucx/inventory/inventorizer.py

@@ -3,12 +3,17 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterator
 from functools import partial
+from itertools import chain
 from typing import Generic, TypeVar
 
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.core import DatabricksError
 from databricks.sdk.service.iam import AccessControlResponse, Group, ObjectPermissions
 from databricks.sdk.service.ml import ModelDatabricks
+from databricks.sdk.service.sql import Alert, Dashboard
+from databricks.sdk.service.sql import GetResponse as SqlPermissions
+from databricks.sdk.service.sql import ObjectTypePlural as SqlRequestObjectType
+from databricks.sdk.service.sql import Query
 from databricks.sdk.service.workspace import (
     AclItem,
     ObjectInfo,
@@ -371,6 +376,75 @@ def inner() -> Iterator[ModelDatabricks]:
     return inner
 
 
+class DBSQLInventorizer(BaseInventorizer[InventoryObject]):
+    def __init__(self, ws: WorkspaceClient):
+        self._ws = ws
+        self._queries: Iterator[Query] = iter([])
+        self._dashboards: Iterator[Dashboard] = iter([])
+        self._alerts: Iterator[Alert] = iter([])
+
+    @property
+    def logical_object_types(self) -> list[LogicalObjectType]:
+        return [LogicalObjectType.ALERT, LogicalObjectType.DASHBOARD, LogicalObjectType.QUERY]
+
+    def preload(self):
+        self._queries = self._ws.queries.list()
+        self._dashboards = self._ws.dashboards.list()
+        self._alerts = self._ws.alerts.list()
+
+    @sleep_and_retry
+    @limits(calls=100, period=1)
+    def _get_dbsql_permissions(
+        self, request_object_type: SqlRequestObjectType, request_object_id: str
+    ) -> SqlPermissions:
+        return self._ws.dbsql_permissions.get(object_type=request_object_type, object_id=request_object_id)
+
+    def _safe_get_dbsql_permissions(
+        self, request_object_type: SqlRequestObjectType, object_id: str
+    ) -> SqlPermissions | None:
+        try:
+            permissions = self._get_dbsql_permissions(request_object_type, object_id)
+            return permissions
+        except DatabricksError as e:
+            if e.error_code in ["RESOURCE_DOES_NOT_EXIST", "RESOURCE_NOT_FOUND", "PERMISSION_DENIED"]:
+                logger.warning(f"Could not get permissions for {request_object_type} {object_id} due to {e.error_code}")
+                return None
+            else:
+                raise e
+
+    def _prepare_permission_item(self, _obj: Alert | Dashboard | Query) -> PermissionsInventoryItem | None:
+        if isinstance(_obj, Alert):
+            logical_type = LogicalObjectType.ALERT
+            request_type = SqlRequestObjectType.ALERTS
+        elif isinstance(_obj, Dashboard):
+            logical_type = LogicalObjectType.DASHBOARD
+            request_type = SqlRequestObjectType.DASHBOARDS
+        elif isinstance(_obj, Query):
+            logical_type = LogicalObjectType.QUERY
+            request_type = SqlRequestObjectType.QUERIES
+        else:
+            logger.warning(f"Unexpected object type {_obj}")
+            return
+
+        _permissions = self._safe_get_dbsql_permissions(request_object_type=request_type, object_id=_obj.id)
+
+        if _permissions:
+            _item = PermissionsInventoryItem(
+                object_id=_obj.id,
+                logical_object_type=logical_type,
+                request_object_type=request_type,
+                raw_object_permissions=json.dumps(_permissions.as_dict()),
+            )
+
+    def inventorize(self) -> list[PermissionsInventoryItem]:
+        chained_objects = chain(self._queries, self._alerts, self._dashboards)
+        executables = [partial(self._prepare_permission_item, _object) for _object in chained_objects]
+        results = ThreadedExecution[PermissionsInventoryItem | None](executables).run()
+        results = [result for result in results if result]  # empty filter
+        logger.info(f"Permissions fetched for {len(results)} DBSQL Objects")
+        return results
+
+
 class Inventorizers:
     @staticmethod
     def provide(ws: WorkspaceClient, migration_state: GroupMigrationState, num_threads: int):

src/databricks/labs/ucx/inventory/permissions.py

@@ -9,6 +9,8 @@
 
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service.iam import AccessControlRequest, Group, ObjectPermissions
+from databricks.sdk.service.sql import GetResponse as SqlPermissions
+from databricks.sdk.service.sql import ObjectTypePlural as SqlRequestObjectType
 from databricks.sdk.service.workspace import AclItem as SdkAclItem
 from ratelimit import limits, sleep_and_retry
 from tenacity import retry, stop_after_attempt, wait_fixed, wait_random
@@ -48,7 +50,19 @@ class RolesAndEntitlementsRequestPayload:
     group_id: str
 
 
-AnyRequestPayload = PermissionRequestPayload | SecretsPermissionRequestPayload | RolesAndEntitlementsRequestPayload
+@dataclass
+class SqlObjectRequestPayload:
+    object_id: str
+    request_object_type: SqlRequestObjectType
+    access_control_list: list
+
+
+AnyRequestPayload = (
+    PermissionRequestPayload
+    | SecretsPermissionRequestPayload
+    | RolesAndEntitlementsRequestPayload
+    | SqlObjectRequestPayload
+)
 
 
 # TODO: this class has too many @staticmethod and they must not be such. write a unit test for this logic.
@@ -155,6 +169,13 @@ def _prepare_request_for_roles_and_entitlements(
         destination_group: Group = getattr(migration_info, destination)
         return RolesAndEntitlementsRequestPayload(payload=item.typed_object_permissions, group_id=destination_group.id)
 
+    def _prepare_request_for_sql_object(
+        self, item: PermissionsInventoryItem, migration_state: GroupMigrationState, destination
+    ) -> SqlObjectRequestPayload:
+        _permissions: SqlPermissions = item.typed_object_permissions
+        # TODO: apply conversion logic
+        raise NotImplementedError()
+
     def _prepare_new_permission_request(
         self,
         item: PermissionsInventoryItem,
@@ -169,6 +190,12 @@ def _prepare_new_permission_request(
             return self._prepare_permission_request_for_secrets_api(item, migration_state, destination)
         elif item.logical_object_type in [LogicalObjectType.ROLES, LogicalObjectType.ENTITLEMENTS]:
             return self._prepare_request_for_roles_and_entitlements(item, migration_state, destination)
+        elif item.logical_object_type in [
+            LogicalObjectType.ALERT,
+            LogicalObjectType.DASHBOARD,
+            LogicalObjectType.QUERY,
+        ]:
+            return self._prepare_request_for_sql_object(item, migration_state, destination)
         else:
             logger.warning(
                 f"Unsupported permissions payload for object {item.object_id} "
@@ -219,6 +246,13 @@ def _standard_permissions_applicator(self, request_payload: PermissionRequestPay
             access_control_list=request_payload.access_control_list,
         )
 
+    def _sql_permissions_applicator(self, request_payload: SqlObjectRequestPayload):
+        self._ws.dbsql_permissions.set(
+            object_type=request_payload.request_object_type,
+            object_id=request_payload.object_id,
+            access_control_list=request_payload.access_control_list,
+        )
+
     def applicator(self, request_payload: AnyRequestPayload):
         if isinstance(request_payload, RolesAndEntitlementsRequestPayload):
             self._apply_roles_and_entitlements(
@@ -230,6 +264,8 @@ def applicator(self, request_payload: AnyRequestPayload):
             self._standard_permissions_applicator(request_payload)
         elif isinstance(request_payload, SecretsPermissionRequestPayload):
             self._scope_permissions_applicator(request_payload)
+        elif isinstance(request_payload, SqlObjectRequestPayload):
+            self._sql_permissions_applicator(request_payload)
         else:
             logger.warning(f"Unsupported payload type {type(request_payload)}")
 

src/databricks/labs/ucx/inventory/types.py

@@ -3,6 +3,8 @@
 
 import pandas as pd
 from databricks.sdk.service.iam import ObjectPermissions
+from databricks.sdk.service.sql import GetResponse as SqlPermissions
+from databricks.sdk.service.sql import ObjectTypePlural as SqlRequestObjectType
 from databricks.sdk.service.workspace import AclItem as SdkAclItem
 from databricks.sdk.service.workspace import AclPermission as SdkAclPermission
 
@@ -48,6 +50,11 @@ class LogicalObjectType(StrEnum):
     INSTANCE_POOL = "INSTANCE_POOL"
     CLUSTER_POLICY = "CLUSTER_POLICY"
 
+    # DBSQL Objects
+    ALERT = "ALERT"
+    DASHBOARD = "DASHBOARD"
+    QUERY = "QUERY"
+
     def __repr__(self):
         return self.value
 
@@ -102,19 +109,25 @@ class RolesAndEntitlements:
 class PermissionsInventoryItem:
     object_id: str
     logical_object_type: LogicalObjectType
-    request_object_type: RequestObjectType | None
+    request_object_type: RequestObjectType | SqlRequestObjectType | None
     raw_object_permissions: str
 
     @property
     def object_permissions(self) -> dict:
         return json.loads(self.raw_object_permissions)
 
     @property
-    def typed_object_permissions(self) -> ObjectPermissions | AclItemsContainer | RolesAndEntitlements:
+    def typed_object_permissions(self) -> ObjectPermissions | AclItemsContainer | RolesAndEntitlements | SqlPermissions:
         if self.logical_object_type == LogicalObjectType.SECRET_SCOPE:
             return AclItemsContainer.from_dict(self.object_permissions)
         elif self.logical_object_type in [LogicalObjectType.ROLES, LogicalObjectType.ENTITLEMENTS]:
             return RolesAndEntitlements(**self.object_permissions)
+        elif self.logical_object_type in [
+            LogicalObjectType.ALERT,
+            LogicalObjectType.DASHBOARD,
+            LogicalObjectType.QUERY,
+        ]:
+            return SqlPermissions.from_dict(self.object_permissions)
         else:
             return ObjectPermissions.from_dict(self.object_permissions)