Connection management support (with optional grpc_gcp dependency)

api_core/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,6 +154,52 @@ def wrap_errors(callable_):
         return _wrap_unary_errors(callable_)
 
 
+def _create_secure_channel(
+        credentials, request, target, ssl_credentials=None, **kwargs):
+    """Creates a secure authorized gRPC channel.
+
+    This overwrites google.auth.transport.grpc.secure_authorized_channel to
+    return a secure channel using grpc_gcp.secure_channel if grpc_gcp is
+    available.
+
+    Args:
+        credentials (google.auth.credentials.Credentials): The credentials to
+            add to requests.
+        request (google.auth.transport.Request): A HTTP transport request
+            object used to refresh credentials as needed. Even though gRPC
+            is a separate transport, there's no way to refresh the credentials
+            without using a standard http transport.
+        target (str): The host and port of the service.
+        ssl_credentials (grpc.ChannelCredentials): Optional SSL channel
+            credentials. This can be used to specify different certificates.
+        kwargs: Additional arguments to pass to
+            :func:`grpc_gcp.secure_channel`.
+
+    Returns:
+        grpc_gcp.Channel: The created gRPC channel.
+    """
+    # 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)
+
+
 def create_channel(target, credentials=None, scopes=None, **kwargs):
     """Create a secure channel with credentials.
 
@@ -161,7 +212,7 @@ def create_channel(target, credentials=None, scopes=None, **kwargs):
             service. These are only used when credentials are not specified and
             are passed to :func:`google.auth.default`.
         kwargs: Additional key-word args passed to
-            :func:`google.auth.transport.grpc.secure_authorized_channel`.
+            :func:`_create_secure_channel`.
 
     Returns:
         grpc.Channel: The created channel.
@@ -174,8 +225,7 @@ 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)
+    return _create_secure_channel(credentials, request, target, **kwargs)
 
 
 _MethodCall = collections.namedtuple(

api_core/nox.py

@@ -65,6 +65,17 @@ def unit(session, py):
     default(session)
 
 
+@nox.session
+@nox.parametrize('py', ['2.7', '3.5', '3.6', '3.7'])
+def unit_with_grpc_gcp(session, py):
+    """Run the unit test suite with grpcio-gcp installed."""
+
+    # Install grpcio-gcp
+    session.install('grpcio-gcp')
+
+    unit(session, py)
+
+
 @nox.session
 def lint(session):
     """Run linters.

api_core/tests/unit/test_grpc_helpers.py

@@ -179,46 +179,64 @@ def test_wrap_errors_streaming(wrap_stream_errors):
 @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('google.api_core.grpc_helpers._create_secure_channel')
+def test_create_channel_implicit(create_secure_channel, default):
     target = 'example.com:443'
 
     channel = grpc_helpers.create_channel(target)
 
-    assert channel is secure_authorized_channel.return_value
+    assert channel is create_secure_channel.return_value
     default.assert_called_once_with(scopes=None)
-    secure_authorized_channel.assert_called_once_with(
+    create_secure_channel.assert_called_once_with(
         mock.sentinel.credentials, mock.ANY, target)
 
 
+@mock.patch('grpc._channel.Channel')
+@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_with_ssl_creds(
+        default, composite, grpc_channel):
+    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.assert_called_once_with(ssl_creds, mock.ANY)
+    grpc_channel.assert_called()
+
+
+@mock.patch(
+    'google.auth.default',
+    return_value=(mock.sentinel.credentials, mock.sentinel.projet))
+@mock.patch('google.api_core.grpc_helpers._create_secure_channel')
 def test_create_channel_implicit_with_scopes(
-        secure_authorized_channel, default):
+        create_secure_channel, default):
     target = 'example.com:443'
 
     channel = grpc_helpers.create_channel(target, scopes=['one', 'two'])
 
-    assert channel is secure_authorized_channel.return_value
+    assert channel is create_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):
+@mock.patch('google.api_core.grpc_helpers._create_secure_channel')
+def test_create_channel_explicit(create_secure_channel):
     target = 'example.com:443'
 
     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(
+    assert channel is create_secure_channel.return_value
+    create_secure_channel.assert_called_once_with(
         mock.sentinel.credentials, mock.ANY, target)
 
 
-@mock.patch('google.auth.transport.grpc.secure_authorized_channel')
-def test_create_channel_explicit_scoped(unused_secure_authorized_channel):
+@mock.patch('google.api_core.grpc_helpers._create_secure_channel')
+def test_create_channel_explicit_scoped(unused_create_secure_channel):
     scopes = ['1', '2']
 
     credentials = mock.create_autospec(
@@ -233,6 +251,44 @@ def test_create_channel_explicit_scoped(unused_secure_authorized_channel):
     credentials.with_scopes.assert_called_once_with(scopes)
 
 
+@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(
+        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)
+
+
 class TestChannelStub(object):
 
     def test_single_response(self):

spanner/MANIFEST.in

@@ -1,4 +1,4 @@
 include README.rst LICENSE
-recursive-include google *.json *.proto
+recursive-include google *.json *.proto *.config
 recursive-include tests *
 global-exclude *.pyc __pycache__

spanner/google/cloud/spanner_v1/gapic/spanner.grpc.config

@@ -0,0 +1,88 @@
+channel_pool: {
+  max_size: 10
+  max_concurrent_streams_low_watermark: 100
+}
+method: {
+  name: "/google.spanner.v1.Spanner/CreateSession"
+  affinity: {
+    command: BIND
+    affinity_key: "name"
+  }
+}
+method: {
+  name: "/google.spanner.v1.Spanner/GetSession"
+  affinity: {
+    command: BOUND
+    affinity_key: "name"
+  }
+}
+method: {
+  name: "/google.spanner.v1.Spanner/DeleteSession"
+  affinity: {
+    command: UNBIND
+    affinity_key: "name"
+  }
+}
+method: {
+  name: "/google.spanner.v1.Spanner/ExecuteSql"
+  affinity: {
+    command: BOUND
+    affinity_key: "session"
+  }
+}
+method: {
+  name: "/google.spanner.v1.Spanner/ExecuteStreamingSql"
+  affinity: {
+    command: BOUND
+    affinity_key: "session"
+  }
+}
+method: {
+  name: "/google.spanner.v1.Spanner/Read"
+  affinity: {
+    command: BOUND
+    affinity_key: "session"
+  }
+}
+method: {
+  name: "/google.spanner.v1.Spanner/StreamingRead"
+  affinity: {
+    command: BOUND
+    affinity_key: "session"
+  }
+}
+method: {
+  name: "/google.spanner.v1.Spanner/BeginTransaction"
+  affinity: {
+    command: BOUND
+    affinity_key: "session"
+  }
+}
+method: {
+  name: "/google.spanner.v1.Spanner/Commit"
+  affinity: {
+    command: BOUND
+    affinity_key: "session"
+  }
+}
+method: {
+  name: "/google.spanner.v1.Spanner/Rollback"
+  affinity: {
+    command: BOUND
+    affinity_key: "session"
+  }
+}
+method: {
+  name: "/google.spanner.v1.Spanner/PartitionQuery"
+  affinity: {
+    command: BOUND
+    affinity_key: "session"
+  }
+}
+method: {
+  name: "/google.spanner.v1.Spanner/PartitionRead"
+  affinity: {
+    command: BOUND
+    affinity_key: "session"
+  }
+}

spanner/google/cloud/spanner_v1/gapic/spanner_client.py

@@ -31,9 +31,17 @@
 from google.cloud.spanner_v1.proto import spanner_pb2
 from google.cloud.spanner_v1.proto import transaction_pb2
 from google.protobuf import struct_pb2
+from google.protobuf import text_format
+
+try:
+    import grpc_gcp
+    HAS_GRPC_GCP = True
+except ImportError:
+    HAS_GRPC_GCP = False
 
 _GAPIC_LIBRARY_VERSION = pkg_resources.get_distribution(
     'google-cloud-spanner', ).version
+_SPANNER_GRPC_CONFIG = 'spanner.grpc.config'
 
 
 class SpannerClient(object):
@@ -113,10 +121,22 @@ def __init__(self,
 
         # Create the channel.
         if channel is None:
+            options = None
+
+            if HAS_GRPC_GCP:
+                # Initialize grpc gcp config for spanner api.
+                grpc_gcp_config = grpc_gcp.proto.grpc_gcp_pb2.ApiConfig()
+                text_format.Merge(
+                    pkg_resources.resource_string(__name__, _SPANNER_GRPC_CONFIG),
+                    grpc_gcp_config
+                )
+                options = [(grpc_gcp.API_CONFIG_CHANNEL_ARG, grpc_gcp_config)]
+
             channel = google.api_core.grpc_helpers.create_channel(
                 self.SERVICE_ADDRESS,
                 credentials=credentials,
                 scopes=self._DEFAULT_SCOPES,
+                options=options,
             )
 
         # Create the gRPC stubs.

spanner/nox.py

@@ -67,6 +67,15 @@ def unit(session, py):
 
     default(session)
 
+@nox.session
+@nox.parametrize('py', ['2.7', '3.4', '3.5', '3.6', '3.7'])
+def unit_with_grpc_gcp(session, py):
+    """Run the unit test suite with grpcio-gcp installed."""
+
+    # Install grpcio-gcp
+    session.install('grpcio-gcp')
+
+    unit(session, py)
 
 @nox.session
 @nox.parametrize('py', ['2.7', '3.6'])