feat(bigquery-storage): send gRPC keepalive pings by default

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(

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