8000 feat: add quota_project_id to service accounts; add with_quota_projec… · akx/google-auth-library-python@b12488c · GitHub
[go: up one dir, main page]
Skip to content

Commit b12488c

Browse files
busunkim96busunkim96
authored
feat: add quota_project_id to service accounts; add with_quota_project methods (googleapis#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.
1 parent 16df8a3 commit b12488c

File tree

4 files changed

+96
-1
lines changed

4 files changed

+96
-1
lines changed

‎google/oauth2/credentials.py

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,13 @@
4848

4949

5050
class Credentials(credentials.ReadOnlyScoped, credentials.Credentials):
51-
"""Credentials using OAuth 2.0 access and refresh tokens."""
51+
"""Credentials using OAuth 2.0 access and refresh tokens.
52+
53+
The credentials are considered immutable. If you want to modify the
54+
quota project, use :meth:`with_quota_project` or ::
55+
56+
credentials = credentials.with_quota_project('myproject-123)
57+
"""
5258

5359
def __init__(
5460
self,
@@ -160,6 +166,27 @@ def requires_scopes(self):
160166
the initial token is requested and can not be changed."""
161167
return False
162168

169+
def with_quota_project(self, quota_project_id):
170+
"""Returns a copy of these credentials with a modified quota project
171+
172+
Args:
173+
quota_project_id (str): The project to use for quota and
174+
billing purposes
175+
176+
Returns:
177+
google.oauth2.credentials.Credentials: A new credentials instance.
178+
"""
179+
return self.__class__(
180+
self.token,
181+
refresh_token=self.refresh_token,
182+
id_token=self.id_token,
183+
token_uri=self.token_uri,
184+
client_id=self.client_id,
185+
client_secret=self.client_secret,
186+
scopes=self.scopes,
187+
quota_project_id=quota_project_id,
188+
)
189+
163190
@_helpers.copy_docstring(credentials.Credentials)
164191
def refresh(self, request):
165192
if (

‎google/oauth2/service_account.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,10 @@ class Credentials(credentials.Signing, credentials.Scoped, credentials.Credentia
112112
113113
scoped_credentials = credentials.with_scopes(['email'])
114114
delegated_credentials = credentials.with_subject(subject)
115+
116+
To add a quota project, use :meth:`with_quota_project`::
117+
118+
credentials = credentials.with_quota_project('myproject-123')
115119
"""
116120

117121
def __init__(
@@ -122,6 +126,7 @@ def __init__(
122126
scopes=None,
123127
subject=None,
124128
project_id=None,
129+
quota_project_id=None,
125130
additional_claims=None,
126131
):
127132
"""
@@ -135,6 +140,8 @@ def __init__(
135140
user to for which to request delegated access.
136141
project_id (str): Project ID associated with the service account
137142
credential.
143+
quota_project_id (Optional[str]): The project ID used for quota and
144+
billing.
138145
additional_claims (Mapping[str, str]): Any additional claims for
139146
the JWT assertion used in the authorization grant.
140147
@@ -150,6 +157,7 @@ def __init__(
150157
self._service_account_email = service_account_email
151158
self._subject = subject
152159
self._project_id = project_id
160+
self._quota_project_id = quota_project_id
153161
self._token_uri = token_uri
154162

155163
if additional_claims is not None:
@@ -229,6 +237,11 @@ def project_id(self):
229237
"""Project ID associated with this credential."""
230238
return self._project_id
231239

240+
@property
241+
def quota_project_id(self):
242+
"""Project ID to use for quota and billing purposes."""
243+
return self._quota_project_id
244+
232245
@property
233246
def requires_scopes(self):
234247
"""Checks if the credentials requires scopes.
@@ -247,6 +260,7 @@ def with_scopes(self, scopes):
247260
token_uri=self._token_uri,
248261
subject=self._subject,
249262
project_id=self._project_id,
263+
quota_project_id=self._quota_project_id,
250264
additional_claims=self._additional_claims.copy(),
251265
)
252266

@@ -267,6 +281,7 @@ def with_subject(self, subject):
267281
token_uri=self._token_uri,
268282
subject=subject,
269283
project_id=self._project_id,
284+
quota_project_id=self._quota_project_id,
270285
additional_claims=self._additional_claims.copy(),
271286
)
272287

@@ -292,9 +307,32 @@ def with_claims(self, additional_claims):
292307
token_uri=self._token_uri,
293308
subject=self._subject,
294309
project_id=self._project_id,
310+
quota_project_id=self._quota_project_id,
295311
additional_claims=new_additional_claims,
296312
)
297313

314+
def with_quota_project(self, quota_project_id):
315+
"""Returns a copy of these credentials with a modified quota project.
316+
317+
Args:
318+
quota_project_id (str): The project to use for quota and
319+
billing purposes
320+
321+
Returns:
322+
google.auth.service_account.Credentials: A new credentials
323+
instance.
324+
"""
325+
return self.__class__(
326+
self._signer,
327+
service_account_email=self._service_account_email,
328+
scopes=self._scopes,
329+
token_uri=self._token_uri,
330+
subject=self._subject,
331+
project_id=self._project_id,
332+
quota_project_id=quota_project_id,
333+
additional_claims=self._additional_claims.copy(),
334+
)
335+
298336
def _make_authorization_grant_assertion(self):
299337
"""Create the OAuth 2.0 assertion.
300338
@@ -335,6 +373,12 @@ def refresh(self, request):
335373
self.token = access_token
336374
self.expiry = expiry
337375

376+
@_helpers.copy_docstring(credentials.Credentials)
377+
def apply(self, headers, token=None):
378+
super(Credentials, self).apply(headers, token=token)
379+
if self.quota_project_id is not None:
380+
headers["x-goog-user-project"] = self.quota_project_id
381+
338382
@_helpers.copy_docstring(credentials.Signing)
339383
def sign_bytes(self, message):
340384
return self._signer.sign(message)

‎tests/oauth2/test_credentials.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -323,6 +323,22 @@ def test_apply_with_no_quota_project_id(self):
323323
creds.apply(headers)
324324
assert "x-goog-user-project" not in headers
325325

326+
def test_with_quota_project(self):
327+
creds = credentials.Credentials(
328+
token="token",
329+
refresh_token=self.REFRESH_TOKEN,
330+
token_uri=self.TOKEN_URI,
331+
client_id=self.CLIENT_ID,
332+
client_secret=self.CLIENT_SECRET,
333+
quota_project_id="quota-project-123",
334+
)
335+
336+
new_creds = creds.with_quota_project("new-project-456")
337+
assert new_creds.quota_project_id == "new-project-456"
338+
headers = {}
339+
creds.apply(headers)
340+
assert "x-goog-user-project" in headers
341+
326342
def test_from_authorized_user_info(self):
327343
info = AUTH_USER_INFO.copy()
328344

‎tests/oauth2/test_service_account.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,14 @@ def test_with_claims(self):
147147
new_credentials = credentials.with_claims({"meep": "moop"})
148148
assert new_credentials._additional_claims == {"meep": "moop"}
149149

150+
def test_with_quota_project(self):
151+
credentials = self.make_credentials()
152+
new_credentials = credentials.with_quota_project("new-project-456")
153+
assert new_credentials.quota_project_id == "new-project-456"
154+
hdrs = {}
155+
new_credentials.apply(hdrs, token="tok")
156+
assert "x-goog-user-project" in hdrs
157+
150158
def test__make_authorization_grant_assertion(self):
151159
credentials = self.make_credentials()
152160
token = credentials._make_authorization_grant_assertion()

0 commit comments

Comments
 (0)
0