github: Rework GraphQL batching and error handling to salvage partial work

revup/forge_utils.py

@@ -125,6 +125,12 @@ async def forge_connection(
         try:
             yield forge
         finally:
+            if endpoint.last_ratelimit_remaining is not None:
+                logging.debug(
+                    "GitHub rate limit: {} points remaining, resets at {}".format(
+                        endpoint.last_ratelimit_remaining, endpoint.last_ratelimit_reset
+                    )
+                )
             await forge.close()
     else:
         raise RevupUsageException(

revup/github/endpoint.py

@@ -7,10 +7,37 @@
 
 from aiohttp import ClientSession, ContentTypeError
 
+from revup.github.graphql import GraphqlResponse
 from revup.types import RevupForgeException, RevupRequestException
 
+# HTTP statuses worth retrying: gateway/timeout (5xx) and secondary-rate-limit (403).
 TRANSIENT_STATUSES = frozenset({500, 502, 503, 504})
-RETRYABLE_GRAPHQL_ERRORS = frozenset({"RESOURCE_LIMITS_EXCEEDED"})
+SECONDARY_LIMIT_STATUS = 403
+# Cap how long we'll auto-sleep waiting for a rate limit to reset. A longer wait
+# (primary budget exhausted) is surfaced to the user instead of hanging silently.
+MAX_BACKOFF_SECONDS = 60.0
+
+
+def _backoff_delay(headers: Any, attempt: int, base_delay: float) -> float:
+    """Seconds to wait before retrying, driven by GitHub's rate-limit headers.
+
+    Precedence: an explicit Retry-After, then waiting out an exhausted budget
+    (remaining == 0) until its reset, else plain exponential backoff.
+    """
+    retry_after = headers.get("retry-after")
+    if retry_after is not None:
+        try:
+            return float(retry_after)
+        except ValueError:
+            pass
+    remaining = headers.get("x-ratelimit-remaining")
+    reset = headers.get("x-ratelimit-reset")
+    if remaining == "0" and reset is not None:
+        try:
+            return max(0.0, int(reset) - time.time())
+        except ValueError:
+            pass
+    return base_delay * (2**attempt)
 
 
 class GitHubEndpoint:
@@ -51,26 +78,20 @@ def __init__(
         self.oauth_token = oauth_token
         self.proxy = proxy
         self.graphql_endpoint = f"https://api.{github_url}/graphql"
+        # Rate-limit budget from the most recent response, for end-of-run reporting.
+        self.last_ratelimit_remaining: Optional[str] = None
+        self.last_ratelimit_reset: Optional[str] = None
 
     async def close(self) -> None:
         if self.session:
             await self.session.close()
 
-    async def _should_retry(
-        self, attempt: int, max_retries: int, base_delay: float, message: str
-    ) -> bool:
-        """Sleep with exponential backoff if retries remain. Returns True to retry."""
-        if attempt >= max_retries - 1:
-            return False
-        delay = base_delay * (2**attempt)
-        logging.warning(
-            "{}, retrying in {}s (attempt {}/{})".format(message, delay, attempt + 1, max_retries)
-        )
-        await asyncio.sleep(delay)
-        return True
-
-    async def _graphql_once(self, query: str, **kwargs: Any) -> Any:
-        """Execute a single GraphQL request. Raises on any error."""
+    async def _post(self, query: str, kwargs: Any) -> Tuple[int, Any, Any]:
+        """POST a GraphQL request. Returns (status, headers, body).
+
+        No policy: never raises for HTTP status or GraphQL errors. body is the
+        parsed JSON, or None if the response wasn't JSON.
+        """
         if self.session is None:
             self.session = ClientSession()
 
@@ -92,56 +113,64 @@ async def _graphql_once(self, query: str, **kwargs: Any) -> Any:
             logging.debug(
                 "Response status: {} took {}".format(resp.status, time.time() - start_time)
             )
-            ratelimit_reset = resp.headers.get("x-ratelimit-reset")
-            if ratelimit_reset is not None:
-                reset_timestamp = datetime.datetime.fromtimestamp(int(ratelimit_reset)).isoformat()
-            else:
-                reset_timestamp = "Unknown"
+            reset = resp.headers.get("x-ratelimit-reset")
+            reset_str = (
+                datetime.datetime.fromtimestamp(int(reset)).isoformat()
+                if reset is not None
+                else "Unknown"
+            )
+            self.last_ratelimit_remaining = resp.headers.get("x-ratelimit-remaining")
+            self.last_ratelimit_reset = reset_str
             logging.debug(
                 "Ratelimit: {} remaining, resets at {}".format(
-                    resp.headers.get("x-ratelimit-remaining"),
-                    reset_timestamp,
+                    self.last_ratelimit_remaining, reset_str
                 )
             )
-
-            if resp.status != 200:
-                try:
-                    r = await resp.json()
-                except (ValueError, ContentTypeError) as exc:
-                    logging.warning("Response body:\n{}".format(await resp.text()))
-                    raise RevupRequestException(resp.status, {}) from exc
-                raise RevupRequestException(resp.status, r)
-
             try:
-                r = await resp.json()
+                body = await resp.json()
+                logging.debug("Response JSON:\n{}".format(json.dumps(body, indent=1)))
             except (ValueError, ContentTypeError):
                 logging.warning("Response body:\n{}".format(await resp.text()))
-                raise
-            else:
-                pretty_json = json.dumps(r, indent=1)
-                logging.debug("Response JSON:\n{}".format(pretty_json))
-
-            if "errors" in r:
-                raise RevupForgeException(r["errors"])
-
-            return r
+                body = None
+            return resp.status, resp.headers, body
 
     async def graphql(
-        self, query: str, *, max_retries: int = 3, base_delay: float = 1.0, **kwargs: Any
-    ) -> Any:
+        self,
+        query: str,
+        *,
+        max_retries: int = 3,
+        base_delay: float = 1.0,
+        **kwargs: Any,
+    ) -> GraphqlResponse:
+        """Execute a GraphQL request, retrying transient failures with backoff.
+
+        Returns a GraphqlResponse carrying partial data and per-field errors — a
+        200 with field errors is NOT raised, so callers can salvage what resolved
+        and re-transact the rest. Raises RevupForgeException for a request error
+        (200 with no data) and RevupRequestException for a non-retryable HTTP error.
+        """
         for attempt in range(max_retries):
-            try:
-                return await self._graphql_once(query, **kwargs)
-            except RevupRequestException as e:
-                if e.status not in TRANSIENT_STATUSES:
-                    raise
-                msg = "GitHub returned {}".format(e.status)
-                if not await self._should_retry(attempt, max_retries, base_delay, msg):
-                    raise
-            except RevupForgeException as e:
-                retryable = set(e.types) & RETRYABLE_GRAPHQL_ERRORS
-                if not retryable:
-                    raise
-                msg = "GitHub GraphQL error ({})".format(", ".join(retryable))
-                if not await self._should_retry(attempt, max_retries, base_delay, msg):
-                    raise
+            status, headers, body = await self._post(query, kwargs)
+
+            if status == 200:
+                if body is None:
+                    raise RevupRequestException(status, {})
+                # A request error (bad query, unresolvable variable) returns no data
+                # and can't be salvaged or retried; fail loudly.
+                if body.get("data") is None and "errors" in body:
+                    raise RevupForgeException(body["errors"])
+                return GraphqlResponse.parse(body)
+
+            retryable = status in TRANSIENT_STATUSES or status == SECONDARY_LIMIT_STATUS
+            if not retryable or attempt >= max_retries - 1:
+                raise RevupRequestException(status, body if body is not None else {})
+
+            delay = min(_backoff_delay(headers, attempt, base_delay), MAX_BACKOFF_SECONDS)
+            logging.warning(
+                "GitHub returned {}, retrying in {}s (attempt {}/{})".format(
+                    status, delay, attempt + 1, max_retries
+                )
+            )
+            await asyncio.sleep(delay)
+
+        raise RevupRequestException(status, body if body is not None else {})