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

Commit e9637cb

Browse files
HemangChothaniHemangChothani
authored
feat: add keep alive timeout (#182)
* feat: add keep alive timeout * feat: override channel settings
1 parent e28e770 commit e9637cb

2 files changed

Lines changed: 84 additions & 13 deletions

File tree

‎google/cloud/bigtable/client.py‎

Lines changed: 57 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,11 @@
3535

3636
from google.cloud import bigtable_v2
3737
from google.cloud import bigtable_admin_v2
38+
from google.cloud.bigtable_v2.gapic.transports import bigtable_grpc_transport
39+
from google.cloud.bigtable_admin_v2.gapic.transports import (
40+
bigtable_table_admin_grpc_transport,
41+
bigtable_instance_admin_grpc_transport,
42+
)
3843

3944
from google.cloud.bigtable import __version__
4045
from google.cloud.bigtable.instance import Instance
@@ -60,13 +65,14 @@
6065
"""Scope for reading table data."""
6166

6267

63-
def _create_gapic_client(client_class, client_options=None):
68+
def _create_gapic_client(client_class, client_options=None, transport=None):
6469
def inner(self):
6570
if self._emulator_host is None:
6671
return client_class(
67-
credentials=self._credentials,
72+
credentials=None,
6873
client_info=self._client_info,
6974
client_options=client_options,
75+
transport=transport,
7076
)
7177
else:
7278
return client_class(
@@ -161,7 +167,13 @@ def __init__(
161167
self._emulator_channel = None
162168

163169
if self._emulator_host is not None:
164-
self._emulator_channel = grpc.insecure_channel(self._emulator_host)
170+
self._emulator_channel = grpc.insecure_channel(
171+
target=self._emulator_host,
172+
options={
173+
"grpc.keepalive_time_ms": 30000,
174+
"grpc.keepalive_timeout_ms": 10000,
175+
}.items(),
176+
)
165177

166178
if channel is not None:
167179
warnings.warn(
@@ -196,6 +208,29 @@ def _get_scopes(self):
196208

197209
return scopes
198210

211+
def _create_gapic_client_channel(self, client_class, grpc_transport):
212+
if self._client_options and self._client_options.api_endpoint:
213+
api_endpoint = self._client_options.api_endpoint
214+
else:
215+
api_endpoint = client_class.SERVICE_ADDRESS
216+
217+
channel = grpc_transport.create_channel(
218+
api_endpoint,
219+
self._credentials,
220+
options={
221+
"grpc.max_send_message_length": -1,
222+
"grpc.max_receive_message_length": -1,
223+
"grpc.keepalive_time_ms": 30000,
224+
"grpc.keepalive_timeout_ms": 10000,
225+
}.items(),
226+
)
227+
transport = grpc_transport(
228+
address=api_endpoint,
229+
channel=channel,
230+
credentials=None,
231+
)
232+
return transport
233+
199234
@property
200235
def project_path(self):
201236
"""Project name to be used with Instance Admin API.
@@ -236,8 +271,14 @@ def table_data_client(self):
236271
:returns: A BigtableClient object.
237272
"""
238273
if self._table_data_client is None:
274+
transport = self._create_gapic_client_channel(
275+
bigtable_v2.BigtableClient,
276+
bigtable_grpc_transport.BigtableGrpcTransport,
277+
)
239278
klass = _create_gapic_client(
240-
bigtable_v2.BigtableClient, client_options=self._client_options
279+
bigtable_v2.BigtableClient,
280+
client_options=self._client_options,
281+
transport=transport,
241282
)
242283
self._table_data_client = klass(self)
243284
return self._table_data_client
@@ -262,9 +303,15 @@ def table_admin_client(self):
262303
if self._table_admin_client is None:
263304
if not self._admin:
264305
raise ValueError("Client is not an admin client.")
306+
307+
transport = self._create_gapic_client_channel(
308+
bigtable_admin_v2.BigtableTableAdminClient,
309+
bigtable_table_admin_grpc_transport.BigtableTableAdminGrpcTransport,
310+
)
265311
klass = _create_gapic_client(
266312
bigtable_admin_v2.BigtableTableAdminClient,
267313
client_options=self._admin_client_options,
314+
transport=transport,
268315
)
269316
self._table_admin_client = klass(self)
270317
return self._table_admin_client
@@ -289,9 +336,15 @@ def instance_admin_client(self):
289336
if self._instance_admin_client is None:
290337
if not self._admin:
291338
raise ValueError("Client is not an admin client.")
339+
340+
transport = self._create_gapic_client_channel(
341+
bigtable_admin_v2.BigtableInstanceAdminClient,
342+
bigtable_instance_admin_grpc_transport.BigtableInstanceAdminGrpcTransport,
343+
)
292344
klass = _create_gapic_client(
293345
bigtable_admin_v2.BigtableInstanceAdminClient,
294346
client_options=self._admin_client_options,
347+
transport=transport,
295348
)
296349
self._instance_admin_client = klass(self)
297350
return self._instance_admin_client

‎tests/unit/test_client.py‎

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -31,14 +31,16 @@ def test_wo_emulator(self):
3131
credentials = _make_credentials()
3232
client = _Client(credentials)
3333
client_info = client._client_info = mock.Mock()
34+
transport = mock.Mock()
3435

35-
result = self._invoke_client_factory(client_class)(client)
36+
result = self._invoke_client_factory(client_class, transport=transport)(client)
3637

3738
self.assertIs(result, client_class.return_value)
3839
client_class.assert_called_once_with(
39-
credentials=client._credentials,
40+
credentials=None,
4041
client_info=client_info,
4142
client_options=None,
43+
transport=transport,
4244
)
4345

4446
def test_wo_emulator_w_client_options(self):
@@ -47,16 +49,18 @@ def test_wo_emulator_w_client_options(self):
4749
client = _Client(credentials)
4850
client_info = client._client_info = mock.Mock()
4951
client_options = mock.Mock()
52+
transport = mock.Mock()
5053

5154
result = self._invoke_client_factory(
52-
client_class, client_options=client_options
55+
client_class, client_options=client_options, transport=transport
5356
)(client)
5457

5558
self.assertIs(result, client_class.return_value)
5659
client_class.assert_called_once_with(
57-
credentials=client._credentials,
60+
credentials=None,
5861
client_info=client_info,
5962
client_options=client_options,
63+
transport=transport,
6064
)
6165

6266
def test_w_emulator(self):
@@ -170,7 +174,13 @@ def test_constructor_with_emulator_host(self):
170174

171175
self.assertEqual(client._emulator_host, emulator_host)
172176
self.assertIs(client._emulator_channel, factory.return_value)
173-
factory.assert_called_once_with(emulator_host)
177+
factory.assert_called_once_with(
178+
target=emulator_host,
179+
options={
180+
"grpc.keepalive_time_ms": 30000,
181+
"grpc.keepalive_timeout_ms": 10000,
182+
}.items(),
183+
)
174184
getenv.assert_called_once_with(BIGTABLE_EMULATOR)
175185

176186
def test__get_scopes_default(self):
@@ -234,7 +244,9 @@ def test_table_data_client_not_initialized_w_client_options(self):
234244
from google.api_core.client_options import ClientOptions
235245

236246
credentials = _make_credentials()
237-
client_options = ClientOptions(quota_project_id="QUOTA-PROJECT")
247+
client_options = ClientOptions(
248+
quota_project_id="QUOTA-PROJECT", api_endpoint="xyz"
249+
)
238250
client = self._make_one(
239251
project=self.PROJECT, credentials=credentials, client_options=client_options
240252
)
@@ -245,9 +257,11 @@ def test_table_data_client_not_initialized_w_client_options(self):
245257

246258
self.assertIs(table_data_client, mocked.return_value)
247259
self.assertIs(client._table_data_client, table_data_client)
260+
248261
mocked.assert_called_once_with(
249262
client_info=client._client_info,
250-
credentials=mock.ANY, # added scopes
263+
credentials=None,
264+
transport=mock.ANY,
251265
client_options=client_options,
252266
)
253267

@@ -308,6 +322,7 @@ def test_table_admin_client_not_initialized_w_client_options(self):
308322
admin_client_options=admin_client_options,
309323
)
310324

325+
client._create_gapic_client_channel = mock.Mock()
311326
patch = mock.patch("google.cloud.bigtable_admin_v2.BigtableTableAdminClient")
312327
with patch as mocked:
313328
table_admin_client = client.table_admin_client
@@ -316,7 +331,8 @@ def test_table_admin_client_not_initialized_w_client_options(self):
316331
self.assertIs(client._table_admin_client, table_admin_client)
317332
mocked.assert_called_once_with(
318333
client_info=client._client_info,
319-
credentials=mock.ANY, # added scopes
334+
credentials=None,
335+
transport=mock.ANY,
320336
client_options=admin_client_options,
321337
)
322338

@@ -377,6 +393,7 @@ def test_instance_admin_client_not_initialized_w_client_options(self):
377393
admin_client_options=admin_client_options,
378394
)
379395

396+
client._create_gapic_client_channel = mock.Mock()
380397
patch = mock.patch("google.cloud.bigtable_admin_v2.BigtableInstanceAdminClient")
381398
with patch as mocked:
382399
instance_admin_client = client.instance_admin_client
@@ -385,7 +402,8 @@ def test_instance_admin_client_not_initialized_w_client_options(self):
385402
self.assertIs(client._instance_admin_client, instance_admin_client)
386403
mocked.assert_called_once_with(
387404
client_info=client._client_info,
388-
credentials=mock.ANY, # added scopes
405+
credentials=None,
406+
transport=mock.ANY,
389407
client_options=admin_client_options,
390408
)
391409

0 commit comments

Comments
 (0)