feat: Read Quota Project from Environment Variable (#1163)

* feat: Read Quota Project from Environment Variable * remove unused code * Update google/auth/credentials.py Co-authored-by: Carl Lundin <[email protected]> * update rt Co-authored-by: Carl Lundin <[email protected]>
googleapis · Oct 31, 2022 · 57b3e42 · 57b3e42
1 parent a12b96d
commit 57b3e42
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 1 deletion.
diff --git a/google/auth/_default.py b/google/auth/_default.py
@@ -479,6 +479,8 @@ def _get_gdch_service_account_credentials(filename, info):
 def _apply_quota_project_id(credentials, quota_project_id):
     if quota_project_id:
         credentials = credentials.with_quota_project(quota_project_id)
+    else:
+        credentials = credentials.with_quota_project_from_environment()
 
     from google.oauth2 import credentials as authorized_user_credentials
 

diff --git a/google/auth/credentials.py b/google/auth/credentials.py
@@ -16,10 +16,11 @@
 """Interfaces for credentials."""
 
 import abc
+import os
 
 import six
 
-from google.auth import _helpers
+from google.auth import _helpers, environment_vars
 
 
 @six.add_metaclass(abc.ABCMeta)
@@ -149,6 +150,12 @@ def with_quota_project(self, quota_project_id):
         """
         raise NotImplementedError("This credential does not support quota project.")
 
+    def with_quota_project_from_environment(self):
+        quota_from_env = os.environ.get(environment_vars.GOOGLE_CLOUD_QUOTA_PROJECT)
+        if quota_from_env:
+            return self.with_quota_project(quota_from_env)
+        return self
+
 
 class CredentialsWithTokenUri(Credentials):
     """Abstract base for credentials supporting ``with_token_uri`` factory"""

diff --git a/google/auth/environment_vars.py b/google/auth/environment_vars.py
@@ -29,6 +29,10 @@
 situations (such as Google App Engine).
 """
 
+GOOGLE_CLOUD_QUOTA_PROJECT = "GOOGLE_CLOUD_QUOTA_PROJECT"
+"""Environment variable defining the project to be used for
+quota and billing."""
+
 CREDENTIALS = "GOOGLE_APPLICATION_CREDENTIALS"
 """Environment variable defining the location of Google application default
 credentials."""

diff --git a/system_tests/secrets.tar.enc b/system_tests/secrets.tar.enc
diff --git a/tests/test__default.py b/tests/test__default.py
@@ -1214,3 +1214,23 @@ def test_default_gdch_service_account_credentials(get_adc_path):
     assert creds._token_uri == "https://service-identity.<Domain>/authenticate"
     assert creds._ca_cert_path == "/path/to/ca/cert"
     assert project == "project_foo"
+
+
+@mock.patch.dict(os.environ)
+@mock.patch(
+    "google.auth._cloud_sdk.get_application_default_credentials_path", autospec=True
+)
+def test_quota_project_from_environment(get_adc_path):
+    get_adc_path.return_value = AUTHORIZED_USER_CLOUD_SDK_WITH_QUOTA_PROJECT_ID_FILE
+
+    credentials, _ = _default.default(quota_project_id=None)
+    assert credentials.quota_project_id == "quota_project_id"
+
+    quota_from_env = "quota_from_env"
+    os.environ[environment_vars.GOOGLE_CLOUD_QUOTA_PROJECT] = quota_from_env
+    credentials, _ = _default.default(quota_project_id=None)
+    assert credentials.quota_project_id == quota_from_env
+
+    explicit_quota = "explicit_quota"
+    credentials, _ = _default.default(quota_project_id=explicit_quota)
+    assert credentials.quota_project_id == explicit_quota