feat: Introduce environment variable for quota project (#1478)

* feat: Introduce environment variable for quota project * remove new line * add awaits * clean up cached cred logic * add explicit quota project * remove explicit value * enhance test Co-authored-by: Daniel Bankhead <danielbankhead@google.com>
googleapis · Nov 4, 2022 · 8706abc · 8706abc
1 parent 6dc4e58
commit 8706abc
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 16 deletions.
diff --git a/src/auth/authclient.ts b/src/auth/authclient.ts
@@ -88,7 +88,11 @@ export abstract class AuthClient
   extends EventEmitter
   implements CredentialsClient
 {
-  protected quotaProjectId?: string;
+  /**
+   * The quota project ID. The quota project can be used by client libraries for the billing purpose.
+   * See {@link https://cloud.google.com/docs/quota| Working with quotas}
+   */
+  quotaProjectId?: string;
   transporter = new DefaultTransporter();
   credentials: Credentials = {};
   projectId?: string | null;

diff --git a/src/auth/googleauth.ts b/src/auth/googleauth.ts
@@ -301,16 +301,18 @@ export class GoogleAuth<T extends AuthClient = JSONClient> {
   private async getApplicationDefaultAsync(
     options: RefreshOptions = {}
   ): Promise<ADCResponse> {
-    // If we've already got a cached credential, just return it.
+    // If we've already got a cached credential, return it.
+    // This will also preserve one's configured quota project, in case they
+    // set one directly on the credential previously.
     if (this.cachedCredential) {
-      return {
-        credential: this.cachedCredential,
-        projectId: await this.getProjectIdOptional(),
-      };
+      return await this.prepareAndCacheADC(this.cachedCredential);
     }
 
+    // Since this is a 'new' ADC to cache we will use the environment variable
+    // if it's available. We prefer this value over the value from ADC.
+    const quotaProjectIdOverride = process.env['GOOGLE_CLOUD_QUOTA_PROJECT'];
+
     let credential: JSONClient | null;
-    let projectId: string | null;
     // Check for the existence of a local environment variable pointing to the
     // location of the credential file. This is typically used in local
     // developer scenarios.
@@ -322,10 +324,8 @@ export class GoogleAuth<T extends AuthClient = JSONClient> {
       } else if (credential instanceof BaseExternalAccountClient) {
         credential.scopes = this.getAnyScopes();
       }
-      this.cachedCredential = credential;
-      projectId = await this.getProjectIdOptional();
 
-      return {credential, projectId};
+      return await this.prepareAndCacheADC(credential, quotaProjectIdOverride);
     }
 
     // Look in the well-known credential file location.
@@ -338,9 +338,7 @@ export class GoogleAuth<T extends AuthClient = JSONClient> {
       } else if (credential instanceof BaseExternalAccountClient) {
         credential.scopes = this.getAnyScopes();
       }
-      this.cachedCredential = credential;
-      projectId = await this.getProjectIdOptional();
-      return {credential, projectId};
+      return await this.prepareAndCacheADC(credential, quotaProjectIdOverride);
     }
 
     // Determine if we're running on GCE.
@@ -365,9 +363,25 @@ export class GoogleAuth<T extends AuthClient = JSONClient> {
     // For GCE, just return a default ComputeClient. It will take care of
     // the rest.
     (options as ComputeOptions).scopes = this.getAnyScopes();
-    this.cachedCredential = new Compute(options);
-    projectId = await this.getProjectIdOptional();
-    return {projectId, credential: this.cachedCredential};
+    return await this.prepareAndCacheADC(
+      new Compute(options),
+      quotaProjectIdOverride
+    );
+  }
+
+  private async prepareAndCacheADC(
+    credential: JSONClient | Impersonated | Compute | T,
+    quotaProjectIdOverride?: string
+  ): Promise<ADCResponse> {
+    const projectId = await this.getProjectIdOptional();
+
+    if (quotaProjectIdOverride) {
+      credential.quotaProjectId = quotaProjectIdOverride;
+    }
+
+    this.cachedCredential = credential;
+
+    return {credential, projectId};
   }
 
   /**

diff --git a/test/test.googleauth.ts b/test/test.googleauth.ts
@@ -127,6 +127,7 @@ describe('googleauth', () => {
         GCLOUD_PROJECT: undefined,
         GOOGLE_APPLICATION_CREDENTIALS: undefined,
         google_application_credentials: undefined,
+        GOOGLE_CLOUD_QUOTA_PROJECT: undefined,
         HOME: path.join('/', 'fake', 'user'),
       });
       sandbox.stub(process, 'env').value(envVars);
@@ -1043,6 +1044,40 @@ describe('googleauth', () => {
       assert.strictEqual(undefined, client.scope);
     });
 
+    it('explicitly set quota project should not be overriden by environment value', async () => {
+      mockLinuxWellKnownFile(
+        './test/fixtures/config-with-quota/.config/gcloud/application_default_credentials.json'
+      );
+      mockEnvVar('GOOGLE_CLOUD_QUOTA_PROJECT', 'quota_from_env');
+      let result = await auth.getApplicationDefault();
+      let client = result.credential as JWT;
+      assert.strictEqual('quota_from_env', client.quotaProjectId);
+
+      client.quotaProjectId = 'explicit_quota';
+      result = await auth.getApplicationDefault();
+      client = result.credential as JWT;
+      assert.strictEqual('explicit_quota', client.quotaProjectId);
+    });
+
+    it('getApplicationDefault should use quota project id from file if environment variable is empty', async () => {
+      mockLinuxWellKnownFile(
+        './test/fixtures/config-with-quota/.config/gcloud/application_default_credentials.json'
+      );
+      mockEnvVar('GOOGLE_CLOUD_QUOTA_PROJECT', '');
+      const result = await auth.getApplicationDefault();
+      const client = result.credential as JWT;
+      assert.strictEqual('my-quota-project', client.quotaProjectId);
+    });
+
+    it('getApplicationDefault should use quota project id from file if environment variable is not set', async () => {
+      mockLinuxWellKnownFile(
+        './test/fixtures/config-with-quota/.config/gcloud/application_default_credentials.json'
+      );
+      const result = await auth.getApplicationDefault();
+      const client = result.credential as JWT;
+      assert.strictEqual('my-quota-project', client.quotaProjectId);
+    });
+
     it('getApplicationDefault should use GCE when well-known file and env const are not set', async () => {
       // Set up the creds.
       // * Environment variable is not set.