feat: sharepoint scalability5 (#8631)

Author: evan-onyx

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

@@ -53,16 +53,18 @@ def sharepoint_group_sync(
 
     msal_app = connector.msal_app
     sp_tenant_domain = connector.sp_tenant_domain
+    sp_domain_suffix = connector.sharepoint_domain_suffix
     for site_descriptor in site_descriptors:
         logger.debug(f"Processing site: {site_descriptor.url}")
 
         ctx = ClientContext(site_descriptor.url).with_access_token(
-            lambda: acquire_token_for_rest(msal_app, sp_tenant_domain)
+            lambda: acquire_token_for_rest(msal_app, sp_tenant_domain, sp_domain_suffix)
         )
 
         external_groups = get_sharepoint_external_groups(
             ctx,
             connector.graph_client,
+            graph_api_base=connector.graph_api_base,
             get_access_token=connector._get_graph_access_token,
             enumerate_all_ad_groups=enumerate_all,
         )

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

@@ -20,7 +20,6 @@
 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
@@ -647,13 +646,14 @@ def add_user_and_group_to_sets(
 def _enumerate_ad_groups_paginated(
     get_access_token: Callable[[], str],
     already_resolved: set[str],
+    graph_api_base: 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_url = f"{graph_api_base}/groups"
     groups_params: dict[str, str] = {"$select": "id,displayName", "$top": "999"}
     total_groups = 0
 
@@ -679,7 +679,7 @@ def _enumerate_ad_groups_paginated(
             continue
 
         member_emails: list[str] = []
-        members_url = f"{GRAPH_API_BASE}/groups/{group_id}/members"
+        members_url = f"{graph_api_base}/groups/{group_id}/members"
         members_params: dict[str, str] = {
             "$select": "userPrincipalName,mail",
             "$top": "999",
@@ -699,6 +699,7 @@ def _enumerate_ad_groups_paginated(
 def get_sharepoint_external_groups(
     client_context: ClientContext,
     graph_client: GraphClient,
+    graph_api_base: str,
     get_access_token: Callable[[], str] | None = None,
     enumerate_all_ad_groups: bool = False,
 ) -> list[ExternalUserGroup]:
@@ -769,7 +770,9 @@ def add_group_to_sets(role_assignments: RoleAssignmentCollection) -> None:
         return external_user_groups
 
     already_resolved = set(groups_and_members.groups_to_emails.keys())
-    for group in _enumerate_ad_groups_paginated(get_access_token, already_resolved):
+    for group in _enumerate_ad_groups_paginated(
+        get_access_token, already_resolved, graph_api_base
+    ):
         external_user_groups.append(group)
 
     return external_user_groups

backend/onyx/connectors/sharepoint/connector.py

@@ -83,7 +83,11 @@
 
 ASPX_EXTENSION = ".aspx"
 
-GRAPH_API_BASE = "https://graph.microsoft.com/v1.0"
+DEFAULT_AUTHORITY_HOST = "https://login.microsoftonline.com"
+DEFAULT_GRAPH_API_HOST = "https://graph.microsoft.com"
+DEFAULT_SHAREPOINT_DOMAIN_SUFFIX = "sharepoint.com"
+
+GRAPH_API_BASE = f"{DEFAULT_GRAPH_API_HOST}/v1.0"
 GRAPH_API_MAX_RETRIES = 5
 GRAPH_API_RETRYABLE_STATUSES = frozenset({429, 500, 502, 503, 504})
 
@@ -285,10 +289,12 @@ def load_certificate_from_pfx(pfx_data: bytes, password: str) -> CertificateData
 
 
 def acquire_token_for_rest(
-    msal_app: msal.ConfidentialClientApplication, sp_tenant_domain: str
+    msal_app: msal.ConfidentialClientApplication,
+    sp_tenant_domain: str,
+    sharepoint_domain_suffix: str,
 ) -> TokenResponse:
     token = msal_app.acquire_token_for_client(
-        scopes=[f"https://{sp_tenant_domain}.sharepoint.com/.default"]
+        scopes=[f"https://{sp_tenant_domain}.{sharepoint_domain_suffix}/.default"]
     )
     return TokenResponse.from_json(token)
 
@@ -403,12 +409,13 @@ def _download_via_graph_api(
     drive_id: str,
     item_id: str,
     bytes_allowed: int,
+    graph_api_base: str,
 ) -> bytes:
     """Download a drive item via the Graph API /content endpoint with a byte cap.
 
     Raises SizeCapExceeded if the cap is exceeded.
     """
-    url = f"{GRAPH_API_BASE}/drives/{drive_id}/items/{item_id}/content"
+    url = f"{graph_api_base}/drives/{drive_id}/items/{item_id}/content"
     headers = {"Authorization": f"Bearer {access_token}"}
     with requests.get(
         url, headers=headers, stream=True, timeout=REQUEST_TIMEOUT_SECONDS
@@ -429,6 +436,7 @@ def _convert_driveitem_to_document_with_permissions(
     drive_name: str,
     ctx: ClientContext | None,
     graph_client: GraphClient,
+    graph_api_base: str,
     include_permissions: bool = False,
     parent_hierarchy_raw_node_id: str | None = None,
     access_token: str | None = None,
@@ -485,6 +493,7 @@ def _convert_driveitem_to_document_with_permissions(
                 driveitem.drive_id,
                 driveitem.id,
                 SHAREPOINT_CONNECTOR_SIZE_THRESHOLD,
+                graph_api_base=graph_api_base,
             )
         except SizeCapExceeded:
             logger.warning(
@@ -804,6 +813,9 @@ def __init__(
         sites: list[str] = [],
         include_site_pages: bool = True,
         include_site_documents: bool = True,
+        authority_host: str = DEFAULT_AUTHORITY_HOST,
+        graph_api_host: str = DEFAULT_GRAPH_API_HOST,
+        sharepoint_domain_suffix: str = DEFAULT_SHAREPOINT_DOMAIN_SUFFIX,
     ) -> None:
         self.batch_size = batch_size
         self.sites = list(sites)
@@ -819,6 +831,10 @@ def __init__(
         self._cached_rest_ctx: ClientContext | None = None
         self._cached_rest_ctx_url: str | None = None
         self._cached_rest_ctx_created_at: float = 0.0
+        self.authority_host = authority_host.rstrip("/")
+        self.graph_api_host = graph_api_host.rstrip("/")
+        self.graph_api_base = f"{self.graph_api_host}/v1.0"
+        self.sharepoint_domain_suffix = sharepoint_domain_suffix
 
     def validate_connector_settings(self) -> None:
         # Validate that at least one content type is enabled
@@ -875,8 +891,9 @@ def _create_rest_client_context(self, site_url: str) -> ClientContext:
 
         msal_app = self.msal_app
         sp_tenant_domain = self.sp_tenant_domain
+        sp_domain_suffix = self.sharepoint_domain_suffix
         self._cached_rest_ctx = ClientContext(site_url).with_access_token(
-            lambda: acquire_token_for_rest(msal_app, sp_tenant_domain)
+            lambda: acquire_token_for_rest(msal_app, sp_tenant_domain, sp_domain_suffix)
         )
         self._cached_rest_ctx_url = site_url
         self._cached_rest_ctx_created_at = time.monotonic()
@@ -1148,7 +1165,7 @@ def _fetch_site_pages(
         site_id = site.id
 
         page_url: str | None = (
-            f"{GRAPH_API_BASE}/sites/{site_id}/pages/microsoft.graph.sitePage"
+            f"{self.graph_api_base}/sites/{site_id}" f"/pages/microsoft.graph.sitePage"
         )
         params: dict[str, str] | None = {"$expand": "canvasLayout"}
         total_yielded = 0
@@ -1175,7 +1192,7 @@ def _acquire_token(self) -> dict[str, Any]:
             raise RuntimeError("MSAL app is not initialized")
 
         token = self.msal_app.acquire_token_for_client(
-            scopes=["https://graph.microsoft.com/.default"]
+            scopes=[f"{self.graph_api_host}/.default"]
         )
         return token
 
@@ -1248,7 +1265,7 @@ def _iter_drive_items_paged(
         Performs BFS folder traversal manually, fetching one page of children
         at a time so that memory usage stays bounded regardless of drive size.
         """
-        base = f"{GRAPH_API_BASE}/drives/{drive_id}"
+        base = f"{self.graph_api_base}/drives/{drive_id}"
         if folder_path:
             start_url = f"{base}/root:/{folder_path}:/children"
         else:
@@ -1308,7 +1325,7 @@ def _iter_drive_items_delta(
         """
         use_timestamp_token = start is not None and start > _EPOCH
 
-        initial_url = f"{GRAPH_API_BASE}/drives/{drive_id}/root/delta"
+        initial_url = f"{self.graph_api_base}/drives/{drive_id}/root/delta"
         if use_timestamp_token:
             assert start is not None  # mypy
             token = quote(start.isoformat(timespec="seconds"))
@@ -1354,7 +1371,7 @@ def _iter_delta_pages(
                         drive_id,
                     )
                     yield from self._iter_delta_pages(
-                        initial_url=f"{GRAPH_API_BASE}/drives/{drive_id}/root/delta",
+                        initial_url=f"{self.graph_api_base}/drives/{drive_id}/root/delta",
                         drive_id=drive_id,
                         start=start,
                         end=end,
@@ -1471,7 +1488,7 @@ def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None
         sp_private_key = credentials.get("sp_private_key")
         sp_certificate_password = credentials.get("sp_certificate_password")
 
-        authority_url = f"https://login.microsoftonline.com/{sp_directory_id}"
+        authority_url = f"{self.authority_host}/{sp_directory_id}"
 
         if auth_method == SharepointAuthMethod.CERTIFICATE.value:
             logger.info("Using certificate authentication")
@@ -1512,7 +1529,7 @@ def _acquire_token_for_graph() -> dict[str, Any]:
                 raise ConnectorValidationError("MSAL app is not initialized")
 
             token = self.msal_app.acquire_token_for_client(
-                scopes=["https://graph.microsoft.com/.default"]
+                scopes=[f"{self.graph_api_host}/.default"]
             )
             if token is None:
                 raise ConnectorValidationError("Failed to acquire token for graph")
@@ -1941,6 +1958,7 @@ def _load_from_checkpoint(
                         self.graph_client,
                         include_permissions=include_permissions,
                         parent_hierarchy_raw_node_id=parent_hierarchy_url,
+                        graph_api_base=self.graph_api_base,
                         access_token=access_token,
                     )
 

backend/onyx/connectors/teams/connector.py

@@ -50,24 +50,31 @@ class TeamsCheckpoint(ConnectorCheckpoint):
     todo_team_ids: list[str] | None = None
 
 
+DEFAULT_AUTHORITY_HOST = "https://login.microsoftonline.com"
+DEFAULT_GRAPH_API_HOST = "https://graph.microsoft.com"
+
+
 class TeamsConnector(
     CheckpointedConnectorWithPermSync[TeamsCheckpoint],
     SlimConnectorWithPermSync,
 ):
     MAX_WORKERS = 10
-    AUTHORITY_URL_PREFIX = "https://login.microsoftonline.com/"
 
     def __init__(
         self,
         # TODO: (chris) move from "Display Names" to IDs, since display names
         # are not necessarily guaranteed to be unique
         teams: list[str] = [],
         max_workers: int = MAX_WORKERS,
+        authority_host: str = DEFAULT_AUTHORITY_HOST,
+        graph_api_host: str = DEFAULT_GRAPH_API_HOST,
     ) -> None:
         self.graph_client: GraphClient | None = None
         self.msal_app: msal.ConfidentialClientApplication | None = None
         self.max_workers = max_workers
         self.requested_team_list: list[str] = teams
+        self.authority_host = authority_host.rstrip("/")
+        self.graph_api_host = graph_api_host.rstrip("/")
 
     # impls for BaseConnector
 
@@ -76,7 +83,7 @@ def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None
         teams_client_secret = credentials["teams_client_secret"]
         teams_directory_id = credentials["teams_directory_id"]
 
-        authority_url = f"{TeamsConnector.AUTHORITY_URL_PREFIX}{teams_directory_id}"
+        authority_url = f"{self.authority_host}/{teams_directory_id}"
         self.msal_app = msal.ConfidentialClientApplication(
             authority=authority_url,
             client_id=teams_client_id,
@@ -91,7 +98,7 @@ def _acquire_token_func() -> dict[str, Any]:
                 raise RuntimeError("MSAL app is not initialized")
 
             token = self.msal_app.acquire_token_for_client(
-                scopes=["https://graph.microsoft.com/.default"]
+                scopes=[f"{self.graph_api_host}/.default"]
             )
 
             if not isinstance(token, dict):

backend/tests/unit/ee/onyx/external_permissions/sharepoint/test_permission_utils.py

@@ -22,6 +22,7 @@
 
 
 MODULE = "ee.onyx.external_permissions.sharepoint.permission_utils"
+GRAPH_API_BASE = "https://graph.microsoft.com/v1.0"
 
 
 # ---------------------------------------------------------------------------
@@ -125,7 +126,11 @@ def test_enumerate_ad_groups_yields_groups(mock_get: MagicMock) -> None:
     }
     mock_get.side_effect = _mock_graph_get_for_enumeration(groups, members)
 
-    results = list(_enumerate_ad_groups_paginated(_fake_token, already_resolved=set()))
+    results = list(
+        _enumerate_ad_groups_paginated(
+            _fake_token, already_resolved=set(), graph_api_base=GRAPH_API_BASE
+        )
+    )
 
     assert len(results) == 2
     eng = next(r for r in results if r.id == "Engineering_g1")
@@ -140,7 +145,11 @@ def test_enumerate_ad_groups_skips_already_resolved(mock_get: MagicMock) -> None
     mock_get.side_effect = _mock_graph_get_for_enumeration(groups, {})
 
     results = list(
-        _enumerate_ad_groups_paginated(_fake_token, already_resolved={"Engineering_g1"})
+        _enumerate_ad_groups_paginated(
+            _fake_token,
+            already_resolved={"Engineering_g1"},
+            graph_api_base=GRAPH_API_BASE,
+        )
     )
     assert results == []
 
@@ -152,7 +161,11 @@ def test_enumerate_ad_groups_circuit_breaker(mock_get: MagicMock) -> None:
     groups = [{"id": f"g{i}", "displayName": f"Group{i}"} for i in range(over_limit)]
     mock_get.side_effect = _mock_graph_get_for_enumeration(groups, {})
 
-    results = list(_enumerate_ad_groups_paginated(_fake_token, already_resolved=set()))
+    results = list(
+        _enumerate_ad_groups_paginated(
+            _fake_token, already_resolved=set(), graph_api_base=GRAPH_API_BASE
+        )
+    )
     assert len(results) <= AD_GROUP_ENUMERATION_THRESHOLD
 
 
@@ -189,6 +202,7 @@ def test_default_skips_ad_enumeration(
     results = get_sharepoint_external_groups(
         client_context=MagicMock(),
         graph_client=MagicMock(),
+        graph_api_base=GRAPH_API_BASE,
     )
 
     assert len(results) == 1
@@ -219,6 +233,7 @@ def test_enumerate_all_includes_ad_groups(
         graph_client=MagicMock(),
         get_access_token=_fake_token,
         enumerate_all_ad_groups=True,
+        graph_api_base=GRAPH_API_BASE,
     )
 
     assert len(results) == 2
@@ -246,6 +261,7 @@ def test_enumerate_all_without_token_skips(
         graph_client=MagicMock(),
         get_access_token=None,
         enumerate_all_ad_groups=True,
+        graph_api_base=GRAPH_API_BASE,
     )
 
     assert results == []

backend/tests/unit/onyx/connectors/sharepoint/test_drive_matching.py

@@ -206,6 +206,7 @@ def fake_convert(
         drive_name: str,
         ctx: Any,  # noqa: ARG001
         graph_client: Any,  # noqa: ARG001
+        graph_api_base: str,  # noqa: ARG001
         include_permissions: bool,  # noqa: ARG001
         parent_hierarchy_raw_node_id: str | None = None,  # noqa: ARG001
         access_token: str | None = None,  # noqa: ARG001