Changing Bigtable stub factories to use beta stubs.

Author: dhermes

gcloud/bigtable/_helpers.py

@@ -19,6 +19,9 @@
 """
 
 
+from grpc.beta import implementations
+
+
 # See https://gist.github.com/dhermes/bbc5b7be1932bfffae77
 # for appropriate values on other systems.
 # NOTE: Even this path is Unix specific.
@@ -66,7 +69,9 @@ def get_certs():
 
 
 def make_stub(client, stub_factory, host, port):
-    """Makes a stub for the an API.
+    """Makes a stub for an RPC service.
+
+    Uses / depends on the beta implementation of gRPC.
 
     :type client: :class:`.client.Client`
     :param client: The client that owns the cluster. Provides authorization and
@@ -85,8 +90,11 @@ def make_stub(client, stub_factory, host, port):
     :rtype: :class:`grpc.beta._stub._AutoIntermediary`
     :returns: The stub object used to make gRPC requests to a given API.
     """
+    root_certificates = get_certs()
+    client_credentials = implementations.ssl_client_credentials(
+        root_certificates, private_key=None, certificate_chain=None)
+    channel = implementations.secure_channel(
+        host, port, client_credentials)
     custom_metadata_transformer = MetadataTransformer(client)
-    return stub_factory(host, port,
-                        metadata_transformer=custom_metadata_transformer,
-                        secure=True,
-                        root_certificates=get_certs())
+    return stub_factory(channel,
+                        metadata_transformer=custom_metadata_transformer)

gcloud/bigtable/client.py

@@ -43,28 +43,27 @@
 from gcloud.credentials import get_credentials
 
 
-TABLE_STUB_FACTORY = (bigtable_table_service_pb2.
-                      early_adopter_create_BigtableTableService_stub)
+TABLE_STUB_FACTORY = (
+    bigtable_table_service_pb2.beta_create_BigtableTableService_stub)
 TABLE_ADMIN_HOST = 'bigtabletableadmin.googleapis.com'
 """Table Admin API request host."""
 TABLE_ADMIN_PORT = 443
 """Table Admin API request port."""
 
-CLUSTER_STUB_FACTORY = (bigtable_cluster_service_pb2.
-                        early_adopter_create_BigtableClusterService_stub)
+CLUSTER_STUB_FACTORY = (
+    bigtable_cluster_service_pb2.beta_create_BigtableClusterService_stub)
 CLUSTER_ADMIN_HOST = 'bigtableclusteradmin.googleapis.com'
 """Cluster Admin API request host."""
 CLUSTER_ADMIN_PORT = 443
 """Cluster Admin API request port."""
 
-DATA_STUB_FACTORY = (bigtable_service_pb2.
-                     early_adopter_create_BigtableService_stub)
+DATA_STUB_FACTORY = bigtable_service_pb2.beta_create_BigtableService_stub
 DATA_API_HOST = 'bigtable.googleapis.com'
 """Data API request host."""
 DATA_API_PORT = 443
 """Data API request port."""
 
-OPERATIONS_STUB_FACTORY = operations_pb2.early_adopter_create_Operations_stub
+OPERATIONS_STUB_FACTORY = operations_pb2.beta_create_Operations_stub
 
 ADMIN_SCOPE = 'https://www.googleapis.com/auth/cloud-bigtable.admin'
 """Scope for interacting with the Cluster Admin and Table Admin APIs."""

gcloud/bigtable/test__helpers.py

@@ -97,10 +97,27 @@ def test_it(self):
         mock_result = object()
         stub_inputs = []
 
-        def mock_stub_factory(host, port, metadata_transformer=None,
-                              secure=None, root_certificates=None):
-            stub_inputs.append((host, port, metadata_transformer,
-                                secure, root_certificates))
+        CLIENT_CREDS = object()
+        CHANNEL = object()
+
+        class _ImplementationsModule(object):
+
+            def __init__(self):
+                self.ssl_client_credentials_args = None
+                self.secure_channel_args = None
+
+            def ssl_client_credentials(self, *args, **kwargs):
+                self.ssl_client_credentials_args = (args, kwargs)
+                return CLIENT_CREDS
+
+            def secure_channel(self, *args, **kwargs):
+                self.secure_channel_args = (args, kwargs)
+                return CHANNEL
+
+        implementations_mod = _ImplementationsModule()
+
+        def mock_stub_factory(channel, metadata_transformer=None):
+            stub_inputs.append((channel, metadata_transformer))
             return mock_result
 
         transformed = object()
@@ -115,12 +132,18 @@ def mock_transformer(client):
         certs = 'FOOBAR'
         client = object()
         with _Monkey(MUT, get_certs=lambda: certs,
+                     implementations=implementations_mod,
                      MetadataTransformer=mock_transformer):
             result = self._callFUT(client, mock_stub_factory, host, port)
 
         self.assertTrue(result is mock_result)
-        self.assertEqual(stub_inputs, [(host, port, transformed, True, certs)])
+        self.assertEqual(stub_inputs, [(CHANNEL, transformed)])
         self.assertEqual(clients, [client])
+        ssl_cli_kwargs = {'private_key': None, 'certificate_chain': None}
+        self.assertEqual(implementations_mod.ssl_client_credentials_args,
+                         ((certs,), ssl_cli_kwargs))
+        self.assertEqual(implementations_mod.secure_channel_args,
+                         ((host, port, CLIENT_CREDS), {}))
 
 
 class _Credentials(object):