From 74424c79e56d941b36b12c67406452d218038c79 Mon Sep 17 00:00:00 2001 From: Bu Sun Kim Date: Tue, 2 Jun 2020 05:02:12 +0000 Subject: [PATCH 1/4] feat: add quota_project_id to service accounts --- google/oauth2/credentials.py | 29 +++++++++++++++++- google/oauth2/service_account.py | 45 ++++++++++++++++++++++++++++ tests/oauth2/test_credentials.py | 14 +++++++++ tests/oauth2/test_service_account.py | 5 ++++ 4 files changed, 92 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..2cbd03d47 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,33 @@ 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 +374,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..0147772d1 100644 --- a/tests/oauth2/test_credentials.py +++ b/tests/oauth2/test_credentials.py @@ -323,6 +323,20 @@ 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" + 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..2c3b76175 100644 --- a/tests/oauth2/test_service_account.py +++ b/tests/oauth2/test_service_account.py @@ -147,6 +147,11 @@ 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" + def test__make_authorization_grant_assertion(self): credentials = self.make_credentials() token = credentials._make_authorization_grant_assertion() From 8a4b2bb6e3fdde1d1f6dd783c29391ae594ee7fd Mon Sep 17 00:00:00 2001 From: Bu Sun Kim Date: Tue, 2 Jun 2020 05:05:57 +0000 Subject: [PATCH 2/4] chore: lint --- google/oauth2/service_account.py | 1 - tests/oauth2/test_credentials.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/google/oauth2/service_account.py b/google/oauth2/service_account.py index 2cbd03d47..54630d34b 100644 --- a/google/oauth2/service_account.py +++ b/google/oauth2/service_account.py @@ -333,7 +333,6 @@ def with_quota_project(self, quota_project_id): additional_claims=self._additional_claims.copy(), ) - def _make_authorization_grant_assertion(self): """Create the OAuth 2.0 assertion. diff --git a/tests/oauth2/test_credentials.py b/tests/oauth2/test_credentials.py index 0147772d1..e85a4dbd0 100644 --- a/tests/oauth2/test_credentials.py +++ b/tests/oauth2/test_credentials.py @@ -330,7 +330,7 @@ def test_with_quota_project(self): token_uri=self.TOKEN_URI, client_id=self.CLIENT_ID, client_secret=self.CLIENT_SECRET, - quota_project_id="quota-project-123" + quota_project_id="quota-project-123", ) new_creds = creds.with_quota_project("new-project-456") From f1f08206112b2c5abecf9f0a7ebf4ee40c5e8b16 Mon Sep 17 00:00:00 2001 From: Dov Shlachter Date: Wed, 10 Jun 2020 10:47:06 -0700 Subject: [PATCH 3/4] Update test_credentials.py Add test coverage --- tests/oauth2/test_credentials.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/oauth2/test_credentials.py b/tests/oauth2/test_credentials.py index e85a4dbd0..78b101252 100644 --- a/tests/oauth2/test_credentials.py +++ b/tests/oauth2/test_credentials.py @@ -334,8 +334,10 @@ def test_with_quota_project(self): ) 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() From bc8b40f24ce66fa780621c853cb84c42c4db21e6 Mon Sep 17 00:00:00 2001 From: Dov Shlachter Date: Wed, 10 Jun 2020 11:44:49 -0700 Subject: [PATCH 4/4] Test coverage --- tests/oauth2/test_service_account.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/oauth2/test_service_account.py b/tests/oauth2/test_service_account.py index 2c3b76175..457d472d7 100644 --- a/tests/oauth2/test_service_account.py +++ b/tests/oauth2/test_service_account.py @@ -151,6 +151,9 @@ 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()