feat: sharepoint scalability4 (#8551)

Author: evan-onyx

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

@@ -1,9 +1,13 @@
 from collections.abc import Generator
 
+from office365.sharepoint.client_context import ClientContext  # type: ignore[import-untyped]
+
 from ee.onyx.db.external_perm import ExternalUserGroup
 from ee.onyx.external_permissions.sharepoint.permission_utils import (
     get_sharepoint_external_groups,
 )
+from onyx.configs.app_configs import SHAREPOINT_EXHAUSTIVE_AD_ENUMERATION
+from onyx.connectors.sharepoint.connector import acquire_token_for_rest
 from onyx.connectors.sharepoint.connector import SharepointConnector
 from onyx.db.models import ConnectorCredentialPair
 from onyx.utils.logger import setup_logger
@@ -43,14 +47,25 @@ def sharepoint_group_sync(
 
     logger.info(f"Processing {len(site_descriptors)} sites for group sync")
 
-    # Process each site
+    enumerate_all = connector_config.get(
+        "exhaustive_ad_enumeration", SHAREPOINT_EXHAUSTIVE_AD_ENUMERATION
+    )
+
+    msal_app = connector.msal_app
+    sp_tenant_domain = connector.sp_tenant_domain
     for site_descriptor in site_descriptors:
         logger.debug(f"Processing site: {site_descriptor.url}")
 
-        ctx = connector._create_rest_client_context(site_descriptor.url)
+        ctx = ClientContext(site_descriptor.url).with_access_token(
+            lambda: acquire_token_for_rest(msal_app, sp_tenant_domain)
+        )
 
-        # Get external groups for this site
-        external_groups = get_sharepoint_external_groups(ctx, connector.graph_client)
+        external_groups = get_sharepoint_external_groups(
+            ctx,
+            connector.graph_client,
+            get_access_token=connector._get_graph_access_token,
+            enumerate_all_ad_groups=enumerate_all,
+        )
 
         # Yield each group
         for group in external_groups:

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

@@ -1,9 +1,13 @@
 import re
+import time
 from collections import deque
+from collections.abc import Callable
+from collections.abc import Generator
 from typing import Any
 from urllib.parse import unquote
 from urllib.parse import urlparse
 
+import requests as _requests
 from office365.graph_client import GraphClient  # type: ignore[import-untyped]
 from office365.onedrive.driveitems.driveItem import DriveItem  # type: ignore[import-untyped]
 from office365.runtime.client_request import ClientRequestException  # type: ignore
@@ -14,7 +18,11 @@
 from ee.onyx.db.external_perm import ExternalUserGroup
 from onyx.access.models import ExternalAccess
 from onyx.access.utils import build_ext_group_name_for_onyx
+from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS
 from onyx.configs.constants import DocumentSource
+from onyx.connectors.sharepoint.connector import GRAPH_API_BASE
+from onyx.connectors.sharepoint.connector import GRAPH_API_MAX_RETRIES
+from onyx.connectors.sharepoint.connector import GRAPH_API_RETRYABLE_STATUSES
 from onyx.connectors.sharepoint.connector import SHARED_DOCUMENTS_MAP_REVERSE
 from onyx.connectors.sharepoint.connector import sleep_and_retry
 from onyx.utils.logger import setup_logger
@@ -33,6 +41,70 @@
 LIMITED_ACCESS_ROLE_NAMES = ["Limited Access", "Web-Only Limited Access"]
 
 
+AD_GROUP_ENUMERATION_THRESHOLD = 100_000
+
+
+def _graph_api_get(
+    url: str,
+    get_access_token: Callable[[], str],
+    params: dict[str, str] | None = None,
+) -> dict[str, Any]:
+    """Authenticated Graph API GET with retry on transient errors."""
+    for attempt in range(GRAPH_API_MAX_RETRIES + 1):
+        access_token = get_access_token()
+        headers = {"Authorization": f"Bearer {access_token}"}
+        try:
+            resp = _requests.get(
+                url, headers=headers, params=params, timeout=REQUEST_TIMEOUT_SECONDS
+            )
+            if (
+                resp.status_code in GRAPH_API_RETRYABLE_STATUSES
+                and attempt < GRAPH_API_MAX_RETRIES
+            ):
+                wait = min(int(resp.headers.get("Retry-After", str(2**attempt))), 60)
+                logger.warning(
+                    f"Graph API {resp.status_code} on attempt {attempt + 1}, "
+                    f"retrying in {wait}s: {url}"
+                )
+                time.sleep(wait)
+                continue
+            resp.raise_for_status()
+            return resp.json()
+        except (_requests.ConnectionError, _requests.Timeout, _requests.HTTPError):
+            if attempt < GRAPH_API_MAX_RETRIES:
+                wait = min(2**attempt, 60)
+                logger.warning(
+                    f"Graph API connection error on attempt {attempt + 1}, "
+                    f"retrying in {wait}s: {url}"
+                )
+                time.sleep(wait)
+                continue
+            raise
+    raise RuntimeError(
+        f"Graph API request failed after {GRAPH_API_MAX_RETRIES + 1} attempts: {url}"
+    )
+
+
+def _iter_graph_collection(
+    initial_url: str,
+    get_access_token: Callable[[], str],
+    params: dict[str, str] | None = None,
+) -> Generator[dict[str, Any], None, None]:
+    """Paginate through a Graph API collection, yielding items one at a time."""
+    url: str | None = initial_url
+    while url:
+        data = _graph_api_get(url, get_access_token, params)
+        params = None
+        yield from data.get("value", [])
+        url = data.get("@odata.nextLink")
+
+
+def _normalize_email(email: str) -> str:
+    if MICROSOFT_DOMAIN in email:
+        return email.replace(MICROSOFT_DOMAIN, "")
+    return email
+
+
 class SharepointGroup(BaseModel):
     model_config = {"frozen": True}
 
@@ -572,8 +644,63 @@ def add_user_and_group_to_sets(
     )
 
 
+def _enumerate_ad_groups_paginated(
+    get_access_token: Callable[[], str],
+    already_resolved: set[str],
+) -> Generator[ExternalUserGroup, None, None]:
+    """Paginate through all Azure AD groups and yield ExternalUserGroup for each.
+
+    Skips groups whose suffixed name is already in *already_resolved*.
+    Stops early if the number of groups exceeds AD_GROUP_ENUMERATION_THRESHOLD.
+    """
+    groups_url = f"{GRAPH_API_BASE}/groups"
+    groups_params: dict[str, str] = {"$select": "id,displayName", "$top": "999"}
+    total_groups = 0
+
+    for group_json in _iter_graph_collection(
+        groups_url, get_access_token, groups_params
+    ):
+        group_id: str = group_json.get("id", "")
+        display_name: str = group_json.get("displayName", "")
+        if not group_id or not display_name:
+            continue
+
+        total_groups += 1
+        if total_groups > AD_GROUP_ENUMERATION_THRESHOLD:
+            logger.warning(
+                f"Azure AD group enumeration exceeded {AD_GROUP_ENUMERATION_THRESHOLD} "
+                "groups — stopping to avoid excessive memory/API usage. "
+                "Remaining groups will be resolved from role assignments only."
+            )
+            return
+
+        name = f"{display_name}_{group_id}"
+        if name in already_resolved:
+            continue
+
+        member_emails: list[str] = []
+        members_url = f"{GRAPH_API_BASE}/groups/{group_id}/members"
+        members_params: dict[str, str] = {
+            "$select": "userPrincipalName,mail",
+            "$top": "999",
+        }
+        for member_json in _iter_graph_collection(
+            members_url, get_access_token, members_params
+        ):
+            email = member_json.get("userPrincipalName") or member_json.get("mail")
+            if email:
+                member_emails.append(_normalize_email(email))
+
+        yield ExternalUserGroup(id=name, user_emails=member_emails)
+
+    logger.info(f"Enumerated {total_groups} Azure AD groups via paginated Graph API")
+
+
 def get_sharepoint_external_groups(
-    client_context: ClientContext, graph_client: GraphClient
+    client_context: ClientContext,
+    graph_client: GraphClient,
+    get_access_token: Callable[[], str] | None = None,
+    enumerate_all_ad_groups: bool = False,
 ) -> list[ExternalUserGroup]:
 
     groups: set[SharepointGroup] = set()
@@ -629,57 +756,20 @@ def add_group_to_sets(role_assignments: RoleAssignmentCollection) -> None:
         client_context, graph_client, groups, is_group_sync=True
     )
 
-    # get all Azure AD groups because if any group is assigned to the drive item, we don't want to miss them
-    # We can't assign sharepoint groups to drive items or drives, so we don't need to get all sharepoint groups
-    azure_ad_groups = sleep_and_retry(
-        graph_client.groups.get_all(page_loaded=lambda _: None),
-        "get_sharepoint_external_groups:get_azure_ad_groups",
-    )
-    logger.info(f"Azure AD Groups: {len(azure_ad_groups)}")
-    identified_groups: set[str] = set(groups_and_members.groups_to_emails.keys())
-    ad_groups_to_emails: dict[str, set[str]] = {}
-    for group in azure_ad_groups:
-        # If the group is already identified, we don't need to get the members
-        if group.display_name in identified_groups:
-            continue
-        # AD groups allows same display name for multiple groups, so we need to add the GUID to the name
-        name = group.display_name
-        name = _get_group_name_with_suffix(group.id, name, graph_client)
-
-        members = sleep_and_retry(
-            group.members.get_all(page_loaded=lambda _: None),
-            "get_sharepoint_external_groups:get_azure_ad_groups:get_members",
-        )
-        for member in members:
-            member_data = member.to_json()
-            user_principal_name = member_data.get("userPrincipalName")
-            mail = member_data.get("mail")
-            if not ad_groups_to_emails.get(name):
-                ad_groups_to_emails[name] = set()
-            if user_principal_name:
-                if MICROSOFT_DOMAIN in user_principal_name:
-                    user_principal_name = user_principal_name.replace(
-                        MICROSOFT_DOMAIN, ""
-                    )
-                ad_groups_to_emails[name].add(user_principal_name)
-            elif mail:
-                if MICROSOFT_DOMAIN in mail:
-                    mail = mail.replace(MICROSOFT_DOMAIN, "")
-                ad_groups_to_emails[name].add(mail)
+    external_user_groups: list[ExternalUserGroup] = [
+        ExternalUserGroup(id=group_name, user_emails=list(emails))
+        for group_name, emails in groups_and_members.groups_to_emails.items()
+    ]
 
-    external_user_groups: list[ExternalUserGroup] = []
-    for group_name, emails in groups_and_members.groups_to_emails.items():
-        external_user_group = ExternalUserGroup(
-            id=group_name,
-            user_emails=list(emails),
+    if not enumerate_all_ad_groups or get_access_token is None:
+        logger.info(
+            "Skipping exhaustive Azure AD group enumeration. "
+            "Only groups found in site role assignments are included."
         )
-        external_user_groups.append(external_user_group)
+        return external_user_groups
 
-    for group_name, emails in ad_groups_to_emails.items():
-        external_user_group = ExternalUserGroup(
-            id=group_name,
-            user_emails=list(emails),
-        )
-        external_user_groups.append(external_user_group)
+    already_resolved = set(groups_and_members.groups_to_emails.keys())
+    for group in _enumerate_ad_groups_paginated(get_access_token, already_resolved):
+        external_user_groups.append(group)
 
     return external_user_groups

backend/onyx/configs/app_configs.py

@@ -637,6 +637,14 @@ def get_current_tz_offset() -> int:
     os.environ.get("SHAREPOINT_CONNECTOR_SIZE_THRESHOLD", 20 * 1024 * 1024)
 )
 
+# When True, group sync enumerates every Azure AD group in the tenant (expensive).
+# When False (default), only groups found in site role assignments are synced.
+# Can be overridden per-connector via the "exhaustive_ad_enumeration" key in
+# connector_specific_config.
+SHAREPOINT_EXHAUSTIVE_AD_ENUMERATION = (
+    os.environ.get("SHAREPOINT_EXHAUSTIVE_AD_ENUMERATION", "").lower() == "true"
+)
+
 BLOB_STORAGE_SIZE_THRESHOLD = int(
     os.environ.get("BLOB_STORAGE_SIZE_THRESHOLD", 20 * 1024 * 1024)
 )