Skip to content

Commit e2d9a7b

Browse files
authored
feat: allow quota project to be passed to create_channel (googleapis#58)
* feat: allow quota project to be passed to create_channel * chore: update test name * chore: lint and increase auth lib version * chore: fix lint * Update setup.py
1 parent 8595f62 commit e2d9a7b

File tree

5 files changed

+73
-6
lines changed

5 files changed

+73
-6
lines changed

google/api_core/grpc_helpers.py

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,12 @@ def wrap_errors(callable_):
176176
return _wrap_unary_errors(callable_)
177177

178178

179-
def _create_composite_credentials(credentials=None, credentials_file=None, scopes=None, ssl_credentials=None):
179+
def _create_composite_credentials(
180+
credentials=None,
181+
credentials_file=None,
182+
scopes=None,
183+
ssl_credentials=None,
184+
quota_project_id=None):
180185
"""Create the composite credentials for secure channels.
181186
182187
Args:
@@ -191,6 +196,7 @@ def _create_composite_credentials(credentials=None, credentials_file=None, scope
191196
are passed to :func:`google.auth.default`.
192197
ssl_credentials (grpc.ChannelCredentials): Optional SSL channel
193198
credentials. This can be used to specify different certificates.
199+
quota_project_id (str): An optional project to use for billing and quota.
194200
195201
Returns:
196202
grpc.ChannelCredentials: The composed channel credentials object.
@@ -210,6 +216,9 @@ def _create_composite_credentials(credentials=None, credentials_file=None, scope
210216
else:
211217
credentials, _ = google.auth.default(scopes=scopes)
212218

219+
if quota_project_id:
220+
credentials = credentials.with_quota_project(quota_project_id)
221+
213222
request = google.auth.transport.requests.Request()
214223

215224
# Create the metadata plugin for inserting the authorization header.
@@ -229,7 +238,14 @@ def _create_composite_credentials(credentials=None, credentials_file=None, scope
229238
)
230239

231240

232-
def create_channel(target, credentials=None, scopes=None, ssl_credentials=None, credentials_file=None, **kwargs):
241+
def create_channel(
242+
target,
243+
credentials=None,
244+
scopes=None,
245+
ssl_credentials=None,
246+
credentials_file=None,
247+
quota_project_id=None,
248+
**kwargs):
233249
"""Create a secure channel with credentials.
234250
235251
Args:
@@ -245,6 +261,7 @@ def create_channel(target, credentials=None, scopes=None, ssl_credentials=None,
245261
credentials_file (str): A file with credentials that can be loaded with
246262
:func:`google.auth.load_credentials_from_file`. This argument is
247263
mutually exclusive with credentials.
264+
quota_project_id (str): An optional project to use for billing and quota.
248265
kwargs: Additional key-word args passed to
249266
:func:`grpc_gcp.secure_channel` or :func:`grpc.secure_channel`.
250267
@@ -259,7 +276,8 @@ def create_channel(target, credentials=None, scopes=None, ssl_credentials=None,
259276
credentials=credentials,
260277
credentials_file=credentials_file,
261278
scopes=scopes,
262-
ssl_credentials=ssl_credentials
279+
ssl_credentials=ssl_credentials,
280+
quota_project_id=quota_project_id,
263281
)
264282

265283
if HAS_GRPC_GCP:

google/api_core/grpc_helpers_async.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -206,7 +206,14 @@ def wrap_errors(callable_):
206206
return _wrap_stream_errors(callable_)
207207

208208

209-
def create_channel(target, credentials=None, scopes=None, ssl_credentials=None, credentials_file=None, **kwargs):
209+
def create_channel(
210+
target,
211+
credentials=None,
212+
scopes=None,
213+
ssl_credentials=None,
214+
credentials_file=None,
215+
quota_project_id=None,
216+
**kwargs):
210217
"""Create an AsyncIO secure channel with credentials.
211218
212219
Args:
@@ -222,6 +229,7 @@ def create_channel(target, credentials=None, scopes=None, ssl_credentials=None,
222229
credentials_file (str): A file with credentials that can be loaded with
223230
:func:`google.auth.load_credentials_from_file`. This argument is
224231
mutually exclusive with credentials.
232+
quota_project_id (str): An optional project to use for billing and quota.
225233
kwargs: Additional key-word args passed to :func:`aio.secure_channel`.
226234
227235
Returns:
@@ -235,7 +243,8 @@ def create_channel(target, credentials=None, scopes=None, ssl_credentials=None,
235243
credentials=credentials,
236244
credentials_file=credentials_file,
237245
scopes=scopes,
238-
ssl_credentials=ssl_credentials
246+
ssl_credentials=ssl_credentials,
247+
quota_project_id=quota_project_id,
239248
)
240249

241250
return aio.secure_channel(target, composite_credentials, **kwargs)

setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
dependencies = [
3232
"googleapis-common-protos >= 1.6.0, < 2.0dev",
3333
"protobuf >= 3.12.0",
34-
"google-auth >= 1.18.0, < 2.0dev",
34+
"google-auth >= 1.19.1, < 2.0dev",
3535
"requests >= 2.18.0, < 3.0.0dev",
3636
"setuptools >= 34.0.0",
3737
"six >= 1.10.0",

tests/asyncio/test_grpc_helpers_async.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -363,6 +363,23 @@ def test_create_channel_explicit_scoped(grpc_secure_channel, composite_creds_cal
363363
grpc_secure_channel.assert_called_once_with(target, composite_creds)
364364

365365

366+
@mock.patch("grpc.composite_channel_credentials")
367+
@mock.patch("grpc.experimental.aio.secure_channel")
368+
def test_create_channel_explicit_with_quota_project(grpc_secure_channel, composite_creds_call):
369+
target = "example.com:443"
370+
composite_creds = composite_creds_call.return_value
371+
372+
credentials = mock.create_autospec(google.auth.credentials.Credentials, instance=True)
373+
374+
channel = grpc_helpers_async.create_channel(
375+
target, credentials=credentials, quota_project_id="project-foo"
376+
)
377+
378+
credentials.with_quota_project.assert_called_once_with("project-foo")
379+
assert channel is grpc_secure_channel.return_value
380+
grpc_secure_channel.assert_called_once_with(target, composite_creds)
381+
382+
366383
@mock.patch("grpc.composite_channel_credentials")
367384
@mock.patch("grpc.experimental.aio.secure_channel")
368385
@mock.patch(

tests/unit/test_grpc_helpers.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -335,6 +335,29 @@ def test_create_channel_explicit_scoped(grpc_secure_channel, composite_creds_cal
335335
grpc_secure_channel.assert_called_once_with(target, composite_creds)
336336

337337

338+
@mock.patch("grpc.composite_channel_credentials")
339+
@mock.patch("grpc.secure_channel")
340+
def test_create_channel_explicit_with_quota_project(grpc_secure_channel, composite_creds_call):
341+
target = "example.com:443"
342+
composite_creds = composite_creds_call.return_value
343+
344+
credentials = mock.create_autospec(google.auth.credentials.Credentials, instance=True)
345+
346+
channel = grpc_helpers.create_channel(
347+
target,
348+
credentials=credentials,
349+
quota_project_id="project-foo"
350+
)
351+
352+
credentials.with_quota_project.assert_called_once_with("project-foo")
353+
354+
assert channel is grpc_secure_channel.return_value
355+
if grpc_helpers.HAS_GRPC_GCP:
356+
grpc_secure_channel.assert_called_once_with(target, composite_creds, None)
357+
else:
358+
grpc_secure_channel.assert_called_once_with(target, composite_creds)
359+
360+
338361
@mock.patch("grpc.composite_channel_credentials")
339362
@mock.patch("grpc.secure_channel")
340363
@mock.patch(

0 commit comments

Comments
 (0)