BED-9434: harden repository GraphQL collection retries

src/openhound_github/auth.py

@@ -2,6 +2,8 @@
 from datetime import datetime, timedelta, timezone
 from threading import Lock
 from typing import Iterator
+from urllib.parse import urlparse
+from weakref import WeakKeyDictionary
 
 import requests
 from dlt.common.configuration import configspec
@@ -17,6 +19,19 @@
 logger = logging.getLogger(__name__)
 
 
+def _normalized_http_origin(url: str) -> tuple[str, str, int | None]:
+    parsed = urlparse(url)
+    scheme = parsed.scheme.lower()
+    if scheme != "https" or not parsed.hostname:
+        raise ValueError("GitHub API URI must be an absolute HTTPS URL")
+
+    port = parsed.port
+    if scheme == "https" and port == 443:
+        port = None
+
+    return scheme, parsed.hostname.lower(), port
+
+
 class AccountConfig(BaseModel):
     id: int
     login: str | None = None
@@ -63,7 +78,8 @@ def __init__(
         private_key_path: str,
         api_uri: str = "https://api.github.com/",
     ):
-        self.api_uri = api_uri
+        _normalized_http_origin(api_uri)
+        self.api_uri = f"{api_uri.rstrip('/')}/"
         self.jwt_issuer = jwt_issuer
         self.private_key_path = private_key_path
         self.client = RESTClient(
@@ -158,11 +174,28 @@ def __init__(
         self,
         installation: GithubInstallation,
         refresh_margin_seconds: int = 300,
+        api_uri: str | None = None,
     ):
         self.installation = installation
         self.refresh_margin_seconds = refresh_margin_seconds
+        installation_api_uri = getattr(
+            installation, "api_uri", "https://api.github.com/"
+        )
+        selected_api_uri = api_uri or installation_api_uri
+        installation_api_origin = _normalized_http_origin(installation_api_uri)
+        selected_api_origin = _normalized_http_origin(selected_api_uri)
+        if selected_api_origin != installation_api_origin:
+            raise ValueError(
+                "GitHub App auth API URI origin must match installation API URI origin"
+            )
+
+        self.api_uri = selected_api_uri
+        self._api_origin = selected_api_origin
         self.access_token: str | None = None
         self.expires_at: datetime | None = None
+        self._response_refreshed_requests: WeakKeyDictionary[
+            requests.PreparedRequest, str
+        ] = WeakKeyDictionary()
         self._token_lock = Lock()
 
     def _should_refresh(self) -> bool:
@@ -172,6 +205,14 @@ def _should_refresh(self) -> bool:
         refresh_at = self.expires_at - timedelta(seconds=self.refresh_margin_seconds)
         return datetime.now(timezone.utc) >= refresh_at
 
+    def _refresh_token(self) -> None:
+        logger.info(
+            f"Refreshing access token for {self.installation.installation_id}"
+        )
+        get_token = self.installation.token
+        self.access_token = get_token.token
+        self.expires_at = get_token.expires_at
+
     def token(self, force_refresh: bool = False) -> str | None:
         if (
             not force_refresh
@@ -182,15 +223,62 @@ def token(self, force_refresh: bool = False) -> str | None:
 
         with self._token_lock:
             if (force_refresh or self._should_refresh()) or self.access_token is None:
-                logger.info(
-                    f"Refreshing access token for {self.installation.installation_id}"
-                )
-                get_token = self.installation.token
-                self.access_token = get_token.token
-                self.expires_at = get_token.expires_at
+                self._refresh_token()
 
         return self.access_token
 
+    def refresh_request(self, request: requests.PreparedRequest) -> bool:
+        """Repair a rejected same-origin request without stampeding token issuance."""
+        try:
+            request_origin = _normalized_http_origin(request.url or "")
+        except ValueError:
+            return False
+
+        if request_origin != self._api_origin:
+            return False
+
+        request_authorization = request.headers.get("Authorization")
+        if (
+            not request_authorization
+            or not request_authorization.startswith("Bearer ")
+            or not request_authorization.removeprefix("Bearer ").strip()
+        ):
+            return False
+
+        with self._token_lock:
+            current_authorization = (
+                f"Bearer {self.access_token}" if self.access_token is not None else None
+            )
+            should_refresh = self._should_refresh()
+            repaired_authorization = self._response_refreshed_requests.get(request)
+            if (
+                not should_refresh
+                and request_authorization == current_authorization
+                and request_authorization == repaired_authorization
+            ):
+                return False
+
+            if (
+                self.access_token is None
+                or should_refresh
+                or request_authorization == current_authorization
+            ):
+                try:
+                    self._refresh_token()
+                except Exception:
+                    logger.warning(
+                        "Failed to refresh GitHub App installation token for "
+                        "installation %s during request retry",
+                        self.installation.installation_id,
+                    )
+                    return False
+
+            replacement_authorization = f"Bearer {self.access_token}"
+            request.headers["Authorization"] = replacement_authorization
+            self._response_refreshed_requests[request] = replacement_authorization
+
+        return True
+
     def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest:
         request.headers["Authorization"] = f"Bearer {self.token()}"
         return request

src/openhound_github/helpers.py

@@ -1,6 +1,7 @@
 import logging
 import time
 from typing import Optional
+from urllib.parse import urlparse
 
 from dlt.common import jsonpath
 from dlt.sources.helpers import requests
@@ -10,6 +11,8 @@
 )
 from requests import Request
 
+from openhound_github.auth import GitHubAppInstallationAuth
+
 logger = logging.getLogger(__name__)
 
 
@@ -151,6 +154,25 @@ def _has_graphql_errors(response: requests.Response) -> bool:
     return isinstance(response_data, dict) and bool(response_data.get("errors"))
 
 
+def _has_invalid_json_body(response: requests.Response) -> bool:
+    try:
+        response.json()
+    except ValueError:
+        return True
+    return False
+
+
+def _is_graphql_response(response: requests.Response) -> bool:
+    if response.headers.get("x-ratelimit-resource") == "graphql":
+        return True
+
+    request = response.request
+    if request is None or not request.url:
+        return False
+
+    return urlparse(request.url).path.rstrip("/").endswith("/graphql")
+
+
 def github_retry_policy(auth: AuthConfigBase):
     def retry_policy(
         response: Optional[requests.Response], exception: Optional[BaseException]
@@ -160,9 +182,32 @@ def retry_policy(
 
         headers = response.headers
         now = int(time.time())
+
+        # DLT retries the same prepared request after long Retry-After sleeps.
+        if (
+            response.status_code == 401
+            and "bad credentials" in _response_message(response).lower()
+            and isinstance(auth, GitHubAppInstallationAuth)
+            and response.request is not None
+        ):
+            if not auth.refresh_request(response.request):
+                return False
+            logger.warning(
+                "GitHub App installation token rejected, retrying request with refreshed token"
+            )
+            return True
+
         if (
             response.status_code == 200
-            and headers.get("x-ratelimit-resource") == "graphql"
+            and _is_graphql_response(response)
+            and _has_invalid_json_body(response)
+        ):
+            logger.warning("GraphQL response body was not valid JSON, retrying request")
+            return True
+
+        if (
+            response.status_code == 200
+            and _is_graphql_response(response)
             and _has_graphql_errors(response)
         ):
             if headers.get("Retry-After"):
@@ -178,10 +223,10 @@ def retry_policy(
                 return True
             return False
 
-        message = _response_message(response).lower()
         if response.status_code not in (403, 429):
             return False
 
+        message = _response_message(response).lower()
         if (
             headers.get("x-ratelimit-remaining") == "0"
             or "api rate limit exceeded" in message

src/openhound_github/models/repository.py

@@ -29,6 +29,7 @@ class GHRepositoryProperties(GHNodeProperties):
         disabled: Whether the repository is disabled.
         visibility: The visibility level: `public`, `private`, or `internal`.
         default_branch: The name of the default branch (e.g., `main`).
+        size: Repository size in kilobytes as reported by GitHub.
         open_issues_count: Number of open issues.
         allow_forking: Whether forking is allowed.
         web_commit_signoff_required: Whether web-based commits require sign-off.
@@ -74,6 +75,7 @@ class GHRepositoryProperties(GHNodeProperties):
     disabled: bool | None = None
     visibility: str | None = None
     default_branch: str | None = None
+    size: int | None = None
     open_issues_count: int | None = None
     allow_forking: bool | None = None
     web_commit_signoff_required: bool | None = None
@@ -195,6 +197,7 @@ class Repository(BaseAsset):
     disabled: bool | None = None
     visibility: str | None = None
     default_branch: str | None = None
+    size: int | None = None
     open_issues_count: int | None = None
     allow_forking: bool | None = None
     web_commit_signoff_required: bool | None = None
@@ -240,6 +243,7 @@ def as_node(self) -> GHNode:
                 disabled=self.disabled,
                 visibility=self.visibility,
                 default_branch=self.default_branch,
+                size=self.size,
                 open_issues_count=self.open_issues_count,
                 allow_forking=self.allow_forking,
                 web_commit_signoff_required=self.web_commit_signoff_required,

src/openhound_github/resources/organization.py

@@ -916,6 +916,8 @@ def repositories_graphql(ctx: SourceContext):
     for org in ctx.organizations:
         org_name = org.org_name
         client = org.client
+        repository_cursor: str | None = None
+        emitted_repositories = 0
         try:
             paginator = GraphQLCursorPaginator(
                 page_info_path="data.organization.repositories.pageInfo",
@@ -939,15 +941,36 @@ def repositories_graphql(ctx: SourceContext):
                 for repo in repos_page["nodes"]:
                     repo_record = {**repo}
                     branch_rulesets = repo_record.pop("branchRulesets", None) or {}
+                    emitted_repositories += 1
                     yield {
                         **repo_record,
                         "branch_ruleset_count": branch_rulesets.get("totalCount"),
                         "org_login": org_name,
                     }
+
+                request_json = getattr(getattr(page_data, "request", None), "json", None)
+                if isinstance(request_json, dict):
+                    variables = request_json.get("variables")
+                    if isinstance(variables, dict):
+                        repository_cursor = variables.get("after")
         except Exception as e:
             logger.error(
-                f"Error in resource 'repositories_graphql' processing organization '{org_name}': {e}",
-                extra={"resource": "repositories_graphql", "phase": "resource_iteration"},
+                "Error in resource 'repositories_graphql' processing organization '%s' "
+                "at repository cursor %r after emitting %d repositories "
+                "(%s): %s",
+                org_name,
+                repository_cursor,
+                emitted_repositories,
+                type(e).__name__,
+                e,
+                extra={
+                    "resource": "repositories_graphql",
+                    "phase": "resource_iteration",
+                    "org_name": org_name,
+                    "repository_cursor": repository_cursor,
+                    "emitted_repositories": emitted_repositories,
+                    "error_type": type(e).__name__,
+                },
             )
             continue
 

src/openhound_github/source.py

@@ -172,19 +172,24 @@ def token_client(token: str) -> RESTClient:
         github_app_session = GithubApp(
             jwt_issuer=jwt_issuer,
             private_key_path=credentials.key_path,
+            api_uri=host,
         )
         for installation in github_app_session.installations:
             if installation.target_type == "Organization":
                 org_installation = GithubInstallation(
                     installation_id=installation.id,
                     jwt_issuer=jwt_issuer,
                     private_key_path=credentials.key_path,
+                    api_uri=host,
                 )
                 ctx.organizations.append(
                     OrgContext(
                         org_name=installation.account.login,
                         client=client(
-                            GitHubAppInstallationAuth(installation=org_installation)
+                            GitHubAppInstallationAuth(
+                                installation=org_installation,
+                                api_uri=host,
+                            )
                         ),
                         enterprise_name=credentials.enterprise_name,
                         github_deployment_id=github_deployment_id,
@@ -196,9 +201,13 @@ def token_client(token: str) -> RESTClient:
                     installation_id=installation.id,
                     jwt_issuer=jwt_issuer,
                     private_key_path=credentials.key_path,
+                    api_uri=host,
                 )
                 ctx.client = client(
-                    GitHubAppInstallationAuth(installation=es_installation)
+                    GitHubAppInstallationAuth(
+                        installation=es_installation,
+                        api_uri=host,
+                    )
                 )
 
         return (*enterprise_resources(ctx), *organization_resources(ctx))
@@ -213,11 +222,17 @@ def token_client(token: str) -> RESTClient:
             installation_id=credentials.install_id,
             jwt_issuer=credentials.client_id,
             private_key_path=credentials.key_path,
+            api_uri=host,
         )
         ctx.organizations.append(
             OrgContext(
                 org_name=credentials.org_name,
-                client=client(GitHubAppInstallationAuth(installation=org_installation)),
+                client=client(
+                    GitHubAppInstallationAuth(
+                        installation=org_installation,
+                        api_uri=host,
+                    )
+                ),
                 github_deployment_id=github_deployment_id,
                 github_web_origin=github_web_origin,
             )