Add enterprise + multi-org collection support

extension/schema.json

@@ -6,6 +6,38 @@
     "namespace": "GH"
   },
   "node_kinds": [
+    {
+      "name": "GH_Enterprise",
+      "display_name": "GitHub Enterprise",
+      "description": "A GitHub Enterprise account that contains organizations, enterprise teams, roles, and managed users",
+      "is_display_kind": true,
+      "icon": "globe",
+      "color": "#6EA8FE"
+    },
+    {
+      "name": "GH_EnterpriseTeam",
+      "display_name": "GitHub Enterprise Team",
+      "description": "A team managed at the GitHub Enterprise level and assignable across organizations",
+      "is_display_kind": true,
+      "icon": "users-between-lines",
+      "color": "#9B8CFF"
+    },
+    {
+      "name": "GH_EnterpriseRole",
+      "display_name": "GitHub Enterprise Role",
+      "description": "The role a user or team has at the GitHub Enterprise level",
+      "is_display_kind": true,
+      "icon": "user-tie",
+      "color": "#B8D7FF"
+    },
+    {
+      "name": "GH_EnterpriseManagedUser",
+      "display_name": "GitHub Enterprise Managed User",
+      "description": "A GitHub Enterprise managed user account linked to an enterprise identity provider",
+      "is_display_kind": true,
+      "icon": "user-lock",
+      "color": "#7DD3FC"
+    },
     {
       "name": "GH_Organization",
       "display_name": "GitHub Organization",
@@ -150,6 +182,30 @@
       "icon": "lock-open",
       "color": "#E89B5C"
     },
+    {
+      "name": "GH_RunnerGroup",
+      "display_name": "GitHub Runner Group",
+      "description": "A GitHub self-hosted runner group that controls runner access and visibility",
+      "is_display_kind": true,
+      "icon": "server",
+      "color": "#94A3B8"
+    },
+    {
+      "name": "GH_OrgRunner",
+      "display_name": "GitHub Org Runner",
+      "description": "An organization-scoped GitHub self-hosted runner available to selected repositories or workflows",
+      "is_display_kind": true,
+      "icon": "microchip",
+      "color": "#22C55E"
+    },
+    {
+      "name": "GH_RepoRunner",
+      "display_name": "GitHub Repo Runner",
+      "description": "A repository-scoped GitHub self-hosted runner available to jobs in a single repository",
+      "is_display_kind": true,
+      "icon": "microchip",
+      "color": "#38BDF8"
+    },
     {
       "name": "GH_EnvironmentVariable",
       "display_name": "GitHub Environment Variable",

justfile

@@ -1,16 +1,16 @@
 set dotenv-load := true
 
-collect +args:
+collect +args='github /tmp/output/raw/':
     @echo "Collecting data"
-    uv run src/main.py collect github {{args}}
+    uv run src/main.py collect {{args}}
 
-preprocess +args:
-    @echo "Collecting data"
-    uv run src/main.py preprocess github {{args}}
+preprocess +args='github /tmp/output/raw/github':
+    @echo "Preprocessing data"
+    uv run openhound preprocess {{args}}
 
-convert +args:
+convert +args='github /tmp/output/raw/github /tmp/output/graph/github':
     @echo "Converting data"
-    uv run src/main.py convert github {{args}}
+    uv run openhound convert {{args}}
 
 sync:
     @echo "Syncing dependencies"

pyproject.toml

@@ -28,7 +28,7 @@ local_scheme = "no-local-version"
 
 [dependency-groups]
 dev = [
-    "openhound",
+    "openhound==0.1.4",
     "pre-commit>=4.5.1",
     "pytest>=9.0.1",
     "ruff>=0.15.5",

src/openhound_github/auth.py

@@ -1,23 +1,28 @@
 """GitHub authentication helpers for JWT-based app authentication."""
 
-import json
 import base64
-import requests
+import json
 from datetime import datetime, timezone
+from typing import Iterator
+
+import requests
 from cryptography.hazmat.primitives import hashes, serialization
 from cryptography.hazmat.primitives.asymmetric import padding
-from typing import Optional
+from dlt.sources.helpers.rest_client.client import RESTClient
+from dlt.sources.helpers.rest_client.paginators import (
+    HeaderLinkPaginator,
+)
 
 
 class GitHubJwtSession:
     """Manages GitHub JWT-based authentication for GitHub Apps."""
 
     def __init__(
         self,
-        org_name: str,
         client_id: str,
         private_key_path: str,
-        app_id: str,
+        app_id: str | None = None,
+        org_name: str | None = None,
         api_uri: str = "https://api.github.com/",
     ):
         """Initialize a GitHub JWT session.
@@ -33,15 +38,13 @@ def __init__(
             FileNotFoundError: If the private key file cannot be found.
             ValueError: If the private key cannot be loaded or JWT cannot be created.
         """
-        self.org_name = org_name
         self.client_id = client_id
         self.private_key_path = private_key_path
         self.app_id = app_id
+        self.org_name = org_name
         self.api_uri = api_uri.rstrip("/") + "/"
 
-        self._jwt_token: Optional[str] = None
-        self._access_token: Optional[str] = None
-        self._token_expires_at: datetime | None = None
+        self._access_tokens: dict[int, tuple[str, datetime | None]] = {}
 
         self._load_private_key()
 
@@ -92,9 +95,7 @@ def _create_jwt(self) -> str:
         except Exception as e:
             raise ValueError(f"Failed to sign JWT: {e}") from e
 
-        jwt_token = f"{header_encoded}.{payload_encoded}.{signature_encoded}"
-        self._jwt_token = jwt_token
-        return jwt_token
+        return f"{header_encoded}.{payload_encoded}.{signature_encoded}"
 
     @staticmethod
     def _base64url_encode(data: str | bytes) -> str:
@@ -117,41 +118,47 @@ def _base64url_encode(data: str | bytes) -> str:
         # Remove padding
         return encoded.rstrip("=")
 
-    def get_access_token(self) -> str:
-        """Get a valid access token for GitHub API requests.
+    @property
+    def jwt_headers(self) -> dict:
+        jwt_token = self._create_jwt()
+        return {
+            "Accept": "application/vnd.github+json",
+            "Authorization": f"Bearer {jwt_token}",
+            "X-GitHub-Api-Version": "2022-11-28",
+        }
 
-        Fetches an installation access token using the JWT token. If a valid
-        access token is already cached and not expired, it returns that.
+    @property
+    def jwt_token(self) -> str:
+        return self._create_jwt()
 
-        Returns:
-            A valid access token string.
+    def _jwt_client(self) -> RESTClient:
+        return RESTClient(
+            base_url=self.api_uri,
+            headers=self.jwt_headers,
+            paginator=HeaderLinkPaginator(),
+        )
 
-        Raises:
-            requests.RequestException: If the API request fails.
-            ValueError: If token cannot be obtained.
-        """
-        # Return cached token if still valid
-        if (
-            self._access_token
-            and self._token_expires_at
-            and datetime.now(timezone.utc) < self._token_expires_at
+    def list_installations(self) -> Iterator[dict]:
+        for page in self._jwt_client().paginate(
+            "/app/installations", params={"per_page": 100}
         ):
-            return self._access_token
+            yield from page
 
-        # Create a fresh JWT
-        jwt_token = self._create_jwt()
+    def installation_id_for_org(self, org_login: str) -> int:
+        response = self._jwt_client().get(f"/orgs/{org_login}/installation").json()
+        return int(response["id"])
 
-        # Exchange JWT for access token
-        headers = {
-            "Accept": "application/vnd.github+json",
-            "Authorization": f"Bearer {jwt_token}",
-            "X-GitHub-Api-Version": "2022-11-28",
-        }
+    def installation_token(self, installation_id: int) -> str:
+        cached = self._access_tokens.get(installation_id)
+        if cached:
+            token, expires_at = cached
+            if expires_at is None or datetime.now(timezone.utc) < expires_at:
+                return token
 
         try:
             response = requests.post(
-                f"{self.api_uri}app/installations/{self.app_id}/access_tokens",
-                headers=headers,
+                f"{self.api_uri}app/installations/{installation_id}/access_tokens",
+                headers=self.jwt_headers,
                 timeout=10,
             )
             response.raise_for_status()
@@ -164,15 +171,41 @@ def get_access_token(self) -> str:
         if "token" not in response_data:
             raise ValueError("No token in GitHub API response. Check app credentials.")
 
-        self._access_token = response_data["token"]
-
-        # Parse expiration time if provided
+        expires_at = None
         if "expires_at" in response_data:
-            self._token_expires_at = datetime.fromisoformat(
+            expires_at = datetime.fromisoformat(
                 response_data["expires_at"].replace("Z", "+00:00")
             )
 
-        return self._access_token
+        token = response_data["token"]
+        self._access_tokens[installation_id] = (token, expires_at)
+        return token
+
+    def get_access_token(
+        self, org_name: str | None = None, installation_id: int | None = None
+    ) -> str:
+        """Get a valid access token for GitHub API requests.
+
+        Fetches an installation access token using the JWT token. If a valid
+        access token is already cached and not expired, it returns that.
+
+        Returns:
+            A valid access token string.
+
+        Raises:
+            requests.RequestException: If the API request fails.
+            ValueError: If token cannot be obtained.
+        """
+        if installation_id is None:
+            org_login = org_name or self.org_name
+            if org_login:
+                installation_id = self.installation_id_for_org(org_login)
+            else:
+                if not self.app_id:
+                    raise ValueError("org_name or installation_id is required")
+                installation_id = int(self.app_id)
+
+        return self.installation_token(installation_id)
 
     def get_headers(self) -> dict:
         """Get HTTP headers for authenticated GitHub API requests.
@@ -189,7 +222,11 @@ def get_headers(self) -> dict:
 
 
 def create_github_jwt_session(
-    org_name: str, client_id: str, private_key_path: str, app_id: str
+    org_name: str | None,
+    client_id: str,
+    private_key_path: str,
+    app_id: str,
+    api_uri: str = "https://api.github.com",
 ) -> GitHubJwtSession:
     """Factory function to create a GitHub JWT session.
 
@@ -217,4 +254,5 @@ def create_github_jwt_session(
         client_id=client_id,
         private_key_path=private_key_path,
         app_id=app_id,
+        api_uri=api_uri,
     )