Skip to content

Commit 4258af4

Browse files
fix: add client id and update trace id population for StreamWriter and JsonWriter (#2389)
* fix: add a bit more message so customers are not going to be scaried by load shedding errors * . * fix:change the client id string, adding client id for stream writer * . * . * . * . * . * . * . * . * 🦉 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>
1 parent e5913cf commit 4258af4

File tree

6 files changed

+44
-37
lines changed

6 files changed

+44
-37
lines changed

google-cloud-bigquerystorage/clirr-ignored-differences.xml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -188,5 +188,10 @@
188188
<className>com/google/cloud/bigquery/storage/v1/StreamWriter$Builder</className>
189189
<method>com.google.cloud.bigquery.storage.v1.StreamWriter$Builder setRetryFirstDelay(org.threeten.bp.Duration)</method>
190190
</difference>
191+
<difference>
192+
<differenceType>7002</differenceType>
193+
<className>com/google/cloud/bigquery/storage/v1/SchemaAwareStreamWriter$Builder</className>
194+
<method>com.google.cloud.bigquery.storage.v1.SchemaAwareStreamWriter$Builder setTraceIdBase(java.lang.String)</method>
195+
</difference>
191196
</differences>
192197

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriter.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
*/
3838
public class JsonStreamWriter implements AutoCloseable {
3939
private final SchemaAwareStreamWriter<Object> schemaAwareStreamWriter;
40+
private static final String CLIENT_ID = "java-jsonwriter";
4041

4142
/**
4243
* Constructs the JsonStreamWriter
@@ -227,8 +228,7 @@ public static final class Builder {
227228
private final SchemaAwareStreamWriter.Builder<Object> schemaAwareStreamWriterBuilder;
228229

229230
private Builder(SchemaAwareStreamWriter.Builder<Object> schemaAwareStreamWriterBuilder) {
230-
this.schemaAwareStreamWriterBuilder =
231-
schemaAwareStreamWriterBuilder.setTraceIdBase("JsonWriter");
231+
this.schemaAwareStreamWriterBuilder = schemaAwareStreamWriterBuilder.setClientId(CLIENT_ID);
232232
}
233233

234234
/**

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/SchemaAwareStreamWriter.java

Lines changed: 6 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -94,14 +94,14 @@ private SchemaAwareStreamWriter(Builder<T> builder)
9494
builder.executorProvider,
9595
builder.endpoint,
9696
builder.flowControlSettings,
97-
builder.traceIdBase,
9897
builder.traceId,
9998
builder.compressorName,
10099
builder.retrySettings);
101100
streamWriterBuilder.setEnableConnectionPool(builder.enableConnectionPool);
102101
streamWriterBuilder.setLocation(builder.location);
103102
streamWriterBuilder.setDefaultMissingValueInterpretation(
104103
builder.defaultMissingValueInterpretation);
104+
streamWriterBuilder.setClientId(builder.clientId);
105105
this.streamWriter = streamWriterBuilder.build();
106106
this.streamName = builder.streamName;
107107
this.tableSchema = builder.tableSchema;
@@ -282,7 +282,6 @@ private void setStreamWriterSettings(
282282
@Nullable ExecutorProvider executorProvider,
283283
@Nullable String endpoint,
284284
@Nullable FlowControlSettings flowControlSettings,
285-
@Nullable String traceIdBase,
286285
@Nullable String traceId,
287286
@Nullable String compressorName,
288287
@Nullable RetrySettings retrySettings) {
@@ -298,18 +297,8 @@ private void setStreamWriterSettings(
298297
if (endpoint != null) {
299298
streamWriterBuilder.setEndpoint(endpoint);
300299
}
301-
if (traceIdBase != null) {
302-
if (traceId != null) {
303-
streamWriterBuilder.setTraceId(traceIdBase + "_" + traceId);
304-
} else {
305-
streamWriterBuilder.setTraceId(traceIdBase + ":null");
306-
}
307-
} else {
308-
if (traceId != null) {
309-
streamWriterBuilder.setTraceId("SchemaAwareStreamWriter_" + traceId);
310-
} else {
311-
streamWriterBuilder.setTraceId("SchemaAwareStreamWriter:null");
312-
}
300+
if (traceId != null) {
301+
streamWriterBuilder.setTraceId(traceId);
313302
}
314303
if (flowControlSettings != null) {
315304
if (flowControlSettings.getMaxOutstandingRequestBytes() != null) {
@@ -445,6 +434,7 @@ public static final class Builder<T> {
445434

446435
private AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation =
447436
MissingValueInterpretation.MISSING_VALUE_INTERPRETATION_UNSPECIFIED;
437+
private String clientId;
448438

449439
private static final String streamPatternString =
450440
"(projects/[^/]+/datasets/[^/]+/tables/[^/]+)/streams/[^/]+";
@@ -581,14 +571,8 @@ public Builder<T> setTraceId(String traceId) {
581571
return this;
582572
}
583573

584-
/**
585-
* Setter for a traceIdBase to help identify traffic origin.
586-
*
587-
* @param traceIdBase
588-
* @return Builder
589-
*/
590-
public Builder<T> setTraceIdBase(String traceIdBase) {
591-
this.traceIdBase = Preconditions.checkNotNull(traceIdBase, "TraceIdBase is null.");
574+
Builder<T> setClientId(String clientId) {
575+
this.clientId = Preconditions.checkNotNull(clientId, "ClientId is null.");
592576
return this;
593577
}
594578

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -229,7 +229,7 @@ private StreamWriter(Builder builder) throws IOException {
229229
builder.maxInflightBytes,
230230
builder.maxRetryDuration,
231231
builder.limitExceededBehavior,
232-
builder.traceId,
232+
builder.getFullTraceId(),
233233
builder.compressorName,
234234
clientSettings,
235235
builder.retrySettings));
@@ -295,7 +295,7 @@ private StreamWriter(Builder builder) throws IOException {
295295
builder.maxInflightBytes,
296296
builder.maxRetryDuration,
297297
builder.limitExceededBehavior,
298-
builder.traceId,
298+
builder.getFullTraceId(),
299299
builder.compressorName,
300300
client.getSettings(),
301301
builder.retrySettings);
@@ -376,12 +376,12 @@ static BigQueryWriteSettings getBigQueryWriteSettings(Builder builder) throws IO
376376
private void validateFetchedConnectonPool(StreamWriter.Builder builder) {
377377
String storedTraceId =
378378
this.singleConnectionOrConnectionPool.connectionWorkerPool().getTraceId();
379-
if (!Objects.equals(storedTraceId, builder.traceId)) {
379+
if (!Objects.equals(storedTraceId, builder.getFullTraceId())) {
380380
throw new IllegalArgumentException(
381381
String.format(
382382
"Trace id used for the same connection pool for the same location must be the same, "
383383
+ "however stored trace id is %s, and expected trace id is %s.",
384-
storedTraceId, builder.traceId));
384+
storedTraceId, builder.getFullTraceId()));
385385
}
386386
FlowController.LimitExceededBehavior storedLimitExceededBehavior =
387387
singleConnectionOrConnectionPool.connectionWorkerPool().limitExceededBehavior();
@@ -629,6 +629,8 @@ public static final class Builder {
629629

630630
private String traceId = null;
631631

632+
private String clientId = "java-streamwriter";
633+
632634
private TableSchema updatedTableSchema = null;
633635

634636
private String location = null;
@@ -730,6 +732,15 @@ public Builder setTraceId(String traceId) {
730732
return this;
731733
}
732734

735+
/**
736+
* Sets the client id of the writer, for example, JsonStreamWriter has the client id of
737+
* "java-jsonwriter".
738+
*/
739+
Builder setClientId(String clientId) {
740+
this.clientId = clientId;
741+
return this;
742+
}
743+
733744
/** Location of the table this stream writer is targeting. */
734745
public Builder setLocation(String location) {
735746
this.location = location;
@@ -811,5 +822,13 @@ public Builder setRetrySettings(RetrySettings retrySettings) {
811822
public StreamWriter build() throws IOException {
812823
return new StreamWriter(this);
813824
}
825+
826+
String getFullTraceId() {
827+
if (traceId == null) {
828+
return clientId;
829+
} else {
830+
return clientId + " " + traceId;
831+
}
832+
}
814833
}
815834
}

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriterTest.java

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -237,7 +237,7 @@ public void testSingleAppendSimpleJson() throws Exception {
237237
.getSerializedRows(0),
238238
expectedProto.toByteString());
239239
assertEquals(
240-
testBigQueryWrite.getAppendRequests().get(0).getTraceId(), "JsonWriter_test:empty");
240+
"java-jsonwriter test:empty", testBigQueryWrite.getAppendRequests().get(0).getTraceId());
241241
}
242242
}
243243

@@ -257,7 +257,7 @@ public void testFlexibleColumnAppend() throws Exception {
257257
jsonArr.put(flexible);
258258

259259
try (JsonStreamWriter writer =
260-
getTestJsonStreamWriterBuilder(TEST_STREAM, tableSchema).setTraceId("test:empty").build()) {
260+
getTestJsonStreamWriterBuilder(TEST_STREAM, tableSchema).build()) {
261261

262262
testBigQueryWrite.addResponse(
263263
AppendRowsResponse.newBuilder()
@@ -284,8 +284,7 @@ public void testFlexibleColumnAppend() throws Exception {
284284
.getRows()
285285
.getSerializedRows(0),
286286
expectedProto.toByteString());
287-
assertEquals(
288-
testBigQueryWrite.getAppendRequests().get(0).getTraceId(), "JsonWriter_test:empty");
287+
assertEquals("java-jsonwriter", testBigQueryWrite.getAppendRequests().get(0).getTraceId());
289288
}
290289
}
291290

@@ -507,7 +506,7 @@ public void testSingleAppendMultipleSimpleJson() throws Exception {
507506
.getProtoRows()
508507
.getRows()
509508
.getSerializedRowsCount());
510-
assertEquals(testBigQueryWrite.getAppendRequests().get(0).getTraceId(), "JsonWriter:null");
509+
assertEquals("java-jsonwriter", testBigQueryWrite.getAppendRequests().get(0).getTraceId());
511510
for (int i = 0; i < 4; i++) {
512511
assertEquals(
513512
testBigQueryWrite

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -286,13 +286,13 @@ private void verifyAppendRequests(long appendCount) {
286286
if (i == 0) {
287287
// First request received by server should have schema and stream name.
288288
assertTrue(serverRequest.getProtoRows().hasWriterSchema());
289-
assertEquals(serverRequest.getWriteStream(), TEST_STREAM_1);
290-
assertEquals(serverRequest.getTraceId(), TEST_TRACE_ID);
289+
assertEquals(TEST_STREAM_1, serverRequest.getWriteStream());
290+
assertEquals("java-streamwriter " + TEST_TRACE_ID, serverRequest.getTraceId());
291291
} else {
292292
// Following request should not have schema and stream name.
293293
assertFalse(serverRequest.getProtoRows().hasWriterSchema());
294-
assertEquals(serverRequest.getWriteStream(), "");
295-
assertEquals(serverRequest.getTraceId(), "");
294+
assertEquals("", serverRequest.getWriteStream());
295+
assertEquals("", serverRequest.getTraceId());
296296
}
297297
}
298298
}

0 commit comments

Comments
 (0)