perf(hub): batch GitHub team sync with GraphQL

Author: MioYuuIH

runtime/hub/core/groups.py

@@ -290,6 +290,194 @@ async def fetch_github_team_members(access_token: str, org_name: str, team_slug:
     return members
 
 
+def _build_github_team_members_graphql_query(team_count: int) -> str:
+    variable_defs = ["$org: String!"]
+    team_fields = []
+    for index in range(team_count):
+        variable_defs.extend([f"$slug{index}: String!", f"$after{index}: String"])
+        team_fields.append(
+            f"""
+      team{index}: team(slug: $slug{index}) {{
+        members(first: 100, after: $after{index}) {{
+          nodes {{
+            login
+          }}
+          pageInfo {{
+            hasNextPage
+            endCursor
+          }}
+        }}
+      }}"""
+        )
+
+    return f"""
+query({", ".join(variable_defs)}) {{
+  organization(login: $org) {{
+{"".join(team_fields)}
+  }}
+}}
+"""
+
+
+_GITHUB_ORG_TEAMS_GRAPHQL_QUERY = """
+query($org: String!, $after: String) {
+  organization(login: $org) {
+    teams(first: 100, after: $after) {
+      nodes {
+        name
+        slug
+      }
+      pageInfo {
+        hasNextPage
+        endCursor
+      }
+    }
+  }
+}
+"""
+
+
+async def fetch_github_org_team_slugs_graphql(access_token: str, org_name: str) -> set[str] | None:
+    """Fetch all team slugs that actually exist in a GitHub organization."""
+    if not access_token or not org_name:
+        return set()
+
+    headers = {
+        "Authorization": f"Bearer {access_token}",
+        "Accept": "application/vnd.github+json",
+        "X-GitHub-Api-Version": "2022-11-28",
+    }
+    team_slugs: set[str] = set()
+    after: str | None = None
+
+    try:
+        async with aiohttp.ClientSession() as session:
+            while True:
+                async with session.post(
+                    "https://api.github.com/graphql",
+                    headers=headers,
+                    json={"query": _GITHUB_ORG_TEAMS_GRAPHQL_QUERY, "variables": {"org": org_name, "after": after}},
+                ) as resp:
+                    if resp.status != 200:
+                        log.warning("GitHub GraphQL API returned status %d when fetching org teams", resp.status)
+                        return None
+                    data = await resp.json()
+
+                if data.get("errors"):
+                    log.warning("GitHub GraphQL API returned errors when fetching org teams: %s", data["errors"])
+                    return None
+
+                organization = data.get("data", {}).get("organization")
+                if organization is None:
+                    log.warning("GitHub GraphQL API did not return organization %s", org_name)
+                    return None
+
+                teams = organization.get("teams") or {}
+                for team in teams.get("nodes") or []:
+                    slug = team.get("slug") if isinstance(team, dict) else None
+                    if isinstance(slug, str) and slug:
+                        team_slugs.add(slug)
+
+                page_info = teams.get("pageInfo") or {}
+                if not page_info.get("hasNextPage"):
+                    break
+                after = page_info.get("endCursor")
+    except Exception as e:
+        log.warning("Error fetching GitHub org teams with GraphQL for org %s: %s", org_name, e)
+        return None
+
+    return team_slugs
+
+
+async def fetch_github_team_members_table_graphql(
+    access_token: str,
+    org_name: str,
+    team_keys: list[str],
+) -> dict[str, list[str]] | None:
+    """Fetch a login -> configured team-key table with batched GraphQL team queries."""
+    if not access_token or not org_name or not team_keys:
+        return {}
+
+    headers = {
+        "Authorization": f"Bearer {access_token}",
+        "Accept": "application/vnd.github+json",
+        "X-GitHub-Api-Version": "2022-11-28",
+    }
+    teams_by_login: dict[str, list[str]] = {}
+    actual_team_slugs = await fetch_github_org_team_slugs_graphql(access_token, org_name)
+    if actual_team_slugs is None:
+        return None
+
+    existing_team_keys = []
+    for team_key in team_keys:
+        api_slug = _github_team_api_slug(team_key)
+        if api_slug in actual_team_slugs:
+            existing_team_keys.append(team_key)
+        else:
+            log.warning("GitHub org %s does not have configured team %s", org_name, api_slug)
+
+    pending = dict.fromkeys(existing_team_keys)
+
+    try:
+        async with aiohttp.ClientSession() as session:
+            while pending:
+                batch_keys = list(pending)
+                variables: dict[str, str | None] = {"org": org_name}
+                for index, team_key in enumerate(batch_keys):
+                    variables[f"slug{index}"] = _github_team_api_slug(team_key)
+                    variables[f"after{index}"] = pending[team_key]
+
+                async with session.post(
+                    "https://api.github.com/graphql",
+                    headers=headers,
+                    json={
+                        "query": _build_github_team_members_graphql_query(len(batch_keys)),
+                        "variables": variables,
+                    },
+                ) as resp:
+                    if resp.status != 200:
+                        log.warning("GitHub GraphQL API returned status %d when fetching team members", resp.status)
+                        return None
+                    data = await resp.json()
+
+                if data.get("errors"):
+                    log.warning("GitHub GraphQL API returned errors when fetching team members: %s", data["errors"])
+                    return None
+
+                organization = data.get("data", {}).get("organization")
+                if organization is None:
+                    log.warning("GitHub GraphQL API did not return organization %s", org_name)
+                    return None
+
+                next_pending: dict[str, str | None] = {}
+                for index, team_key in enumerate(batch_keys):
+                    team = organization.get(f"team{index}")
+                    if team is None:
+                        log.warning(
+                            "GitHub GraphQL API did not return team %s in org %s",
+                            _github_team_api_slug(team_key),
+                            org_name,
+                        )
+                        continue
+
+                    members = team.get("members") or {}
+                    for member in members.get("nodes") or []:
+                        login = member.get("login") if isinstance(member, dict) else None
+                        if isinstance(login, str) and login:
+                            teams_by_login.setdefault(login.lower(), []).append(team_key)
+
+                    page_info = members.get("pageInfo") or {}
+                    if page_info.get("hasNextPage"):
+                        next_pending[team_key] = page_info.get("endCursor")
+
+                pending = next_pending
+    except Exception as e:
+        log.warning("Error fetching GitHub team members with GraphQL for org %s: %s", org_name, e)
+        return None
+
+    return teams_by_login
+
+
 async def fetch_github_team_members_table(
     app_id: str,
     installation_id: str,
@@ -329,13 +517,13 @@ async def fetch_github_team_members_table(
         if not installation_token:
             return None
 
-        teams_by_login: dict[str, list[str]] = {}
-        for team_key in github_team_keys:
-            members = await fetch_github_team_members(installation_token, org_name, _github_team_api_slug(team_key))
-            if members is None:
-                return None
-            for login in members:
-                teams_by_login.setdefault(login, []).append(team_key)
+        teams_by_login = await fetch_github_team_members_table_graphql(
+            installation_token,
+            org_name,
+            github_team_keys,
+        )
+        if teams_by_login is None:
+            return None
 
         _GITHUB_TEAM_MEMBERS_CACHE[cache_key] = (now, teams_by_login)
         return {login: list(teams) for login, teams in teams_by_login.items()}

runtime/hub/tests/test_groups.py

@@ -124,6 +124,7 @@ class DummyClientSession:
     created = 0
     get_calls = 0
     post_calls = 0
+    graphql_calls = 0
 
     def __init__(self):
         type(self).created += 1
@@ -144,17 +145,55 @@ def get(self, url, headers=None):
             return DummyResponse(200, [{"login": "OctoUser"}])
         return DummyResponse(200, {"repositories": []})
 
-    def post(self, url, headers=None):
+    def post(self, url, headers=None, json=None):
         type(self).post_calls += 1
         if url.endswith("/app/installations/12345/access_tokens"):
             return DummyResponse(201, {"token": "cached-token", "expires_at": "2099-01-01T00:00:00Z"})
+        if url == "https://api.github.com/graphql":
+            type(self).graphql_calls += 1
+            if "teams(first: 100" in (json or {}).get("query", ""):
+                return DummyResponse(
+                    200,
+                    {
+                        "data": {
+                            "organization": {
+                                "teams": {
+                                    "nodes": [{"name": "AUP", "slug": "aup"}],
+                                    "pageInfo": {"hasNextPage": False, "endCursor": None},
+                                }
+                            }
+                        }
+                    },
+                )
+
+            organization = {}
+            variables = (json or {}).get("variables", {})
+            for key, value in variables.items():
+                if not key.startswith("slug"):
+                    continue
+                index = key.removeprefix("slug")
+                if value == "missing-team":
+                    organization[f"team{index}"] = None
+                elif value == "aup":
+                    organization[f"team{index}"] = {
+                        "members": {
+                            "nodes": [{"login": "OctoUser"}],
+                            "pageInfo": {"hasNextPage": False, "endCursor": None},
+                        }
+                    }
+                else:
+                    organization[f"team{index}"] = {
+                        "members": {"nodes": [], "pageInfo": {"hasNextPage": False, "endCursor": None}}
+                    }
+            return DummyResponse(200, {"data": {"organization": organization}})
         return DummyResponse(500, {})
 
 
 def _reset_dummy_client_session():
     DummyClientSession.created = 0
     DummyClientSession.get_calls = 0
     DummyClientSession.post_calls = 0
+    DummyClientSession.graphql_calls = 0
     groups._GITHUB_APP_INSTALLATION_ID_CACHE.clear()
     groups._GITHUB_APP_INSTALLATION_TOKEN.clear()
     groups._GITHUB_TEAM_MEMBERS_CACHE.clear()
@@ -228,6 +267,30 @@ def test_fetch_github_team_members_table_uses_api_slug_and_preserves_group_key(m
     )
 
     assert teams_by_login == {"octouser": ["AUP"]}
+    assert DummyClientSession.graphql_calls == 2
+
+
+def test_fetch_github_team_members_table_skips_missing_graphql_team(monkeypatch, caplog):
+    _reset_dummy_client_session()
+    monkeypatch.setattr(groups.aiohttp, "ClientSession", DummyClientSession)
+    monkeypatch.setattr(groups.jwt, "encode", lambda payload, private_key, algorithm: "jwt-token")
+
+    caplog.set_level("WARNING", logger="jupyterhub.groups")
+    teams_by_login = asyncio.run(
+        fetch_github_team_members_table(
+            "app-123",
+            "",
+            "dummy-private-key",
+            "",
+            "test-org",
+            {"AUP", "missing-team"},
+            force=True,
+        )
+    )
+
+    assert teams_by_login == {"octouser": ["AUP"]}
+    assert "configured team missing-team" in caplog.text
+    assert DummyClientSession.graphql_calls == 2
 
 
 def test_resolve_resources_for_user_uses_group_mapping():