feat: efficient sharepoint role fetching (#13813)

Author: evan-onyx

backend/ee/onyx/external_permissions/sharepoint/permission_utils.py

@@ -30,6 +30,11 @@
     SHARED_DOCUMENTS_MAP_REVERSE,
     sleep_and_retry,
 )
+from onyx.connectors.sharepoint.connector_utils import (
+    SharepointGroup,
+    SharepointGroupExpansion,
+    SharepointPermissionCache,
+)
 from onyx.db.enums import HierarchyNodeType
 from onyx.utils.logger import setup_logger
 from onyx.utils.retry_after import parse_retry_after_seconds
@@ -44,6 +49,7 @@
 SHAREPOINT_GROUP_PRINCIPAL_TYPE = 8  # SharePoint site groups (local to the site)
 MICROSOFT_DOMAIN = ".onmicrosoft"
 SHAREPOINT_GROUP_SCOPE_SEPARATOR = "::"
+GROUP_CACHE_KEY_SEPARATOR = ":"
 # PnP RoleType defines Guest=1 and RestrictedGuest=9:
 # https://github.com/pnp/pnpcore/blob/4e4f58fcac797f2957bfcd14fedcecd690dfe7ee/src/sdk/PnP.Core/Model/SharePoint/Core/Public/Enums/RoleType.cs
 LIMITED_ACCESS_ROLE_TYPES = frozenset({1, 9})
@@ -145,19 +151,16 @@ def _normalize_email(email: str) -> str:
     return email
 
 
-class SharepointGroup(BaseModel):
-    model_config = {"frozen": True}
-
-    name: str
-    login_name: str
-    principal_type: int
-
-
 class GroupsResult(BaseModel):
     groups_to_emails: dict[str, set[str]]
     found_public_group: bool
 
 
+class DocumentGroupsResult(BaseModel):
+    group_ids: set[str]
+    found_public_group: bool
+
+
 def _get_azuread_group_guid_by_name(
     graph_client: GraphClient, group_name: str
 ) -> str | None:
@@ -515,10 +518,88 @@ def _get_groups_and_members_recursively(
     )
 
 
+def _group_cache_key(
+    client_context: ClientContext,
+    group: SharepointGroup,
+) -> str:
+    identity = group.login_name
+    if group.principal_type == SHAREPOINT_GROUP_PRINCIPAL_TYPE:
+        identity = _get_site_scoped_group_name(client_context, identity)
+    elif guid := _extract_guid_from_claims_token(identity):
+        identity = guid
+    return f"{group.principal_type}{GROUP_CACHE_KEY_SEPARATOR}{identity}"
+
+
+def _get_cached_group_expansion(
+    client_context: ClientContext,
+    graph_client: GraphClient,
+    group: SharepointGroup,
+    permission_cache: SharepointPermissionCache,
+) -> SharepointGroupExpansion:
+    cache_key = _group_cache_key(client_context, group)
+    cached_expansion = permission_cache.group_expansions.get(cache_key)
+    if cached_expansion is not None:
+        return cached_expansion
+
+    try:
+        if group.principal_type == SHAREPOINT_GROUP_PRINCIPAL_TYPE:
+            nested_groups, _ = _get_sharepoint_groups(
+                client_context, group.login_name, graph_client
+            )
+        else:
+            nested_groups, _ = _get_azuread_groups(graph_client, group.login_name)
+    except ClientRequestException as e:
+        if (
+            group.principal_type != AZURE_AD_GROUP_PRINCIPAL_TYPE
+            or e.response is None
+            or e.response.status_code != 404
+        ):
+            raise
+        logger.warning("Group %s not found", group.login_name)
+        nested_groups = set()
+
+    expansion = SharepointGroupExpansion(nested_groups=nested_groups)
+    permission_cache.group_expansions[cache_key] = expansion
+    return expansion
+
+
+def _resolve_document_groups(
+    client_context: ClientContext,
+    graph_client: GraphClient,
+    groups: set[SharepointGroup],
+    permission_cache: SharepointPermissionCache,
+) -> DocumentGroupsResult:
+    group_queue: deque[SharepointGroup] = deque(groups)
+    visited_group_keys: set[str] = set()
+    group_ids: set[str] = set()
+
+    while group_queue:
+        group = group_queue.popleft()
+        if _is_public_login_name(group.login_name):
+            return DocumentGroupsResult(group_ids=set(), found_public_group=True)
+
+        group_ids.add(group.name)
+        cache_key = _group_cache_key(client_context, group)
+        if cache_key in visited_group_keys:
+            continue
+        visited_group_keys.add(cache_key)
+
+        expansion = _get_cached_group_expansion(
+            client_context, graph_client, group, permission_cache
+        )
+        group_queue.extend(expansion.nested_groups)
+
+    return DocumentGroupsResult(
+        group_ids=group_ids,
+        found_public_group=False,
+    )
+
+
 def _get_external_access_from_securable_object(
     client_context: ClientContext,
     graph_client: GraphClient,
     securable_object: SecurableObject,
+    permission_cache: SharepointPermissionCache,
     add_prefix: bool = False,
 ) -> ExternalAccess:
     groups: set[SharepointGroup] = set()
@@ -578,17 +659,20 @@ def add_user_and_group_to_sets(
         "get_external_access_from_sharepoint",
     )
 
-    groups_and_members = _get_groups_and_members_recursively(
-        client_context, graph_client, groups
+    resolved_groups = _resolve_document_groups(
+        client_context,
+        graph_client,
+        groups,
+        permission_cache,
     )
-    if groups_and_members.found_public_group:
+    if resolved_groups.found_public_group:
         return ExternalAccess(
             external_user_emails=set(),
             external_user_group_ids=set(),
             is_public=True,
         )
 
-    for group_name in groups_and_members.groups_to_emails:
+    for group_name in resolved_groups.group_ids:
         if add_prefix:
             group_name = build_ext_group_name_for_onyx(
                 group_name, DocumentSource.SHAREPOINT
@@ -612,7 +696,9 @@ def get_external_access_from_sharepoint(
     site_page: dict[str, Any] | None,
     add_prefix: bool = False,
     treat_sharing_link_as_public: bool = False,
+    permission_cache: SharepointPermissionCache | None = None,
 ) -> ExternalAccess:
+    permission_cache = permission_cache or SharepointPermissionCache()
     if drive_item and drive_name:
         is_public = _is_public_item(drive_item, treat_sharing_link_as_public)
         if is_public:
@@ -655,6 +741,7 @@ def get_external_access_from_sharepoint(
         client_context,
         graph_client,
         item,
+        permission_cache,
         add_prefix,
     )
 
@@ -665,7 +752,9 @@ def get_hierarchy_node_external_access_from_sharepoint(
     node_type: HierarchyNodeType,
     drive_name: str | None,
     folder_url: str | None,
+    permission_cache: SharepointPermissionCache | None = None,
 ) -> ExternalAccess:
+    permission_cache = permission_cache or SharepointPermissionCache()
     if node_type == HierarchyNodeType.SITE:
         securable_object = client_context.web
     elif node_type == HierarchyNodeType.DRIVE and drive_name:
@@ -683,6 +772,7 @@ def get_hierarchy_node_external_access_from_sharepoint(
         client_context,
         graph_client,
         securable_object,
+        permission_cache,
         add_prefix=True,
     )
 

backend/onyx/connectors/sharepoint/connector.py

@@ -69,6 +69,7 @@
     TextSection,
 )
 from onyx.connectors.sharepoint.connector_utils import (
+    SharepointPermissionCache,
     get_sharepoint_external_access,
     get_sharepoint_hierarchy_node_external_access,
 )
@@ -464,6 +465,9 @@ class SharepointConnectorCheckpoint(ConnectorCheckpoint):
     # Track yielded document IDs to avoid processing the same document twice.
     # The Microsoft Graph delta API can return the same item on multiple pages.
     seen_document_ids: set[str] = Field(default_factory=set)
+    permission_cache: SharepointPermissionCache = Field(
+        default_factory=SharepointPermissionCache
+    )
 
 
 class SharepointAuthMethod(Enum):
@@ -799,12 +803,14 @@ def _convert_driveitem_to_document_with_permissions(
     access_token: str | None = None,
     treat_sharing_link_as_public: bool = False,
     raw_file_callback: RawFileCallback | None = None,
+    permission_cache: SharepointPermissionCache | None = None,
 ) -> Document | ConnectorFailure | None:
     if not driveitem.name or not driveitem.id:
         raise ValueError("DriveItem name/id is required")
 
     if include_permissions and ctx is None:
         raise ValueError("ClientContext is required for permissions")
+    permission_cache = permission_cache or SharepointPermissionCache()
 
     mime_type = driveitem.mime_type
     if not mime_type or mime_type in OnyxMimeTypes.EXCLUDED_IMAGE_TYPES:
@@ -939,6 +945,7 @@ def _convert_driveitem_to_document_with_permissions(
         external_access = get_sharepoint_external_access(
             ctx=ctx,
             graph_client=graph_client,
+            permission_cache=permission_cache,
             drive_item=sdk_item,
             drive_name=drive_name,
             add_prefix=True,
@@ -981,6 +988,7 @@ def _convert_sitepage_to_document(
     site_name: str | None,
     ctx: ClientContext | None,
     graph_client: GraphClient,
+    permission_cache: SharepointPermissionCache,
     include_permissions: bool = False,
     parent_hierarchy_raw_node_id: str | None = None,
     treat_sharing_link_as_public: bool = False,
@@ -1111,6 +1119,7 @@ def _convert_sitepage_to_document(
         external_access = get_sharepoint_external_access(
             ctx=ctx,  # ty: ignore[invalid-argument-type]
             graph_client=graph_client,
+            permission_cache=permission_cache,
             site_page=site_page,
             add_prefix=True,
             treat_sharing_link_as_public=treat_sharing_link_as_public,
@@ -1155,6 +1164,7 @@ def _convert_driveitem_to_slim_document(
     drive_name: str,
     ctx: ClientContext,
     graph_client: GraphClient,
+    permission_cache: SharepointPermissionCache,
     parent_hierarchy_raw_node_id: str | None = None,
     treat_sharing_link_as_public: bool = False,
 ) -> SlimDocument:
@@ -1165,6 +1175,7 @@ def _convert_driveitem_to_slim_document(
     external_access = get_sharepoint_external_access(
         ctx=ctx,
         graph_client=graph_client,
+        permission_cache=permission_cache,
         drive_item=sdk_item,
         drive_name=drive_name,
         treat_sharing_link_as_public=treat_sharing_link_as_public,
@@ -1186,6 +1197,7 @@ def _convert_sitepage_to_slim_document(
     site_page: dict[str, Any],
     ctx: ClientContext | None,
     graph_client: GraphClient,
+    permission_cache: SharepointPermissionCache,
     parent_hierarchy_raw_node_id: str | None = None,
     treat_sharing_link_as_public: bool = False,
 ) -> SlimDocument:
@@ -1197,6 +1209,7 @@ def _convert_sitepage_to_slim_document(
     external_access = get_sharepoint_external_access(
         ctx=ctx,  # ty: ignore[invalid-argument-type]
         graph_client=graph_client,
+        permission_cache=permission_cache,
         site_page=site_page,
         treat_sharing_link_as_public=treat_sharing_link_as_public,
     )
@@ -2323,6 +2336,7 @@ def _fetch_slim_documents_from_sharepoint(
                                     drive_name,
                                     ctx,
                                     self.graph_client,
+                                    temp_checkpoint.permission_cache,
                                     parent_hierarchy_raw_node_id=parent_hierarchy_url,
                                     treat_sharing_link_as_public=self.treat_sharing_link_as_public,
                                 )
@@ -2372,6 +2386,7 @@ def _fetch_slim_documents_from_sharepoint(
                                         site_page,
                                         ctx,
                                         self.graph_client,
+                                        temp_checkpoint.permission_cache,
                                         parent_hierarchy_raw_node_id=site_descriptor.url,
                                         treat_sharing_link_as_public=self.treat_sharing_link_as_public,
                                     )
@@ -2570,6 +2585,7 @@ def _yield_site_hierarchy_node(
             external_access = get_sharepoint_hierarchy_node_external_access(
                 ctx,
                 self.graph_client,
+                checkpoint.permission_cache,
                 HierarchyNodeType.SITE,
             )
 
@@ -2604,6 +2620,7 @@ def _yield_drive_hierarchy_node(
             external_access = get_sharepoint_hierarchy_node_external_access(
                 ctx,
                 self.graph_client,
+                checkpoint.permission_cache,
                 HierarchyNodeType.DRIVE,
                 drive_name=drive_name,
             )
@@ -2657,6 +2674,7 @@ def _yield_folder_hierarchy_nodes(
                 external_access = get_sharepoint_hierarchy_node_external_access(
                     ctx,
                     self.graph_client,
+                    checkpoint.permission_cache,
                     HierarchyNodeType.FOLDER,
                     folder_url=folder_url,
                 )
@@ -2794,6 +2812,7 @@ def _process_drive_item(
                 drive_name,
                 ctx,
                 self.graph_client,
+                permission_cache=checkpoint.permission_cache,
                 include_permissions=include_permissions,
                 parent_hierarchy_raw_node_id=parent_hierarchy_url,
                 graph_api_base=self.graph_api_base,
@@ -3187,6 +3206,7 @@ def _load_from_checkpoint(
                                 site_descriptor.drive_name,
                                 client_ctx,
                                 self.graph_client,
+                                permission_cache=checkpoint.permission_cache,
                                 include_permissions=include_permissions,
                                 # Site pages have the site as their parent
                                 parent_hierarchy_raw_node_id=site_descriptor.url,
@@ -3449,6 +3469,7 @@ def _reindex_site_page(
             site_descriptor.drive_name,
             ctx,
             self.graph_client,
+            permission_cache=dedup.permission_cache,
             include_permissions=include_permissions,
             parent_hierarchy_raw_node_id=site_descriptor.url,
             treat_sharing_link_as_public=self.treat_sharing_link_as_public,