Allow setting BatchSize and MaxBufferingDuration in JDBC IO WriteWithResults method

Author: tvalentyn

sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java

@@ -1949,6 +1949,8 @@ public <V extends JdbcWriteResult> WriteWithResults<T, V> withWriteResults(
           .setStatement(inner.getStatement())
           .setTable(inner.getTable())
           .setAutoSharding(inner.getAutoSharding())
+          .setBatchSize(inner.getBatchSize())
+          .setMaxBatchBufferingDuration(inner.getMaxBatchBufferingDuration())
           .build();
     }
 
@@ -2055,6 +2057,10 @@ public abstract static class WriteWithResults<T, V extends JdbcWriteResult>
 
     abstract @Nullable RowMapper<V> getRowMapper();
 
+    abstract @Nullable Long getBatchSize();
+
+    abstract @Nullable Long getMaxBatchBufferingDuration();
+
     abstract Builder<T, V> toBuilder();
 
     @AutoValue.Builder
@@ -2064,6 +2070,10 @@ abstract Builder<T, V> setDataSourceProviderFn(
 
       abstract Builder<T, V> setAutoSharding(@Nullable Boolean autoSharding);
 
+      abstract Builder<T, V> setBatchSize(@Nullable Long batchSize);
+
+      abstract Builder<T, V> setMaxBatchBufferingDuration(@Nullable Long maxBatchBufferingDuration);
+
       abstract Builder<T, V> setStatement(@Nullable ValueProvider<String> statement);
 
       abstract Builder<T, V> setPreparedStatementSetter(
@@ -2080,6 +2090,19 @@ abstract Builder<T, V> setPreparedStatementSetter(
       abstract WriteWithResults<T, V> build();
     }
 
+    public WriteWithResults<T, V> withBatchSize(long batchSize) {
+      checkArgument(batchSize > 0, "batchSize must be > 0, but was %s", batchSize);
+      return toBuilder().setBatchSize(batchSize).build();
+    }
+
+    public WriteWithResults<T, V> withMaxBatchBufferingDuration(long maxBatchBufferingDuration) {
+      checkArgument(
+          maxBatchBufferingDuration > 0,
+          "maxBatchBufferingDuration must be > 0, but was %s",
+          maxBatchBufferingDuration);
+      return toBuilder().setMaxBatchBufferingDuration(maxBatchBufferingDuration).build();
+    }
+
     public WriteWithResults<T, V> withDataSourceConfiguration(DataSourceConfiguration config) {
       return withDataSourceProviderFn(new DataSourceProviderFromDataSourceConfiguration(config));
     }
@@ -2173,9 +2196,16 @@ public PCollection<V> expand(PCollection<T> input) {
           autoSharding == null || (autoSharding && input.isBounded() != IsBounded.UNBOUNDED),
           "Autosharding is only supported for streaming pipelines.");
 
+      Long batchSizeAsLong = getBatchSize();
+      long batchSize = batchSizeAsLong == null ? DEFAULT_BATCH_SIZE : batchSizeAsLong;
+      Long maxBufferingDurationAsLong = getMaxBatchBufferingDuration();
+      long maxBufferingDuration =
+          maxBufferingDurationAsLong == null
+              ? DEFAULT_MAX_BATCH_BUFFERING_DURATION
+              : maxBufferingDurationAsLong;
+
       PCollection<Iterable<T>> iterables =
-          JdbcIO.<T>batchElements(
-              input, autoSharding, DEFAULT_BATCH_SIZE, DEFAULT_MAX_BATCH_BUFFERING_DURATION);
+          JdbcIO.<T>batchElements(input, autoSharding, batchSize, maxBufferingDuration);
       return iterables.apply(
           ParDo.of(
               new WriteFn<T, V>(
@@ -2187,8 +2217,8 @@ public PCollection<V> expand(PCollection<T> input) {
                       .setStatement(getStatement())
                       .setRetryConfiguration(getRetryConfiguration())
                       .setReturnResults(true)
-                      .setBatchSize(1L)
-                      .setMaxBatchBufferingDuration(DEFAULT_MAX_BATCH_BUFFERING_DURATION)
+                      .setBatchSize(1L) // We are writing iterables 1 at a time.
+                      .setMaxBatchBufferingDuration(maxBufferingDuration)
                       .build())));
     }
   }