From 0e749d9ee0201bc4653735ac34638452502a1c26 Mon Sep 17 00:00:00 2001
From: Yiru Tang <yiru@google.com>
Date: Wed, 8 Sep 2021 15:00:26 -0700
Subject: [PATCH] feat: add trace id support to JsonWriter and add default
 trace id to help identify json writer users. (#1302)

* fix: update code comment to reflect max size change

* feat: add setTraceId to JsonWriter to annotate writes from json writer
---
 .../storage/v1beta2/JsonStreamWriter.java     | 23 +++++++++++++++++--
 .../storage/v1beta2/JsonStreamWriterTest.java |  8 ++++++-
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/JsonStreamWriter.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/JsonStreamWriter.java
index 2b2e531174..ed5524a892 100644
--- a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/JsonStreamWriter.java
+++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/JsonStreamWriter.java
@@ -72,7 +72,8 @@ private JsonStreamWriter(Builder builder)
         builder.channelProvider,
         builder.credentialsProvider,
         builder.endpoint,
-        builder.flowControlSettings);
+        builder.flowControlSettings,
+        builder.traceId);
     this.streamWriter = streamWriterBuilder.build();
     this.streamName = builder.streamName;
   }
@@ -156,7 +157,8 @@ private void setStreamWriterSettings(
       @Nullable TransportChannelProvider channelProvider,
       @Nullable CredentialsProvider credentialsProvider,
       @Nullable String endpoint,
-      @Nullable FlowControlSettings flowControlSettings) {
+      @Nullable FlowControlSettings flowControlSettings,
+      @Nullable String traceId) {
     if (channelProvider != null) {
       streamWriterBuilder.setChannelProvider(channelProvider);
     }
@@ -166,6 +168,11 @@ private void setStreamWriterSettings(
     if (endpoint != null) {
       streamWriterBuilder.setEndpoint(endpoint);
     }
+    if (traceId != null) {
+      streamWriterBuilder.setTraceId("JsonWriterBeta_" + traceId);
+    } else {
+      streamWriterBuilder.setTraceId("JsonWriterBeta:null");
+    }
     if (flowControlSettings != null) {
       if (flowControlSettings.getMaxOutstandingRequestBytes() != null) {
         streamWriterBuilder.setMaxInflightBytes(
@@ -246,6 +253,7 @@ public static final class Builder {
     private FlowControlSettings flowControlSettings;
     private String endpoint;
     private boolean createDefaultStream = false;
+    private String traceId;
 
     private static String streamPatternString =
         "(projects/[^/]+/datasets/[^/]+/tables/[^/]+)/streams/[^/]+";
@@ -336,6 +344,17 @@ public Builder setEndpoint(String endpoint) {
       return this;
     }
 
+    /**
+     * Setter for a traceId to help identify traffic origin.
+     *
+     * @param traceId
+     * @return Builder
+     */
+    public Builder setTraceId(String traceId) {
+      this.traceId = Preconditions.checkNotNull(traceId, "TraceId is null.");
+      return this;
+    }
+
     /**
      * Builds JsonStreamWriter
      *
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/JsonStreamWriterTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/JsonStreamWriterTest.java
index 0e888ef267..4bd9ea91dc 100644
--- a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/JsonStreamWriterTest.java
+++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/JsonStreamWriterTest.java
@@ -253,7 +253,9 @@ public void testSingleAppendSimpleJson() throws Exception {
     jsonArr.put(foo);
 
     try (JsonStreamWriter writer =
-        getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA).build()) {
+        getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA)
+            .setTraceId("test:empty")
+            .build()) {
 
       testBigQueryWrite.addResponse(
           AppendRowsResponse.newBuilder()
@@ -280,6 +282,8 @@ public void testSingleAppendSimpleJson() throws Exception {
               .getRows()
               .getSerializedRows(0),
           expectedProto.toByteString());
+      assertEquals(
+          testBigQueryWrite.getAppendRequests().get(0).getTraceId(), "JsonWriterBeta_test:empty");
     }
   }
 
@@ -320,6 +324,8 @@ public void testSingleAppendMultipleSimpleJson() throws Exception {
               .getProtoRows()
               .getRows()
               .getSerializedRowsCount());
+      assertEquals(
+          testBigQueryWrite.getAppendRequests().get(0).getTraceId(), "JsonWriterBeta:null");
       for (int i = 0; i < 4; i++) {
         assertEquals(
             testBigQueryWrite