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
googleapis · Sep 8, 2021 · 0e749d9 · 0e749d9
1 parent b3ef5ee
commit 0e749d9
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 3 deletions.
diff --git a/...uerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/JsonStreamWriter.java b/...uerystorage/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/...storage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/JsonStreamWriterTest.java b/...storage/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