From ce7d433335919a8bf57db72e356ca0a6e649234c Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 20 Oct 2026 20:37:40 +0000 Subject: [PATCH] feat(bigquery-storage): send gRPC keepalive pings by default A ReadRows stream can sit idle for a long time, for example while BigQuery finishes a query job before the first row. Without keepalive pings, a NAT or load balancer can drop the idle connection without telling either side. BigQueryReadClient now uses a gRPC channel with keepalive pings when no transport is passed. A transport passed by the caller is used as given. Fixes #18784 --- .../cloud/bigquery_storage_v1/client.py | 31 ++++++++++++ .../tests/unit/test_read_client_v1.py | 47 +++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py index 791be3f2a0dc..36a7ef4a51f3 100644 --- a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py +++ b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py @@ -37,16 +37,47 @@ client_library_version=package_version.__version__ ) +# A ReadRows stream can sit idle for a long time, for example while BigQuery +# finishes a query job before the first row. Without keepalive pings, a NAT or +# load balancer can drop the idle connection without telling either side. +_KEEPALIVE_TIME_MS = 60 * 1000 +_KEEPALIVE_TIMEOUT_MS = 60 * 1000 + + +def _create_keepalive_transport(**kwargs): + """Create the default gRPC transport with a channel that sends keepalive + pings. + + ``BigQueryReadClient`` calls this with the arguments it passes to its + default transport. + """ + transport_class = big_query_read.BigQueryReadClient.get_transport_class("grpc") + + def create_channel(*args, options=(), **channel_kwargs): + options = list(options) + [ + ("grpc.keepalive_time_ms", _KEEPALIVE_TIME_MS), + ("grpc.keepalive_timeout_ms", _KEEPALIVE_TIMEOUT_MS), + ] + return transport_class.create_channel(*args, options=options, **channel_kwargs) + + return transport_class(channel=create_channel, **kwargs) + class BigQueryReadClient(big_query_read.BigQueryReadClient): """Client for interacting with BigQuery Storage API. The BigQuery storage API can be used to read data stored in BigQuery. + + By default, the client's gRPC channel sends keepalive pings, so idle + ``ReadRows`` streams aren't dropped. A ``transport`` passed to the + constructor is used as given. """ def __init__(self, **kwargs): if "client_info" not in kwargs: kwargs["client_info"] = VENEER_CLIENT_INFO + if kwargs.get("transport") in (None, "grpc"): + kwargs["transport"] = _create_keepalive_transport super().__init__(**kwargs) def read_rows( diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py b/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py index bca9647bdffd..65e69a7df96f 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py @@ -222,3 +222,50 @@ def test_init_default_client_info(module_under_test): client_info = call.kwargs.get("client_info") assert client_info is not None assert expected_client_info in client_info.to_user_agent() + + +_KEEPALIVE_OPTIONS = [ + ("grpc.keepalive_time_ms", 60 * 1000), + ("grpc.keepalive_timeout_ms", 60 * 1000), +] + + +@pytest.mark.parametrize("transport", [None, "grpc"]) +def test_constructor_default_transport_sends_keepalive(transport): + from google.cloud import bigquery_storage + + kwargs = {} if transport is None else {"transport": transport} + with mock.patch( + "google.api_core.grpc_helpers.create_channel", autospec=True + ) as create_channel: + client = bigquery_storage.BigQueryReadClient( + credentials=credentials.AnonymousCredentials(), **kwargs + ) + + create_channel.assert_called_once() + options = create_channel.call_args.kwargs["options"] + for option in _KEEPALIVE_OPTIONS: + assert option in options + # The transport's own channel options are kept. + assert ("grpc.max_receive_message_length", -1) in options + assert client._transport.grpc_channel is create_channel.return_value + + +def test_constructor_keeps_given_transport(mock_transport): + from google.cloud import bigquery_storage + + client = bigquery_storage.BigQueryReadClient(transport=mock_transport) + + assert client._transport is mock_transport + + +def test_constructor_keeps_given_transport_callable(mock_transport): + from google.cloud import bigquery_storage + + create_transport = mock.Mock(return_value=mock_transport) + client = bigquery_storage.BigQueryReadClient( + credentials=credentials.AnonymousCredentials(), transport=create_transport + ) + + create_transport.assert_called_once() + assert client._transport is mock_transport