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