Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
feat: allow quota project to be passed to create_channel (#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
  • Loading branch information
busunkim96 committed Jul 20, 2020
1 parent 8595f62 commit e2d9a7b
Show file tree
Hide file tree
Showing 5 changed files with 73 additions and 6 deletions.
24 changes: 21 additions & 3 deletions google/api_core/grpc_helpers.py
Expand Up @@ -176,7 +176,12 @@ def wrap_errors(callable_):
return _wrap_unary_errors(callable_)


def _create_composite_credentials(credentials=None, credentials_file=None, scopes=None, ssl_credentials=None):
def _create_composite_credentials(
credentials=None,
credentials_file=None,
scopes=None,
ssl_credentials=None,
quota_project_id=None):
"""Create the composite credentials for secure channels.
Args:
Expand All @@ -191,6 +196,7 @@ def _create_composite_credentials(credentials=None, credentials_file=None, scope
are passed to :func:`google.auth.default`.
ssl_credentials (grpc.ChannelCredentials): Optional SSL channel
credentials. This can be used to specify different certificates.
quota_project_id (str): An optional project to use for billing and quota.
Returns:
grpc.ChannelCredentials: The composed channel credentials object.
Expand All @@ -210,6 +216,9 @@ def _create_composite_credentials(credentials=None, credentials_file=None, scope
else:
credentials, _ = google.auth.default(scopes=scopes)

if quota_project_id:
credentials = credentials.with_quota_project(quota_project_id)

request = google.auth.transport.requests.Request()

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


def create_channel(target, credentials=None, scopes=None, ssl_credentials=None, credentials_file=None, **kwargs):
def create_channel(
target,
credentials=None,
scopes=None,
ssl_credentials=None,
credentials_file=None,
quota_project_id=None,
**kwargs):
"""Create a secure channel with credentials.
Args:
Expand All @@ -245,6 +261,7 @@ def create_channel(target, credentials=None, scopes=None, ssl_credentials=None,
credentials_file (str): A file with credentials that can be loaded with
:func:`google.auth.load_credentials_from_file`. This argument is
mutually exclusive with credentials.
quota_project_id (str): An optional project to use for billing and quota.
kwargs: Additional key-word args passed to
:func:`grpc_gcp.secure_channel` or :func:`grpc.secure_channel`.
Expand All @@ -259,7 +276,8 @@ def create_channel(target, credentials=None, scopes=None, ssl_credentials=None,
credentials=credentials,
credentials_file=credentials_file,
scopes=scopes,
ssl_credentials=ssl_credentials
ssl_credentials=ssl_credentials,
quota_project_id=quota_project_id,
)

if HAS_GRPC_GCP:
Expand Down
13 changes: 11 additions & 2 deletions google/api_core/grpc_helpers_async.py
Expand Up @@ -206,7 +206,14 @@ def wrap_errors(callable_):
return _wrap_stream_errors(callable_)


def create_channel(target, credentials=None, scopes=None, ssl_credentials=None, credentials_file=None, **kwargs):
def create_channel(
target,
credentials=None,
scopes=None,
ssl_credentials=None,
credentials_file=None,
quota_project_id=None,
**kwargs):
"""Create an AsyncIO secure channel with credentials.
Args:
Expand All @@ -222,6 +229,7 @@ def create_channel(target, credentials=None, scopes=None, ssl_credentials=None,
credentials_file (str): A file with credentials that can be loaded with
:func:`google.auth.load_credentials_from_file`. This argument is
mutually exclusive with credentials.
quota_project_id (str): An optional project to use for billing and quota.
kwargs: Additional key-word args passed to :func:`aio.secure_channel`.
Returns:
Expand All @@ -235,7 +243,8 @@ def create_channel(target, credentials=None, scopes=None, ssl_credentials=None,
credentials=credentials,
credentials_file=credentials_file,
scopes=scopes,
ssl_credentials=ssl_credentials
ssl_credentials=ssl_credentials,
quota_project_id=quota_project_id,
)

return aio.secure_channel(target, composite_credentials, **kwargs)
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Expand Up @@ -31,7 +31,7 @@
dependencies = [
"googleapis-common-protos >= 1.6.0, < 2.0dev",
"protobuf >= 3.12.0",
"google-auth >= 1.18.0, < 2.0dev",
"google-auth >= 1.19.1, < 2.0dev",
"requests >= 2.18.0, < 3.0.0dev",
"setuptools >= 34.0.0",
"six >= 1.10.0",
Expand Down
17 changes: 17 additions & 0 deletions tests/asyncio/test_grpc_helpers_async.py
Expand Up @@ -363,6 +363,23 @@ def test_create_channel_explicit_scoped(grpc_secure_channel, composite_creds_cal
grpc_secure_channel.assert_called_once_with(target, composite_creds)


@mock.patch("grpc.composite_channel_credentials")
@mock.patch("grpc.experimental.aio.secure_channel")
def test_create_channel_explicit_with_quota_project(grpc_secure_channel, composite_creds_call):
target = "example.com:443"
composite_creds = composite_creds_call.return_value

credentials = mock.create_autospec(google.auth.credentials.Credentials, instance=True)

channel = grpc_helpers_async.create_channel(
target, credentials=credentials, quota_project_id="project-foo"
)

credentials.with_quota_project.assert_called_once_with("project-foo")
assert channel is grpc_secure_channel.return_value
grpc_secure_channel.assert_called_once_with(target, composite_creds)


@mock.patch("grpc.composite_channel_credentials")
@mock.patch("grpc.experimental.aio.secure_channel")
@mock.patch(
Expand Down
23 changes: 23 additions & 0 deletions tests/unit/test_grpc_helpers.py
Expand Up @@ -335,6 +335,29 @@ def test_create_channel_explicit_scoped(grpc_secure_channel, composite_creds_cal
grpc_secure_channel.assert_called_once_with(target, composite_creds)


@mock.patch("grpc.composite_channel_credentials")
@mock.patch("grpc.secure_channel")
def test_create_channel_explicit_with_quota_project(grpc_secure_channel, composite_creds_call):
target = "example.com:443"
composite_creds = composite_creds_call.return_value

credentials = mock.create_autospec(google.auth.credentials.Credentials, instance=True)

channel = grpc_helpers.create_channel(
target,
credentials=credentials,
quota_project_id="project-foo"
)

credentials.with_quota_project.assert_called_once_with("project-foo")

assert channel is grpc_secure_channel.return_value
if grpc_helpers.HAS_GRPC_GCP:
grpc_secure_channel.assert_called_once_with(target, composite_creds, None)
else:
grpc_secure_channel.assert_called_once_with(target, composite_creds)


@mock.patch("grpc.composite_channel_credentials")
@mock.patch("grpc.secure_channel")
@mock.patch(
Expand Down

0 comments on commit e2d9a7b

Please sign in to comment.