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

Author: WeiranFang

packages/google-cloud-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__

packages/google-cloud-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"
+  }
+}

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

@@ -32,8 +32,15 @@
 from google.cloud.spanner_v1.proto import transaction_pb2
 from google.protobuf import struct_pb2
 
+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 +120,20 @@ 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.api_config_from_text_pb(
+                    pkg_resources.resource_string(__name__,
+                                                  _SPANNER_GRPC_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.

packages/google-cloud-spanner/nox.py

@@ -69,20 +69,23 @@ def unit(session, py):
 
 
 @nox.session
-@nox.parametrize('py', ['2.7', '3.6'])
-def system(session, py):
-    """Run the system test suite."""
-
-    # Sanity check: Only run system tests if the environment variable is set.
-    if not os.environ.get('GOOGLE_APPLICATION_CREDENTIALS', ''):
-        session.skip('Credentials must be set via environment variable.')
+@nox.parametrize('py', ['2.7', '3.4', '3.5', '3.6', '3.7'])
+def unit_grpc_gcp(session, py):
+    """Run the unit test suite with grpcio-gcp installed."""
 
-    # Run the system tests against latest Python 2 and Python 3 only.
+    # Run unit tests against all supported versions of Python.
     session.interpreter = 'python{}'.format(py)
 
     # Set the virtualenv dirname.
-    session.virtualenv_dirname = 'sys-' + py
+    session.virtualenv_dirname = 'unit-grpc-gcp-' + py
+
+    # Install grpcio-gcp
+    session.install('grpcio-gcp')
 
+    default(session)
+
+
+def system_common(session):
     # Use pre-release gRPC for system tests.
     session.install('--pre', 'grpcio')
 
@@ -96,6 +99,45 @@ def system(session, py):
     session.run('py.test', '--quiet', 'tests/system', *session.posargs)
 
 
+@nox.session
+@nox.parametrize('py', ['2.7', '3.6'])
+def system(session, py):
+    """Run the system test suite."""
+
+    # Sanity check: Only run system tests if the environment variable is set.
+    if not os.environ.get('GOOGLE_APPLICATION_CREDENTIALS', ''):
+        session.skip('Credentials must be set via environment variable.')
+
+    # Run the system tests against latest Python 2 and Python 3 only.
+    session.interpreter = 'python{}'.format(py)
+
+    # Set the virtualenv dirname.
+    session.virtualenv_dirname = 'sys-' + py
+
+    system_common(session)
+
+
+@nox.session
+@nox.parametrize('py', ['2.7', '3.6'])
+def system_grpc_gcp(session, py):
+    """Run the system test suite with grpcio-gcp installed."""
+
+    # Sanity check: Only run system tests if the environment variable is set.
+    if not os.environ.get('GOOGLE_APPLICATION_CREDENTIALS', ''):
+        session.skip('Credentials must be set via environment variable.')
+
+    # Run the system tests against latest Python 2 and Python 3 only.
+    session.interpreter = 'python{}'.format(py)
+
+    # Set the virtualenv dirname.
+    session.virtualenv_dirname = 'sys-grpc-gcp-' + py
+
+    # Install grpcio-gcp
+    session.install('grpcio-gcp')
+
+    system_common(session)
+
+
 @nox.session
 def lint(session):
     """Run linters.

packages/google-cloud-spanner/tests/unit/gapic/v1/test_spanner_client_v1.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Unit tests."""
 
+import mock
 import pytest
 
 # Manual edit to auto-generated import because we do not expose the
@@ -554,3 +555,24 @@ def test_partition_read_exception(self):
 
         with pytest.raises(CustomException):
             client.partition_read(session, table, key_set)
+
+    @pytest.mark.skipif(not spanner_v1.HAS_GRPC_GCP,
+                        reason='grpc_gcp module not available')
+    @mock.patch(
+        'google.auth.default',
+        return_value=(mock.sentinel.credentials, mock.sentinel.projet))
+    @mock.patch('google.protobuf.text_format.Merge')
+    @mock.patch('grpc_gcp.proto.grpc_gcp_pb2.ApiConfig',
+                return_value=mock.sentinel.api_config)
+    @mock.patch('grpc_gcp.secure_channel')
+    def test_client_with_grpc_gcp_channel(self,
+                                          grpc_gcp_secure_channel,
+                                          api_config,
+                                          merge,
+                                          auth_default):
+        spanner_target = 'spanner.googleapis.com:443'
+        client = spanner_v1.SpannerClient()
+        merge.assert_called_once_with(mock.ANY, mock.sentinel.api_config)
+        options = [('grpc_gcp.api_config', mock.sentinel.api_config)]
+        grpc_gcp_secure_channel.assert_called_once_with(
+            spanner_target, mock.ANY, options=options)