feat(bigtable): add 'client_options' / 'admin_client_options' to Client (#9517)

Toward #8475.

Author: tseaver

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):

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(