Bigtable: Add transport as wrapper for channel

bigtable/google/cloud/bigtable/client.py

@@ -35,6 +35,11 @@
 
 from google.cloud import bigtable_v2
 from google.cloud import bigtable_admin_v2
+from google.cloud.bigtable_v2.gapic.transports import bigtable_grpc_transport
+from google.cloud.bigtable_admin_v2.gapic.transports import (
+    bigtable_instance_admin_grpc_transport,
+    bigtable_table_admin_grpc_transport,
+)
 
 from google.cloud.bigtable import __version__
 from google.cloud.bigtable.instance import Instance
@@ -61,11 +66,29 @@
 
 
 def _create_gapic_client(client_class):
-    def inner(self):
+    def inner(self, transport_class=None):
         if self._emulator_host is None:
-            return client_class(
-                credentials=self._credentials, client_info=self._client_info
-            )
+            if transport_class is not None:
+                if self._channel is not None:
+                    # Use transport with channel
+                    transport_obj = _transport(
+                        channel=self._channel, default_class=transport_class
+                    )
+                    return client_class(
+                        transport=transport_obj, client_info=self._client_info
+                    )
+                else:
+                    # Use transport as pointer to function with credentials
+                    return client_class(
+                        transport=_transport,
+                        credentials=self._credentials,
+                        client_info=self._client_info,
+                    )
+            else:
+                # Use credentials
+                return client_class(
+                    credentials=self._credentials, client_info=self._client_info
+                )
         else:
             return client_class(
                 channel=self._emulator_channel, client_info=self._client_info
@@ -74,6 +97,29 @@ def inner(self):
     return inner
 
 
+def _transport(channel=None, credentials=None, default_class=None):
+    """ Returns a object for gRPC class passed in argument default_class.
+    Args:
+
+        channel (grpc.Channel): A ``Channel`` instance through
+                which to make calls. If not passed, it's None.
+
+        credentials (google.auth.credentials.Credentials): The
+                authorization credentials to attach to requests. These
+                credentials identify this application to the service. If not passed, it's None.
+
+        default_class (GrpcTransport): A transport
+                instance, responsible for actually making the API calls.
+                The default transport uses the gRPC protocol.. Defaults to None.
+    """
+    if default_class is not None:
+        if channel is not None:
+            return default_class(channel=channel)
+        elif credentials is not None:
+            return default_class(credentials=credentials)
+    return None
+
+
 class Client(ClientWithProject):
     """Client for interacting with Google Cloud Bigtable API.
 
@@ -122,6 +168,8 @@ class Client(ClientWithProject):
     _table_data_client = None
     _table_admin_client = None
     _instance_admin_client = None
+    _admin_transport = None
+    _data_transport = None
 
     def __init__(
         self,
@@ -140,6 +188,7 @@ def __init__(
         # NOTE: We set the scopes **before** calling the parent constructor.
         #       It **may** use those scopes in ``with_scopes_if_required``.
         self._read_only = bool(read_only)
+        self._credentials = credentials
         self._admin = bool(admin)
         self._client_info = client_info
         self._emulator_host = os.getenv(BIGTABLE_EMULATOR)
@@ -213,8 +262,13 @@ def table_data_client(self):
         :returns: A BigtableClient object.
         """
         if self._table_data_client is None:
+            if self._data_transport is not None:
+                transport_class = bigtable_grpc_transport.BigtableGrpcTransport
+            else:
+                transport_class = None
+
             self._table_data_client = _create_gapic_client(bigtable_v2.BigtableClient)(
-                self
+                self, transport_class=transport_class
             )
         return self._table_data_client
 
@@ -237,9 +291,16 @@ def table_admin_client(self):
         if self._table_admin_client is None:
             if not self._admin:
                 raise ValueError("Client is not an admin client.")
+            if self._admin_transport is not None:
+                transport_class = (
+                    bigtable_table_admin_grpc_transport.BigtableTableAdminGrpcTransport
+                )
+            else:
+                transport_class = None
+
             self._table_admin_client = _create_gapic_client(
                 bigtable_admin_v2.BigtableTableAdminClient
-            )(self)
+            )(self, transport_class=transport_class)
         return self._table_admin_client
 
     @property
@@ -261,11 +322,35 @@ def instance_admin_client(self):
         if self._instance_admin_client is None:
             if not self._admin:
                 raise ValueError("Client is not an admin client.")
+
+            if self._admin_transport is not None:
+                transport_class = (
+                    bigtable_instance_admin_grpc_transport.BigtableInstanceAdminGrpcTransport
+                )
+            else:
+                transport_class = None
             self._instance_admin_client = _create_gapic_client(
                 bigtable_admin_v2.BigtableInstanceAdminClient
-            )(self)
+            )(self, transport_class=transport_class)
         return self._instance_admin_client
 
+    def data_transport(self, channel=None):
+        if self._data_transport is None:
+            self._data_transport = _transport
+            self._channel = channel
+
+        return self._data_transport
+
+    def admin_transport(self, channel=None):
+        if self._admin_transport is None:
+            if not self._admin:
+                raise ValueError("Client is not an admin client.")
+
+            self._admin_transport = _transport
+            self._channel = channel
+
+        return self._admin_transport
+
     def instance(self, instance_id, display_name=None, instance_type=None, labels=None):
         """Factory to create a instance associated with this client.
 

bigtable/tests/unit/test_client.py

@@ -210,6 +210,61 @@ def test_table_data_client_not_initialized_w_client_info(self):
         self.assertIs(table_data_client._client_info, client_info)
         self.assertIs(client._table_data_client, table_data_client)
 
+    def test_data_transport_not_initialized_w_admin_flag(self):
+        client = self._make_one(project=self.PROJECT, admin=True)
+
+        self.assertTrue(callable(client.data_transport))
+        data_transport = client.data_transport()
+        self.assertTrue(callable(data_transport))
+
+        self.assertIsNone(data_transport())
+
+    def test_table_data_client_not_initialized_w_data_transport_w_channel(self):
+        from google.cloud.bigtable.client import _CLIENT_INFO
+        from google.cloud.bigtable_v2 import BigtableClient
+        from google.cloud.bigtable_v2.gapic.transports import bigtable_grpc_transport
+
+        client = self._make_one(project=self.PROJECT, admin=True)
+
+        channel = mock.Mock()
+
+        data_transport = client.data_transport(channel=channel)
+
+        self.assertTrue(callable(data_transport))
+
+        data_transport_obj = data_transport(
+            channel=channel, default_class=bigtable_grpc_transport.BigtableGrpcTransport
+        )
+
+        self.assertIsInstance(
+            data_transport_obj, bigtable_grpc_transport.BigtableGrpcTransport
+        )
+
+        table_data_client = client.table_data_client
+        self.assertIsInstance(table_data_client, BigtableClient)
+        self.assertIs(table_data_client._client_info, _CLIENT_INFO)
+        self.assertIs(client._table_data_client, table_data_client)
+
+    def test_table_data_client_not_initialized_w_data_transport_w_credentials(self):
+        from google.cloud.bigtable.client import _CLIENT_INFO
+        from google.cloud.bigtable_v2 import BigtableClient
+
+        credentials = _make_credentials()
+        client = self._make_one(
+            project=self.PROJECT, credentials=credentials, admin=True
+        )
+
+        data_transport = client.data_transport()
+        self.assertTrue(callable(data_transport))
+
+        self.assertIsNotNone(client._credentials)
+
+        table_data_client = client.table_data_client
+
+        self.assertIsInstance(table_data_client, BigtableClient)
+        self.assertIs(table_data_client._client_info, _CLIENT_INFO)
+        self.assertIs(client._table_data_client, table_data_client)
+
     def test_table_data_client_initialized(self):
         credentials = _make_credentials()
         client = self._make_one(
@@ -226,7 +281,46 @@ def test_table_admin_client_not_initialized_no_admin_flag(self):
         with self.assertRaises(ValueError):
             client.table_admin_client()
 
-    def test_table_admin_client_not_initialized_w_admin_flag(self):
+    def test_admin_transport_not_initialized_w_admin_flag(self):
+        client = self._make_one(project=self.PROJECT, admin=True)
+
+        self.assertTrue(callable(client.admin_transport))
+        admin_transport = client.admin_transport()
+        self.assertTrue(callable(admin_transport))
+
+        self.assertIsNone(admin_transport())
+
+    def test_table_admin_client_not_initialized_w_admin_transport_w_channel(self):
+        from google.cloud.bigtable.client import _CLIENT_INFO
+        from google.cloud.bigtable_admin_v2 import BigtableTableAdminClient
+        from google.cloud.bigtable_admin_v2.gapic.transports import (
+            bigtable_table_admin_grpc_transport,
+        )
+
+        client = self._make_one(project=self.PROJECT, admin=True)
+
+        channel = mock.Mock()
+
+        admin_transport = client.admin_transport(channel=channel)
+
+        self.assertTrue(callable(admin_transport))
+
+        admin_transport_obj = admin_transport(
+            channel=channel,
+            default_class=bigtable_table_admin_grpc_transport.BigtableTableAdminGrpcTransport,
+        )
+
+        self.assertIsInstance(
+            admin_transport_obj,
+            bigtable_table_admin_grpc_transport.BigtableTableAdminGrpcTransport,
+        )
+
+        table_admin_client = client.table_admin_client
+        self.assertIsInstance(table_admin_client, BigtableTableAdminClient)
+        self.assertIs(table_admin_client._client_info, _CLIENT_INFO)
+        self.assertIs(client._table_admin_client, table_admin_client)
+
+    def test_table_admin_client_not_initialized_w_admin_transport_w_credentials(self):
         from google.cloud.bigtable.client import _CLIENT_INFO
         from google.cloud.bigtable_admin_v2 import BigtableTableAdminClient
 
@@ -235,7 +329,13 @@ def test_table_admin_client_not_initialized_w_admin_flag(self):
             project=self.PROJECT, credentials=credentials, admin=True
         )
 
+        admin_transport = client.admin_transport()
+        self.assertTrue(callable(admin_transport))
+
+        self.assertIsNotNone(client._credentials)
+
         table_admin_client = client.table_admin_client
+
         self.assertIsInstance(table_admin_client, BigtableTableAdminClient)
         self.assertIs(table_admin_client._client_info, _CLIENT_INFO)
         self.assertIs(client._table_admin_client, table_admin_client)
@@ -287,6 +387,58 @@ def test_instance_admin_client_not_initialized_w_admin_flag(self):
         self.assertIs(instance_admin_client._client_info, _CLIENT_INFO)
         self.assertIs(client._instance_admin_client, instance_admin_client)
 
+    def test_instance_admin_client_not_initialized_w_admin_transport_w_channel(self):
+        from google.cloud.bigtable.client import _CLIENT_INFO
+        from google.cloud.bigtable_admin_v2 import BigtableInstanceAdminClient
+        from google.cloud.bigtable_admin_v2.gapic.transports import (
+            bigtable_instance_admin_grpc_transport,
+        )
+
+        client = self._make_one(project=self.PROJECT, admin=True)
+
+        channel = mock.Mock()
+
+        admin_transport = client.admin_transport(channel=channel)
+
+        self.assertTrue(callable(admin_transport))
+
+        admin_transport_obj = admin_transport(
+            channel=channel,
+            default_class=bigtable_instance_admin_grpc_transport.BigtableInstanceAdminGrpcTransport,
+        )
+
+        self.assertIsInstance(
+            admin_transport_obj,
+            bigtable_instance_admin_grpc_transport.BigtableInstanceAdminGrpcTransport,
+        )
+
+        instance_admin_client = client.instance_admin_client
+        self.assertIsInstance(instance_admin_client, BigtableInstanceAdminClient)
+        self.assertIs(instance_admin_client._client_info, _CLIENT_INFO)
+        self.assertIs(client._instance_admin_client, instance_admin_client)
+
+    def test_instance_admin_client_not_initialized_w_admin_transport_w_credentials(
+        self
+    ):
+        from google.cloud.bigtable.client import _CLIENT_INFO
+        from google.cloud.bigtable_admin_v2 import BigtableInstanceAdminClient
+
+        credentials = _make_credentials()
+        client = self._make_one(
+            project=self.PROJECT, credentials=credentials, admin=True
+        )
+
+        admin_transport = client.admin_transport()
+        self.assertTrue(callable(admin_transport))
+
+        self.assertIsNotNone(client._credentials)
+
+        instance_admin_client = client.instance_admin_client
+
+        self.assertIsInstance(instance_admin_client, BigtableInstanceAdminClient)
+        self.assertIs(instance_admin_client._client_info, _CLIENT_INFO)
+        self.assertIs(client._instance_admin_client, instance_admin_client)
+
     def test_instance_admin_client_not_initialized_w_admin_and_client_info(self):
         from google.cloud.bigtable_admin_v2 import BigtableInstanceAdminClient