Skip to content
This repository was archived by the owner on Apr 1, 2026. It is now read-only.

Commit 406ae29

Browse files
gkevinzhenggkevinzheng
committed
feat: initializing the new data client
1 parent 263332a commit 406ae29

7 files changed

Lines changed: 278 additions & 10 deletions

File tree

‎google/cloud/bigtable/client.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,11 @@
5454
from google.cloud.bigtable.cluster import _CLUSTER_NAME_RE
5555
from google.cloud.environment_vars import BIGTABLE_EMULATOR # type: ignore
5656

57+
from google.cloud.bigtable.data import BigtableDataClient
58+
from google.cloud.bigtable.data._sync_autogen.client import (
59+
_DEFAULT_BIGTABLE_EMULATOR_CLIENT,
60+
)
61+
5762

5863
INSTANCE_TYPE_PRODUCTION = instance.Instance.Type.PRODUCTION
5964
INSTANCE_TYPE_DEVELOPMENT = instance.Instance.Type.DEVELOPMENT
@@ -66,7 +71,6 @@
6671
READ_ONLY_SCOPE = "https://www.googleapis.com/auth/bigtable.data.readonly"
6772
"""Scope for reading table data."""
6873

69-
_DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator"
7074
_GRPC_CHANNEL_OPTIONS = (
7175
("grpc.max_send_message_length", -1),
7276
("grpc.max_receive_message_length", -1),
@@ -146,6 +150,7 @@ class Client(ClientWithProject):
146150
_table_data_client = None
147151
_table_admin_client = None
148152
_instance_admin_client = None
153+
_new_table_data_client = None
149154

150155
def __init__(
151156
self,
@@ -369,6 +374,23 @@ def instance_admin_client(self):
369374
self._instance_admin_client = klass(self)
370375
return self._instance_admin_client
371376

377+
@property
378+
def new_table_data_client(self):
379+
"""Getter for the new Data Table API.
380+
381+
TODO: Replace table_data_client with this implementation
382+
when shimming legacy client is finished.
383+
"""
384+
if self._new_table_data_client is None:
385+
self._new_table_data_client = BigtableDataClient(
386+
project=self.project,
387+
credentials=self._credentials,
388+
client_options=self._client_options,
389+
client_info=self._client_info,
390+
disable_background_channel_refresh=True,
391+
)
392+
return self._new_table_data_client
393+
372394
def instance(self, instance_id, display_name=None, instance_type=None, labels=None):
373395
"""Factory to create a instance associated with this client.
374396

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,6 @@
6464
import google.auth.credentials
6565
import google.auth._default
6666
from google.api_core import client_options as client_options_lib
67-
from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT
6867
from google.cloud.bigtable.data.row import Row
6968
from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
7069
from google.cloud.bigtable.data.exceptions import FailedQueryShardError
@@ -133,6 +132,7 @@
133132

134133

135134
__CROSS_SYNC_OUTPUT__ = "google.cloud.bigtable.data._sync_autogen.client"
135+
_DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator"
136136

137137

138138
@CrossSync.convert_class(
@@ -185,7 +185,7 @@ def __init__(
185185
if "pool_size" in kwargs:
186186
warnings.warn("pool_size no longer supported")
187187
# set up client info headers for veneer library
188-
self.client_info = DEFAULT_CLIENT_INFO
188+
self.client_info = kwargs.get("client_info") or DEFAULT_CLIENT_INFO
189189
self.client_info.client_library_version = self._client_version()
190190
# parse client options
191191
if type(client_options) is dict:
@@ -236,6 +236,9 @@ def __init__(
236236
"is the default."
237237
)
238238
self._is_closed = CrossSync.Event()
239+
self._disable_background_channel_refresh = bool(
240+
kwargs.get("disable_background_channel_refresh", False)
241+
)
239242
self.transport = cast(TransportType, self._gapic_client.transport)
240243
# keep track of active instances to for warmup on channel refresh
241244
self._active_instances: Set[_WarmedInstanceKey] = set()
@@ -329,6 +332,7 @@ def _start_background_channel_refresh(self) -> None:
329332
not self._channel_refresh_task
330333
and not self._emulator_host
331334
and not self._is_closed.is_set()
335+
and not self._disable_background_channel_refresh
332336
):
333337
# raise error if not in an event loop in async client
334338
CrossSync.verify_async_event_loop()

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,6 @@
5454
import google.auth.credentials
5555
import google.auth._default
5656
from google.api_core import client_options as client_options_lib
57-
from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT
5857
from google.cloud.bigtable.data.row import Row
5958
from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
6059
from google.cloud.bigtable.data.exceptions import FailedQueryShardError
@@ -93,6 +92,7 @@
9392
from google.cloud.bigtable.data.execute_query._sync_autogen.execute_query_iterator import (
9493
ExecuteQueryIterator,
9594
)
95+
_DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator"
9696

9797

9898
@CrossSync._Sync_Impl.add_mapping_decorator("DataClient")
@@ -127,7 +127,7 @@ def __init__(
127127
"""
128128
if "pool_size" in kwargs:
129129
warnings.warn("pool_size no longer supported")
130-
self.client_info = DEFAULT_CLIENT_INFO
130+
self.client_info = kwargs.get("client_info") or DEFAULT_CLIENT_INFO
131131
self.client_info.client_library_version = self._client_version()
132132
if type(client_options) is dict:
133133
client_options = client_options_lib.from_dict(client_options)
@@ -168,6 +168,9 @@ def __init__(
168168
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."
169169
)
170170
self._is_closed = CrossSync._Sync_Impl.Event()
171+
self._disable_background_channel_refresh = bool(
172+
kwargs.get("disable_background_channel_refresh", False)
173+
)
171174
self.transport = cast(TransportType, self._gapic_client.transport)
172175
self._active_instances: Set[_WarmedInstanceKey] = set()
173176
self._instance_owners: dict[_WarmedInstanceKey, Set[int]] = {}
@@ -238,6 +241,7 @@ def _start_background_channel_refresh(self) -> None:
238241
not self._channel_refresh_task
239242
and (not self._emulator_host)
240243
and (not self._is_closed.is_set())
244+
and (not self._disable_background_channel_refresh)
241245
):
242246
CrossSync._Sync_Impl.verify_async_event_loop()
243247
self._channel_refresh_task = CrossSync._Sync_Impl.create_task(

‎tests/unit/data/_async/test_client.py‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,11 @@ async def test_ctor(self):
120120
async def test_ctor_super_inits(self):
121121
from google.cloud.client import ClientWithProject
122122
from google.api_core import client_options as client_options_lib
123+
from google.cloud.bigtable_v2.services.bigtable.transports.base import (
124+
DEFAULT_CLIENT_INFO,
125+
)
126+
127+
import copy
123128

124129
project = "project-id"
125130
credentials = AnonymousCredentials()
@@ -147,13 +152,72 @@ async def test_ctor_super_inits(self):
147152
kwargs = bigtable_client_init.call_args[1]
148153
assert kwargs["credentials"] == credentials
149154
assert kwargs["client_options"] == options_parsed
155+
156+
expected_client_info = copy.copy(DEFAULT_CLIENT_INFO)
157+
expected_client_info.client_library_version = (
158+
CrossSync.DataClient._client_version()
159+
)
160+
assert (
161+
kwargs["client_info"].to_user_agent()
162+
== expected_client_info.to_user_agent()
163+
)
164+
assert (
165+
kwargs["client_info"].to_grpc_metadata()
166+
== expected_client_info.to_grpc_metadata()
167+
)
168+
150169
# test mixin superclass init was called
151170
assert client_project_init.call_count == 1
152171
kwargs = client_project_init.call_args[1]
153172
assert kwargs["project"] == project
154173
assert kwargs["credentials"] == credentials
155174
assert kwargs["client_options"] == options_parsed
156175

176+
@CrossSync.pytest
177+
async def test_ctor_client_info(self):
178+
from google.api_core import client_options as client_options_lib
179+
from google.api_core.gapic_v1.client_info import ClientInfo
180+
181+
import copy
182+
183+
project = "project-id"
184+
credentials = AnonymousCredentials()
185+
client_info = ClientInfo(gapic_version="1.2.3", user_agent="test-client-")
186+
client_options = {"api_endpoint": "foo.bar:1234"}
187+
options_parsed = client_options_lib.from_dict(client_options)
188+
with mock.patch.object(
189+
CrossSync.GapicClient, "__init__"
190+
) as bigtable_client_init:
191+
try:
192+
self._make_client(
193+
project=project,
194+
credentials=credentials,
195+
client_info=client_info,
196+
client_options=options_parsed,
197+
use_emulator=False,
198+
)
199+
except TypeError:
200+
pass
201+
202+
# test gapic superclass init was called with the right arguments
203+
assert bigtable_client_init.call_count == 1
204+
kwargs = bigtable_client_init.call_args[1]
205+
assert kwargs["credentials"] == credentials
206+
assert kwargs["client_options"] == options_parsed
207+
208+
expected_client_info = copy.copy(client_info)
209+
expected_client_info.client_library_version = (
210+
CrossSync.DataClient._client_version()
211+
)
212+
assert (
213+
kwargs["client_info"].to_user_agent()
214+
== expected_client_info.to_user_agent()
215+
)
216+
assert (
217+
kwargs["client_info"].to_grpc_metadata()
218+
== expected_client_info.to_grpc_metadata()
219+
)
220+
157221
@CrossSync.pytest
158222
async def test_ctor_dict_options(self):
159223
from google.api_core.client_options import ClientOptions
@@ -245,6 +309,19 @@ async def test__start_background_channel_refresh(self):
245309
assert ping_and_warm.call_count == 1
246310
await client.close()
247311

312+
@CrossSync.pytest
313+
async def test__start_background_channel_refresh_disable_refresh(self):
314+
client = self._make_client(
315+
project="project-id",
316+
disable_background_channel_refresh=True,
317+
)
318+
# should create background tasks for each channel
319+
with mock.patch.object(client, "_ping_and_warm_instances", CrossSync.Mock()):
320+
client._emulator_host = None
321+
client.transport._grpc_channel = CrossSync.SwappableChannel(mock.Mock)
322+
client._start_background_channel_refresh()
323+
assert client._channel_refresh_task is None
324+
248325
@CrossSync.drop
249326
@CrossSync.pytest
250327
@pytest.mark.skipif(

‎tests/unit/data/_sync_autogen/test_client.py‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,10 @@ def test_ctor(self):
9090
def test_ctor_super_inits(self):
9191
from google.cloud.client import ClientWithProject
9292
from google.api_core import client_options as client_options_lib
93+
from google.cloud.bigtable_v2.services.bigtable.transports.base import (
94+
DEFAULT_CLIENT_INFO,
95+
)
96+
import copy
9397

9498
project = "project-id"
9599
credentials = AnonymousCredentials()
@@ -116,12 +120,64 @@ def test_ctor_super_inits(self):
116120
kwargs = bigtable_client_init.call_args[1]
117121
assert kwargs["credentials"] == credentials
118122
assert kwargs["client_options"] == options_parsed
123+
expected_client_info = copy.copy(DEFAULT_CLIENT_INFO)
124+
expected_client_info.client_library_version = (
125+
CrossSync._Sync_Impl.DataClient._client_version()
126+
)
127+
assert (
128+
kwargs["client_info"].to_user_agent()
129+
== expected_client_info.to_user_agent()
130+
)
131+
assert (
132+
kwargs["client_info"].to_grpc_metadata()
133+
== expected_client_info.to_grpc_metadata()
134+
)
119135
assert client_project_init.call_count == 1
120136
kwargs = client_project_init.call_args[1]
121137
assert kwargs["project"] == project
122138
assert kwargs["credentials"] == credentials
123139
assert kwargs["client_options"] == options_parsed
124140

141+
def test_ctor_client_info(self):
142+
from google.api_core import client_options as client_options_lib
143+
from google.api_core.gapic_v1.client_info import ClientInfo
144+
import copy
145+
146+
project = "project-id"
147+
credentials = AnonymousCredentials()
148+
client_info = ClientInfo(gapic_version="1.2.3", user_agent="test-client-")
149+
client_options = {"api_endpoint": "foo.bar:1234"}
150+
options_parsed = client_options_lib.from_dict(client_options)
151+
with mock.patch.object(
152+
CrossSync._Sync_Impl.GapicClient, "__init__"
153+
) as bigtable_client_init:
154+
try:
155+
self._make_client(
156+
project=project,
157+
credentials=credentials,
158+
client_info=client_info,
159+
client_options=options_parsed,
160+
use_emulator=False,
161+
)
162+
except TypeError:
163+
pass
164+
assert bigtable_client_init.call_count == 1
165+
kwargs = bigtable_client_init.call_args[1]
166+
assert kwargs["credentials"] == credentials
167+
assert kwargs["client_options"] == options_parsed
168+
expected_client_info = copy.copy(client_info)
169+
expected_client_info.client_library_version = (
170+
CrossSync._Sync_Impl.DataClient._client_version()
171+
)
172+
assert (
173+
kwargs["client_info"].to_user_agent()
174+
== expected_client_info.to_user_agent()
175+
)
176+
assert (
177+
kwargs["client_info"].to_grpc_metadata()
178+
== expected_client_info.to_grpc_metadata()
179+
)
180+
125181
def test_ctor_dict_options(self):
126182
from google.api_core.client_options import ClientOptions
127183

@@ -194,6 +250,20 @@ def test__start_background_channel_refresh(self):
194250
assert ping_and_warm.call_count == 1
195251
client.close()
196252

253+
def test__start_background_channel_refresh_disable_refresh(self):
254+
client = self._make_client(
255+
project="project-id", disable_background_channel_refresh=True
256+
)
257+
with mock.patch.object(
258+
client, "_ping_and_warm_instances", CrossSync._Sync_Impl.Mock()
259+
):
260+
client._emulator_host = None
261+
client.transport._grpc_channel = CrossSync._Sync_Impl.SwappableChannel(
262+
mock.Mock
263+
)
264+
client._start_background_channel_refresh()
265+
assert client._channel_refresh_task is None
266+
197267
def test__ping_and_warm_instances(self):
198268
"""test ping and warm with mocked asyncio.gather"""
199269
client_mock = mock.Mock()

‎tests/unit/v2_client/_testing.py‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,11 +27,22 @@ def __init__(self, *results):
2727

2828

2929
def _make_credentials():
30+
from google.cloud.bigtable_v2 import BigtableClient
3031
import google.auth.credentials
3132

32-
class _CredentialsWithScopes(
33-
google.auth.credentials.Credentials, google.auth.credentials.Scoped
33+
class _CredentialsWithScopesAndQuotaProject(
34+
google.auth.credentials.Scoped,
35+
google.auth.credentials.CredentialsWithQuotaProject,
36+
google.auth.credentials.Credentials,
3437
):
3538
pass
3639

37-
return mock.Mock(spec=_CredentialsWithScopes)
40+
credentials = mock.Mock(spec=_CredentialsWithScopesAndQuotaProject)
41+
42+
# Needed to mock universe domain and quota project for new data client tests
43+
credentials_with_scopes = mock.Mock(spec=_CredentialsWithScopesAndQuotaProject)
44+
credentials_with_scopes.universe_domain = BigtableClient._DEFAULT_UNIVERSE
45+
credentials_with_scopes.with_quota_project.return_value = credentials_with_scopes
46+
credentials.with_scopes.return_value = credentials_with_scopes
47+
48+
return credentials

0 commit comments

Comments
 (0)