Addressed code review feedback + test fixes

Author: gkevinzheng

google/cloud/bigtable/client.py

@@ -34,9 +34,7 @@
 from google.api_core.gapic_v1 import client_info as client_info_lib
 from google.auth.credentials import AnonymousCredentials  # type: ignore
 
-from google.cloud import bigtable_v2
 from google.cloud import bigtable_admin_v2
-from google.cloud.bigtable_v2.services.bigtable.transports import BigtableGrpcTransport
 from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin.transports import (
     BigtableInstanceAdminGrpcTransport,
 )
@@ -150,7 +148,6 @@ class Client(ClientWithProject):
     _table_data_client = None
     _table_admin_client = None
     _instance_admin_client = None
-    _new_table_data_client = None
 
     def __init__(
         self,
@@ -295,18 +292,7 @@ def table_data_client(self):
         :rtype: :class:`.bigtable_v2.BigtableClient`
         :returns: A BigtableClient object.
         """
-        if self._table_data_client is None:
-            transport = self._create_gapic_client_channel(
-                bigtable_v2.BigtableClient,
-                BigtableGrpcTransport,
-            )
-            klass = _create_gapic_client(
-                bigtable_v2.BigtableClient,
-                client_options=self._client_options,
-                transport=transport,
-            )
-            self._table_data_client = klass(self)
-        return self._table_data_client
+        return self._data_client._gapic_client
 
     @property
     def table_admin_client(self):
@@ -375,21 +361,17 @@ def instance_admin_client(self):
         return self._instance_admin_client
 
     @property
-    def new_table_data_client(self):
-        """Getter for the new Data Table API.
-
-        TODO: Replace table_data_client with this implementation
-        when shimming legacy client is finished.
-        """
-        if self._new_table_data_client is None:
-            self._new_table_data_client = BigtableDataClient(
+    def _data_client(self):
+        """Getter for the new Data Table API."""
+        if self._table_data_client is None:
+            self._table_data_client = BigtableDataClient(
                 project=self.project,
                 credentials=self._credentials,
                 client_options=self._client_options,
-                client_info=self._client_info,
-                disable_background_channel_refresh=True,
+                _client_info=self._client_info,
+                _disable_background_channel_refresh=True,
             )
-        return self._new_table_data_client
+        return self._table_data_client
 
     def instance(self, instance_id, display_name=None, instance_type=None, labels=None):
         """Factory to create a instance associated with this client.

google/cloud/bigtable/data/_async/client.py

@@ -184,9 +184,15 @@ def __init__(
         """
         if "pool_size" in kwargs:
             warnings.warn("pool_size no longer supported")
-        # set up client info headers for veneer library
-        self.client_info = kwargs.get("client_info") or DEFAULT_CLIENT_INFO
-        self.client_info.client_library_version = self._client_version()
+        if "_client_info" in kwargs:
+            # use client_info passed in from legacy client. For internal use only, for the legacy
+            # client shim.
+            self.client_info = kwargs["_client_info"]
+        else:
+            # set up client info headers for veneer library.
+            self.client_info = kwargs.get("_client_info", DEFAULT_CLIENT_INFO)
+            self.client_info.client_library_version = self._client_version()
+
         # parse client options
         if type(client_options) is dict:
             client_options = client_options_lib.from_dict(client_options)
@@ -237,7 +243,7 @@ def __init__(
             )
         self._is_closed = CrossSync.Event()
         self._disable_background_channel_refresh = bool(
-            kwargs.get("disable_background_channel_refresh", False)
+            kwargs.get("_disable_background_channel_refresh", False)
         )
         self.transport = cast(TransportType, self._gapic_client.transport)
         # keep track of active instances to for warmup on channel refresh

google/cloud/bigtable/data/_sync_autogen/client.py

@@ -127,8 +127,11 @@ def __init__(
         """
         if "pool_size" in kwargs:
             warnings.warn("pool_size no longer supported")
-        self.client_info = kwargs.get("client_info") or DEFAULT_CLIENT_INFO
-        self.client_info.client_library_version = self._client_version()
+        if "_client_info" in kwargs:
+            self.client_info = kwargs["_client_info"]
+        else:
+            self.client_info = kwargs.get("_client_info", DEFAULT_CLIENT_INFO)
+            self.client_info.client_library_version = self._client_version()
         if type(client_options) is dict:
             client_options = client_options_lib.from_dict(client_options)
         client_options = cast(
@@ -169,7 +172,7 @@ def __init__(
             )
         self._is_closed = CrossSync._Sync_Impl.Event()
         self._disable_background_channel_refresh = bool(
-            kwargs.get("disable_background_channel_refresh", False)
+            kwargs.get("_disable_background_channel_refresh", False)
         )
         self.transport = cast(TransportType, self._gapic_client.transport)
         self._active_instances: Set[_WarmedInstanceKey] = set()

google/cloud/bigtable/table.py

@@ -132,6 +132,15 @@ def __init__(self, table_id, instance, mutation_timeout=None, app_profile_id=Non
         self._app_profile_id = app_profile_id
         self.mutation_timeout = mutation_timeout
 
+        self._table_impl = (
+            self._instance._client._data_client.get_table(
+                self._instance.instance_id,
+                self.table_id,
+            )
+            if self._instance
+            else None
+        )
+
     @property
     def name(self):
         """Table name used in requests.

tests/unit/data/_async/test_client.py

@@ -178,8 +178,6 @@ async def test_ctor_client_info(self):
         from google.api_core import client_options as client_options_lib
         from google.api_core.gapic_v1.client_info import ClientInfo
 
-        import copy
-
         project = "project-id"
         credentials = AnonymousCredentials()
         client_info = ClientInfo(gapic_version="1.2.3", user_agent="test-client-")
@@ -192,9 +190,9 @@ async def test_ctor_client_info(self):
                 self._make_client(
                     project=project,
                     credentials=credentials,
-                    client_info=client_info,
                     client_options=options_parsed,
                     use_emulator=False,
+                    _client_info=client_info,
                 )
             except TypeError:
                 pass
@@ -205,18 +203,7 @@ async def test_ctor_client_info(self):
             assert kwargs["credentials"] == credentials
             assert kwargs["client_options"] == options_parsed
 
-            expected_client_info = copy.copy(client_info)
-            expected_client_info.client_library_version = (
-                CrossSync.DataClient._client_version()
-            )
-            assert (
-                kwargs["client_info"].to_user_agent()
-                == expected_client_info.to_user_agent()
-            )
-            assert (
-                kwargs["client_info"].to_grpc_metadata()
-                == expected_client_info.to_grpc_metadata()
-            )
+        kwargs["client_info"] == client_info
 
     @CrossSync.pytest
     async def test_ctor_dict_options(self):
@@ -313,14 +300,17 @@ async def test__start_background_channel_refresh(self):
     async def test__start_background_channel_refresh_disable_refresh(self):
         client = self._make_client(
             project="project-id",
-            disable_background_channel_refresh=True,
+            _disable_background_channel_refresh=True,
         )
         # should create background tasks for each channel
-        with mock.patch.object(client, "_ping_and_warm_instances", CrossSync.Mock()):
+        with mock.patch.object(
+            client, "_ping_and_warm_instances", CrossSync.Mock()
+        ) as ping_and_warm:
             client._emulator_host = None
             client.transport._grpc_channel = CrossSync.SwappableChannel(mock.Mock)
             client._start_background_channel_refresh()
             assert client._channel_refresh_task is None
+            ping_and_warm.assert_not_called()
 
     @CrossSync.drop
     @CrossSync.pytest

tests/unit/data/_sync_autogen/test_client.py

@@ -141,7 +141,6 @@ def test_ctor_super_inits(self):
     def test_ctor_client_info(self):
         from google.api_core import client_options as client_options_lib
         from google.api_core.gapic_v1.client_info import ClientInfo
-        import copy
 
         project = "project-id"
         credentials = AnonymousCredentials()
@@ -155,28 +154,17 @@ def test_ctor_client_info(self):
                 self._make_client(
                     project=project,
                     credentials=credentials,
-                    client_info=client_info,
                     client_options=options_parsed,
                     use_emulator=False,
+                    _client_info=client_info,
                 )
             except TypeError:
                 pass
             assert bigtable_client_init.call_count == 1
             kwargs = bigtable_client_init.call_args[1]
             assert kwargs["credentials"] == credentials
             assert kwargs["client_options"] == options_parsed
-            expected_client_info = copy.copy(client_info)
-            expected_client_info.client_library_version = (
-                CrossSync._Sync_Impl.DataClient._client_version()
-            )
-            assert (
-                kwargs["client_info"].to_user_agent()
-                == expected_client_info.to_user_agent()
-            )
-            assert (
-                kwargs["client_info"].to_grpc_metadata()
-                == expected_client_info.to_grpc_metadata()
-            )
+        kwargs["client_info"] == client_info
 
     def test_ctor_dict_options(self):
         from google.api_core.client_options import ClientOptions
@@ -252,17 +240,18 @@ def test__start_background_channel_refresh(self):
 
     def test__start_background_channel_refresh_disable_refresh(self):
         client = self._make_client(
-            project="project-id", disable_background_channel_refresh=True
+            project="project-id", _disable_background_channel_refresh=True
         )
         with mock.patch.object(
             client, "_ping_and_warm_instances", CrossSync._Sync_Impl.Mock()
-        ):
+        ) as ping_and_warm:
             client._emulator_host = None
             client.transport._grpc_channel = CrossSync._Sync_Impl.SwappableChannel(
                 mock.Mock
             )
             client._start_background_channel_refresh()
             assert client._channel_refresh_task is None
+            ping_and_warm.assert_not_called()
 
     def test__ping_and_warm_instances(self):
         """test ping and warm with mocked asyncio.gather"""