feat: add quota project to base credentials class (#546)

.gitignore

@@ -13,6 +13,8 @@ docs/_build
 .tox/
 .cache/
 .pytest_cache/
+cert_path
+key_path
 
 # Django test database
 db.sqlite3

google/auth/_default.py

@@ -69,7 +69,7 @@ def _warn_about_problematic_credentials(credentials):
         warnings.warn(_CLOUD_SDK_CREDENTIALS_WARNING)
 
 
-def load_credentials_from_file(filename, scopes=None):
+def load_credentials_from_file(filename, scopes=None, quota_project_id=None):
     """Loads Google credentials from a file.
 
     The credentials file must be a service account key or stored authorized
@@ -79,7 +79,9 @@ def load_credentials_from_file(filename, scopes=None):
         filename (str): The full path to the credentials file.
         scopes (Optional[Sequence[str]]): The list of scopes for the credentials. If
             specified, the credentials will automatically be scoped if
-            necessary.
+            necessary
+        quota_project_id (Optional[str]):  The project ID used for
+                quota and billing.
 
     Returns:
         Tuple[google.auth.credentials.Credentials, Optional[str]]: Loaded
@@ -114,7 +116,7 @@ def load_credentials_from_file(filename, scopes=None):
         try:
             credentials = credentials.Credentials.from_authorized_user_info(
                 info, scopes=scopes
-            )
+            ).with_quota_project(quota_project_id)
         except ValueError as caught_exc:
             msg = "Failed to load authorized user credentials from {}".format(filename)
             new_exc = exceptions.DefaultCredentialsError(msg, caught_exc)
@@ -129,7 +131,7 @@ def load_credentials_from_file(filename, scopes=None):
         try:
             credentials = service_account.Credentials.from_service_account_info(
                 info, scopes=scopes
-            )
+            ).with_quota_project(quota_project_id)
         except ValueError as caught_exc:
             msg = "Failed to load service account credentials from {}".format(filename)
             new_exc = exceptions.DefaultCredentialsError(msg, caught_exc)
@@ -226,7 +228,7 @@ def _get_gce_credentials(request=None):
         return None, None
 
 
-def default(scopes=None, request=None):
+def default(scopes=None, request=None, quota_project_id=None):
     """Gets the default credentials for the current environment.
 
     `Application Default Credentials`_ provides an easy way to obtain
@@ -286,7 +288,8 @@ def default(scopes=None, request=None):
             HTTP requests. This is used to detect whether the application
             is running on Compute Engine. If not specified, then it will
             use the standard library http client to make requests.
-
+        quota_project_id (Optional[str]):  The project ID used for
+            quota and billing.
     Returns:
         Tuple[~google.auth.credentials.Credentials, Optional[str]]:
             the current environment's credentials and project ID. Project ID
@@ -314,7 +317,9 @@ def default(scopes=None, request=None):
     for checker in checkers:
         credentials, project_id = checker()
         if credentials is not None:
-            credentials = with_scopes_if_required(credentials, scopes)
+            credentials = with_scopes_if_required(
+                credentials, scopes
+            ).with_quota_project(quota_project_id)
             effective_project_id = explicit_project_id or project_id
             if not effective_project_id:
                 _LOGGER.warning(

google/auth/app_engine.py

@@ -84,7 +84,7 @@ class Credentials(credentials.Scoped, credentials.Signing, credentials.Credentia
     tokens.
     """
 
-    def __init__(self, scopes=None, service_account_id=None):
+    def __init__(self, scopes=None, service_account_id=None, quota_project_id=None):
         """
         Args:
             scopes (Sequence[str]): Scopes to request from the App Identity
@@ -93,6 +93,8 @@ def __init__(self, scopes=None, service_account_id=None):
                 :func:`google.appengine.api.app_identity.get_access_token`.
                 If not specified, the default application service account
                 ID will be used.
+            quota_project_id (Optional[str]): The project ID used for quota
+                and billing.
 
         Raises:
             EnvironmentError: If the App Engine APIs are unavailable.
@@ -107,6 +109,7 @@ def __init__(self, scopes=None, service_account_id=None):
         self._scopes = scopes
         self._service_account_id = service_account_id
         self._signer = Signer()
+        self._quota_project_id = quota_project_id
 
     @_helpers.copy_docstring(credentials.Credentials)
     def refresh(self, request):
@@ -137,7 +140,17 @@ def requires_scopes(self):
     @_helpers.copy_docstring(credentials.Scoped)
     def with_scopes(self, scopes):
         return self.__class__(
-            scopes=scopes, service_account_id=self._service_account_id
+            scopes=scopes,
+            service_account_id=self._service_account_id,
+            quota_project_id=self.quota_project_id,
+        )
+
+    @_helpers.copy_docstring(credentials.Credentials)
+    def with_quota_project(self, quota_project_id):
+        return self.__class__(
+            scopes=self._scopes,
+            service_account_id=self._service_account_id,
+            quota_project_id=quota_project_id,
         )
 
     @_helpers.copy_docstring(credentials.Signing)

google/auth/compute_engine/credentials.py

@@ -54,15 +54,18 @@ class Credentials(credentials.ReadOnlyScoped, credentials.Credentials):
         https://cloud.google.com/compute/docs/authentication#using
     """
 
-    def __init__(self, service_account_email="default"):
+    def __init__(self, service_account_email="default", quota_project_id=None):
         """
         Args:
             service_account_email (str): The service account email to use, or
                 'default'. A Compute Engine instance may have multiple service
                 accounts.
+            quota_project_id (Optional[str]): The project ID used for quota and
+                billing.
         """
         super(Credentials, self).__init__()
         self._service_account_email = service_account_email
+        self._quota_project_id = quota_project_id
 
     def _retrieve_info(self, request):
         """Retrieve information about the service account.
@@ -115,6 +118,13 @@ def requires_scopes(self):
         """False: Compute Engine credentials can not be scoped."""
         return False
 
+    @_helpers.copy_docstring(credentials.Credentials)
+    def with_quota_project(self, quota_project_id):
+        return self.__class__(
+            service_account_email=self._service_account_email,
+            quota_project_id=quota_project_id,
+        )
+
 
 _DEFAULT_TOKEN_LIFETIME_SECS = 3600  # 1 hour in seconds
 _DEFAULT_TOKEN_URI = "https://www.googleapis.com/oauth2/v4/token"
@@ -143,6 +153,7 @@ def __init__(
         service_account_email=None,
         signer=None,
         use_metadata_identity_endpoint=False,
+        quota_project_id=None,
     ):
         """
         Args:
@@ -165,6 +176,8 @@ def __init__(
                 is False. If set to True, ``token_uri``, ``additional_claims``,
                 ``service_account_email``, ``signer`` argument should not be set;
                 otherwise ValueError will be raised.
+            quota_project_id (Optional[str]): The project ID used for quota and
+                billing.
 
         Raises:
             ValueError:
@@ -174,6 +187,7 @@ def __init__(
         """
         super(IDTokenCredentials, self).__init__()
 
+        self._quota_project_id = quota_project_id
         self._use_metadata_identity_endpoint = use_metadata_identity_endpoint
         self._target_audience = target_audience
 
@@ -226,6 +240,7 @@ def with_target_audience(self, target_audience):
                 None,
                 target_audience=target_audience,
                 use_metadata_identity_endpoint=True,
+                quota_project_id=self._quota_project_id,
             )
         else:
             return self.__class__(
@@ -236,6 +251,31 @@ def with_target_audience(self, target_audience):
                 additional_claims=self._additional_claims.copy(),
                 signer=self.signer,
                 use_metadata_identity_endpoint=False,
+                quota_project_id=self._quota_project_id,
+            )
+
+    @_helpers.copy_docstring(credentials.Credentials)
+    def with_quota_project(self, quota_project_id):
+
+        # since the signer is already instantiated,
+        # the request is not needed
+        if self._use_metadata_identity_endpoint:
+            return self.__class__(
+                None,
+                target_audience=self._target_audience,
+                use_metadata_identity_endpoint=True,
+                quota_project_id=quota_project_id,
+            )
+        else:
+            return self.__class__(
+                None,
+                service_account_email=self._service_account_email,
+                token_uri=self._token_uri,
+                target_audience=self._target_audience,
+                additional_claims=self._additional_claims.copy(),
+                signer=self.signer,
+                use_metadata_identity_endpoint=False,
+                quota_project_id=quota_project_id,
             )
 
     def _make_authorization_grant_assertion(self):

google/auth/credentials.py

@@ -49,6 +49,8 @@ def __init__(self):
         self.expiry = None
         """Optional[datetime]: When the token expires and is no longer valid.
         If this is None, the token is assumed to never expire."""
+        self._quota_project_id = None
+        """Optional[str]: Project to use for quota and billing purposes."""
 
     @property
     def expired(self):
@@ -75,6 +77,11 @@ def valid(self):
         """
         return self.token is not None and not self.expired
 
+    @property
+    def quota_project_id(self):
+        """Project to use for quota and billing purposes."""
+        return self._quota_project_id
+
     @abc.abstractmethod
     def refresh(self, request):
         """Refreshes the access token.
@@ -102,6 +109,8 @@ def apply(self, headers, token=None):
         headers["authorization"] = "Bearer {}".format(
             _helpers.from_bytes(token or self.token)
         )
+        if self.quota_project_id:
+            headers["x-goog-user-project"] = self.quota_project_id
 
     def before_request(self, request, method, url, headers):
         """Performs credential-specific before request logic.
@@ -124,6 +133,18 @@ def before_request(self, request, method, url, headers):
             self.refresh(request)
         self.apply(headers)
 
+    def with_quota_project(self, quota_project_id):
+        """Returns a copy of these credentials with a modified quota project
+
+        Args:
+            quota_project_id (str): The project to use for quota and
+                billing purposes
+
+        Returns:
+            google.oauth2.credentials.Credentials: A new credentials instance.
+        """
+        raise NotImplementedError("This class does not support quota project.")
+
 
 class AnonymousCredentials(Credentials):
     """Credentials that do not provide any authentication information.
@@ -161,6 +182,9 @@ def apply(self, headers, token=None):
     def before_request(self, request, method, url, headers):
         """Anonymous credentials do nothing to the request."""
 
+    def with_quota_project(self, quota_project_id):
+        raise ValueError("Anonymous credentials don't support quota project.")
+
 
 @six.add_metaclass(abc.ABCMeta)
 class ReadOnlyScoped(object):

google/auth/impersonated_credentials.py

@@ -184,6 +184,7 @@ def __init__(
         target_scopes,
         delegates=None,
         lifetime=_DEFAULT_TOKEN_LIFETIME_SECS,
+        quota_project_id=None,
     ):
         """
         Args:
@@ -205,6 +206,9 @@ def __init__(
                 target_principal.
             lifetime (int): Number of seconds the delegated credential should
                 be valid for (upto 3600).
+            quota_project_id (Optional[str]): The project ID used for quota and billing.
+                This project may be different from the project used to
+                create the credentials.
         """
 
         super(Credentials, self).__init__()
@@ -221,6 +225,7 @@ def __init__(
         self._lifetime = lifetime
         self.token = None
         self.expiry = _helpers.utcnow()
+        self._quota_project_id = quota_project_id
 
     @_helpers.copy_docstring(credentials.Credentials)
     def refresh(self, request):
@@ -288,19 +293,38 @@ def service_account_email(self):
     def signer(self):
         return self
 
+    @_helpers.copy_docstring(credentials.Credentials)
+    def with_quota_project(self, quota_project_id):
+        return self.__class__(
+            self._source_credentials,
+            target_principal=self._target_principal,
+            target_scopes=self._target_scopes,
+            delegates=self._delegates,
+            lifetime=self._lifetime,
+            quota_project_id=quota_project_id,
+        )
+
 
 class IDTokenCredentials(credentials.Credentials):
     """Open ID Connect ID Token-based service account credentials.
 
     """
 
-    def __init__(self, target_credentials, target_audience=None, include_email=False):
+    def __init__(
+        self,
+        target_credentials,
+        target_audience=None,
+        include_email=False,
+        quota_project_id=None,
+    ):
         """
         Args:
             target_credentials (google.auth.Credentials): The target
                 credential used as to acquire the id tokens for.
             target_audience (string): Audience to issue the token for.
             include_email (bool): Include email in IdToken
+            quota_project_id (Optional[str]):  The project ID used for
+                quota and billing.
         """
         super(IDTokenCredentials, self).__init__()
 
@@ -311,22 +335,37 @@ def __init__(self, target_credentials, target_audience=None, include_email=False
         self._target_credentials = target_credentials
         self._target_audience = target_audience
         self._include_email = include_email
+        self._quota_project_id = quota_project_id
 
     def from_credentials(self, target_credentials, target_audience=None):
         return self.__class__(
-            target_credentials=self._target_credentials, target_audience=target_audience
+            target_credentials=self._target_credentials,
+            target_audience=target_audience,
+            quota_project_id=self._quota_project_id,
         )
 
     def with_target_audience(self, target_audience):
         return self.__class__(
-            target_credentials=self._target_credentials, target_audience=target_audience
+            target_credentials=self._target_credentials,
+            target_audience=target_audience,
+            quota_project_id=self._quota_project_id,
         )
 
     def with_include_email(self, include_email):
         return self.__class__(
             target_credentials=self._target_credentials,
             target_audience=self._target_audience,
             include_email=include_email,
+            quota_project_id=self._quota_project_id,
+        )
+
+    @_helpers.copy_docstring(credentials.Credentials)
+    def with_quota_project(self, quota_project_id):
+        return self.__class__(
+            target_credentials=self._target_credentials,
+            target_audience=self._target_audience,
+            include_email=self._include_email,
+            quota_project_id=quota_project_id,
         )
 
     @_helpers.copy_docstring(credentials.Credentials)