[Java] Honor withMaxRetryJobs for bounded BigQueryIO file loads

CHANGES.md

@@ -94,6 +94,7 @@
 * (Java) KafkaIO dynamic reads no longer require the obsolete `beam_fn_api` experiment ([#29998](https://github.com/apache/beam/issues/29998)).
 * (Prism) Self-checkpointing splittable DoFns now resume after their requested delay instead of immediately, so polling SDFs no longer busy-spin ([#39848](https://github.com/apache/beam/issues/39848)).
 * (Java) MongoDbIO read splitting now preserves non-ObjectId `_id` types (e.g. string ids) instead of failing to parse the generated range filters ([#39900](https://github.com/apache/beam/issues/39900)).
+* (Java) `BigQueryIO.Write.withMaxRetryJobs` is now honored for bounded (batch) pipelines using `FILE_LOADS`, which previously always retried failed load jobs 3 times. Pipelines that never call it keep their existing defaults ([#28281](https://github.com/apache/beam/issues/28281)).
 
 ## Security Fixes
 

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java

@@ -135,7 +135,11 @@ class BatchLoads<DestinationT, ElementT>
   // It sets to {@code Integer.MAX_VALUE} to block until the BigQuery job finishes.
   static final int LOAD_JOB_POLL_MAX_RETRIES = Integer.MAX_VALUE;
 
+  // The number of times a failed load or copy job is retried in place before the bundle is failed.
+  // A bounded pipeline keeps this low because the runner can just rerun the failed bundle; an
+  // unbounded pipeline retries far more, because failing a bundle in streaming is expensive.
   static final int DEFAULT_MAX_RETRY_JOBS = 3;
+  static final int DEFAULT_MAX_RETRY_JOBS_UNBOUNDED = 1000;
 
   private BigQueryServices bigQueryServices;
   private final WriteDisposition writeDisposition;

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java

@@ -2523,7 +2523,6 @@ public static <T> Write<T> write() {
         .setPropagateSuccessful(true)
         .setAutoSchemaUpdate(false)
         .setDeterministicRecordIdFn(null)
-        .setMaxRetryJobs(1000)
         .setPropagateSuccessfulStorageApiWrites(false)
         .setPropagateSuccessfulStorageApiWritesPredicate(Predicates.alwaysTrue())
         .setDirectWriteProtos(true)
@@ -2785,7 +2784,7 @@ public enum Method {
 
     abstract boolean getIgnoreInsertIds();
 
-    abstract int getMaxRetryJobs();
+    abstract @Nullable Integer getMaxRetryJobs();
 
     abstract @Nullable String getKmsKey();
 
@@ -2917,7 +2916,7 @@ abstract Builder<T> setDefaultMissingValueInterpretation(
 
       abstract Builder<T> setAutoSharding(boolean autoSharding);
 
-      abstract Builder<T> setMaxRetryJobs(int maxRetryJobs);
+      abstract Builder<T> setMaxRetryJobs(@Nullable Integer maxRetryJobs);
 
       abstract Builder<T> setPropagateSuccessful(boolean propagateSuccessful);
 
@@ -3568,7 +3567,17 @@ public Write<T> withAutoSharding() {
       return toBuilder().setAutoSharding(true).build();
     }
 
-    /** If set, this will set the max number of retry of batch load jobs. */
+    /**
+     * Sets the maximum number of times a failed BigQuery load or copy job is retried before the
+     * write fails.
+     *
+     * <p>Only applies when the write method is {@link Method#FILE_LOADS}. The streaming insert and
+     * Storage Write API methods retry at the row level and ignore this setting.
+     *
+     * <p>If this is not called, a bounded (batch) pipeline retries 3 times and an unbounded
+     * (streaming) pipeline retries 1000 times. Streaming defaults higher because failing a bundle
+     * in streaming is far more expensive than retrying the load job.
+     */
     public Write<T> withMaxRetryJobs(int maxRetryJobs) {
       return toBuilder().setMaxRetryJobs(maxRetryJobs).build();
     }
@@ -4227,10 +4236,15 @@ private <DestinationT> WriteResult continueExpandTyped(
         batchLoads.setMaxFilesPerPartition(getMaxFilesPerPartition());
         batchLoads.setMaxBytesPerPartition(getMaxBytesPerPartition());
 
-        // When running in streaming (unbounded mode) we want to retry failed load jobs
-        // indefinitely. Failing the bundle is expensive, so we set a fairly high limit on retries.
-        if (IsBounded.UNBOUNDED.equals(input.isBounded())) {
-          batchLoads.setMaxRetryJobs(getMaxRetryJobs());
+        // an explicit withMaxRetryJobs applies to batch and streaming alike. left unset, streaming
+        // retries a failed load job far more often than batch does: failing the bundle in streaming
+        // is expensive, so we would rather keep retrying the job than hand the work back to the
+        // runner. batch leaves BatchLoads on its own lower default
+        Integer maxRetryJobs = getMaxRetryJobs();
+        if (maxRetryJobs != null) {
+          batchLoads.setMaxRetryJobs(maxRetryJobs);
+        } else if (IsBounded.UNBOUNDED.equals(input.isBounded())) {
+          batchLoads.setMaxRetryJobs(BatchLoads.DEFAULT_MAX_RETRY_JOBS_UNBOUNDED);
         }
         batchLoads.setTriggeringFrequency(getTriggeringFrequency());
         if (getAutoSharding()) {

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java

@@ -830,8 +830,18 @@ public Write<?> fromConfigRow(Row configRow, PipelineOptions options) {
         if (ignoreInsertIds != null) {
           builder = builder.setIgnoreInsertIds(ignoreInsertIds);
         }
+        // before 2.77.0 BigQueryIO.write() always seeded maxRetryJobs with 1000, whether or not the
+        // pipeline ever called withMaxRetryJobs, and a bounded write ignored the value and used
+        // BatchLoads' own default of 3. so a row from one of those versions that holds exactly 1000
+        // tells us nothing about what the user asked for. leaving it unset makes both modes fall
+        // back to the same numbers they used before the upgrade, 3 for bounded and 1000 for
+        // unbounded. any other value was chosen deliberately and is carried over
         Integer maxRetryJobs = configRow.getInt32("max_retry_jobs");
-        if (maxRetryJobs != null) {
+        boolean isPreservedLegacyDefault =
+            TransformUpgrader.compareVersions(updateCompatibilityBeamVersion, "2.77.0") < 0
+                && Integer.valueOf(BatchLoads.DEFAULT_MAX_RETRY_JOBS_UNBOUNDED)
+                    .equals(maxRetryJobs);
+        if (maxRetryJobs != null && !isPreservedLegacyDefault) {
           builder = builder.setMaxRetryJobs(maxRetryJobs);
         }
         String kmsKey = configRow.getString("kms_key");

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java

@@ -20,6 +20,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
 import com.google.api.services.bigquery.model.Clustering;
@@ -276,6 +277,70 @@ public void testReCreateWriteTransformFromRowTable() {
         writeTransformFromRow.getJsonClustering().get());
   }
 
+  @Test
+  public void testReCreateWriteTransformDropsLegacyMaxRetryJobsDefault() {
+    // an SDK older than 2.77.0 put 1000 in every config row, so this is what a pipeline that never
+    // called withMaxRetryJobs looks like once one of those versions has serialized it
+    BigQueryIO.Write<?> writeTransform =
+        BigQueryIO.write()
+            .to("dummyproject:dummydataset.dummytable")
+            .withMaxRetryJobs(BatchLoads.DEFAULT_MAX_RETRY_JOBS_UNBOUNDED);
+
+    BigQueryIOTranslation.BigQueryIOWriteTranslator translator =
+        new BigQueryIOTranslation.BigQueryIOWriteTranslator();
+    Row row = translator.toConfigRow(writeTransform);
+
+    PipelineOptions options = PipelineOptionsFactory.create();
+    options.as(StreamingOptions.class).setUpdateCompatibilityVersion("2.76.0");
+    BigQueryIO.Write<?> writeTransformFromRow =
+        (BigQueryIO.Write<?>) translator.fromConfigRow(row, options);
+
+    // unset, so a bounded write falls back to BatchLoads' default of 3 as it did before the
+    // upgrade, rather than jumping to 1000
+    assertNull(writeTransformFromRow.getMaxRetryJobs());
+  }
+
+  @Test
+  public void testReCreateWriteTransformKeepsLegacyMaxRetryJobsOtherThanDefault() {
+    BigQueryIO.Write<?> writeTransform =
+        BigQueryIO.write().to("dummyproject:dummydataset.dummytable").withMaxRetryJobs(7);
+
+    BigQueryIOTranslation.BigQueryIOWriteTranslator translator =
+        new BigQueryIOTranslation.BigQueryIOWriteTranslator();
+    Row row = translator.toConfigRow(writeTransform);
+
+    PipelineOptions options = PipelineOptionsFactory.create();
+    options.as(StreamingOptions.class).setUpdateCompatibilityVersion("2.76.0");
+    BigQueryIO.Write<?> writeTransformFromRow =
+        (BigQueryIO.Write<?>) translator.fromConfigRow(row, options);
+
+    // 7 was never a default in any version, so the pipeline must have asked for it
+    assertEquals(Integer.valueOf(7), writeTransformFromRow.getMaxRetryJobs());
+  }
+
+  @Test
+  public void testReCreateWriteTransformKeepsMaxRetryJobsFromCurrentVersion() {
+    BigQueryIO.Write<?> writeTransform =
+        BigQueryIO.write()
+            .to("dummyproject:dummydataset.dummytable")
+            .withMaxRetryJobs(BatchLoads.DEFAULT_MAX_RETRY_JOBS_UNBOUNDED);
+
+    BigQueryIOTranslation.BigQueryIOWriteTranslator translator =
+        new BigQueryIOTranslation.BigQueryIOWriteTranslator();
+    Row row = translator.toConfigRow(writeTransform);
+
+    PipelineOptions options = PipelineOptionsFactory.create();
+    options.as(StreamingOptions.class).setUpdateCompatibilityVersion("2.77.0");
+    BigQueryIO.Write<?> writeTransformFromRow =
+        (BigQueryIO.Write<?>) translator.fromConfigRow(row, options);
+
+    // 2.77.0 and later only write this field when the pipeline set it, so 1000 is a real choice
+    // here
+    assertEquals(
+        Integer.valueOf(BatchLoads.DEFAULT_MAX_RETRY_JOBS_UNBOUNDED),
+        writeTransformFromRow.getMaxRetryJobs());
+  }
+
   @Test
   public void testWriteTransformRowIncludesAllFields() {
     // These fields do not represent properties of the transform.

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java

@@ -2154,12 +2154,40 @@ public void testWriteFailedJobs() throws Exception {
 
     thrown.expect(RuntimeException.class);
     thrown.expectMessage("Failed to create job with prefix");
-    thrown.expectMessage("reached max retries");
+    // a bounded write that never calls withMaxRetryJobs keeps BatchLoads' own default of 3
+    thrown.expectMessage("reached max retries: 3");
     thrown.expectMessage("last failed job");
 
     p.run();
   }
 
+  @Test
+  public void testWriteFailedJobsRespectsMaxRetryJobsWhenBounded() throws Exception {
+    assumeTrue(!useStorageApi);
+    assumeTrue(!useStreaming);
+    p.apply(
+            Create.of(
+                    new TableRow().set("name", "a").set("number", 1),
+                    new TableRow().set("name", "b").set("number", 2),
+                    new TableRow().set("name", "c").set("number", 3))
+                .withCoder(TableRowJsonCoder.of()))
+        .apply(
+            BigQueryIO.writeTableRows()
+                .to("dataset-id.table-id")
+                .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_NEVER)
+                .withMaxRetryJobs(1)
+                .withTestServices(fakeBqServices)
+                .withoutValidation());
+
+    thrown.expect(RuntimeException.class);
+    // the load job fails on every attempt because the destination table does not exist, and the
+    // failure message reports the retry limit that was actually applied. a bounded pipeline used to
+    // drop withMaxRetryJobs on the floor, so this used to read "reached max retries: 3"
+    thrown.expectMessage("reached max retries: 1");
+
+    p.run();
+  }
+
   @Test
   public void testWriteWithMissingSchemaFromView() throws Exception {
     // Because no messages
@@ -3019,7 +3047,7 @@ public void testMaxRetryJobs() {
             .withSchemaUpdateOptions(
                 EnumSet.of(BigQueryIO.Write.SchemaUpdateOption.ALLOW_FIELD_ADDITION))
             .withMaxRetryJobs(500);
-    assertEquals(500, write.getMaxRetryJobs());
+    assertEquals(Integer.valueOf(500), write.getMaxRetryJobs());
   }
 
   @Test