diff --git a/google-cloud-bigquerystorage/clirr-ignored-differences.xml b/google-cloud-bigquerystorage/clirr-ignored-differences.xml index 6d99d6a031..0add32e9f1 100644 --- a/google-cloud-bigquerystorage/clirr-ignored-differences.xml +++ b/google-cloud-bigquerystorage/clirr-ignored-differences.xml @@ -188,5 +188,10 @@ com/google/cloud/bigquery/storage/v1/StreamWriter$Builder com.google.cloud.bigquery.storage.v1.StreamWriter$Builder setRetryFirstDelay(org.threeten.bp.Duration) + + 7002 + com/google/cloud/bigquery/storage/v1/SchemaAwareStreamWriter$Builder + com.google.cloud.bigquery.storage.v1.SchemaAwareStreamWriter$Builder setTraceIdBase(java.lang.String) + 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 index 64d683438c..0e8d669764 100644 --- 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 @@ -37,6 +37,7 @@ */ public class JsonStreamWriter implements AutoCloseable { private final SchemaAwareStreamWriter schemaAwareStreamWriter; + private static final String CLIENT_ID = "java-jsonwriter"; /** * Constructs the JsonStreamWriter @@ -227,8 +228,7 @@ public static final class Builder { private final SchemaAwareStreamWriter.Builder schemaAwareStreamWriterBuilder; private Builder(SchemaAwareStreamWriter.Builder schemaAwareStreamWriterBuilder) { - this.schemaAwareStreamWriterBuilder = - schemaAwareStreamWriterBuilder.setTraceIdBase("JsonWriter"); + this.schemaAwareStreamWriterBuilder = schemaAwareStreamWriterBuilder.setClientId(CLIENT_ID); } /** diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/SchemaAwareStreamWriter.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/SchemaAwareStreamWriter.java index c95ee72ae8..51df76a06c 100644 --- a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/SchemaAwareStreamWriter.java +++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/SchemaAwareStreamWriter.java @@ -94,7 +94,6 @@ private SchemaAwareStreamWriter(Builder builder) builder.executorProvider, builder.endpoint, builder.flowControlSettings, - builder.traceIdBase, builder.traceId, builder.compressorName, builder.retrySettings); @@ -102,6 +101,7 @@ private SchemaAwareStreamWriter(Builder builder) streamWriterBuilder.setLocation(builder.location); streamWriterBuilder.setDefaultMissingValueInterpretation( builder.defaultMissingValueInterpretation); + streamWriterBuilder.setClientId(builder.clientId); this.streamWriter = streamWriterBuilder.build(); this.streamName = builder.streamName; this.tableSchema = builder.tableSchema; @@ -282,7 +282,6 @@ private void setStreamWriterSettings( @Nullable ExecutorProvider executorProvider, @Nullable String endpoint, @Nullable FlowControlSettings flowControlSettings, - @Nullable String traceIdBase, @Nullable String traceId, @Nullable String compressorName, @Nullable RetrySettings retrySettings) { @@ -298,18 +297,8 @@ private void setStreamWriterSettings( if (endpoint != null) { streamWriterBuilder.setEndpoint(endpoint); } - if (traceIdBase != null) { - if (traceId != null) { - streamWriterBuilder.setTraceId(traceIdBase + "_" + traceId); - } else { - streamWriterBuilder.setTraceId(traceIdBase + ":null"); - } - } else { - if (traceId != null) { - streamWriterBuilder.setTraceId("SchemaAwareStreamWriter_" + traceId); - } else { - streamWriterBuilder.setTraceId("SchemaAwareStreamWriter:null"); - } + if (traceId != null) { + streamWriterBuilder.setTraceId(traceId); } if (flowControlSettings != null) { if (flowControlSettings.getMaxOutstandingRequestBytes() != null) { @@ -445,6 +434,7 @@ public static final class Builder { private AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation = MissingValueInterpretation.MISSING_VALUE_INTERPRETATION_UNSPECIFIED; + private String clientId; private static final String streamPatternString = "(projects/[^/]+/datasets/[^/]+/tables/[^/]+)/streams/[^/]+"; @@ -581,14 +571,8 @@ public Builder setTraceId(String traceId) { return this; } - /** - * Setter for a traceIdBase to help identify traffic origin. - * - * @param traceIdBase - * @return Builder - */ - public Builder setTraceIdBase(String traceIdBase) { - this.traceIdBase = Preconditions.checkNotNull(traceIdBase, "TraceIdBase is null."); + Builder setClientId(String clientId) { + this.clientId = Preconditions.checkNotNull(clientId, "ClientId is null."); return this; } diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java index 09988c4f1b..e7107785df 100644 --- a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java +++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java @@ -229,7 +229,7 @@ private StreamWriter(Builder builder) throws IOException { builder.maxInflightBytes, builder.maxRetryDuration, builder.limitExceededBehavior, - builder.traceId, + builder.getFullTraceId(), builder.compressorName, clientSettings, builder.retrySettings)); @@ -295,7 +295,7 @@ private StreamWriter(Builder builder) throws IOException { builder.maxInflightBytes, builder.maxRetryDuration, builder.limitExceededBehavior, - builder.traceId, + builder.getFullTraceId(), builder.compressorName, client.getSettings(), builder.retrySettings); @@ -376,12 +376,12 @@ static BigQueryWriteSettings getBigQueryWriteSettings(Builder builder) throws IO private void validateFetchedConnectonPool(StreamWriter.Builder builder) { String storedTraceId = this.singleConnectionOrConnectionPool.connectionWorkerPool().getTraceId(); - if (!Objects.equals(storedTraceId, builder.traceId)) { + if (!Objects.equals(storedTraceId, builder.getFullTraceId())) { throw new IllegalArgumentException( String.format( "Trace id used for the same connection pool for the same location must be the same, " + "however stored trace id is %s, and expected trace id is %s.", - storedTraceId, builder.traceId)); + storedTraceId, builder.getFullTraceId())); } FlowController.LimitExceededBehavior storedLimitExceededBehavior = singleConnectionOrConnectionPool.connectionWorkerPool().limitExceededBehavior(); @@ -629,6 +629,8 @@ public static final class Builder { private String traceId = null; + private String clientId = "java-streamwriter"; + private TableSchema updatedTableSchema = null; private String location = null; @@ -730,6 +732,15 @@ public Builder setTraceId(String traceId) { return this; } + /** + * Sets the client id of the writer, for example, JsonStreamWriter has the client id of + * "java-jsonwriter". + */ + Builder setClientId(String clientId) { + this.clientId = clientId; + return this; + } + /** Location of the table this stream writer is targeting. */ public Builder setLocation(String location) { this.location = location; @@ -811,5 +822,13 @@ public Builder setRetrySettings(RetrySettings retrySettings) { public StreamWriter build() throws IOException { return new StreamWriter(this); } + + String getFullTraceId() { + if (traceId == null) { + return clientId; + } else { + return clientId + " " + traceId; + } + } } } 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 index 3662330bae..d3a25510f8 100644 --- 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 @@ -237,7 +237,7 @@ public void testSingleAppendSimpleJson() throws Exception { .getSerializedRows(0), expectedProto.toByteString()); assertEquals( - testBigQueryWrite.getAppendRequests().get(0).getTraceId(), "JsonWriter_test:empty"); + "java-jsonwriter test:empty", testBigQueryWrite.getAppendRequests().get(0).getTraceId()); } } @@ -257,7 +257,7 @@ public void testFlexibleColumnAppend() throws Exception { jsonArr.put(flexible); try (JsonStreamWriter writer = - getTestJsonStreamWriterBuilder(TEST_STREAM, tableSchema).setTraceId("test:empty").build()) { + getTestJsonStreamWriterBuilder(TEST_STREAM, tableSchema).build()) { testBigQueryWrite.addResponse( AppendRowsResponse.newBuilder() @@ -284,8 +284,7 @@ public void testFlexibleColumnAppend() throws Exception { .getRows() .getSerializedRows(0), expectedProto.toByteString()); - assertEquals( - testBigQueryWrite.getAppendRequests().get(0).getTraceId(), "JsonWriter_test:empty"); + assertEquals("java-jsonwriter", testBigQueryWrite.getAppendRequests().get(0).getTraceId()); } } @@ -507,7 +506,7 @@ public void testSingleAppendMultipleSimpleJson() throws Exception { .getProtoRows() .getRows() .getSerializedRowsCount()); - assertEquals(testBigQueryWrite.getAppendRequests().get(0).getTraceId(), "JsonWriter:null"); + assertEquals("java-jsonwriter", testBigQueryWrite.getAppendRequests().get(0).getTraceId()); for (int i = 0; i < 4; i++) { assertEquals( testBigQueryWrite diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java index ce51601394..d727cf4fd5 100644 --- a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java +++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java @@ -286,13 +286,13 @@ private void verifyAppendRequests(long appendCount) { if (i == 0) { // First request received by server should have schema and stream name. assertTrue(serverRequest.getProtoRows().hasWriterSchema()); - assertEquals(serverRequest.getWriteStream(), TEST_STREAM_1); - assertEquals(serverRequest.getTraceId(), TEST_TRACE_ID); + assertEquals(TEST_STREAM_1, serverRequest.getWriteStream()); + assertEquals("java-streamwriter " + TEST_TRACE_ID, serverRequest.getTraceId()); } else { // Following request should not have schema and stream name. assertFalse(serverRequest.getProtoRows().hasWriterSchema()); - assertEquals(serverRequest.getWriteStream(), ""); - assertEquals(serverRequest.getTraceId(), ""); + assertEquals("", serverRequest.getWriteStream()); + assertEquals("", serverRequest.getTraceId()); } } }