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_table_admin_grpc_transport,
+    bigtable_instance_admin_grpc_transport,
+)
 
 from google.cloud.bigtable import __version__
 from google.cloud.bigtable.instance import Instance
@@ -61,11 +66,14 @@
 
 
 def _create_gapic_client(client_class):
-    def inner(self):
+    def inner(self, transport=None):
         if self._emulator_host is None:
-            return client_class(
-                credentials=self._credentials, client_info=self._client_info
-            )
+            if transport is not None:
+                return client_class(transport=transport, client_info=self._client_info)
+            else:
+                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 +82,18 @@ def inner(self):
     return inner
 
 
+def transport(channel=None):
+    def inner(credentials=None, default_class=None):
+        if channel is not None:
+            return default_class(channel=channel)
+        elif credentials is not None:
+            return default_class(credentials=credentials)
+        else:
+            return None
+
+    return inner
+
+
 class Client(ClientWithProject):
     """Client for interacting with Google Cloud Bigtable API.
 
@@ -122,6 +142,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 +162,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 +236,15 @@ def table_data_client(self):
         :returns: A BigtableClient object.
         """
         if self._table_data_client is None:
+            if self._data_transport is not None:
+                transport = self.data_transport(
+                    default_class=bigtable_grpc_transport.BigtableGrpcTransport
+                )
+            else:
+                transport = None
+
             self._table_data_client = _create_gapic_client(bigtable_v2.BigtableClient)(
-                self
+                self, transport=transport
             )
         return self._table_data_client
 
@@ -237,9 +267,15 @@ 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 = self.admin_transport(
+                    default_class=bigtable_table_admin_grpc_transport.BigtableTableAdminGrpcTransport
+                )
+            else:
+                transport = None
             self._table_admin_client = _create_gapic_client(
                 bigtable_admin_v2.BigtableTableAdminClient
-            )(self)
+            )(self, transport=transport)
         return self._table_admin_client
 
     @property
@@ -261,11 +297,44 @@ 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 = self.admin_transport(
+                    default_class=bigtable_instance_admin_grpc_transport.BigtableInstanceAdminGrpcTransport
+                )
+            else:
+                transport = None
             self._instance_admin_client = _create_gapic_client(
                 bigtable_admin_v2.BigtableInstanceAdminClient
-            )(self)
+            )(self, transport=transport)
         return self._instance_admin_client
 
+    @property
+    def data_transport(self):
+        return self._data_transport
+
+    @data_transport.setter
+    def data_transport(self, channel=None):
+        if self._data_transport is None:
+            if channel is not None:
+                self._data_transport = transport(channel=channel)
+            elif self._credentials is not None:
+                self._data_transport = transport()
+
+    @property
+    def admin_transport(self):
+        return self._admin_transport
+
+    @admin_transport.setter
+    def admin_transport(self, channel=None):
+
+        if self._admin_transport is None:
+            if not self._admin:
+                raise ValueError("Client is not an admin client.")
+            if channel is not None:
+                self._admin_transport = transport(channel=channel)
+            elif self._credentials is not None:
+                self._admin_transport = 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,79 @@ 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):
+        from google.cloud.bigtable_v2.gapic.transports import bigtable_grpc_transport
+
+        channel = mock.Mock()
+
+        client = self._make_one(project=self.PROJECT, admin=True)
+
+        client.data_transport = channel
+        self.assertTrue(callable(client.data_transport))
+        data_transport = client.data_transport(
+            default_class=bigtable_grpc_transport.BigtableGrpcTransport
+        )
+        self.assertIsInstance(
+            data_transport, bigtable_grpc_transport.BigtableGrpcTransport
+        )
+
+        self.assertIsInstance(
+            client._data_transport(
+                default_class=bigtable_grpc_transport.BigtableGrpcTransport
+            ),
+            bigtable_grpc_transport.BigtableGrpcTransport,
+        )
+
+    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()
+
+        client.data_transport = channel
+
+        self.assertTrue(callable(client.data_transport))
+        data_transport = client.data_transport(
+            default_class=bigtable_grpc_transport.BigtableGrpcTransport
+        )
+
+        table_data_client = client.table_data_client
+        self.assertIsInstance(
+            data_transport, bigtable_grpc_transport.BigtableGrpcTransport
+        )
+
+        self.assertIsInstance(
+            table_data_client.transport, bigtable_grpc_transport.BigtableGrpcTransport
+        )
+
+        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
+        )
+
+        client.data_transport = None
+        data_transport = client.data_transport
+        self.assertTrue(callable(data_transport))
+
+        self.assertIsNot(client._credentials, None)
+
+        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 +299,66 @@ 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):
+        from google.cloud.bigtable_admin_v2.gapic.transports import (
+            bigtable_table_admin_grpc_transport,
+        )
+
+        channel = mock.Mock()
+
+        client = self._make_one(project=self.PROJECT, admin=True)
+
+        client.admin_transport = channel
+        self.assertTrue(callable(client.admin_transport))
+        admin_transport = client.admin_transport(
+            default_class=bigtable_table_admin_grpc_transport.BigtableTableAdminGrpcTransport
+        )
+        self.assertIsInstance(
+            admin_transport,
+            bigtable_table_admin_grpc_transport.BigtableTableAdminGrpcTransport,
+        )
+
+        self.assertIsInstance(
+            client._admin_transport(
+                default_class=bigtable_table_admin_grpc_transport.BigtableTableAdminGrpcTransport
+            ),
+            bigtable_table_admin_grpc_transport.BigtableTableAdminGrpcTransport,
+        )
+
+    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()
+
+        client.admin_transport = channel
+
+        self.assertTrue(callable(client.admin_transport))
+        admin_transport = client.admin_transport(
+            default_class=bigtable_table_admin_grpc_transport.BigtableTableAdminGrpcTransport
+        )
+
+        table_admin_client = client.table_admin_client
+        self.assertIsInstance(
+            admin_transport,
+            bigtable_table_admin_grpc_transport.BigtableTableAdminGrpcTransport,
+        )
+
+        self.assertIsInstance(
+            table_admin_client.transport,
+            bigtable_table_admin_grpc_transport.BigtableTableAdminGrpcTransport,
+        )
+
+        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 +367,14 @@ def test_table_admin_client_not_initialized_w_admin_flag(self):
             project=self.PROJECT, credentials=credentials, admin=True
         )
 
+        client.admin_transport = None
+        admin_transport = client.admin_transport
+        self.assertTrue(callable(admin_transport))
+
+        self.assertIsNot(client._credentials, None)
+
         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 +426,54 @@ 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_flag_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()
+
+        client.admin_transport = channel
+        admin_transport = client.admin_transport(
+            default_class=bigtable_instance_admin_grpc_transport.BigtableInstanceAdminGrpcTransport
+        )
+        self.assertIsInstance(
+            admin_transport,
+            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
+        )
+
+        client.admin_transport = None
+        admin_transport = client.admin_transport
+        self.assertTrue(callable(admin_transport))
+
+        self.assertIsNot(client._credentials, None)
+
+        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