diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/TemplateUtils.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/TemplateUtils.java index 53f79ae2b6..59ed9bd908 100644 --- a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/TemplateUtils.java +++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/TemplateUtils.java @@ -106,6 +106,12 @@ public static CloudBigtableScanConfiguration buildExportConfig(ExportOptions opt BigtableOptionsFactory.BIGTABLE_READ_RPC_ATTEMPT_TIMEOUT_MS_KEY, options.getBigtableReadRpcAttemptTimeoutMs()); } + if (options.getBigtableMaxAttempts() != null) { + configBuilder.withConfiguration( + BigtableOptionsFactory.MAX_SCAN_TIMEOUT_RETRIES, + ValueProvider.NestedValueProvider.of( + options.getBigtableMaxAttempts(), value -> String.valueOf(value))); + } return configBuilder.build(); } } diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/sequencefiles/ExportJob.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/sequencefiles/ExportJob.java index b70eef5c7d..9d5806b70d 100644 --- a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/sequencefiles/ExportJob.java +++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/sequencefiles/ExportJob.java @@ -153,6 +153,12 @@ public interface ExportOptions extends GcpOptions, GcsOptions { @SuppressWarnings("unused") void setBigtableFilter(ValueProvider filter); + @Description("The maximum number of retry attempts for the Bigtable client.") + ValueProvider getBigtableMaxAttempts(); + + @SuppressWarnings("unused") + void setBigtableMaxAttempts(ValueProvider maxAttempts); + @Description("The destination directory") ValueProvider getDestinationPath(); diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/TemplateUtilsTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/TemplateUtilsTest.java new file mode 100644 index 0000000000..b0b830d7af --- /dev/null +++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/TemplateUtilsTest.java @@ -0,0 +1,60 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigtable.beam; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; + +import com.google.cloud.bigtable.beam.sequencefiles.ExportJob.ExportOptions; +import com.google.cloud.bigtable.hbase.BigtableOptionsFactory; +import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; +import org.apache.hadoop.conf.Configuration; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mockito; + +@RunWith(JUnit4.class) +public class TemplateUtilsTest { + + @Test + public void testBuildExportConfig_propagatesMaxAttempts() { + ExportOptions options = Mockito.mock(ExportOptions.class); + Mockito.when(options.getBigtableProject()).thenReturn(StaticValueProvider.of("my-project")); + Mockito.when(options.getBigtableInstanceId()).thenReturn(StaticValueProvider.of("my-instance")); + Mockito.when(options.getBigtableTableId()).thenReturn(StaticValueProvider.of("my-table")); + Mockito.when(options.getBigtableMaxAttempts()).thenReturn(StaticValueProvider.of(15)); + Mockito.when(options.getBigtableAppProfileId()).thenReturn(StaticValueProvider.of("my-app-profile")); + + CloudBigtableScanConfiguration config = TemplateUtils.buildExportConfig(options); + Configuration hbaseConfig = config.toHBaseConfig(); + assertEquals("15", hbaseConfig.get(BigtableOptionsFactory.MAX_SCAN_TIMEOUT_RETRIES)); + } + + @Test + public void testBuildExportConfig_defaultMaxAttempts() { + ExportOptions options = Mockito.mock(ExportOptions.class); + Mockito.when(options.getBigtableProject()).thenReturn(StaticValueProvider.of("my-project")); + Mockito.when(options.getBigtableInstanceId()).thenReturn(StaticValueProvider.of("my-instance")); + Mockito.when(options.getBigtableTableId()).thenReturn(StaticValueProvider.of("my-table")); + Mockito.when(options.getBigtableAppProfileId()).thenReturn(StaticValueProvider.of((String) null)); + Mockito.when(options.getBigtableMaxAttempts()).thenReturn(null); + + CloudBigtableScanConfiguration config = TemplateUtils.buildExportConfig(options); + Configuration hbaseConfig = config.toHBaseConfig(); + assertNull(hbaseConfig.get(BigtableOptionsFactory.MAX_SCAN_TIMEOUT_RETRIES)); + } +}