Skip to content

Commit

Permalink
feat: add quota_project_id to service accounts; add with_quota_projec…
Browse files Browse the repository at this point in the history
…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.
  • Loading branch information
busunkim96 committed Jun 10, 2020
1 parent 16df8a3 commit b12488c
Show file tree
Hide file tree
Showing 4 changed files with 96 additions and 1 deletion.
29 changes: 28 additions & 1 deletion google/oauth2/credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 (
Expand Down
44 changes: 44 additions & 0 deletions google/oauth2/service_account.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__(
Expand All @@ -122,6 +126,7 @@ def __init__(
scopes=None,
subject=None,
project_id=None,
quota_project_id=None,
additional_claims=None,
):
"""
Expand All @@ -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.
Expand All @@ -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:
Expand Down Expand Up @@ -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.
Expand All @@ -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(),
)

Expand All @@ -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(),
)

Expand All @@ -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.
Expand Down Expand Up @@ -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)
Expand Down
16 changes: 16 additions & 0 deletions tests/oauth2/test_credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
8 changes: 8 additions & 0 deletions tests/oauth2/test_service_account.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down

0 comments on commit b12488c

Please sign in to comment.