INT-5579 Decouple iterateApi into forEachPage and withErrorHandling

src/google-cloud/client.test.ts

@@ -1,13 +1,13 @@
-import { google } from 'googleapis';
-import { GoogleAuth, GoogleAuthOptions } from 'google-auth-library';
-import { getMockIntegrationConfig } from '../../test/config';
-import { Client, withErrorHandling } from './client';
-import { parseServiceAccountKeyFile } from '../utils/parseServiceAccountKeyFile';
 import {
   IntegrationProviderAPIError,
   IntegrationProviderAuthorizationError,
 } from '@jupiterone/integration-sdk-core';
+import { GoogleAuth, GoogleAuthOptions } from 'google-auth-library';
+import { google } from 'googleapis';
 import { IntegrationConfig } from '..';
+import { getMockIntegrationConfig } from '../../test/config';
+import { parseServiceAccountKeyFile } from '../utils/parseServiceAccountKeyFile';
+import { Client } from './client';
 
 describe('#getAuthenticatedServiceClient', () => {
   let googleAuthSpy: jest.SpyInstance<
@@ -75,6 +75,19 @@ describe('#getAuthenticatedServiceClient', () => {
 describe('withErrorHandling', () => {
   // Specific error handling for this method is tested in the index.test.ts files where the errors were seen. Ex: src/steps/compute/index.test.ts
 
+  const config = {
+    projectId: 'projectId',
+    serviceAccountKeyConfig: { project_id: 'serviceAccountProjectId' },
+  } as unknown as IntegrationConfig;
+
+  let client;
+  let onRetry;
+
+  beforeEach(() => {
+    onRetry = jest.fn();
+    client = new Client({ config, onRetry: onRetry });
+  });
+
   [IntegrationProviderAuthorizationError, IntegrationProviderAPIError].forEach(
     (J1Error) => {
       test('should forward on errors that have already been handled', async () => {
@@ -86,45 +99,37 @@ describe('withErrorHandling', () => {
         const executionHandler = jest
           .fn()
           .mockRejectedValue(mockForbiddenError);
-        const handledFunction = withErrorHandling(executionHandler);
-        await expect(handledFunction()).rejects.toThrow(J1Error);
+        const handledFunction = client.withErrorHandling(executionHandler);
+        await expect(handledFunction).rejects.toThrow(J1Error);
       });
     },
   );
 
   test('should handle errors of unknown format', async () => {
     const mockUnknownError = new Error() as any;
     const executionHandler = jest.fn().mockRejectedValue(mockUnknownError);
-    const handledFunction = withErrorHandling(executionHandler);
-    await expect(handledFunction()).rejects.toThrow(
-      IntegrationProviderAPIError,
-    );
+    const handledFunction = client.withErrorHandling(executionHandler);
+    await expect(handledFunction).rejects.toThrow(IntegrationProviderAPIError);
   });
 
   test('should throw an IntegrationProviderAPIError on all unknown errors', async () => {
     const executionHandler = jest
       .fn()
       .mockRejectedValue(new Error('Something esploded'));
 
-    const onRetry = jest.fn();
-    const handledFunction = withErrorHandling(executionHandler, { onRetry });
+    const handledFunction = client.withErrorHandling(executionHandler);
 
-    await expect(handledFunction()).rejects.toThrow(
-      IntegrationProviderAPIError,
-    );
+    await expect(handledFunction).rejects.toThrow(IntegrationProviderAPIError);
 
     expect(onRetry).toHaveBeenCalledTimes(0);
   });
 
   test('should pass parameters to the wrapped function return the result if no errors', async () => {
     const executionHandler = jest
       .fn()
-      .mockImplementation((...params) => Promise.resolve(params));
-    const handledFunction = withErrorHandling(executionHandler);
-    await expect(handledFunction('param1', 'param2')).resolves.toEqual([
-      'param1',
-      'param2',
-    ]);
+      .mockImplementation(() => Promise.resolve(['param1', 'param2']));
+    const response = await client.withErrorHandling(executionHandler);
+    expect(response).toEqual(['param1', 'param2']);
   });
 
   test('should retry if quota error received with 403 status code', async () => {
@@ -136,15 +141,11 @@ describe('withErrorHandling', () => {
     const executionHandler = jest
       .fn()
       .mockRejectedValueOnce(err)
-      .mockImplementationOnce((...params) => Promise.resolve(params));
+      .mockImplementationOnce(() => Promise.resolve(['param1', 'param2']));
 
-    const onRetry = jest.fn();
-    const handledFunction = withErrorHandling(executionHandler, { onRetry });
+    const response = await client.withErrorHandling(executionHandler);
 
-    await expect(handledFunction('param1', 'param2')).resolves.toEqual([
-      'param1',
-      'param2',
-    ]);
+    expect(response).toEqual(['param1', 'param2']);
 
     expect(onRetry).toHaveBeenCalledTimes(1);
     expect(onRetry).toHaveBeenCalledWith(err);
@@ -157,15 +158,11 @@ describe('withErrorHandling', () => {
     const executionHandler = jest
       .fn()
       .mockRejectedValueOnce(err)
-      .mockImplementationOnce((...params) => Promise.resolve(params));
+      .mockImplementationOnce(() => Promise.resolve(['param1', 'param2']));
 
-    const onRetry = jest.fn();
-    const handledFunction = withErrorHandling(executionHandler, { onRetry });
+    const response = await client.withErrorHandling(executionHandler);
 
-    await expect(handledFunction('param1', 'param2')).resolves.toEqual([
-      'param1',
-      'param2',
-    ]);
+    expect(response).toEqual(['param1', 'param2']);
 
     expect(onRetry).toHaveBeenCalledTimes(1);
     expect(onRetry).toHaveBeenCalledWith(err);

src/google-cloud/client.ts

@@ -1,14 +1,14 @@
-import { IntegrationConfig } from '../types';
-import { google } from 'googleapis';
-import { CredentialBody, BaseExternalAccountClient } from 'google-auth-library';
-import { GaxiosResponse } from 'gaxios';
 import {
-  IntegrationProviderAuthorizationError,
-  IntegrationProviderAPIError,
   IntegrationError,
+  IntegrationProviderAPIError,
+  IntegrationProviderAuthorizationError,
 } from '@jupiterone/integration-sdk-core';
-import { createErrorProps } from './utils/createErrorProps';
 import { retry } from '@lifeomic/attempt';
+import { GaxiosResponse } from 'gaxios';
+import { BaseExternalAccountClient, CredentialBody } from 'google-auth-library';
+import { google } from 'googleapis';
+import { IntegrationConfig } from '../types';
+import { createErrorProps } from './utils/createErrorProps';
 // import { GoogleCloudServiceApiDisabledError } from './errors';
 
 export interface ClientOptions {
@@ -39,21 +39,6 @@ export type IterateApiOptions = {
   onRetry?: (err: any) => void;
 };
 
-export async function iterateApi<T>(
-  fn: (nextPageToken?: string) => Promise<PageableGaxiosResponse<T>>,
-  callback: (data: T) => Promise<void>,
-  options?: IterateApiOptions,
-) {
-  let nextPageToken: string | undefined;
-
-  do {
-    const wrappedFn = withErrorHandling(fn, options);
-    const result = await wrappedFn(nextPageToken);
-    nextPageToken = result.data.nextPageToken || undefined;
-    await callback(result.data);
-  } while (nextPageToken);
-}
-
 export class Client {
   readonly projectId: string;
   readonly organizationId?: string;
@@ -100,24 +85,31 @@ export class Client {
     fn: (nextPageToken?: string) => Promise<PageableGaxiosResponse<T>>,
     callback: (data: T) => Promise<void>,
   ) {
-    return iterateApi(fn, callback, {
-      onRetry: this.onRetry,
+    return this.forEachPage(async (nextPageToken) => {
+      const result = await this.withErrorHandling(() => fn(nextPageToken));
+      await callback(result.data);
+
+      return result;
     });
   }
-}
 
-export type WithErrorHandlingOptions = {
-  onRetry?: (err: any) => void;
-};
+  async forEachPage<T>(
+    cb: (nextToken: string | undefined) => Promise<PageableGaxiosResponse<T>>,
+  ): Promise<any> {
+    let nextToken: string | undefined;
+    do {
+      const response = await cb(nextToken);
+      nextToken = response.data.nextPageToken
+        ? response.data.nextPageToken
+        : undefined;
+    } while (nextToken);
+  }
 
-export function withErrorHandling<T extends (...params: any) => any>(
-  fn: T,
-  options?: WithErrorHandlingOptions,
-) {
-  return async (...params: any) => {
+  withErrorHandling<T>(fn: () => Promise<T>) {
+    const onRetry = this.onRetry;
     return retry(
       async () => {
-        return await fn(...params);
+        return await fn();
       },
       {
         delay: 2_000,
@@ -130,13 +122,13 @@ export function withErrorHandling<T extends (...params: any) => any>(
           if (!newError.retryable) {
             ctx.abort();
             throw newError;
-          } else if (options?.onRetry) {
-            options.onRetry(err);
+          } else if (onRetry) {
+            onRetry(err);
           }
         },
       },
     );
-  };
+  }
 }
 
 /**

src/steps/cloud-run/client.ts

@@ -1,5 +1,5 @@
-import { Client } from '../../google-cloud/client';
 import { google, run_v1 } from 'googleapis';
+import { Client } from '../../google-cloud/client';
 
 export class CloudRunClient extends Client {
   private client = google.run({ version: 'v1', retry: false });
@@ -9,61 +9,51 @@ export class CloudRunClient extends Client {
   ): Promise<void> {
     const auth = await this.getAuthenticatedServiceClient();
 
-    await this.iterateApi(
-      async () => {
-        // Doesn't support pageToken
-        return this.client.namespaces.services.list({
-          parent: `namespaces/${this.projectId}`,
-          auth,
-        });
-      },
-      async (data: run_v1.Schema$ListServicesResponse) => {
-        for (const service of data.items || []) {
-          await callback(service);
-        }
-      },
+    const response = await this.withErrorHandling(async () =>
+      this.client.namespaces.services.list({
+        parent: `namespaces/${this.projectId}`,
+        auth,
+      }),
     );
+
+    for (const service of response.data.items || []) {
+      await callback(service);
+    }
   }
 
   async iterateCloudRunRoutes(
     callback: (data: run_v1.Schema$Route) => Promise<void>,
   ): Promise<void> {
     const auth = await this.getAuthenticatedServiceClient();
 
-    await this.iterateApi(
-      async () => {
-        // Doesn't support pageToken
-        return this.client.namespaces.routes.list({
+    const response = await this.withErrorHandling(
+      async () =>
+        await this.client.namespaces.routes.list({
           parent: `namespaces/${this.projectId}`,
           auth,
-        });
-      },
-      async (data: run_v1.Schema$ListRoutesResponse) => {
-        for (const route of data.items || []) {
-          await callback(route);
-        }
-      },
+        }),
     );
+
+    for (const route of response.data.items || []) {
+      await callback(route);
+    }
   }
 
   async iterateCloudRunConfigurations(
     callback: (data: run_v1.Schema$Configuration) => Promise<void>,
   ): Promise<void> {
     const auth = await this.getAuthenticatedServiceClient();
 
-    await this.iterateApi(
-      async () => {
-        // Doesn't support pageToken
-        return this.client.namespaces.configurations.list({
+    const response = await this.withErrorHandling(
+      async () =>
+        await this.client.namespaces.configurations.list({
           parent: `namespaces/${this.projectId}`,
           auth,
-        });
-      },
-      async (data: run_v1.Schema$ListConfigurationsResponse) => {
-        for (const configuration of data.items || []) {
-          await callback(configuration);
-        }
-      },
+        }),
     );
+
+    for (const configuration of response.data.items || []) {
+      await callback(configuration);
+    }
   }
 }

src/steps/compute/client.ts

@@ -1,7 +1,7 @@
-import { google, compute_v1 } from 'googleapis';
+import { BaseExternalAccountClient } from 'google-auth-library';
+import { compute_v1, google } from 'googleapis';
 import { Client, PageableGaxiosResponse } from '../../google-cloud/client';
 import { iterateRegions, iterateRegionZones } from '../../google-cloud/regions';
-import { BaseExternalAccountClient } from 'google-auth-library';
 
 export class ComputeClient extends Client {
   private client = google.compute({ version: 'v1', retry: false });

src/steps/sql-admin/client.ts

@@ -17,7 +17,7 @@ export class SQLAdminClient extends Client {
           pageToken: nextPageToken,
         });
       },
-      async (data: sqladmin_v1beta4.Schema$DatabasesListResponse) => {
+      async (data: sqladmin_v1beta4.Schema$InstancesListResponse) => {
         for (const sqlInstance of data.items || []) {
           await callback(sqlInstance);
         }