Skip to content
Permalink
Browse files
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
  • Loading branch information
yirutang committed Sep 8, 2021
1 parent b3ef5ee commit 0e749d9ee0201bc4653735ac34638452502a1c26
@@ -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
*
@@ -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

0 comments on commit 0e749d9

Please sign in to comment.