Visitar URL original
feat(bigquery-storage): send gRPC keepalive pings by default (#18792) · googleapis/google-cloud-python@31d9365 · GitHub
Skip to content

Commit 31d9365

Browse files
ohmayrohmayr
authored
feat(bigquery-storage): send gRPC keepalive pings by default (#18792)
Makes `BigQueryReadClient` send gRPC keepalive pings by default. `ReadRows` streams can sit idle for a long time, for example while BigQuery finishes a query job, and without pings a NAT or load balancer can drop the idle connection without telling either side. google-cloud-bigquery works around this with its own transport for Arrow query results (#18652). This moves the setting here, so every long read gets it. - Uses a gRPC channel with `grpc.keepalive_time_ms` and `grpc.keepalive_timeout_ms` set to 60s when no `transport` is passed (or `transport="grpc"`). - Leaves a caller's transport instance or callable unchanged. - Adds unit tests for the default channel options and for caller-provided transports. This changes the default channel for every bigquery-storage user, but normal table reads shouldn't notice: pings only go out while a call is open and the connection has been quiet for 60s (`keepalive_permit_without_calls` isn't set). A read on a dropped connection now fails with `UNAVAILABLE` after about 2 minutes instead of hanging, and `ReadRowsStream` already reconnects on that. google-cloud-bigtable and google-cloud-firestore already send keepalive pings by default, every 30s. Fixes #18784
1 parent 89cbc50 commit 31d9365

2 files changed

Lines changed: 78 additions & 0 deletions

File tree

‎packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,16 +37,47 @@
3737
client_library_version=package_version.__version__
3838
)
3939

40+
# A ReadRows stream can sit idle for a long time, for example while BigQuery
41+
# finishes a query job before the first row. Without keepalive pings, a NAT or
42+
# load balancer can drop the idle connection without telling either side.
43+
_KEEPALIVE_TIME_MS = 60 * 1000
44+
_KEEPALIVE_TIMEOUT_MS = 60 * 1000
45+
46+
47+
def _create_keepalive_transport(**kwargs):
48+
"""Create the default gRPC transport with a channel that sends keepalive
49+
pings.
50+
51+
``BigQueryReadClient`` calls this with the arguments it passes to its
52+
default transport.
53+
"""
54+
transport_class = big_query_read.BigQueryReadClient.get_transport_class("grpc")
55+
56+
def create_channel(*args, options=(), **channel_kwargs):
57+
options = list(options) + [
58+
("grpc.keepalive_time_ms", _KEEPALIVE_TIME_MS),
59+
("grpc.keepalive_timeout_ms", _KEEPALIVE_TIMEOUT_MS),
60+
]
61+
return transport_class.create_channel(*args, options=options, **channel_kwargs)
62+
63+
return transport_class(channel=create_channel, **kwargs)
64+
4065

4166
class BigQueryReadClient(big_query_read.BigQueryReadClient):
4267
"""Client for interacting with BigQuery Storage API.
4368
4469
The BigQuery storage API can be used to read data stored in BigQuery.
70+
71+
By default, the client's gRPC channel sends keepalive pings, so idle
72+
``ReadRows`` streams aren't dropped. A ``transport`` passed to the
73+
constructor is used as given.
4574
"""
4675

4776
def __init__(self, **kwargs):
4877
if "client_info" not in kwargs:
4978
kwargs["client_info"] = VENEER_CLIENT_INFO
79+
if kwargs.get("transport") in (None, "grpc"):
80+
kwargs["transport"] = _create_keepalive_transport
5081
super().__init__(**kwargs)
5182

5283
def read_rows(

‎packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -222,3 +222,50 @@ def test_init_default_client_info(module_under_test):
222222
client_info = call.kwargs.get("client_info")
223223
assert client_info is not None
224224
assert expected_client_info in client_info.to_user_agent()
225+
226+
227+
_KEEPALIVE_OPTIONS = [
228+
("grpc.keepalive_time_ms", 60 * 1000),
229+
("grpc.keepalive_timeout_ms", 60 * 1000),
230+
]
231+
232+
233+
@pytest.mark.parametrize("transport", [None, "grpc"])
234+
def test_constructor_default_transport_sends_keepalive(transport):
235+
from google.cloud import bigquery_storage
236+
237+
kwargs = {} if transport is None else {"transport": transport}
238+
with mock.patch(
239+
"google.api_core.grpc_helpers.create_channel", autospec=True
240+
) as create_channel:
241+
client = bigquery_storage.BigQueryReadClient(
242+
credentials=credentials.AnonymousCredentials(), **kwargs
243+
)
244+
245+
create_channel.assert_called_once()
246+
options = create_channel.call_args.kwargs["options"]
247+
for option in _KEEPALIVE_OPTIONS:
248+
assert option in options
249+
# The transport's own channel options are kept.
250+
assert ("grpc.max_receive_message_length", -1) in options
251+
assert client._transport.grpc_channel is create_channel.return_value
252+
253+
254+
def test_constructor_keeps_given_transport(mock_transport):
255+
from google.cloud import bigquery_storage
256+
257+
client = bigquery_storage.BigQueryReadClient(transport=mock_transport)
258+
259+
assert client._transport is mock_transport
260+
261+
262+
def test_constructor_keeps_given_transport_callable(mock_transport):
263+
from google.cloud import bigquery_storage
264+
265+
create_transport = mock.Mock(return_value=mock_transport)
266+
client = bigquery_storage.BigQueryReadClient(
267+
credentials=credentials.AnonymousCredentials(), transport=create_transport
268+
)
269+
270+
create_transport.assert_called_once()
271+
assert client._transport is mock_transport

0 commit comments

Comments
 (0)