samples: Update WriteBatch sample to use a batcher (#2135)

* samples: Update WriteBatch sample to use a batcher * remove unused imports * format * address comments * handle exception * handle exception * fix format * fix test
googleapis · Feb 27, 2024 · ee740f8 · ee740f8
1 parent 899c50b
commit ee740f8
Showing 1 changed file with 42 additions and 28 deletions.
diff --git a/samples/snippets/src/main/java/com/example/bigtable/WriteBatch.java b/samples/snippets/src/main/java/com/example/bigtable/WriteBatch.java
@@ -18,10 +18,15 @@
 
 // [START bigtable_writes_batch]
 
+import com.google.api.core.ApiFuture;
+import com.google.api.gax.batching.Batcher;
+import com.google.api.gax.batching.BatchingException;
 import com.google.cloud.bigtable.data.v2.BigtableDataClient;
-import com.google.cloud.bigtable.data.v2.models.BulkMutation;
-import com.google.cloud.bigtable.data.v2.models.Mutation;
+import com.google.cloud.bigtable.data.v2.models.RowMutationEntry;
 import com.google.protobuf.ByteString;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
 
 public class WriteBatch {
  private static final String COLUMN_FAMILY_NAME = "stats_summary";
@@ -32,34 +37,43 @@ public static void writeBatch(String projectId, String instanceId, String tableI
  // String tableId = "mobile-time-series";
 
  try (BigtableDataClient dataClient = BigtableDataClient.create(projectId, instanceId)) {
- long timestamp = System.currentTimeMillis() * 1000;
+ List<ApiFuture<Void>> batchFutures = new ArrayList<>();
+ try (Batcher<RowMutationEntry, Void> batcher = dataClient.newBulkMutationBatcher(tableId)) {
+ long timestamp = System.currentTimeMillis() * 1000;
+ batchFutures.add(
+ batcher.add(
+ RowMutationEntry.create("tablet#a0b81f74#20190501")
+ .setCell(
+ COLUMN_FAMILY_NAME, ByteString.copyFromUtf8("connected_wifi"), timestamp, 1)
+ .setCell(COLUMN_FAMILY_NAME, "os_build", timestamp, "12155.0.0-rc1")));
+ batchFutures.add(
+ batcher.add(
+ RowMutationEntry.create("tablet#a0b81f74#20190502")
+ .setCell(
+ COLUMN_FAMILY_NAME, ByteString.copyFromUtf8("connected_wifi"), timestamp, 1)
+ .setCell(COLUMN_FAMILY_NAME, "os_build", timestamp, "12155.0.0-rc6")));
 
- BulkMutation bulkMutation =
- BulkMutation.create(tableId)
- .add(
- "tablet#a0b81f74#20190501",
- Mutation.create()
- .setCell(
- COLUMN_FAMILY_NAME,
- ByteString.copyFrom("connected_wifi".getBytes()),
- timestamp,
- 1)
- .setCell(COLUMN_FAMILY_NAME, "os_build", timestamp, "12155.0.0-rc1"))
- .add(
- "tablet#a0b81f74#20190502",
- Mutation.create()
- .setCell(
- COLUMN_FAMILY_NAME,
- ByteString.copyFrom("connected_wifi".getBytes()),
- timestamp,
- 1)
- .setCell(COLUMN_FAMILY_NAME, "os_build", timestamp, "12155.0.0-rc6"));
-
- dataClient.bulkMutateRows(bulkMutation);
-
- System.out.print("Successfully wrote 2 rows");
+ // Blocks until mutations are applied on all submitted row entries.
+ // flush will be called automatically when a batch is full.
+ batcher.flush();
+ // Before batcher is closed, all remaining (if any) mutations are applied.
+ } catch (BatchingException batchingException) {
+ System.out.println(
+ "At least one entry failed to apply. Summary of the errors: \n" + batchingException);
+ // get individual entry error details
+ for (ApiFuture<Void> future : batchFutures) {
+ try {
+ future.get();
+ } catch (ExecutionException entryException) {
+ System.out.println("Entry failure: " + entryException.getCause());
+ } catch (InterruptedException e) {
+ // handle interrupted exception
+ }
+ }
+ }
+ System.out.println("Successfully wrote 2 rows");
  } catch (Exception e) {
- System.out.println("Error during WriteBatch: \n" + e.toString());
+ System.out.println("Error during WriteBatch: \n" + e);
  }
  }
 }