From 2beaa95214d8f5cb18c5d57fb6f8ed2fd326741f Mon Sep 17 00:00:00 2001 From: Tres Seaver Date: Thu, 24 Oct 2019 14:33:48 -0400 Subject: [PATCH] feat(bigtable): add 'client_options' / 'admin_client_options' to Client (#9517) Toward #8475. --- bigtable/google/cloud/bigtable/client.py | 42 ++++++++--- bigtable/tests/unit/test_client.py | 95 ++++++++++++++++++++++-- 2 files changed, 121 insertions(+), 16 deletions(-) diff --git a/bigtable/google/cloud/bigtable/client.py b/bigtable/google/cloud/bigtable/client.py index f9a625b15843..8a8315623cae 100644 --- a/bigtable/google/cloud/bigtable/client.py +++ b/bigtable/google/cloud/bigtable/client.py @@ -60,11 +60,13 @@ """Scope for reading table data.""" -def _create_gapic_client(client_class): +def _create_gapic_client(client_class, client_options=None): def inner(self): if self._emulator_host is None: return client_class( - credentials=self._credentials, client_info=self._client_info + credentials=self._credentials, + client_info=self._client_info, + client_options=client_options, ) else: return client_class( @@ -109,6 +111,17 @@ class Client(ClientWithProject): you only need to set this if you're developing your own library or partner tool. + :type client_options: :class:`~google.api_core.client_options.ClientOptions` + or :class:`dict` + :param client_options: (Optional) Client options used to set user options + on the client. API Endpoint should be set through client_options. + + :type admin_client_options: + :class:`~google.api_core.client_options.ClientOptions` or :class:`dict` + :param admin_client_options: (Optional) Client options used to set user + options on the client. API Endpoint for admin operations should be set + through admin_client_options. + :type channel: :instance: grpc.Channel :param channel (grpc.Channel): (Optional) DEPRECATED: A ``Channel`` instance through which to make calls. @@ -130,6 +143,8 @@ def __init__( read_only=False, admin=False, client_info=_CLIENT_INFO, + client_options=None, + admin_client_options=None, channel=None, ): if read_only and admin: @@ -155,6 +170,8 @@ def __init__( stacklevel=2, ) + self._client_options = client_options + self._admin_client_options = admin_client_options self._channel = channel self.SCOPE = self._get_scopes() super(Client, self).__init__(project=project, credentials=credentials) @@ -213,9 +230,10 @@ def table_data_client(self): :returns: A BigtableClient object. """ if self._table_data_client is None: - self._table_data_client = _create_gapic_client(bigtable_v2.BigtableClient)( - self + klass = _create_gapic_client( + bigtable_v2.BigtableClient, client_options=self._client_options ) + self._table_data_client = klass(self) return self._table_data_client @property @@ -237,9 +255,11 @@ def table_admin_client(self): if self._table_admin_client is None: if not self._admin: raise ValueError("Client is not an admin client.") - self._table_admin_client = _create_gapic_client( - bigtable_admin_v2.BigtableTableAdminClient - )(self) + klass = _create_gapic_client( + bigtable_admin_v2.BigtableTableAdminClient, + client_options=self._admin_client_options, + ) + self._table_admin_client = klass(self) return self._table_admin_client @property @@ -261,9 +281,11 @@ def instance_admin_client(self): if self._instance_admin_client is None: if not self._admin: raise ValueError("Client is not an admin client.") - self._instance_admin_client = _create_gapic_client( - bigtable_admin_v2.BigtableInstanceAdminClient - )(self) + klass = _create_gapic_client( + bigtable_admin_v2.BigtableInstanceAdminClient, + client_options=self._admin_client_options, + ) + self._instance_admin_client = klass(self) return self._instance_admin_client def instance(self, instance_id, display_name=None, instance_type=None, labels=None): diff --git a/bigtable/tests/unit/test_client.py b/bigtable/tests/unit/test_client.py index 05a017d898af..8a2ef3c64b56 100644 --- a/bigtable/tests/unit/test_client.py +++ b/bigtable/tests/unit/test_client.py @@ -21,12 +21,12 @@ class Test__create_gapic_client(unittest.TestCase): - def _invoke_client_factory(self, client_class): + def _invoke_client_factory(self, client_class, **kw): from google.cloud.bigtable.client import _create_gapic_client - return _create_gapic_client(client_class) + return _create_gapic_client(client_class, **kw) - def test_without_emulator(self): + def test_wo_emulator(self): client_class = mock.Mock() credentials = _make_credentials() client = _Client(credentials) @@ -36,10 +36,30 @@ def test_without_emulator(self): self.assertIs(result, client_class.return_value) client_class.assert_called_once_with( - credentials=client._credentials, client_info=client_info + credentials=client._credentials, + client_info=client_info, + client_options=None, ) - def test_with_emulator(self): + def test_wo_emulator_w_client_options(self): + client_class = mock.Mock() + credentials = _make_credentials() + client = _Client(credentials) + client_info = client._client_info = mock.Mock() + client_options = mock.Mock() + + result = self._invoke_client_factory( + client_class, client_options=client_options + )(client) + + self.assertIs(result, client_class.return_value) + client_class.assert_called_once_with( + credentials=client._credentials, + client_info=client_info, + client_options=client_options, + ) + + def test_w_emulator(self): client_class = mock.Mock() emulator_host = emulator_channel = object() credentials = _make_credentials() @@ -210,6 +230,25 @@ 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_table_data_client_not_initialized_w_client_options(self): + credentials = _make_credentials() + client_options = mock.Mock() + client = self._make_one( + project=self.PROJECT, credentials=credentials, client_options=client_options + ) + + patch = mock.patch("google.cloud.bigtable_v2.BigtableClient") + with patch as mocked: + table_data_client = client.table_data_client + + self.assertIs(table_data_client, mocked.return_value) + self.assertIs(client._table_data_client, table_data_client) + mocked.assert_called_once_with( + client_info=client._client_info, + credentials=mock.ANY, # added scopes + client_options=client_options, + ) + def test_table_data_client_initialized(self): credentials = _make_credentials() client = self._make_one( @@ -257,6 +296,28 @@ def test_table_admin_client_not_initialized_w_client_info(self): 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_client_options(self): + credentials = _make_credentials() + admin_client_options = mock.Mock() + client = self._make_one( + project=self.PROJECT, + credentials=credentials, + admin=True, + admin_client_options=admin_client_options, + ) + + patch = mock.patch("google.cloud.bigtable_admin_v2.BigtableTableAdminClient") + with patch as mocked: + table_admin_client = client.table_admin_client + + self.assertIs(table_admin_client, mocked.return_value) + self.assertIs(client._table_admin_client, table_admin_client) + mocked.assert_called_once_with( + client_info=client._client_info, + credentials=mock.ANY, # added scopes + client_options=admin_client_options, + ) + def test_table_admin_client_initialized(self): credentials = _make_credentials() client = self._make_one( @@ -287,7 +348,7 @@ 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_and_client_info(self): + def test_instance_admin_client_not_initialized_w_client_info(self): from google.cloud.bigtable_admin_v2 import BigtableInstanceAdminClient credentials = _make_credentials() @@ -304,6 +365,28 @@ def test_instance_admin_client_not_initialized_w_admin_and_client_info(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_client_options(self): + credentials = _make_credentials() + admin_client_options = mock.Mock() + client = self._make_one( + project=self.PROJECT, + credentials=credentials, + admin=True, + admin_client_options=admin_client_options, + ) + + patch = mock.patch("google.cloud.bigtable_admin_v2.BigtableInstanceAdminClient") + with patch as mocked: + instance_admin_client = client.instance_admin_client + + self.assertIs(instance_admin_client, mocked.return_value) + self.assertIs(client._instance_admin_client, instance_admin_client) + mocked.assert_called_once_with( + client_info=client._client_info, + credentials=mock.ANY, # added scopes + client_options=admin_client_options, + ) + def test_instance_admin_client_initialized(self): credentials = _make_credentials() client = self._make_one(