Add support for gRPC connection management (available when using optional grpc_gcp dependency) (#5553)

3 files changed, 165 insertions, 22 deletions

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.credentials
 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 @@ def wrap_errors(callable_):
         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 @@ def create_channel(target, credentials=None, scopes=None, **kwargs):
         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 @@ def create_channel(target, credentials=None, scopes=None, **kwargs):
 
     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 @@ def unit(session, py):
 
 
 @nox.session
+@nox.parametrize('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)
+
+
+@nox.session
 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 @@ def test_wrap_errors_streaming(wrap_stream_errors):
     wrap_stream_errors.assert_called_once_with(callable_)
 
 
+@mock.patch('grpc.composite_channel_credentials')
 @mock.patch(
     'google.auth.default',
     return_value=(mock.sentinel.credentials, mock.sentinel.projet))
-@mock.patch('google.auth.transport.grpc.secure_authorized_channel')
-def test_create_channel_implicit(secure_authorized_channel, default):
+@mock.patch('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)
 
 
+@mock.patch('grpc.composite_channel_credentials')
 @mock.patch(
     'google.auth.default',
     return_value=(mock.sentinel.credentials, mock.sentinel.projet))
-@mock.patch('google.auth.transport.grpc.secure_authorized_channel')
+@mock.patch('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)
+
+
+@mock.patch('grpc.composite_channel_credentials')
+@mock.patch(
+    'google.auth.default',
+    return_value=(mock.sentinel.credentials, mock.sentinel.projet))
+@mock.patch('grpc.secure_channel')
 def test_create_channel_implicit_with_scopes(
-        secure_authorized_channel, default):
+        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'])
-
-
-@mock.patch('google.auth.transport.grpc.secure_authorized_channel')
-def test_create_channel_explicit(secure_authorized_channel):
+    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('google.auth.credentials.with_scopes_if_required')
+@mock.patch('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)
 
 
-@mock.patch('google.auth.transport.grpc.secure_authorized_channel')
-def test_create_channel_explicit_scoped(unused_secure_authorized_channel):
+@mock.patch('grpc.composite_channel_credentials')
+@mock.patch('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)
+
+
+@pytest.mark.skipif(not grpc_helpers.HAS_GRPC_GCP,
+                    reason='grpc_gcp module not available')
+@mock.patch('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)
 
+
+@pytest.mark.skipif(grpc_helpers.HAS_GRPC_GCP,
+                    reason='grpc_gcp module not available')
+@mock.patch('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)