Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: allow quota project to be passed to create_channel #58

Merged
merged 7 commits into from Jul 20, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
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