Add Storage API streaming max retries parameter for BigQueryOptions (#…

…31683)
apache · Jun 26, 2024 · b482337 · b482337
1 parent 90d3f8a
commit b482337
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 7 deletions.
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
@@ -133,6 +133,14 @@ public interface BigQueryOptions
 
   void setStorageWriteApiTriggeringFrequencySec(Integer value);
 
+  @Description(
+      "Maximum number of retries for Storage Write API writes. "
+          + "Currently it is only applicable for streaming pipeline.")
+  @Default.Integer(500)
+  Integer getStorageWriteApiMaxRetries();
+
+  void setStorageWriteApiMaxRetries(Integer value);
+
   @Description(
       "When auto-sharding is used, the maximum duration in milliseconds the input records are"
           + " allowed to be buffered before being written to BigQuery.")

diff --git a/...src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java b/...src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java
@@ -108,7 +108,8 @@ public PCollectionTuple expand(PCollection<KV<DestinationT, StorageApiWritePaylo
                         createDisposition,
                         kmsKey,
                         usesCdc,
-                        defaultMissingValueInterpretation))
+                        defaultMissingValueInterpretation,
+                        bigQueryOptions.getStorageWriteApiMaxRetries()))
                 .withOutputTags(finalizeTag, tupleTagList)
                 .withSideInputs(dynamicDestinations.getSideInputs()));
     result.get(failedRowsTag).setCoder(failedRowsCoder);

diff --git a/...rm/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java b/...rm/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java
@@ -217,7 +217,8 @@ public PCollectionTuple expand(PCollection<KV<DestinationT, StorageApiWritePaylo
                         createDisposition,
                         kmsKey,
                         usesCdc,
-                        defaultMissingValueInterpretation))
+                        defaultMissingValueInterpretation,
+                        options.getStorageWriteApiMaxRetries()))
                 .withOutputTags(finalizeTag, tupleTagList)
                 .withSideInputs(dynamicDestinations.getSideInputs()));
 
@@ -889,6 +890,7 @@ void postFlush() {
     private int numPendingRecordBytes = 0;
     private final int flushThresholdBytes;
     private final int flushThresholdCount;
+    private final int maxRetries;
     private final StorageApiDynamicDestinations<ElementT, DestinationT> dynamicDestinations;
     private final BigQueryServices bqServices;
     private final boolean useDefaultStream;
@@ -910,7 +912,8 @@ void postFlush() {
         BigQueryIO.Write.CreateDisposition createDisposition,
         @Nullable String kmsKey,
         boolean usesCdc,
-        AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation) {
+        AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation,
+        int maxRetries) {
       this.messageConverters = new TwoLevelMessageConverterCache<>(operationName);
       this.dynamicDestinations = dynamicDestinations;
       this.bqServices = bqServices;
@@ -927,6 +930,7 @@ void postFlush() {
       this.kmsKey = kmsKey;
       this.usesCdc = usesCdc;
       this.defaultMissingValueInterpretation = defaultMissingValueInterpretation;
+      this.maxRetries = maxRetries;
     }
 
     boolean shouldFlush() {
@@ -960,7 +964,7 @@ void flushAll(
             new RetryManager<>(
                 Duration.standardSeconds(1),
                 Duration.standardSeconds(20),
-                500,
+                maxRetries,
                 BigQuerySinkMetrics.throttledTimeCounter(
                     BigQuerySinkMetrics.RpcMethod.APPEND_ROWS));
         retryManagers.add(retryManager);

diff --git a/...orm/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java b/...orm/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java
@@ -242,6 +242,7 @@ public PCollectionTuple expand(
     BigQueryOptions bigQueryOptions = input.getPipeline().getOptions().as(BigQueryOptions.class);
     final long splitSize = bigQueryOptions.getStorageApiAppendThresholdBytes();
     final long maxRequestSize = bigQueryOptions.getStorageWriteApiMaxRequestSize();
+    final int maxRetries = bigQueryOptions.getStorageWriteApiMaxRetries();
 
     String operationName = input.getName() + "/" + getName();
     TupleTagList tupleTagList = TupleTagList.of(failedRowsTag);
@@ -252,7 +253,9 @@ public PCollectionTuple expand(
     PCollectionTuple writeRecordsResult =
         input.apply(
             "Write Records",
-            ParDo.of(new WriteRecordsDoFn(operationName, streamIdleTime, splitSize, maxRequestSize))
+            ParDo.of(
+                    new WriteRecordsDoFn(
+                        operationName, streamIdleTime, splitSize, maxRequestSize, maxRetries))
                 .withSideInputs(dynamicDestinations.getSideInputs())
                 .withOutputTags(flushTag, tupleTagList));
 
@@ -340,13 +343,19 @@ class WriteRecordsDoFn
     private final Duration streamIdleTime;
     private final long splitSize;
     private final long maxRequestSize;
+    private final int maxRetries;
 
     public WriteRecordsDoFn(
-        String operationName, Duration streamIdleTime, long splitSize, long maxRequestSize) {
+        String operationName,
+        Duration streamIdleTime,
+        long splitSize,
+        long maxRequestSize,
+        int maxRetries) {
       this.messageConverters = new TwoLevelMessageConverterCache<>(operationName);
       this.streamIdleTime = streamIdleTime;
       this.splitSize = splitSize;
       this.maxRequestSize = maxRequestSize;
+      this.maxRetries = maxRetries;
     }
 
     @StartBundle
@@ -819,7 +828,7 @@ public void process(
           new RetryManager<>(
               Duration.standardSeconds(1),
               Duration.standardSeconds(20),
-              500,
+              maxRetries,
               BigQuerySinkMetrics.throttledTimeCounter(BigQuerySinkMetrics.RpcMethod.APPEND_ROWS));
       int numAppends = 0;
       for (SplittingIterable.Value splitValue : messages) {