Add support for setting an HTTP read timeout for BigQueryIO (#32118)

Shamelessly stolen from #7097
apache · Aug 10, 2024 · 4889969 · 4889969
1 parent fc5a71d
commit 4889969
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/...gle-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java b/...gle-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java
@@ -39,7 +39,14 @@ public interface BigQueryOptions
   void setTempDatasetId(String value);
 
   @Description(
-      "Timeout for HTTP requests to BigQuery service in milliseconds. Set to 0 to disable.")
+      "Timeout for HTTP read requests to BigQuery service in milliseconds. Set to 0 to disable.")
+  @Default.Integer(80 * 1000)
+  Integer getHTTPReadTimeout();
+
+  void setHTTPReadTimeout(Integer timeout);
+
+  @Description(
+      "Timeout for HTTP write requests to BigQuery service in milliseconds. Set to 0 to disable.")
   @Default.Integer(900 * 1000)
   Integer getHTTPWriteTimeout();
 

diff --git a/...loud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java b/...loud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java
@@ -1584,6 +1584,7 @@ private static Bigquery.Builder newBigQueryClient(BigQueryOptions options) {
     RetryHttpRequestInitializer httpRequestInitializer =
         new RetryHttpRequestInitializer(ImmutableList.of(404));
     httpRequestInitializer.setCustomErrors(createBigQueryClientCustomErrors());
+    httpRequestInitializer.setReadTimeout(options.getHTTPReadTimeout());
     httpRequestInitializer.setWriteTimeout(options.getHTTPWriteTimeout());
     ImmutableList.Builder<HttpRequestInitializer> initBuilder = ImmutableList.builder();
     Credentials credential = options.getGcpCredential();