Merge pull request #642 from JupiterOne/INT-10459-4

NT-10459: add retry to client get function calls
JupiterOne-Archives · Feb 23, 2024 · 46e7b4e · 46e7b4e
2 parents 685222a + d460f61
commit 46e7b4e
Show file tree

Hide file tree

Showing 16 changed files with 162 additions and 115 deletions.
diff --git a/src/google-cloud/client.ts b/src/google-cloud/client.ts
@@ -77,8 +77,10 @@ export class Client {
       scopes: ['https://www.googleapis.com/auth/cloud-platform'],
     });
 
-    const client = (await auth.getClient()) as BaseExternalAccountClient;
-    await client.getAccessToken();
+    const client = (await this.withErrorHandling(() =>
+      auth.getClient(),
+    )) as BaseExternalAccountClient;
+    await this.withErrorHandling(() => client.getAccessToken());
     return client;
   }
 

diff --git a/src/steps/api-gateway/client.ts b/src/steps/api-gateway/client.ts
@@ -23,11 +23,12 @@ export class ApiGatewayClient extends Client {
   ): Promise<apigateway_v1.Schema$ApigatewayPolicy> {
     const auth = await this.getAuthenticatedServiceClient();
 
-    const result =
-      await this.client.projects.locations.apis.configs.getIamPolicy({
+    const result = await this.withErrorHandling(() =>
+      this.client.projects.locations.apis.configs.getIamPolicy({
         resource: `projects/${this.projectId}/locations/global/apis/${apiId}/configs/${configId}`,
         auth,
-      });
+      }),
+    );
 
     return result.data;
   }
@@ -37,10 +38,12 @@ export class ApiGatewayClient extends Client {
   ): Promise<apigateway_v1.Schema$ApigatewayPolicy> {
     const auth = await this.getAuthenticatedServiceClient();
 
-    const result = await this.client.projects.locations.gateways.getIamPolicy({
-      resource: `projects/${this.projectId}/locations/global/gateways/${gatewayId}`,
-      auth,
-    });
+    const result = await this.withErrorHandling(() =>
+      this.client.projects.locations.gateways.getIamPolicy({
+        resource: `projects/${this.projectId}/locations/global/gateways/${gatewayId}`,
+        auth,
+      }),
+    );
 
     return result.data;
   }

diff --git a/src/steps/app-engine/client.ts b/src/steps/app-engine/client.ts
@@ -7,10 +7,12 @@ export class AppEngineClient extends Client {
   async getAppEngineApplication() {
     const auth = await this.getAuthenticatedServiceClient();
 
-    const response = await this.client.apps.get({
-      appsId: this.projectId,
-      auth,
-    });
+    const response = await this.withErrorHandling(() =>
+      this.client.apps.get({
+        appsId: this.projectId,
+        auth,
+      }),
+    );
 
     return response.data;
   }

diff --git a/src/steps/big-query/client.ts b/src/steps/big-query/client.ts
@@ -57,10 +57,12 @@ export class BigQueryClient extends Client {
     if (!projectId || !datasetId || !tableId) {
       return undefined;
     }
-    const policyResponse = await this.client.tables.getIamPolicy({
-      auth,
-      resource: `projects/${projectId}/datasets/${datasetId}/tables/${tableId}`,
-    });
+    const policyResponse = await this.withErrorHandling(() =>
+      this.client.tables.getIamPolicy({
+        auth,
+        resource: `projects/${projectId}/datasets/${datasetId}/tables/${tableId}`,
+      }),
+    );
     return policyResponse?.data;
   }
 
@@ -69,12 +71,14 @@ export class BigQueryClient extends Client {
   ): Promise<bigquery_v2.Schema$Table> {
     const auth = await this.getAuthenticatedServiceClient();
 
-    const resp = await this.client.tables.get({
-      auth,
-      projectId: data.tableReference?.projectId!,
-      datasetId: data.tableReference?.datasetId!,
-      tableId: data.tableReference?.tableId!,
-    });
+    const resp = await this.withErrorHandling(() =>
+      this.client.tables.get({
+        auth,
+        projectId: data.tableReference?.projectId!,
+        datasetId: data.tableReference?.datasetId!,
+        tableId: data.tableReference?.tableId!,
+      }),
+    );
 
     return resp.data;
   }
@@ -95,11 +99,13 @@ export class BigQueryClient extends Client {
       async (data: bigquery_v2.Schema$DatasetList) => {
         for (const datasetRef of data.datasets || []) {
           if (datasetRef?.datasetReference?.datasetId) {
-            const dataset = await this.client.datasets.get({
-              auth,
-              projectId: this.projectId,
-              datasetId: datasetRef.datasetReference?.datasetId,
-            });
+            const dataset = await this.withErrorHandling(() =>
+              this.client.datasets.get({
+                auth,
+                projectId: this.projectId,
+                datasetId: datasetRef.datasetReference?.datasetId as string,
+              }),
+            );
 
             await callback(dataset.data);
           }
@@ -126,12 +132,14 @@ export class BigQueryClient extends Client {
       async (data: bigquery_v2.Schema$ListModelsResponse) => {
         for (const modelRef of data.models || []) {
           if (modelRef.modelReference?.modelId) {
-            const model = await this.client.models.get({
-              auth,
-              projectId: this.projectId,
-              datasetId,
-              modelId: modelRef.modelReference.modelId,
-            });
+            const model = await this.withErrorHandling(() =>
+              this.client.models.get({
+                auth,
+                projectId: this.projectId,
+                datasetId: datasetId,
+                modelId: modelRef.modelReference?.modelId as string,
+              }),
+            );
             await callback(model.data);
           }
         }

diff --git a/src/steps/binary-authorization/client.ts b/src/steps/binary-authorization/client.ts
@@ -11,10 +11,12 @@ export class BinaryAuthorizationClient extends Client {
   async fetchPolicy(): Promise<binaryauthorization_v1.Schema$Policy> {
     const auth = await this.getAuthenticatedServiceClient();
 
-    const result: PolicyResponse = await this.client.projects.getPolicy({
-      auth,
-      name: `projects/${this.projectId}/policy`,
-    });
+    const result: PolicyResponse = await this.withErrorHandling(() =>
+      this.client.projects.getPolicy({
+        auth,
+        name: `projects/${this.projectId}/policy`,
+      }),
+    );
 
     return result.data;
   }

diff --git a/src/steps/cloud-build/client.ts b/src/steps/cloud-build/client.ts
@@ -86,10 +86,12 @@ export class CloudBuildClient extends Client {
     const auth = await this.getAuthenticatedServiceClient();
 
     try {
-      const res = await this.client.projects.githubEnterpriseConfigs.list({
-        auth,
-        parent: `projects/${this.projectId}`,
-      });
+      const res = await this.withErrorHandling(() =>
+        this.client.projects.githubEnterpriseConfigs.list({
+          auth,
+          parent: `projects/${this.projectId}`,
+        }),
+      );
 
       if (res.data?.configs) {
         for (const config of res.data.configs) {

diff --git a/src/steps/compute/client.ts b/src/steps/compute/client.ts
@@ -141,24 +141,28 @@ export class ComputeClient extends Client {
   async fetchComputeImagePolicy(name: string) {
     const auth = await this.getAuthenticatedServiceClient();
 
-    const resp = await this.client.images.getIamPolicy({
-      auth,
-      project: this.projectId,
-      resource: name,
-    });
+    const resp = await this.withErrorHandling(() =>
+      this.client.images.getIamPolicy({
+        auth,
+        project: this.projectId,
+        resource: name,
+      }),
+    );
 
     return resp.data;
   }
 
   async fetchComputeImage(name: string, projectId: string) {
     const auth = await this.getAuthenticatedServiceClient();
 
-    const resp = await this.client.images.get({
-      auth,
-      image: name,
-      // allow us to use the same method for both custom and public images
-      project: projectId,
-    });
+    const resp = await this.withErrorHandling(() =>
+      this.client.images.get({
+        auth,
+        image: name,
+        // allow us to use the same method for both custom and public images
+        project: projectId,
+      }),
+    );
 
     return resp.data;
   }
@@ -231,10 +235,12 @@ export class ComputeClient extends Client {
   async fetchComputeProject(): Promise<compute_v1.Schema$Project> {
     const auth = await this.getAuthenticatedServiceClient();
 
-    const computeProjectResponse = await this.client.projects.get({
-      auth: auth,
-      project: this.projectId,
-    });
+    const computeProjectResponse = await this.withErrorHandling(() =>
+      this.client.projects.get({
+        auth: auth,
+        project: this.projectId,
+      }),
+    );
 
     return computeProjectResponse.data;
   }
@@ -270,11 +276,12 @@ export class ComputeClient extends Client {
   ): Promise<osconfig_v1.Schema$Inventory> {
     const auth = await this.getAuthenticatedServiceClient();
 
-    const resp =
-      await this.osConfigClient.projects.locations.instances.inventories.get({
+    const resp = await this.withErrorHandling(() =>
+      this.osConfigClient.projects.locations.instances.inventories.get({
         auth,
         name: `projects/${this.projectId}/locations/${location}/instances/${instanceId}/inventory`,
-      });
+      }),
+    );
 
     return resp.data;
   }

diff --git a/src/steps/compute/steps/fetch-compute-instances.ts b/src/steps/compute/steps/fetch-compute-instances.ts
@@ -134,7 +134,7 @@ export async function fetchComputeInstances(
       }
     } catch (e) {
       // Do not make this inventory call if api is disabled and customer is not using this feature.
-      if (e.response.status === 403) {
+      if (e.code === 403) {
         inventoryApiDisabled = true;
       }
     }

diff --git a/src/steps/containers/client.ts b/src/steps/containers/client.ts
@@ -9,10 +9,12 @@ export class ContainerClient extends Client {
   ): Promise<void> {
     const auth = await this.getAuthenticatedServiceClient();
 
-    const result = await this.client.projects.locations.clusters.list({
-      auth,
-      parent: `projects/${this.projectId}/locations/-`,
-    });
+    const result = await this.withErrorHandling(() =>
+      this.client.projects.locations.clusters.list({
+        auth,
+        parent: `projects/${this.projectId}/locations/-`,
+      }),
+    );
 
     for (const cluster of result.data.clusters || []) {
       await callback(cluster);

diff --git a/src/steps/iam/client.ts b/src/steps/iam/client.ts
@@ -101,10 +101,12 @@ export class IamClient extends Client {
     callback: (data: iam_v1.Schema$ServiceAccountKey) => Promise<void>,
   ): Promise<void> {
     const auth = await this.getAuthenticatedServiceClient();
-    const response = await this.client.projects.serviceAccounts.keys.list({
-      auth,
-      name: serviceAccountName,
-    });
+    const response = await this.withErrorHandling(() =>
+      this.client.projects.serviceAccounts.keys.list({
+        auth,
+        name: serviceAccountName,
+      }),
+    );
 
     for (const k of response.data.keys || []) {
       await callback(k);
@@ -116,12 +118,14 @@ export class IamClient extends Client {
   ): Promise<void> {
     const auth = await this.getAuthenticatedServiceClient();
 
-    const response = await this.client.iamPolicies.queryAuditableServices({
-      auth,
-      requestBody: {
-        fullResourceName: `//cloudresourcemanager.googleapis.com/projects/${this.projectId}`,
-      },
-    });
+    const response = await this.withErrorHandling(() =>
+      this.client.iamPolicies.queryAuditableServices({
+        auth,
+        requestBody: {
+          fullResourceName: `//cloudresourcemanager.googleapis.com/projects/${this.projectId}`,
+        },
+      }),
+    );
 
     for (const service of response.data.services || []) {
       const name = service.name;

diff --git a/src/steps/kms/client.ts b/src/steps/kms/client.ts
@@ -93,11 +93,12 @@ export class CloudKmsClient extends Client {
   async fetchCryptoKeyPolicy(resource: string) {
     const auth = await this.getAuthenticatedServiceClient();
 
-    const result =
-      await this.client.projects.locations.keyRings.cryptoKeys.getIamPolicy({
+    const result = await this.withErrorHandling(() =>
+      this.client.projects.locations.keyRings.cryptoKeys.getIamPolicy({
         auth,
         resource,
-      });
+      }),
+    );
 
     return result.data;
   }

diff --git a/src/steps/orgpolicy/client.ts b/src/steps/orgpolicy/client.ts
@@ -8,10 +8,12 @@ export class OrgPolicyClient extends Client {
     boolean | undefined
   > {
     const auth = await this.getAuthenticatedServiceClient();
-    const resp = await this.client.projects.policies.getEffectivePolicy({
-      name: `projects/${this.projectId}/policies/storage.publicAccessPrevention`,
-      auth,
-    });
+    const resp = await this.withErrorHandling(() =>
+      this.client.projects.policies.getEffectivePolicy({
+        name: `projects/${this.projectId}/policies/storage.publicAccessPrevention`,
+        auth,
+      }),
+    );
 
     if (resp.data && resp.data.spec?.rules) {
       return resp.data.spec?.rules[0].enforce as boolean;

diff --git a/src/steps/privateca/client.ts b/src/steps/privateca/client.ts
@@ -10,10 +10,12 @@ export class PrivateCaClient extends Client {
   ): Promise<privateca_v1.Schema$Policy> {
     const auth = await this.getAuthenticatedServiceClient();
 
-    const result = await this.client.projects.locations.caPools.getIamPolicy({
-      resource: `projects/${this.projectId}/locations/${location}/caPools/${caPoolId}`,
-      auth,
-    });
+    const result = await this.withErrorHandling(() =>
+      this.client.projects.locations.caPools.getIamPolicy({
+        resource: `projects/${this.projectId}/locations/${location}/caPools/${caPoolId}`,
+        auth,
+      }),
+    );
 
     return result.data;
   }