Add support for gRPC connection management (available when using optional grpc_gcp dependency) (#5553)
diff --git a/google/api_core/grpc_helpers.py b/google/api_core/grpc_helpers.py
index 6a2f052..b4ac9e0 100644
--- a/google/api_core/grpc_helpers.py
+++ b/google/api_core/grpc_helpers.py
@@ -26,6 +26,11 @@
import google.auth.transport.grpc
import google.auth.transport.requests
+try:
+ import grpc_gcp
+ HAS_GRPC_GCP = True
+except ImportError:
+ HAS_GRPC_GCP = False
# The list of gRPC Callable interfaces that return iterators.
_STREAM_WRAP_CLASSES = (
@@ -149,7 +154,11 @@
return _wrap_unary_errors(callable_)
-def create_channel(target, credentials=None, scopes=None, **kwargs):
+def create_channel(target,
+ credentials=None,
+ scopes=None,
+ ssl_credentials=None,
+ **kwargs):
"""Create a secure channel with credentials.
Args:
@@ -160,8 +169,10 @@
scopes (Sequence[str]): A optional list of scopes needed for this
service. These are only used when credentials are not specified and
are passed to :func:`google.auth.default`.
+ ssl_credentials (grpc.ChannelCredentials): Optional SSL channel
+ credentials. This can be used to specify different certificates.
kwargs: Additional key-word args passed to
- :func:`google.auth.transport.grpc.secure_authorized_channel`.
+ :func:`grpc_gcp.secure_channel` or :func:`grpc.secure_channel`.
Returns:
grpc.Channel: The created channel.
@@ -174,8 +185,26 @@
request = google.auth.transport.requests.Request()
- return google.auth.transport.grpc.secure_authorized_channel(
- credentials, request, target, **kwargs)
+ # Create the metadata plugin for inserting the authorization header.
+ metadata_plugin = google.auth.transport.grpc.AuthMetadataPlugin(
+ credentials, request)
+
+ # Create a set of grpc.CallCredentials using the metadata plugin.
+ google_auth_credentials = grpc.metadata_call_credentials(metadata_plugin)
+
+ if ssl_credentials is None:
+ ssl_credentials = grpc.ssl_channel_credentials()
+
+ # Combine the ssl credentials and the authorization credentials.
+ composite_credentials = grpc.composite_channel_credentials(
+ ssl_credentials, google_auth_credentials)
+
+ if HAS_GRPC_GCP:
+ # If grpc_gcp module is available use grpc_gcp.secure_channel,
+ # otherwise, use grpc.secure_channel to create grpc channel.
+ return grpc_gcp.secure_channel(target, composite_credentials, **kwargs)
+ else:
+ return grpc.secure_channel(target, composite_credentials, **kwargs)
_MethodCall = collections.namedtuple(
diff --git a/nox.py b/nox.py
index 7b40c68..dacfbb5 100644
--- a/nox.py
+++ b/nox.py
@@ -66,6 +66,23 @@
@nox.session
[email protected]('py', ['2.7', '3.5', '3.6', '3.7'])
+def unit_grpc_gcp(session, py):
+ """Run the unit test suite with grpcio-gcp installed."""
+
+ # Run unit tests against all supported versions of Python.
+ session.interpreter = 'python{}'.format(py)
+
+ # Set the virtualenv dirname.
+ session.virtualenv_dirname = 'unit-grpc-gcp-' + py
+
+ # Install grpcio-gcp
+ session.install('grpcio-gcp')
+
+ default(session)
+
+
[email protected]
def lint(session):
"""Run linters.
diff --git a/tests/unit/test_grpc_helpers.py b/tests/unit/test_grpc_helpers.py
index e5e4311..b91847c 100644
--- a/tests/unit/test_grpc_helpers.py
+++ b/tests/unit/test_grpc_helpers.py
@@ -176,60 +176,157 @@
wrap_stream_errors.assert_called_once_with(callable_)
[email protected]('grpc.composite_channel_credentials')
@mock.patch(
'google.auth.default',
return_value=(mock.sentinel.credentials, mock.sentinel.projet))
[email protected]('google.auth.transport.grpc.secure_authorized_channel')
-def test_create_channel_implicit(secure_authorized_channel, default):
[email protected]('grpc.secure_channel')
+def test_create_channel_implicit(
+ grpc_secure_channel, default, composite_creds_call):
target = 'example.com:443'
+ composite_creds = composite_creds_call.return_value
channel = grpc_helpers.create_channel(target)
- assert channel is secure_authorized_channel.return_value
+ assert channel is grpc_secure_channel.return_value
default.assert_called_once_with(scopes=None)
- secure_authorized_channel.assert_called_once_with(
- mock.sentinel.credentials, mock.ANY, target)
+ 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)
[email protected]('grpc.composite_channel_credentials')
@mock.patch(
'google.auth.default',
return_value=(mock.sentinel.credentials, mock.sentinel.projet))
[email protected]('google.auth.transport.grpc.secure_authorized_channel')
-def test_create_channel_implicit_with_scopes(
- secure_authorized_channel, default):
[email protected]('grpc.secure_channel')
+def test_create_channel_implicit_with_ssl_creds(
+ grpc_secure_channel, default, composite_creds_call):
target = 'example.com:443'
+ ssl_creds = grpc.ssl_channel_credentials()
+
+ grpc_helpers.create_channel(target, ssl_credentials=ssl_creds)
+
+ default.assert_called_once_with(scopes=None)
+ composite_creds_call.assert_called_once_with(ssl_creds, mock.ANY)
+ composite_creds = composite_creds_call.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)
+
+
[email protected]('grpc.composite_channel_credentials')
[email protected](
+ 'google.auth.default',
+ return_value=(mock.sentinel.credentials, mock.sentinel.projet))
[email protected]('grpc.secure_channel')
+def test_create_channel_implicit_with_scopes(
+ grpc_secure_channel, default, composite_creds_call):
+ target = 'example.com:443'
+ composite_creds = composite_creds_call.return_value
+
channel = grpc_helpers.create_channel(target, scopes=['one', 'two'])
- assert channel is secure_authorized_channel.return_value
+ assert channel is grpc_secure_channel.return_value
default.assert_called_once_with(scopes=['one', 'two'])
+ 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)
[email protected]('google.auth.transport.grpc.secure_authorized_channel')
-def test_create_channel_explicit(secure_authorized_channel):
[email protected]('grpc.composite_channel_credentials')
[email protected]('google.auth.credentials.with_scopes_if_required')
[email protected]('grpc.secure_channel')
+def test_create_channel_explicit(
+ grpc_secure_channel, auth_creds, composite_creds_call):
target = 'example.com:443'
+ composite_creds = composite_creds_call.return_value
channel = grpc_helpers.create_channel(
target, credentials=mock.sentinel.credentials)
- assert channel is secure_authorized_channel.return_value
- secure_authorized_channel.assert_called_once_with(
- mock.sentinel.credentials, mock.ANY, target)
+ auth_creds.assert_called_once_with(mock.sentinel.credentials, None)
+ 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)
[email protected]('google.auth.transport.grpc.secure_authorized_channel')
-def test_create_channel_explicit_scoped(unused_secure_authorized_channel):
[email protected]('grpc.composite_channel_credentials')
[email protected]('grpc.secure_channel')
+def test_create_channel_explicit_scoped(
+ grpc_secure_channel, composite_creds_call):
+ target = 'example.com:443'
scopes = ['1', '2']
+ composite_creds = composite_creds_call.return_value
+
+ credentials = mock.create_autospec(
+ google.auth.credentials.Scoped, instance=True)
+ credentials.requires_scopes = True
+
+ channel = grpc_helpers.create_channel(
+ target,
+ credentials=credentials,
+ scopes=scopes)
+
+ credentials.with_scopes.assert_called_once_with(scopes)
+ 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)
+
+
[email protected](not grpc_helpers.HAS_GRPC_GCP,
+ reason='grpc_gcp module not available')
[email protected]('grpc_gcp.secure_channel')
+def test_create_channel_with_grpc_gcp(grpc_gcp_secure_channel):
+ target = 'example.com:443'
+ scopes = ['test_scope']
credentials = mock.create_autospec(
google.auth.credentials.Scoped, instance=True)
credentials.requires_scopes = True
grpc_helpers.create_channel(
- mock.sentinel.target,
+ target,
credentials=credentials,
scopes=scopes)
+ grpc_gcp_secure_channel.assert_called()
+ credentials.with_scopes.assert_called_once_with(scopes)
+
[email protected](grpc_helpers.HAS_GRPC_GCP,
+ reason='grpc_gcp module not available')
[email protected]('grpc.secure_channel')
+def test_create_channel_without_grpc_gcp(grpc_secure_channel):
+ target = 'example.com:443'
+ scopes = ['test_scope']
+
+ credentials = mock.create_autospec(
+ google.auth.credentials.Scoped, instance=True)
+ credentials.requires_scopes = True
+
+ grpc_helpers.create_channel(
+ target,
+ credentials=credentials,
+ scopes=scopes)
+ grpc_secure_channel.assert_called()
credentials.with_scopes.assert_called_once_with(scopes)
