From b12488cf552888299425c8009ea075511627cf08 Mon Sep 17 00:00:00 2001 From: Bu Sun Kim <8822365+busunkim96@users.noreply.github.com> Date: Wed, 10 Jun 2020 13:44:07 -0700 Subject: [PATCH] feat: add quota_project_id to service accounts; add with_quota_project 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 | 29 +++++++++++++++++- google/oauth2/service_account.py | 44 ++++++++++++++++++++++++++++ tests/oauth2/test_credentials.py | 16 ++++++++++ tests/oauth2/test_service_account.py | 8 +++++ 4 files changed, 96 insertions(+), 1 deletion(-) diff --git a/google/oauth2/credentials.py b/google/oauth2/credentials.py index baf3cf7f4..d39848b41 100644 --- a/google/oauth2/credentials.py +++ b/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 ( diff --git a/google/oauth2/service_account.py b/google/oauth2/service_account.py index af86588d5..54630d34b 100644 --- a/google/oauth2/service_account.py +++ b/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) diff --git a/tests/oauth2/test_credentials.py b/tests/oauth2/test_credentials.py index 76aa463cb..78b101252 100644 --- a/tests/oauth2/test_credentials.py +++ b/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() diff --git a/tests/oauth2/test_service_account.py b/tests/oauth2/test_service_account.py index 897374a6f..457d472d7 100644 --- a/tests/oauth2/test_service_account.py +++ b/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()