BED-9434: harden repository GraphQL collection retries

src/openhound_github/auth.py

@@ -19,11 +19,22 @@
 logger = logging.getLogger(__name__)
 
 
-def _normalized_http_origin(url: str) -> tuple[str, str, int | None]:
+def _normalized_http_origin(
+    url: str, *, allow_query: bool = False
+) -> 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")
+    if (
+        parsed.username is not None
+        or parsed.password is not None
+        or (parsed.query and not allow_query)
+        or parsed.fragment
+    ):
+        raise ValueError(
+            "GitHub API URI must not contain user-info, query strings, or fragments"
+        )
 
     port = parsed.port
     if scheme == "https" and port == 443:
@@ -230,7 +241,9 @@ def token(self, force_refresh: bool = False) -> str | None:
     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 "")
+            request_origin = _normalized_http_origin(
+                request.url or "", allow_query=True
+            )
         except ValueError:
             return False
 

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
@@ -153,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]
@@ -162,12 +182,11 @@ def retry_policy(
 
         headers = response.headers
         now = int(time.time())
-        message = _response_message(response).lower()
 
         # DLT retries the same prepared request after long Retry-After sleeps.
         if (
             response.status_code == 401
-            and "bad credentials" in message
+            and "bad credentials" in _response_message(response).lower()
             and isinstance(auth, GitHubAppInstallationAuth)
             and response.request is not None
         ):
@@ -180,7 +199,15 @@ def retry_policy(
 
         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"):
@@ -199,6 +226,7 @@ def retry_policy(
         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,34 @@ 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,
                     }
+
+                page_info = repos_page.get("pageInfo") or {}
+                if isinstance(page_info, dict):
+                    repository_cursor = page_info.get("endCursor")
         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
 

tests/test_app_auth.py

@@ -120,6 +120,26 @@ def test_github_session_rejects_plaintext_api_uri() -> None:
         )
 
 
+@pytest.mark.parametrize(
+    "api_uri",
+    (
+        "https://user:password@ghe.example/api/v3/",
+        "https://ghe.example/api/v3/?token=secret",
+        "https://ghe.example/api/v3/#fragment",
+    ),
+)
+def test_github_session_rejects_api_uri_with_unsafe_components(api_uri: str) -> None:
+    with pytest.raises(
+        ValueError,
+        match="must not contain user-info, query strings, or fragments",
+    ):
+        GithubSession(
+            jwt_issuer="123456",
+            private_key_path="/tmp/github-app.pem",
+            api_uri=api_uri,
+        )
+
+
 def test_enterprise_source_reuses_selected_issuer_for_installation_tokens(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:

tests/test_github_app_retry.py

@@ -99,6 +99,23 @@ def test_refresh_request_refreshes_rejected_current_token() -> None:
     assert installation.token_calls == 1
 
 
+def test_refresh_request_allows_same_origin_request_query_string() -> None:
+    installation = FakeInstallation("new-token")
+    auth = GitHubAppInstallationAuth(installation=installation)
+    auth.access_token = "old-token"
+    auth.expires_at = datetime.now(timezone.utc) + timedelta(hours=1)
+    request = prepared_request(
+        "old-token",
+        url="https://api.github.com/repos/example/repo?page=2",
+    )
+
+    repaired = auth.refresh_request(request)
+
+    assert repaired is True
+    assert request.headers["Authorization"] == "Bearer new-token"
+    assert installation.token_calls == 1
+
+
 def test_refresh_request_reuses_token_refreshed_by_another_request() -> None:
     installation = BlockingFakeInstallation("new-token")
     auth = GitHubAppInstallationAuth(installation=installation)

tests/test_helpers.py

@@ -0,0 +1,102 @@
+import json
+
+import requests
+from dlt.sources.helpers.rest_client.auth import BearerTokenAuth
+from dlt.sources.helpers.requests.retry import Client
+from dlt.sources.helpers.requests.session import Session
+
+from openhound_github.helpers import github_retry_policy
+
+
+def graphql_response(
+    *,
+    headers: dict[str, str] | None = None,
+    body: dict[str, object] | None = None,
+    text: str | None = None,
+    url: str = "https://ghe.example/api/v3/graphql",
+) -> requests.Response:
+    response = requests.Response()
+    response.status_code = 200
+    response.headers.update(headers or {})
+    if text is not None:
+        response._content = text.encode("utf-8")
+    else:
+        response._content = json.dumps(body or {}).encode("utf-8")
+    response.request = requests.Request("POST", url).prepare()
+    return response
+
+
+def test_retry_policy_recognizes_graphql_endpoint_without_resource_header() -> None:
+    response = graphql_response(
+        headers={"Retry-After": "0"},
+        body={"errors": [{"message": "temporary GraphQL failure"}]},
+    )
+
+    should_retry = github_retry_policy(BearerTokenAuth(token="static-token"))(
+        response,
+        None,
+    )
+
+    assert should_retry is True
+
+
+def test_retry_policy_retries_malformed_graphql_json() -> None:
+    response = graphql_response(text='{"data":{"organization":')
+
+    should_retry = github_retry_policy(BearerTokenAuth(token="static-token"))(
+        response,
+        None,
+    )
+
+    assert should_retry is True
+
+
+def test_retry_policy_does_not_retry_malformed_non_graphql_json() -> None:
+    response = graphql_response(
+        text='{"data":{"organization":',
+        url="https://ghe.example/api/v3/repos/example/repo",
+    )
+
+    should_retry = github_retry_policy(BearerTokenAuth(token="static-token"))(
+        response,
+        None,
+    )
+
+    assert should_retry is False
+
+
+def test_retry_client_recovers_from_malformed_graphql_json(monkeypatch) -> None:
+    responses = [
+        graphql_response(text='{"data":{"organization":'),
+        graphql_response(body={"data": {"organization": {"repositories": {}}}}),
+    ]
+    requests_seen: list[requests.PreparedRequest] = []
+
+    def fake_send(
+        _session: Session,
+        request: requests.PreparedRequest,
+        **_kwargs,
+    ) -> requests.Response:
+        response = responses[len(requests_seen)]
+        response.request = request
+        requests_seen.append(request)
+        return response
+
+    monkeypatch.setattr(Session, "send", fake_send)
+    session = Client(
+        raise_for_status=False,
+        status_codes=(),
+        exceptions=(),
+        request_max_attempts=2,
+        request_backoff_factor=0,
+        retry_condition=github_retry_policy(BearerTokenAuth(token="static-token")),
+    ).session
+    request = requests.Request(
+        "POST",
+        "https://ghe.example/api/v3/graphql",
+    ).prepare()
+
+    response = session.send(request)
+
+    assert response.json() == {"data": {"organization": {"repositories": {}}}}
+    assert len(requests_seen) == 2