feat(bigtable): initialize internal data client in classic client (#18080)

Migrating over @gkevinzheng PR from bigtable monorepo
googleapis/python-bigtable#1238


Original description:
> Changes:
> - Added a new property ~`new_table_data_client`~ `_veneer_data_client`
in the classic client for the new table data client
> - Added two hidden kwargs (`client_info` and
`disable_background_channel_refresh`) to the new data client constructor
to support the classic client.
>
> Fixes #15298 


Note to reviewers: This PR has already been reviewed and merged to a
staging branch, with the intention of doing a single merge to main. We
are now planning to slowly rollout these changes back to the main
branch. Minimal re-review should be necessary

---------

Co-authored-by: Kevin Zheng <147537668+gkevinzheng@users.noreply.github.com>
Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: 3 people

packages/google-cloud-bigtable/google/cloud/bigtable/client.py

@@ -28,6 +28,7 @@
   :class:`~google.cloud.bigtable.row.Row` (and all the cells in the row)
 """
 
+import copy
 import os
 import warnings
 
@@ -37,7 +38,7 @@
 from google.cloud.client import ClientWithProject  # type: ignore
 from google.cloud.environment_vars import BIGTABLE_EMULATOR  # type: ignore
 
-from google.cloud import bigtable, bigtable_admin_v2, bigtable_v2
+from google.cloud import bigtable, bigtable_admin_v2
 from google.cloud.bigtable.cluster import _CLUSTER_NAME_RE, Cluster
 from google.cloud.bigtable.instance import Instance
 from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin.transports import (
@@ -47,7 +48,6 @@
     BigtableTableAdminGrpcTransport,
 )
 from google.cloud.bigtable_admin_v2.types import instance
-from google.cloud.bigtable_v2.services.bigtable.transports import BigtableGrpcTransport
 
 INSTANCE_TYPE_PRODUCTION = instance.Instance.Type.PRODUCTION
 INSTANCE_TYPE_DEVELOPMENT = instance.Instance.Type.DEVELOPMENT
@@ -60,14 +60,15 @@
 READ_ONLY_SCOPE = "https://www.googleapis.com/auth/bigtable.data.readonly"
 """Scope for reading table data."""
 
-_DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator"
 _GRPC_CHANNEL_OPTIONS = (
     ("grpc.max_send_message_length", -1),
     ("grpc.max_receive_message_length", -1),
     ("grpc.keepalive_time_ms", 30000),
     ("grpc.keepalive_timeout_ms", 10000),
 )
 
+_DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator"
+
 
 def _create_gapic_client(client_class, client_options=None, transport=None):
     def inner(self):
@@ -284,18 +285,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._veneer_data_client._gapic_client
 
     @property
     def table_admin_client(self):
@@ -363,6 +353,23 @@ def instance_admin_client(self):
             self._instance_admin_client = klass(self)
         return self._instance_admin_client
 
+    @property
+    def _veneer_data_client(self):
+        """Getter for the new Data Table API."""
+        if self._table_data_client is None:
+            from google.cloud.bigtable.data import BigtableDataClient
+
+            client_info = copy.copy(self._client_info)
+            client_info.client_library_version = f"{bigtable.__version__}-data-shim"
+            self._table_data_client = BigtableDataClient(
+                project=self.project,
+                credentials=self._credentials,
+                client_options=self._client_options,
+                _client_info=client_info,
+                _disable_background_refresh=True,
+            )
+        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.
 

packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py

@@ -23,7 +23,7 @@
 import google.cloud.bigtable_v2.types.bigtable as types_pb
 from google.cloud.bigtable.data._cross_sync import CrossSync
 from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
-from google.cloud.bigtable.data._metrics import tracked_retry
+from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry
 
 # mutate_rows requests are limited to this number of mutations
 from google.cloud.bigtable.data.mutations import (

packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py

@@ -25,7 +25,7 @@
 from google.cloud.bigtable.data._helpers import (
     _attempt_timeout_generator,
 )
-from google.cloud.bigtable.data._metrics import tracked_retry
+from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry
 from google.cloud.bigtable.data.exceptions import (
     InvalidChunk,
     _ResetRow,

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

@@ -50,10 +50,10 @@
 from google.protobuf.message import Message
 from grpc import Channel
 
-from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT
 from google.cloud.bigtable.data._cross_sync import CrossSync
 from google.cloud.bigtable.data._helpers import (
     _CONCURRENCY_LIMIT,
+    _DEFAULT_BIGTABLE_EMULATOR_CLIENT,
     TABLE_DEFAULT,
     _align_timeouts,
     _attempt_timeout_generator,
@@ -68,8 +68,8 @@
     ActiveOperationMetric,
     BigtableClientSideMetricsController,
     OperationType,
-    tracked_retry,
 )
+from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry
 from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler
 from google.cloud.bigtable.data._metrics.handlers.gcp_exporter import (
     BigtableMetricsExporter,
@@ -218,9 +218,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 = DEFAULT_CLIENT_INFO
-        self.client_info.client_library_version = self._client_version()
+
+        # set up client info headers for veneer library. _client_info is for internal use only,
+        # for the legacy client shim.
+        if kwargs.get("_client_info"):
+            self.client_info = kwargs["_client_info"]
+        else:
+            self.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)
@@ -271,6 +277,10 @@ def __init__(
                 "is the default."
             )
         self._is_closed = CrossSync.Event()
+        # Private argument, for internal use only
+        self._disable_background_refresh = bool(
+            kwargs.get("_disable_background_refresh", False)
+        )
         handlers: list[MetricsHandler] = []
         if self._emulator_host is None:
             try:
@@ -402,6 +412,7 @@ def _start_background_channel_refresh(self) -> None:
             not self._channel_refresh_task
             and not self._emulator_host
             and not self._is_closed.is_set()
+            and not self._disable_background_refresh
         ):
             # raise error if not in an event loop in async client
             CrossSync.verify_async_event_loop()

packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py

@@ -47,6 +47,9 @@
 # used by read_rows_sharded to limit how many requests are attempted in parallel
 _CONCURRENCY_LIMIT = 10
 
+# used by every data client as a default project name for testing on Bigtable emulator.
+_DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator"
+
 # used to identify an active bigtable resource that needs to be warmed through PingAndWarm
 # each instance/app_profile_id pair needs to be individually tracked
 _WarmedInstanceKey = namedtuple(

packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/handlers/opentelemetry.py

@@ -17,7 +17,7 @@
 import socket
 import uuid
 
-from google.cloud.bigtable import __version__ as bigtable_version
+from google.cloud.bigtable.gapic_version import __version__ as bigtable_version
 from google.cloud.bigtable.data._metrics.data_model import (
     DEFAULT_CLUSTER_ID,
     DEFAULT_ZONE,

packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/tracked_retry.py

@@ -29,7 +29,10 @@
 from grpc import StatusCode
 
 from google.cloud.bigtable.data._helpers import _retry_exception_factory
-from google.cloud.bigtable.data._metrics import ActiveOperationMetric, OperationState
+from google.cloud.bigtable.data._metrics.data_model import (
+    ActiveOperationMetric,
+    OperationState,
+)
 from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete
 
 T = TypeVar("T")

packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py

@@ -26,7 +26,7 @@
 import google.cloud.bigtable_v2.types.bigtable as types_pb
 from google.cloud.bigtable.data._cross_sync import CrossSync
 from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
-from google.cloud.bigtable.data._metrics import tracked_retry
+from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry
 from google.cloud.bigtable.data.mutations import (
     _MUTATE_ROWS_REQUEST_MUTATION_LIMIT,
     _EntryWithProto,

packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py

@@ -26,7 +26,7 @@
 
 from google.cloud.bigtable.data._cross_sync import CrossSync
 from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
-from google.cloud.bigtable.data._metrics import tracked_retry
+from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry
 from google.cloud.bigtable.data.exceptions import (
     InvalidChunk,
     _ResetRow,

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

@@ -44,10 +44,10 @@
 from google.protobuf.message import Message
 from grpc import Channel, insecure_channel, intercept_channel
 
-from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT
 from google.cloud.bigtable.data._cross_sync import CrossSync
 from google.cloud.bigtable.data._helpers import (
     _CONCURRENCY_LIMIT,
+    _DEFAULT_BIGTABLE_EMULATOR_CLIENT,
     TABLE_DEFAULT,
     _align_timeouts,
     _attempt_timeout_generator,
@@ -62,8 +62,8 @@
     ActiveOperationMetric,
     BigtableClientSideMetricsController,
     OperationType,
-    tracked_retry,
 )
+from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry
 from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler
 from google.cloud.bigtable.data._metrics.handlers.gcp_exporter import (
     BigtableMetricsExporter,
@@ -158,8 +158,11 @@ def __init__(
         """
         if "pool_size" in kwargs:
             warnings.warn("pool_size no longer supported")
-        self.client_info = DEFAULT_CLIENT_INFO
-        self.client_info.client_library_version = self._client_version()
+        if kwargs.get("_client_info"):
+            self.client_info = kwargs["_client_info"]
+        else:
+            self.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(
@@ -201,6 +204,9 @@ def __init__(
                 f"The configured universe domain ({self.universe_domain}) does not match the universe domain found in the credentials ({self._credentials.universe_domain}). If you haven't configured the universe domain explicitly, `googleapis.com` is the default."
             )
         self._is_closed = CrossSync._Sync_Impl.Event()
+        self._disable_background_refresh = bool(
+            kwargs.get("_disable_background_refresh", False)
+        )
         handlers: list[MetricsHandler] = []
         if self._emulator_host is None:
             try:
@@ -297,6 +303,7 @@ def _start_background_channel_refresh(self) -> None:
             not self._channel_refresh_task
             and (not self._emulator_host)
             and (not self._is_closed.is_set())
+            and (not self._disable_background_refresh)
         ):
             CrossSync._Sync_Impl.verify_async_event_loop()
             self._channel_refresh_task = CrossSync._Sync_Impl.create_task(