feat: Add isDone to JsonWriter to indicate a JsonWriter is no longer usable and needs to be recreated. (#1978)

yirutang · gcf-owl-bot[bot] · web-flow · commit dc07ed847780 · 2023-02-07T14:52:01.000-08:00
* . * . * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriter.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriter.java
@@ -382,6 +382,15 @@ public void close() {
     this.streamWriter.close();
   }
 
+  /**
+   * @return if a Json writer can no longer be used for writing. It is due to either the
+   *     JsonStreamWriter is explicitly closed or the underlying connection is broken when
+   *     connection pool is not used. Client should recreate JsonStreamWriter in this case.
+   */
+  public boolean isDone() {
+    return this.streamWriter.isDone();
+  }
+
   public static final class Builder {
     private String streamName;
     private BigQueryWriteClient client;
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriterTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriterTest.java
@@ -1139,6 +1139,14 @@ public void testWriterId()
     Assert.assertNotEquals(writer1.getWriterId(), writer2.getWriterId());
   }
 
+  @Test
+  public void testIsDone() throws DescriptorValidationException, IOException, InterruptedException {
+    JsonStreamWriter writer1 = getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA).build();
+    Assert.assertFalse(writer1.isDone());
+    writer1.close();
+    Assert.assertTrue(writer1.isDone());
+  }
+
   private AppendRowsResponse createAppendResponse(long offset) {
     return AppendRowsResponse.newBuilder()
         .setAppendResult(
