From c0d566a780d198acc2fe197fdbe80f2c0811be65 Mon Sep 17 00:00:00 2001 From: Jonathan Lemos Date: Thu, 8 Aug 2024 15:39:31 -0400 Subject: [PATCH] Add support for setting an HTTP read timeout for BigQueryIO Shamelessly stolen from https://github.com/apache/beam/pull/7097 --- .../apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java | 9 ++++++++- .../beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java | 1 + 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java index ba76f483f774..faa252e79b2d 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java +++ b/sdks/java/io/google-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/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java index b87b6a222a4d..115875c59411 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java +++ b/sdks/java/io/google-cloud-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 initBuilder = ImmutableList.builder(); Credentials credential = options.getGcpCredential();