chore: clean up BigtableIO client creation code (#29538)

* chore: clean up BigtableIO client creation code * update logging format
apache · Nov 28, 2023 · df324fd · df324fd
1 parent cbddd4c
commit df324fd
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 28 deletions.
diff --git a/...-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableConfigTranslator.java b/...-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableConfigTranslator.java
@@ -256,6 +256,13 @@ private static BigtableDataSettings configureReadSettings(
           Duration.ofMillis(readOptions.getOperationTimeout().getMillis()));
     }
 
+    if (readOptions.getWaitTimeout() != null) {
+      settings
+          .stubSettings()
+          .readRowsSettings()
+          .setWaitTimeout(Duration.ofMillis(readOptions.getWaitTimeout().getMillis()));
+    }
+
     settings.stubSettings().readRowsSettings().setRetrySettings(retrySettings.build());
 
     return settings.build();

diff --git a/...ud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceFactory.java b/...ud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceFactory.java
@@ -128,12 +128,7 @@ BigtableServiceEntry getServiceForReading(
         BigtableDataSettings.enableBuiltinMetrics();
       }
 
-      BigtableService service;
-      if (opts.getWaitTimeout() != null) {
-        service = new BigtableServiceImpl(settings, opts.getWaitTimeout());
-      } else {
-        service = new BigtableServiceImpl(settings);
-      }
+      BigtableService service = new BigtableServiceImpl(settings);
       entry = BigtableServiceEntry.create(configId, service);
       entries.put(configId.id(), entry);
       refCounts.put(configId.id(), new AtomicInteger(1));

diff --git a/...cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java b/...cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java
@@ -101,33 +101,13 @@ class BigtableServiceImpl implements BigtableService {
   private static final long MIN_BYTE_BUFFER_SIZE = 100 * 1024 * 1024; // 100MB
 
   BigtableServiceImpl(BigtableDataSettings settings) throws IOException {
-    this(settings, null);
-  }
-
-  // TODO remove this constructor once https://github.com/googleapis/gapic-generator-java/pull/1473
-  // is resolved. readWaitTimeout is a hack to workaround incorrect mapping from attempt timeout to
-  // Watchdog's wait timeout.
-  BigtableServiceImpl(BigtableDataSettings settings, Duration readWaitTimeout) throws IOException {
     this.projectId = settings.getProjectId();
     this.instanceId = settings.getInstanceId();
     RetrySettings retry = settings.getStubSettings().readRowsSettings().getRetrySettings();
     this.readAttemptTimeout = Duration.millis(retry.getInitialRpcTimeout().toMillis());
     this.readOperationTimeout = Duration.millis(retry.getTotalTimeout().toMillis());
-    BigtableDataSettings.Builder builder = settings.toBuilder();
-    if (readWaitTimeout != null) {
-      builder
-          .stubSettings()
-          .readRowsSettings()
-          .setRetrySettings(
-              retry
-                  .toBuilder()
-                  .setInitialRpcTimeout(
-                      org.threeten.bp.Duration.ofMillis(readWaitTimeout.getMillis()))
-                  .setMaxRpcTimeout(org.threeten.bp.Duration.ofMillis(readWaitTimeout.getMillis()))
-                  .build());
-    }
-    LOG.info("Started Bigtable service with settings " + builder.build());
-    this.client = BigtableDataClient.create(builder.build());
+    LOG.info("Started Bigtable service with settings {}", settings);
+    this.client = BigtableDataClient.create(settings);
   }
 
   private final BigtableDataClient client;