feat: provide append() methods that accept com.google.gson objects (#2985)

agrawal-siddharth · web-flow · commit b8dff570a27b · 2025-05-28T13:36:25.000-07:00
diff --git a/google-cloud-bigquerystorage/pom.xml b/google-cloud-bigquerystorage/pom.xml
@@ -117,6 +117,11 @@
       <groupId>com.google.api.grpc</groupId>
       <artifactId>proto-google-cloud-bigquerystorage-v1</artifactId>
     </dependency>
+    <dependency>
+      <groupId>com.google.code.gson</groupId>
+      <artifactId>gson</artifactId>
+      <version>2.12.1</version>
+    </dependency>
     <dependency>
       <groupId>com.google.guava</groupId>
       <artifactId>guava</artifactId>
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
@@ -21,6 +21,7 @@
 import com.google.api.gax.core.ExecutorProvider;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.rpc.TransportChannelProvider;
+import com.google.gson.JsonArray;
 import com.google.protobuf.Descriptors;
 import java.io.IOException;
 import java.time.Duration;
@@ -83,6 +84,41 @@ public ApiFuture<AppendRowsResponse> append(JSONArray jsonArr, long offset)
     return this.schemaAwareStreamWriter.append(jsonArr, offset);
   }
 
+  private JSONArray gsonToOrgJSON(JsonArray jsonArr) {
+    return new JSONArray(jsonArr.toString());
+  }
+
+  /**
+   * Writes a JsonArray that contains JsonObjects to the BigQuery table by first converting the JSON
+   * data to protobuf messages, then using StreamWriter's append() to write the data at current end
+   * of stream. If there is a schema update, the current StreamWriter is closed. A new StreamWriter
+   * is created with the updated TableSchema.
+   *
+   * @param jsonArr The JSON array that contains JsonObjects to be written
+   * @return {@code ApiFuture<AppendRowsResponse>} returns an AppendRowsResponse message wrapped in
+   *     an ApiFuture
+   */
+  public ApiFuture<AppendRowsResponse> append(JsonArray jsonArr)
+      throws IOException, Descriptors.DescriptorValidationException {
+    return this.append(jsonArr, -1);
+  }
+
+  /**
+   * Writes a JsonArray that contains JsonObjects to the BigQuery table by first converting the JSON
+   * data to protobuf messages, then using StreamWriter's append() to write the data at the
+   * specified offset. If there is a schema update, the current StreamWriter is closed. A new
+   * StreamWriter is created with the updated TableSchema.
+   *
+   * @param jsonArr The JSON array that contains JSONObjects to be written
+   * @param offset Offset for deduplication
+   * @return {@code ApiFuture<AppendRowsResponse>} returns an AppendRowsResponse message wrapped in
+   *     an ApiFuture
+   */
+  public ApiFuture<AppendRowsResponse> append(JsonArray jsonArr, long offset)
+      throws IOException, Descriptors.DescriptorValidationException {
+    return this.append(gsonToOrgJSON(jsonArr), offset);
+  }
+
   public String getStreamName() {
     return this.schemaAwareStreamWriter.getStreamName();
   }
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
@@ -38,6 +38,8 @@
 import com.google.cloud.bigquery.storage.v1.Exceptions.AppendSerializationError;
 import com.google.cloud.bigquery.storage.v1.TableFieldSchema.Mode;
 import com.google.common.collect.ImmutableMap;
+import com.google.gson.JsonArray;
+import com.google.gson.JsonObject;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.Descriptors.DescriptorValidationException;
 import com.google.protobuf.Int64Value;
@@ -251,6 +253,49 @@ public void testSingleAppendSimpleJson() throws Exception {
     }
   }
 
+  @Test
+  public void testSingleAppendSimpleGson() throws Exception {
+    FooType expectedProto = FooType.newBuilder().setFoo("allen").build();
+    JsonObject foo = new JsonObject();
+    foo.addProperty("foo", "allen");
+    JsonArray jsonArr = new JsonArray();
+    jsonArr.add(foo);
+
+    try (JsonStreamWriter writer =
+        getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA)
+            .setTraceId("test:empty")
+            .build()) {
+
+      testBigQueryWrite.addResponse(
+          AppendRowsResponse.newBuilder()
+              .setAppendResult(
+                  AppendRowsResponse.AppendResult.newBuilder().setOffset(Int64Value.of(0)).build())
+              .build());
+
+      ApiFuture<AppendRowsResponse> appendFuture = writer.append(jsonArr);
+      assertEquals(0L, appendFuture.get().getAppendResult().getOffset().getValue());
+      appendFuture.get();
+      assertEquals(
+          1,
+          testBigQueryWrite
+              .getAppendRequests()
+              .get(0)
+              .getProtoRows()
+              .getRows()
+              .getSerializedRowsCount());
+      assertEquals(
+          testBigQueryWrite
+              .getAppendRequests()
+              .get(0)
+              .getProtoRows()
+              .getRows()
+              .getSerializedRows(0),
+          expectedProto.toByteString());
+      assertEquals(
+          "java-jsonwriter test:empty", testBigQueryWrite.getAppendRequests().get(0).getTraceId());
+    }
+  }
+
   @Test
   public void testFlexibleColumnAppend() throws Exception {
     TableFieldSchema field =
