feat: add quota_project_id to service accounts; add with_quota_projec…

…t methods (#519)

Adds quota_project_id to service account credentials, making it possible to set quota_project_id on OAuth2 credentials and service account credentials.

This PR also adds the method with_quota_project to both classes.

google/oauth2/credentials.py

@@ -48,7 +48,13 @@
 
 
 class Credentials(credentials.ReadOnlyScoped, credentials.Credentials):
-    """Credentials using OAuth 2.0 access and refresh tokens."""
+    """Credentials using OAuth 2.0 access and refresh tokens.
+
+    The credentials are considered immutable. If you want to modify the
+    quota project, use :meth:`with_quota_project` or ::
+
+        credentials = credentials.with_quota_project('myproject-123)
+    """
 
     def __init__(
         self,
@@ -160,6 +166,27 @@ def requires_scopes(self):
         the initial token is requested and can not be changed."""
         return False
 
+    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.
+        """
+        return self.__class__(
+            self.token,
+            refresh_token=self.refresh_token,
+            id_token=self.id_token,
+            token_uri=self.token_uri,
+            client_id=self.client_id,
+            client_secret=self.client_secret,
+            scopes=self.scopes,
+            quota_project_id=quota_project_id,
+        )
+
     @_helpers.copy_docstring(credentials.Credentials)
     def refresh(self, request):
         if (

google/oauth2/service_account.py

@@ -112,6 +112,10 @@ class Credentials(credentials.Signing, credentials.Scoped, credentials.Credentia
 
         scoped_credentials = credentials.with_scopes(['email'])
         delegated_credentials = credentials.with_subject(subject)
+
+    To add a quota project, use :meth:`with_quota_project`::
+
+        credentials = credentials.with_quota_project('myproject-123')
     """
 
     def __init__(
@@ -122,6 +126,7 @@ def __init__(
         scopes=None,
         subject=None,
         project_id=None,
+        quota_project_id=None,
         additional_claims=None,
     ):
         """
@@ -135,6 +140,8 @@ def __init__(
                 user to for which to request delegated access.
             project_id  (str): Project ID associated with the service account
                 credential.
+            quota_project_id (Optional[str]): The project ID used for quota and
+                billing.
             additional_claims (Mapping[str, str]): Any additional claims for
                 the JWT assertion used in the authorization grant.
 
@@ -150,6 +157,7 @@ def __init__(
         self._service_account_email = service_account_email
         self._subject = subject
         self._project_id = project_id
+        self._quota_project_id = quota_project_id
         self._token_uri = token_uri
 
         if additional_claims is not None:
@@ -229,6 +237,11 @@ def project_id(self):
         """Project ID associated with this credential."""
         return self._project_id
 
+    @property
+    def quota_project_id(self):
+        """Project ID to use for quota and billing purposes."""
+        return self._quota_project_id
+
     @property
     def requires_scopes(self):
         """Checks if the credentials requires scopes.
@@ -247,6 +260,7 @@ def with_scopes(self, scopes):
             token_uri=self._token_uri,
             subject=self._subject,
             project_id=self._project_id,
+            quota_project_id=self._quota_project_id,
             additional_claims=self._additional_claims.copy(),
         )
 
@@ -267,6 +281,7 @@ def with_subject(self, subject):
             token_uri=self._token_uri,
             subject=subject,
             project_id=self._project_id,
+            quota_project_id=self._quota_project_id,
             additional_claims=self._additional_claims.copy(),
         )
 
@@ -292,9 +307,32 @@ def with_claims(self, additional_claims):
             token_uri=self._token_uri,
             subject=self._subject,
             project_id=self._project_id,
+            quota_project_id=self._quota_project_id,
             additional_claims=new_additional_claims,
         )
 
+    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.auth.service_account.Credentials: A new credentials
+                instance.
+        """
+        return self.__class__(
+            self._signer,
+            service_account_email=self._service_account_email,
+            scopes=self._scopes,
+            token_uri=self._token_uri,
+            subject=self._subject,
+            project_id=self._project_id,
+            quota_project_id=quota_project_id,
+            additional_claims=self._additional_claims.copy(),
+        )
+
     def _make_authorization_grant_assertion(self):
         """Create the OAuth 2.0 assertion.
 
@@ -335,6 +373,12 @@ def refresh(self, request):
         self.token = access_token
         self.expiry = expiry
 
+    @_helpers.copy_docstring(credentials.Credentials)
+    def apply(self, headers, token=None):
+        super(Credentials, self).apply(headers, token=token)
+        if self.quota_project_id is not None:
+            headers["x-goog-user-project"] = self.quota_project_id
+
     @_helpers.copy_docstring(credentials.Signing)
     def sign_bytes(self, message):
         return self._signer.sign(message)

tests/oauth2/test_credentials.py

@@ -323,6 +323,22 @@ def test_apply_with_no_quota_project_id(self):
         creds.apply(headers)
         assert "x-goog-user-project" not in headers
 
+    def test_with_quota_project(self):
+        creds = credentials.Credentials(
+            token="token",
+            refresh_token=self.REFRESH_TOKEN,
+            token_uri=self.TOKEN_URI,
+            client_id=self.CLIENT_ID,
+            client_secret=self.CLIENT_SECRET,
+            quota_project_id="quota-project-123",
+        )
+
+        new_creds = creds.with_quota_project("new-project-456")
+        assert new_creds.quota_project_id == "new-project-456"
+        headers = {}
+        creds.apply(headers)
+        assert "x-goog-user-project" in headers
+
     def test_from_authorized_user_info(self):
         info = AUTH_USER_INFO.copy()
 

tests/oauth2/test_service_account.py

@@ -147,6 +147,14 @@ def test_with_claims(self):
         new_credentials = credentials.with_claims({"meep": "moop"})
         assert new_credentials._additional_claims == {"meep": "moop"}
 
+    def test_with_quota_project(self):
+        credentials = self.make_credentials()
+        new_credentials = credentials.with_quota_project("new-project-456")
+        assert new_credentials.quota_project_id == "new-project-456"
+        hdrs = {}
+        new_credentials.apply(hdrs, token="tok")
+        assert "x-goog-user-project" in hdrs
+
     def test__make_authorization_grant_assertion(self):
         credentials = self.make_credentials()
         token = credentials._make_authorization_grant_assertion()