Fix issue with missing users and temp groups after workspace-local groups migration and skip table when crawling table size if it does not exist anymore (#770)

Author: mwojtyczka

src/databricks/labs/ucx/framework/dashboards.py

@@ -47,7 +47,7 @@ def viz_args(self) -> dict:
 class VizColumn:
     name: str
     title: str
-    type: str = "string"  # noqa: A003
+    type: str = "string"
     imageUrlTemplate: str = "{{ @ }}"  # noqa: N815
     imageTitleTemplate: str = "{{ @ }}"  # noqa: N815
     linkUrlTemplate: str = "{{ @ }}"  # noqa: N815

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

@@ -42,7 +42,10 @@ def _crawl(self) -> Iterable[TableSize]:
                 continue
             if not table.is_dbfs_root:
                 continue
-            size_in_bytes = self.get_table_size(table.key)
+            size_in_bytes = self._safe_get_table_size(table.key)
+            if size_in_bytes is None:
+                continue  # table does not exist anymore
+
             yield TableSize(
                 catalog=table.catalog, database=table.database, name=table.name, size_in_bytes=size_in_bytes
             )
@@ -55,12 +58,19 @@ def _try_load(self) -> Iterable[TableSize]:
     def snapshot(self) -> list[TableSize]:
         """
         Takes a snapshot of tables in the specified catalog and database.
+        Return None if the table cannot be found anymore.
 
         Returns:
             list[Table]: A list of Table objects representing the snapshot of tables.
         """
         return self._snapshot(partial(self._try_load), partial(self._crawl))
 
-    def get_table_size(self, table_full_name: str) -> int:
+    def _safe_get_table_size(self, table_full_name: str) -> int | None:
         logger.debug(f"Evaluating {table_full_name} table size.")
-        return self._spark._jsparkSession.table(table_full_name).queryExecution().analyzed().stats().sizeInBytes()
+        try:
+            return self._spark._jsparkSession.table(table_full_name).queryExecution().analyzed().stats().sizeInBytes()
+        except Exception as e:
+            if "[TABLE_OR_VIEW_NOT_FOUND]" in str(e):
+                logger.warning(f"Failed to evaluate {table_full_name} table size. Table not found.")
+                return None
+            raise RuntimeError(str(e)) from e

src/databricks/labs/ucx/mixins/fixtures.py

@@ -17,7 +17,7 @@
 from databricks.sdk.core import DatabricksError
 from databricks.sdk.errors import ResourceConflict
 from databricks.sdk.retries import retried
-from databricks.sdk.service import compute, iam, jobs, pipelines, workspace
+from databricks.sdk.service import compute, iam, jobs, pipelines, sql, workspace
 from databricks.sdk.service.catalog import (
     CatalogInfo,
     DataSourceFormat,
@@ -27,6 +27,8 @@
 )
 from databricks.sdk.service.sql import (
     CreateWarehouseRequestWarehouseType,
+    GetResponse,
+    ObjectTypePlural,
     Query,
     QueryInfo,
 )
@@ -271,6 +273,47 @@ def _path(ws, path):
     ]
 
 
+def _redash_permissions_mapping():
+    def _simple(_, object_id):
+        return object_id
+
+    return [
+        (
+            "query",
+            ObjectTypePlural.QUERIES,
+            [
+                sql.PermissionLevel.CAN_VIEW,
+                sql.PermissionLevel.CAN_RUN,
+                sql.PermissionLevel.CAN_MANAGE,
+                sql.PermissionLevel.CAN_EDIT,
+            ],
+            _simple,
+        ),
+        (
+            "alert",
+            ObjectTypePlural.ALERTS,
+            [
+                sql.PermissionLevel.CAN_VIEW,
+                sql.PermissionLevel.CAN_RUN,
+                sql.PermissionLevel.CAN_MANAGE,
+                sql.PermissionLevel.CAN_EDIT,
+            ],
+            _simple,
+        ),
+        (
+            "dashboard",
+            ObjectTypePlural.DASHBOARDS,
+            [
+                sql.PermissionLevel.CAN_VIEW,
+                sql.PermissionLevel.CAN_RUN,
+                sql.PermissionLevel.CAN_MANAGE,
+                sql.PermissionLevel.CAN_EDIT,
+            ],
+            _simple,
+        ),
+    ]
+
+
 class _PermissionsChange:
     def __init__(self, object_id: str, before: list[iam.AccessControlRequest], after: list[iam.AccessControlRequest]):
         self._object_id = object_id
@@ -293,6 +336,26 @@ def __repr__(self):
         return f"{self._object_id} [{self._list(self._before)}] -> [{self._list(self._after)}]"
 
 
+class _RedashPermissionsChange:
+    def __init__(self, object_id: str, before: list[sql.AccessControl], after: list[sql.AccessControl]):
+        self._object_id = object_id
+        self._before = before
+        self._after = after
+
+    @staticmethod
+    def _principal(acr: sql.AccessControl) -> str:
+        if acr.user_name is not None:
+            return f"user_name {acr.user_name}"
+        else:
+            return f"group_name {acr.group_name}"
+
+    def _list(self, acl: list[sql.AccessControl]):
+        return ", ".join(f"{self._principal(_)} {_.permission_level.value}" for _ in acl)
+
+    def __repr__(self):
+        return f"{self._object_id} [{self._list(self._before)}] -> [{self._list(self._after)}]"
+
+
 def _make_permissions_factory(name, resource_type, levels, id_retriever):
     def _non_inherited(x: iam.ObjectPermissions):
         out: list[iam.AccessControlRequest] = []
@@ -337,14 +400,29 @@ def create(
                     names = ", ".join(_.value for _ in levels)
                     msg = f"invalid permission level: {permission_level.value}. Valid levels: {names}"
                     raise ValueError(msg)
-                access_control_list = [
-                    iam.AccessControlRequest(
-                        group_name=group_name,
-                        user_name=user_name,
-                        service_principal_name=service_principal_name,
-                        permission_level=permission_level,
+
+                access_control_list = []
+                if group_name is not None:
+                    access_control_list.append(
+                        iam.AccessControlRequest(
+                            group_name=group_name,
+                            permission_level=permission_level,
+                        )
+                    )
+                if user_name is not None:
+                    access_control_list.append(
+                        iam.AccessControlRequest(
+                            user_name=user_name,
+                            permission_level=permission_level,
+                        )
+                    )
+                if service_principal_name is not None:
+                    access_control_list.append(
+                        iam.AccessControlRequest(
+                            service_principal_name=service_principal_name,
+                            permission_level=permission_level,
+                        )
                     )
-                ]
             ws.permissions.update(resource_type, object_id, access_control_list=access_control_list)
             return _PermissionsChange(object_id, initial, access_control_list)
 
@@ -356,12 +434,86 @@ def remove(change: _PermissionsChange):
     return _make_permissions
 
 
+def _make_redash_permissions_factory(name, resource_type, levels, id_retriever):
+    def _non_inherited(x: GetResponse):
+        out: list[sql.AccessControl] = []
+        assert x.access_control_list is not None
+        for access_control in x.access_control_list:
+            out.append(
+                sql.AccessControl(
+                    permission_level=access_control.permission_level,
+                    group_name=access_control.group_name,
+                    user_name=access_control.user_name,
+                )
+            )
+        return out
+
+    def _make_permissions(ws):
+        def create(
+            *,
+            object_id: str,
+            permission_level: sql.PermissionLevel | None = None,
+            group_name: str | None = None,
+            user_name: str | None = None,
+            access_control_list: Optional["list[sql.AccessControl]"] = None,
+        ):
+            nothing_specified = permission_level is None and access_control_list is None
+            both_specified = permission_level is not None and access_control_list is not None
+            if nothing_specified or both_specified:
+                msg = "either permission_level or access_control_list has to be specified"
+                raise ValueError(msg)
+
+            object_id = id_retriever(ws, object_id)
+            initial = _non_inherited(ws.dbsql_permissions.get(resource_type, object_id))
+
+            if access_control_list is None:
+                if permission_level not in levels:
+                    assert permission_level is not None
+                    names = ", ".join(_.value for _ in levels)
+                    msg = f"invalid permission level: {permission_level.value}. Valid levels: {names}"
+                    raise ValueError(msg)
+
+                access_control_list = []
+                if group_name is not None:
+                    access_control_list.append(
+                        sql.AccessControl(
+                            group_name=group_name,
+                            permission_level=permission_level,
+                        )
+                    )
+                if user_name is not None:
+                    access_control_list.append(
+                        sql.AccessControl(
+                            user_name=user_name,
+                            permission_level=permission_level,
+                        )
+                    )
+
+            ws.dbsql_permissions.set(resource_type, object_id, access_control_list=access_control_list)
+            return _RedashPermissionsChange(object_id, initial, access_control_list)
+
+        def remove(change: _RedashPermissionsChange):
+            ws.dbsql_permissions.set(
+                sql.ObjectTypePlural(resource_type), change._object_id, access_control_list=change._before
+            )
+
+        yield from factory(f"{name} permissions", create, remove)
+
+    return _make_permissions
+
+
 for name, resource_type, levels, id_retriever in _permissions_mapping():
     # wrap function factory, otherwise loop scope sticks the wrong way
     locals()[f"make_{name}_permissions"] = pytest.fixture(
         _make_permissions_factory(name, resource_type, levels, id_retriever)
     )
 
+for name, resource_type, levels, id_retriever in _redash_permissions_mapping():
+    # wrap function factory, otherwise loop scope sticks the wrong way
+    locals()[f"make_{name}_permissions"] = pytest.fixture(
+        _make_redash_permissions_factory(name, resource_type, levels, id_retriever)
+    )
+
 
 @pytest.fixture
 def make_secret_scope(ws, make_random):

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

@@ -68,6 +68,12 @@ def get_target_principal(self, name: str) -> str | None:
             return None
         return mg.name_in_account
 
+    def get_temp_principal(self, name: str) -> str | None:
+        mg = self._name_to_group.get(name)
+        if mg is None:
+            return None
+        return mg.temporary_name
+
     def is_in_scope(self, name: str) -> bool:
         if name is None:
             return False

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

@@ -41,10 +41,20 @@ def __iter__(self):
 
 class RedashPermissionsSupport(AclSupport):
     def __init__(
-        self, ws: WorkspaceClient, listings: list[Listing], verify_timeout: timedelta | None = timedelta(minutes=1)
+        self,
+        ws: WorkspaceClient,
+        listings: list[Listing],
+        set_permissions_timeout: timedelta | None = timedelta(minutes=1),
+        # Group information in Redash are cached for up to 10 minutes causing inconsistencies.
+        # If a group is renamed, the old name may still be returned by the dbsql get permissions api.
+        # Note that the update/set API is strongly consistent and is not affected by this behaviour.
+        # The validation step should keep retrying for at least 10 mins until the get api returns the new group name.
+        # More details here: https://databricks.atlassian.net/browse/ES-992619
+        verify_timeout: timedelta | None = timedelta(minutes=11),
     ):
         self._ws = ws
         self._listings = listings
+        self._set_permissions_timeout = set_permissions_timeout
         self._verify_timeout = verify_timeout
 
     @staticmethod
@@ -85,6 +95,41 @@ def _safe_get_dbsql_permissions(self, object_type: sql.ObjectTypePlural, object_
             logger.warning(f"removed on backend: {object_type} {object_id}")
             return None
 
+    def _load_as_request(self, object_type: sql.ObjectTypePlural, object_id: str) -> list[sql.AccessControl]:
+        loaded = self._safe_get_dbsql_permissions(object_type, object_id)
+        if loaded is None:
+            return []
+        acl: list[sql.AccessControl] = []
+        if not loaded.access_control_list:
+            return acl
+
+        for v in loaded.access_control_list:
+            acl.append(
+                sql.AccessControl(
+                    permission_level=v.permission_level,
+                    group_name=v.group_name,
+                    user_name=v.user_name,
+                )
+            )
+        # sort to return deterministic results
+        return sorted(acl, key=lambda v: f"{v.group_name}:{v.user_name}")
+
+    def load_as_dict(self, object_type: sql.ObjectTypePlural, object_id: str) -> dict[str, sql.PermissionLevel]:
+        result = {}
+        for acl in self._load_as_request(object_type, object_id):
+            if not acl.permission_level:
+                continue
+            result[self._key_for_acl_dict(acl)] = acl.permission_level
+        return result
+
+    @staticmethod
+    def _key_for_acl_dict(acl: sql.AccessControl) -> str:
+        if acl.group_name is not None:
+            return acl.group_name
+        if acl.user_name is not None:
+            return acl.user_name
+        return "UNKNOWN"
+
     @rate_limited(max_requests=100)
     def _crawler_task(self, object_id: str, object_type: sql.ObjectTypePlural) -> Permissions | None:
         permissions = self._safe_get_dbsql_permissions(object_type=object_type, object_id=object_id)
@@ -120,7 +165,7 @@ def _applier_task(self, object_type: sql.ObjectTypePlural, object_id: str, acl:
         This affects the way how we prepare the new ACL request.
         """
 
-        set_retry_on_value_error = retried(on=[InternalError, ValueError], timeout=self._verify_timeout)
+        set_retry_on_value_error = retried(on=[InternalError, ValueError], timeout=self._set_permissions_timeout)
         set_retried_check = set_retry_on_value_error(self._safe_set_permissions)
         set_retried_check(object_type, object_id, acl)
 
@@ -133,22 +178,45 @@ def _prepare_new_acl(
     ) -> list[sql.AccessControl]:
         """
         Please note the comment above on how we apply these permissions.
+        Permissions are set/replaced and not updated/patched, therefore all existing ACLs need to be collected
+        including users and temp/backup groups.
         """
         acl_requests: list[sql.AccessControl] = []
         for access_control in acl:
+            if access_control.user_name:
+                logger.debug(f"Including redash permissions acl for user: `{access_control.user_name}`")
+                acl_requests.append(access_control)
+                continue
+
             if not access_control.group_name:
                 continue
+
             if not migration_state.is_in_scope(access_control.group_name):
-                logger.debug(f"Skipping redash item for `{access_control.group_name}`: not in scope")
+                logger.debug(
+                    f"Including redash permissions acl for group not in the scope: `{access_control.group_name}`"
+                )
                 acl_requests.append(access_control)
                 continue
+
             target_principal = migration_state.get_target_principal(access_control.group_name)
-            if target_principal is None:
-                logger.debug(f"Skipping redash item for `{access_control.group_name}`: no target principal")
+            if not target_principal:
+                logger.debug(
+                    f"Including redash permissions acl for group without target principal: "
+                    f"`{access_control.group_name}`"
+                )
                 acl_requests.append(access_control)
                 continue
+
+            logger.debug(f"Including redash permissions acl for target group `{target_principal}`")
             new_acl_request = dataclasses.replace(access_control, group_name=target_principal)
             acl_requests.append(new_acl_request)
+
+            temp_principal = migration_state.get_temp_principal(access_control.group_name)
+            if temp_principal is not None:
+                logger.debug(f"Including redash permissions acl for temp group `{temp_principal}`")
+                temp_group_new_acl_request = dataclasses.replace(access_control, group_name=temp_principal)
+                acl_requests.append(temp_group_new_acl_request)
+
         return acl_requests
 
     @rate_limited(burst_period_seconds=30)