Merge pull request #1531 from github/handle-rate-limiting-in-graphql-queries

offbyone · web-flow · commit 27cf3c3f12f3 · 2026-04-01T10:19:17.000-07:00
Catch and retry rate limited responses in read-only GraphQL queries
diff --git a/src/Octoshift/Services/GithubApi.cs b/src/Octoshift/Services/GithubApi.cs
@@ -561,14 +561,14 @@ ... on Migration {
         {
             return await _retryPolicy.Retry(async () =>
             {
-                var data = await _client.PostGraphQLAsync(url, payload);
+                var data = await _client.PostGraphQLWithRetryAsync(url, payload);
 
                 return (
-                    State: (string)data["data"]["node"]["state"],
-                    RepositoryName: (string)data["data"]["node"]["repositoryName"],
-                    WarningsCount: (int)data["data"]["node"]["warningsCount"],
-                    FailureReason: (string)data["data"]["node"]["failureReason"],
-                    MigrationLogUrl: (string)data["data"]["node"]["migrationLogUrl"]);
+                        State: (string)data["data"]["node"]["state"],
+                        RepositoryName: (string)data["data"]["node"]["repositoryName"],
+                        WarningsCount: (int)data["data"]["node"]["warningsCount"],
+                        FailureReason: (string)data["data"]["node"]["failureReason"],
+                        MigrationLogUrl: (string)data["data"]["node"]["migrationLogUrl"]);
             });
         }
         catch (Exception ex)
diff --git a/src/Octoshift/Services/GithubClient.cs b/src/Octoshift/Services/GithubClient.cs
@@ -25,6 +25,9 @@ public class GithubClient
     private const int SECONDARY_RATE_LIMIT_MAX_RETRIES = 3;
     private const int SECONDARY_RATE_LIMIT_DEFAULT_DELAY = 60; // 60 seconds default delay
 
+    internal int _secondaryRateLimitMaxRetries = SECONDARY_RATE_LIMIT_MAX_RETRIES; // Exposed for testing purposes
+    internal int _secondaryRateLimitDefaultDelay = SECONDARY_RATE_LIMIT_DEFAULT_DELAY; // Exposed for testing purposes
+
     public GithubClient(OctoLogger log, HttpClient httpClient, IVersionProvider versionProvider, RetryPolicy retryPolicy, DateTimeProvider dateTimeProvider, string personalAccessToken)
     {
         _log = log;
@@ -93,6 +96,36 @@ public virtual async Task<JToken> PostGraphQLAsync(
         return data;
     }
 
+    public virtual async Task<JToken> PostGraphQLWithRetryAsync(
+        string url,
+        object body,
+        Dictionary<string, string> customHeaders = null,
+        int retryCount = 0)
+    {
+        var currentRetryCount = retryCount;
+
+        while (true)
+        {
+            var (response, headers) = await PostWithRetry(url, body, customHeaders);
+
+            if (IsGraphQLSecondaryRateLimit(response))
+            {
+                (response, _) = await HandleSecondaryRateLimit(HttpMethod.Post, url, body, HttpStatusCode.OK, customHeaders, headers, currentRetryCount);
+
+                if (IsGraphQLSecondaryRateLimit(response))
+                {
+                    currentRetryCount++;
+                    continue;
+                }
+            }
+
+            var data = JObject.Parse(response);
+            EnsureSuccessGraphQLResponse(data);
+
+            return data;
+        }
+    }
+
     public virtual async IAsyncEnumerable<JToken> PostGraphQLWithPaginationAsync(
         string url,
         object body,
@@ -156,6 +189,13 @@ public virtual async Task<string> PatchAsync(string url, object body, Dictionary
         HttpStatusCode expectedStatus = HttpStatusCode.OK) =>
         await _retryPolicy.Retry(async () => await SendAsync(HttpMethod.Get, url, customHeaders: customHeaders, expectedStatus: expectedStatus));
 
+    private async Task<(string Content, KeyValuePair<string, IEnumerable<string>>[] ResponseHeaders)> PostWithRetry(
+        string url,
+        object body,
+        Dictionary<string, string> customHeaders = null,
+        HttpStatusCode expectedStatus = HttpStatusCode.OK) =>
+        await _retryPolicy.Retry(async () => await SendAsync(HttpMethod.Post, url, body, customHeaders: customHeaders, expectedStatus: expectedStatus));
+
     private async Task<(string Content, KeyValuePair<string, IEnumerable<string>>[] ResponseHeaders)> SendAsync(
         HttpMethod httpMethod,
         string url,
@@ -312,6 +352,22 @@ private bool IsSecondaryRateLimit(HttpStatusCode statusCode, string content)
                statusCode == HttpStatusCode.TooManyRequests;
     }
 
+    private bool IsGraphQLSecondaryRateLimit(string content)
+    {
+        var response = JObject.Parse(content);
+
+        if (response.TryGetValue("errors", out var jErrors) && jErrors is JArray { Count: > 0 } errors)
+        {
+            var error = (JObject)errors[0];
+            if (error.TryGetValue("type", out var jType) && jType.Type == JTokenType.String)
+            {
+                return ((string)jType) == "RATE_LIMIT";
+            }
+        }
+
+        return false;
+    }
+
     private async Task<(string Content, KeyValuePair<string, IEnumerable<string>>[] ResponseHeaders)> HandleSecondaryRateLimit(
         HttpMethod httpMethod,
         string url,
@@ -321,7 +377,7 @@ private bool IsSecondaryRateLimit(HttpStatusCode statusCode, string content)
         KeyValuePair<string, IEnumerable<string>>[] headers,
         int retryCount = 0)
     {
-        if (retryCount >= SECONDARY_RATE_LIMIT_MAX_RETRIES)
+        if (retryCount >= _secondaryRateLimitMaxRetries)
         {
             throw new OctoshiftCliException($"Secondary rate limit exceeded. Maximum retries ({SECONDARY_RATE_LIMIT_MAX_RETRIES}) reached. Please wait before retrying your request.");
         }
@@ -359,6 +415,6 @@ private int GetSecondaryRateLimitDelay(KeyValuePair<string, IEnumerable<string>>
         }
 
         // Otherwise use exponential backoff: 1m → 2m → 4m
-        return SECONDARY_RATE_LIMIT_DEFAULT_DELAY * (int)Math.Pow(2, retryCount);
+        return _secondaryRateLimitDefaultDelay * (int)Math.Pow(2, retryCount);
     }
 }
diff --git a/src/OctoshiftCLI.Tests/Octoshift/Services/GithubApiTests.cs b/src/OctoshiftCLI.Tests/Octoshift/Services/GithubApiTests.cs
@@ -1282,7 +1282,7 @@ ... on Migration {
             }}");
 
         _githubClientMock
-            .Setup(m => m.PostGraphQLAsync(url, It.Is<object>(x => x.ToJson() == payload.ToJson()), null))
+            .Setup(m => m.PostGraphQLWithRetryAsync(url, It.Is<object>(x => x.ToJson() == payload.ToJson()), null, 0))
             .ReturnsAsync(response);
 
         // Act
@@ -1341,7 +1341,7 @@ ... on Migration {
             }}");
 
         _githubClientMock
-            .SetupSequence(m => m.PostGraphQLAsync(url, It.Is<object>(x => x.ToJson() == payload.ToJson()), null))
+            .SetupSequence(m => m.PostGraphQLWithRetryAsync(url, It.Is<object>(x => x.ToJson() == payload.ToJson()), null, 0))
             .Throws(new HttpRequestException(null, null, statusCode: HttpStatusCode.BadGateway))
             .Throws(new HttpRequestException(null, null, statusCode: HttpStatusCode.BadGateway))
             .ReturnsAsync(response);
@@ -1399,7 +1399,7 @@ ... on Migration {
             }}");
 
         _githubClientMock
-            .SetupSequence(m => m.PostGraphQLAsync(url, It.Is<object>(x => x.ToJson() == payload.ToJson()), null))
+            .SetupSequence(m => m.PostGraphQLWithRetryAsync(url, It.Is<object>(x => x.ToJson() == payload.ToJson()), null, 0))
             .ReturnsAsync(GQL_ERROR_RESPONSE)
             .ReturnsAsync(GQL_ERROR_RESPONSE)
             .ReturnsAsync(response);
@@ -1457,7 +1457,7 @@ ... on Migration {
             }}");
 
         _githubClientMock
-            .Setup(m => m.PostGraphQLAsync(url, It.Is<object>(x => x.ToJson() == payload.ToJson()), null))
+            .Setup(m => m.PostGraphQLWithRetryAsync(url, It.Is<object>(x => x.ToJson() == payload.ToJson()), null, 0))
             .ReturnsAsync(response);
 
         // Act
diff --git a/src/OctoshiftCLI.Tests/Octoshift/Services/GithubClientTests.cs b/src/OctoshiftCLI.Tests/Octoshift/Services/GithubClientTests.cs
@@ -2160,15 +2160,19 @@ public async Task SendAsync_Uses_Exponential_Backoff_When_No_Retry_Headers()
             .ReturnsAsync(CreateHttpResponseFactory(content: "SUCCESS_RESPONSE")());
 
         using var httpClient = new HttpClient(handlerMock.Object);
-        var githubClient = new GithubClient(_mockOctoLogger.Object, httpClient, null, _retryPolicy, _dateTimeProvider.Object, PERSONAL_ACCESS_TOKEN);
+        var githubClient = new GithubClient(_mockOctoLogger.Object, httpClient, null, _retryPolicy, _dateTimeProvider.Object, PERSONAL_ACCESS_TOKEN)
+        {
+            // Set short default delay to speed up the test. This is reflected in the logs below.
+            _secondaryRateLimitDefaultDelay = 0
+        };
 
         // Act
         var result = await githubClient.PatchAsync("http://example.com", _rawRequestBody);
 
         // Assert
         result.Should().Be("SUCCESS_RESPONSE");
-        _mockOctoLogger.Verify(m => m.LogWarning("Secondary rate limit detected (attempt 1/3). Waiting 60 seconds before retrying..."), Times.Once);
-        _mockOctoLogger.Verify(m => m.LogWarning("Secondary rate limit detected (attempt 2/3). Waiting 120 seconds before retrying..."), Times.Once);
+        _mockOctoLogger.Verify(m => m.LogWarning("Secondary rate limit detected (attempt 1/3). Waiting 0 seconds before retrying..."), Times.Once);
+        _mockOctoLogger.Verify(m => m.LogWarning("Secondary rate limit detected (attempt 2/3). Waiting 0 seconds before retrying..."), Times.Once);
     }
 
     [Fact]
@@ -2187,7 +2191,11 @@ public async Task SendAsync_Throws_Exception_After_Max_Secondary_Rate_Limit_Retr
                 content: "Too many requests")());
 
         using var httpClient = new HttpClient(handlerMock.Object);
-        var githubClient = new GithubClient(_mockOctoLogger.Object, httpClient, null, _retryPolicy, _dateTimeProvider.Object, PERSONAL_ACCESS_TOKEN);
+        var githubClient = new GithubClient(_mockOctoLogger.Object, httpClient, null, _retryPolicy, _dateTimeProvider.Object, PERSONAL_ACCESS_TOKEN)
+        {
+            // Set short default delay to speed up the test. This is reflected in the logs below.
+            _secondaryRateLimitDefaultDelay = 0
+        };
 
         // Act & Assert
         await FluentActions
@@ -2197,9 +2205,9 @@ await FluentActions
             .WithMessage("Secondary rate limit exceeded. Maximum retries (3) reached. Please wait before retrying your request.");
 
         // Verify all retry attempts were logged
-        _mockOctoLogger.Verify(m => m.LogWarning("Secondary rate limit detected (attempt 1/3). Waiting 60 seconds before retrying..."), Times.Once);
-        _mockOctoLogger.Verify(m => m.LogWarning("Secondary rate limit detected (attempt 2/3). Waiting 120 seconds before retrying..."), Times.Once);
-        _mockOctoLogger.Verify(m => m.LogWarning("Secondary rate limit detected (attempt 3/3). Waiting 240 seconds before retrying..."), Times.Once);
+        _mockOctoLogger.Verify(m => m.LogWarning("Secondary rate limit detected (attempt 1/3). Waiting 0 seconds before retrying..."), Times.Once);
+        _mockOctoLogger.Verify(m => m.LogWarning("Secondary rate limit detected (attempt 2/3). Waiting 0 seconds before retrying..."), Times.Once);
+        _mockOctoLogger.Verify(m => m.LogWarning("Secondary rate limit detected (attempt 3/3). Waiting 0 seconds before retrying..."), Times.Once);
     }
 
     private object CreateRepositoryMigration(string migrationId = null, string state = RepositoryMigrationStatus.Succeeded) => new
