diff --git a/bigtable/google/cloud/bigtable/client.py b/bigtable/google/cloud/bigtable/client.py index f9a625b15843..d1d9a5dbbcff 100644 --- a/bigtable/google/cloud/bigtable/client.py +++ b/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. diff --git a/bigtable/tests/unit/test_client.py b/bigtable/tests/unit/test_client.py index 05a017d898af..c7c530d4c283 100644 --- a/bigtable/tests/unit/test_client.py +++ b/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