Merge pull request #28085 from sonikro/add-throttling-and-rate-limit-github-entity-provider

feat(GitHubOrgEntityProvider): adds throttling and rate limit handling to GithubOrgEntityProvider

plugins/catalog-backend-module-github/package.json

@@ -61,7 +61,9 @@
     "@backstage/plugin-catalog-common": "workspace:^",
     "@backstage/plugin-catalog-node": "workspace:^",
     "@backstage/plugin-events-node": "workspace:^",
-    "@octokit/graphql": "^5.0.0",
+    "@octokit/core": "^5.2.0",
+    "@octokit/graphql": "^7.0.2",
+    "@octokit/plugin-throttling": "^8.1.3",
     "@octokit/rest": "^19.0.3",
     "git-url-parse": "^15.0.0",
     "lodash": "^4.17.21",

plugins/catalog-backend-module-github/src/lib/github.test.ts

@@ -14,7 +14,10 @@
  * limitations under the License.
  */
 
-import { registerMswTestHooks } from '@backstage/backend-test-utils';
+import {
+  mockServices,
+  registerMswTestHooks,
+} from '@backstage/backend-test-utils';
 import { GroupEntity, UserEntity } from '@backstage/catalog-model';
 import { graphql as graphqlOctokit } from '@octokit/graphql';
 import { graphql as graphqlMsw, HttpResponse } from 'msw';
@@ -32,7 +35,17 @@ import {
   createAddEntitiesOperation,
   createRemoveEntitiesOperation,
   createReplaceEntitiesOperation,
+  createGraphqlClient,
 } from './github';
+import { Octokit } from '@octokit/core';
+import { throttling } from '@octokit/plugin-throttling';
+
+jest.mock('@octokit/core', () => ({
+  ...jest.requireActual('@octokit/core'),
+  Octokit: {
+    plugin: jest.fn().mockReturnValue({ defaults: jest.fn() }),
+  },
+}));
 
 describe('github', () => {
   const server = setupServer();
@@ -699,4 +712,81 @@ describe('github', () => {
       });
     });
   });
+
+  describe('createGraphqlClient', () => {
+    const headers = {};
+
+    const baseUrl = 'https://api.github.com';
+
+    const logger = mockServices.rootLogger();
+
+    const mockClient = jest.fn().mockImplementation();
+
+    const graphqlDefaults = jest.fn().mockReturnValue(mockClient);
+    const mockedOctokit = jest.fn().mockImplementation(() => ({
+      graphql: {
+        defaults: graphqlDefaults,
+      },
+    }));
+    (Octokit.plugin as jest.Mock).mockReturnValue(mockedOctokit);
+
+    const rateLimitOptions = {
+      method: 'POST',
+      url: '/graphql',
+    };
+    const client = createGraphqlClient({
+      headers,
+      baseUrl,
+      logger,
+    });
+    it('should return a graphql client with throttling', async () => {
+      expect(client).toBeDefined();
+      expect(Octokit.plugin).toHaveBeenCalledWith(throttling);
+    });
+
+    it('should return a graphql client with the correct options', async () => {
+      expect(graphqlDefaults).toHaveBeenCalledWith({
+        baseUrl,
+        headers,
+      });
+    });
+
+    describe('onRateLimit', () => {
+      it.each([
+        { retryCount: 0, expectedResult: true },
+        { retryCount: 1, expectedResult: true },
+        { retryCount: 2, expectedResult: false },
+      ])('should return %s', async ({ retryCount, expectedResult }) => {
+        const throttleOptions = mockedOctokit.mock.calls[0][0].throttle;
+
+        const result = throttleOptions.onRateLimit(
+          60,
+          rateLimitOptions,
+          undefined,
+          retryCount,
+        );
+
+        expect(result).toBe(expectedResult);
+      });
+    });
+
+    describe('onSecondaryRateLimit', () => {
+      it.each([
+        { retryCount: 0, expectedResult: true },
+        { retryCount: 1, expectedResult: true },
+        { retryCount: 2, expectedResult: false },
+      ])('should return %s', async ({ retryCount, expectedResult }) => {
+        const throttleOptions = mockedOctokit.mock.calls[0][0].throttle;
+
+        const result = throttleOptions.onSecondaryRateLimit(
+          60,
+          rateLimitOptions,
+          undefined,
+          retryCount,
+        );
+
+        expect(result).toBe(expectedResult);
+      });
+    });
+  });
 });

plugins/catalog-backend-module-github/src/lib/github.ts

@@ -27,7 +27,9 @@ import {
 import { withLocations } from './withLocations';
 
 import { DeferredEntity } from '@backstage/plugin-catalog-node';
-
+import { Octokit } from '@octokit/core';
+import { LoggerService } from '@backstage/backend-plugin-api';
+import { throttling } from '@octokit/plugin-throttling';
 // Graphql types
 
 export type QueryResponse = {
@@ -710,3 +712,58 @@ export const createReplaceEntitiesOperation =
       added: entitiesToReplace,
     };
   };
+
+/**
+ * Creates a GraphQL Client with Throttling
+ */
+export const createGraphqlClient = (args: {
+  headers:
+    | {
+        [name: string]: string;
+      }
+    | undefined;
+  baseUrl: string;
+  logger: LoggerService;
+}): typeof graphql => {
+  const { headers, baseUrl, logger } = args;
+  const ThrottledOctokit = Octokit.plugin(throttling);
+  const octokit = new ThrottledOctokit({
+    throttle: {
+      onRateLimit: (retryAfter, rateLimitData, _, retryCount) => {
+        logger.warn(
+          `Request quota exhausted for request ${rateLimitData?.method} ${rateLimitData?.url}`,
+        );
+
+        if (retryCount < 2) {
+          logger.warn(
+            `Retrying after ${retryAfter} seconds for the ${retryCount} time due to Rate Limit!`,
+          );
+          return true;
+        }
+
+        return false;
+      },
+      onSecondaryRateLimit: (retryAfter, rateLimitData, _, retryCount) => {
+        logger.warn(
+          `Secondary Rate Limit Exhausted for request ${rateLimitData?.method} ${rateLimitData?.url}`,
+        );
+
+        if (retryCount < 2) {
+          logger.warn(
+            `Retrying after ${retryAfter} seconds for the ${retryCount} time due to Secondary Rate Limit!`,
+          );
+          return true;
+        }
+
+        return false;
+      },
+    },
+  });
+
+  const client = octokit.graphql.defaults({
+    headers,
+    baseUrl,
+  });
+
+  return client;
+};

plugins/catalog-backend-module-github/src/providers/GithubOrgEntityProvider.test.ts

@@ -14,32 +14,37 @@
  * limitations under the License.
  */
 
+import { mockServices } from '@backstage/backend-test-utils';
 import { GroupEntity, UserEntity } from '@backstage/catalog-model';
 import {
   GithubCredentialsProvider,
   GithubIntegrationConfig,
 } from '@backstage/integration';
 import { EntityProviderConnection } from '@backstage/plugin-catalog-node';
-import { graphql } from '@octokit/graphql';
 import {
   DefaultEventsService,
   EventParams,
 } from '@backstage/plugin-events-node';
-import { GithubOrgEntityProvider } from './GithubOrgEntityProvider';
+import { graphql } from '@octokit/graphql';
+import { createGraphqlClient } from '../lib/github';
 import { withLocations } from '../lib/withLocations';
-import { mockServices } from '@backstage/backend-test-utils';
+import { GithubOrgEntityProvider } from './GithubOrgEntityProvider';
 
 jest.mock('@octokit/graphql');
+jest.mock('../lib/github', () => ({
+  ...jest.requireActual('../lib/github'),
+  createGraphqlClient: jest.fn(),
+}));
 
 describe('GithubOrgEntityProvider', () => {
   describe('read', () => {
-    let mockClient;
+    let mockClient: any;
     let entityProviderConnection: EntityProviderConnection;
     let entityProvider: GithubOrgEntityProvider;
 
     const setupMocks = (response: ((...args: any) => any) | undefined) => {
       mockClient = jest.fn().mockImplementation(response);
-      (graphql.defaults as jest.Mock).mockReturnValue(mockClient);
+      (createGraphqlClient as jest.Mock).mockReturnValue(mockClient);
     };
 
     beforeEach(() => {

plugins/catalog-backend-module-github/src/providers/GithubOrgEntityProvider.ts

@@ -14,7 +14,10 @@
  * limitations under the License.
  */
 
-import { SchedulerServiceTaskRunner } from '@backstage/backend-plugin-api';
+import {
+  LoggerService,
+  SchedulerServiceTaskRunner,
+} from '@backstage/backend-plugin-api';
 import { Entity, isGroupEntity } from '@backstage/catalog-model';
 import { Config } from '@backstage/config';
 import {
@@ -47,6 +50,7 @@ import {
 } from '../lib/defaultTransformers';
 import {
   createAddEntitiesOperation,
+  createGraphqlClient,
   createRemoveEntitiesOperation,
   createReplaceEntitiesOperation,
   DeferredEntitiesBuilder,
@@ -56,11 +60,10 @@ import {
   getOrganizationUsers,
   GithubTeam,
 } from '../lib/github';
+import { areGroupEntities, areUserEntities } from '../lib/guards';
 import { assignGroupsToUsers, buildOrgHierarchy } from '../lib/org';
 import { parseGithubOrgUrl } from '../lib/util';
 import { withLocations } from '../lib/withLocations';
-import { areGroupEntities, areUserEntities } from '../lib/guards';
-import { LoggerService } from '@backstage/backend-plugin-api';
 
 const EVENT_TOPICS = [
   'github.membership',
@@ -220,9 +223,11 @@ export class GithubOrgEntityProvider implements EntityProvider {
       await this.credentialsProvider.getCredentials({
         url: this.options.orgUrl,
       });
-    const client = graphql.defaults({
-      baseUrl: this.options.gitHubConfig.apiBaseUrl,
+
+    const client = createGraphqlClient({
       headers,
+      baseUrl: this.options.gitHubConfig.apiBaseUrl!,
+      logger,
     });
 
     const { org } = parseGithubOrgUrl(this.options.orgUrl);