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

Commit f44b36b

Browse files
daniel-sanchedaniel-sanche
authored
fix: grpc channel refresh (#1087)
* added failing test * prevent _start_bg_task from running * let bg task run * invalidate transport stubs after channel refresh * added sync implementation
1 parent 89b8da8 commit f44b36b

4 files changed

Lines changed: 70 additions & 6 deletions

File tree

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -164,8 +164,8 @@ def __init__(
164164
if "pool_size" in kwargs:
165165
warnings.warn("pool_size no longer supported")
166166
# set up client info headers for veneer library
167-
client_info = DEFAULT_CLIENT_INFO
168-
client_info.client_library_version = self._client_version()
167+
self.client_info = DEFAULT_CLIENT_INFO
168+
self.client_info.client_library_version = self._client_version()
169169
# parse client options
170170
if type(client_options) is dict:
171171
client_options = client_options_lib.from_dict(client_options)
@@ -196,7 +196,7 @@ def __init__(
196196
self._gapic_client = CrossSync.GapicClient(
197197
credentials=credentials,
198198
client_options=client_options,
199-
client_info=client_info,
199+
client_info=self.client_info,
200200
transport=lambda *args, **kwargs: TransportType(
201201
*args, **kwargs, channel=custom_channel
202202
),
@@ -371,6 +371,9 @@ async def _manage_channel(
371371
await self._ping_and_warm_instances(channel=new_channel)
372372
# cycle channel out of use, with long grace window before closure
373373
self.transport._grpc_channel = new_channel
374+
# invalidate caches
375+
self.transport._stubs = {}
376+
self.transport._prep_wrapped_messages(self.client_info)
374377
# give old_channel a chance to complete existing rpcs
375378
if CrossSync.is_async:
376379
await old_channel.close(grace_period)

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

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -114,8 +114,8 @@ def __init__(
114114
"""
115115
if "pool_size" in kwargs:
116116
warnings.warn("pool_size no longer supported")
117-
client_info = DEFAULT_CLIENT_INFO
118-
client_info.client_library_version = self._client_version()
117+
self.client_info = DEFAULT_CLIENT_INFO
118+
self.client_info.client_library_version = self._client_version()
119119
if type(client_options) is dict:
120120
client_options = client_options_lib.from_dict(client_options)
121121
client_options = cast(
@@ -143,7 +143,7 @@ def __init__(
143143
self._gapic_client = CrossSync._Sync_Impl.GapicClient(
144144
credentials=credentials,
145145
client_options=client_options,
146-
client_info=client_info,
146+
client_info=self.client_info,
147147
transport=lambda *args, **kwargs: TransportType(
148148
*args, **kwargs, channel=custom_channel
149149
),
@@ -284,6 +284,8 @@ def _manage_channel(
284284
new_channel = self.transport.create_channel()
285285
self._ping_and_warm_instances(channel=new_channel)
286286
self.transport._grpc_channel = new_channel
287+
self.transport._stubs = {}
288+
self.transport._prep_wrapped_messages(self.client_info)
287289
if grace_period:
288290
self._is_closed.wait(grace_period)
289291
old_channel.close()

‎tests/system/data/test_system_async.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,38 @@ async def test_ping_and_warm(self, client, table):
207207
assert len(results) == 1
208208
assert results[0] is None
209209

210+
@CrossSync.pytest
211+
async def test_channel_refresh(self, table_id, instance_id, temp_rows):
212+
"""
213+
change grpc channel to refresh after 1 second. Schedule a read_rows call after refresh,
214+
to ensure new channel works
215+
"""
216+
await temp_rows.add_row(b"row_key_1")
217+
await temp_rows.add_row(b"row_key_2")
218+
project = os.getenv("GOOGLE_CLOUD_PROJECT") or None
219+
client = CrossSync.DataClient(project=project)
220+
# start custom refresh task
221+
try:
222+
client._channel_refresh_task = CrossSync.create_task(
223+
client._manage_channel,
224+
refresh_interval_min=1,
225+
refresh_interval_max=1,
226+
sync_executor=client._executor,
227+
)
228+
# let task run
229+
await CrossSync.yield_to_event_loop()
230+
async with client.get_table(instance_id, table_id) as table:
231+
rows = await table.read_rows({})
232+
first_channel = client.transport.grpc_channel
233+
assert len(rows) == 2
234+
await CrossSync.sleep(2)
235+
rows_after_refresh = await table.read_rows({})
236+
assert len(rows_after_refresh) == 2
237+
assert client.transport.grpc_channel is not first_channel
238+
print(table)
239+
finally:
240+
await client.close()
241+
210242
@CrossSync.pytest
211243
@pytest.mark.usefixtures("table")
212244
@CrossSync.Retry(

‎tests/system/data/test_system_autogen.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,33 @@ def test_ping_and_warm(self, client, table):
165165
assert len(results) == 1
166166
assert results[0] is None
167167

168+
def test_channel_refresh(self, table_id, instance_id, temp_rows):
169+
"""change grpc channel to refresh after 1 second. Schedule a read_rows call after refresh,
170+
to ensure new channel works"""
171+
temp_rows.add_row(b"row_key_1")
172+
temp_rows.add_row(b"row_key_2")
173+
project = os.getenv("GOOGLE_CLOUD_PROJECT") or None
174+
client = CrossSync._Sync_Impl.DataClient(project=project)
175+
try:
176+
client._channel_refresh_task = CrossSync._Sync_Impl.create_task(
177+
client._manage_channel,
178+
refresh_interval_min=1,
179+
refresh_interval_max=1,
180+
sync_executor=client._executor,
181+
)
182+
CrossSync._Sync_Impl.yield_to_event_loop()
183+
with client.get_table(instance_id, table_id) as table:
184+
rows = table.read_rows({})
185+
first_channel = client.transport.grpc_channel
186+
assert len(rows) == 2
187+
CrossSync._Sync_Impl.sleep(2)
188+
rows_after_refresh = table.read_rows({})
189+
assert len(rows_after_refresh) == 2
190+
assert client.transport.grpc_channel is not first_channel
191+
print(table)
192+
finally:
193+
client.close()
194+
168195
@pytest.mark.usefixtures("table")
169196
@CrossSync._Sync_Impl.Retry(
170197
predicate=retry.if_exception_type(ClientError), initial=1, maximum=5

0 commit comments

Comments
 (0)