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

1 files changed, 33 insertions, 4 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(