diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2.java index 51296e9096..33abf53bf4 100644 --- a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2.java +++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2.java @@ -600,8 +600,16 @@ public Builder setCredentialsProvider(CredentialsProvider credentialsProvider) { return this; } - /** TraceId for debuging purpose. */ + /** + * Sets traceId for debuging purpose. TraceId must follow the format of + * CustomerDomain:DebugString, e.g. DATAFLOW:job_id_x. + */ public Builder setTraceId(String traceId) { + int colonIndex = traceId.indexOf(':'); + if (colonIndex == -1 || colonIndex == 0 || colonIndex == traceId.length() - 1) { + throw new IllegalArgumentException( + "TraceId must follow the format of A:B. Actual:" + traceId); + } this.traceId = traceId; return this; } diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2Test.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2Test.java index fed7ca778d..69aa4341a0 100644 --- a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2Test.java +++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2Test.java @@ -54,7 +54,7 @@ public class StreamWriterV2Test { private static final Logger log = Logger.getLogger(StreamWriterV2Test.class.getName()); private static final String TEST_STREAM = "projects/p/datasets/d/tables/t/streams/s"; - private static final String TEST_TRACE_ID = "test trace id"; + private static final String TEST_TRACE_ID = "DATAFLOW:job_id"; private FakeScheduledExecutorService fakeExecutor; private FakeBigQueryWrite testBigQueryWrite; private static MockServiceHelper serviceHelper; @@ -212,7 +212,10 @@ public void testBuildBigQueryWriteClientInWriter() throws Exception { @Test public void testAppendWithRowsSuccess() throws Exception { StreamWriterV2 writer = - StreamWriterV2.newBuilder(TEST_STREAM, client).setWriterSchema(createProtoSchema()).build(); + StreamWriterV2.newBuilder(TEST_STREAM, client) + .setWriterSchema(createProtoSchema()) + .setTraceId(TEST_TRACE_ID) + .build(); long appendCount = 100; for (int i = 0; i < appendCount; i++) { @@ -272,6 +275,34 @@ public void run() throws Throwable { assertTrue(ex.getStatus().getDescription().contains("Writer schema must be provided")); } + @Test + public void testInvalidTraceId() throws Exception { + assertThrows( + IllegalArgumentException.class, + new ThrowingRunnable() { + @Override + public void run() throws Throwable { + StreamWriterV2.newBuilder(TEST_STREAM).setTraceId("abc"); + } + }); + assertThrows( + IllegalArgumentException.class, + new ThrowingRunnable() { + @Override + public void run() throws Throwable { + StreamWriterV2.newBuilder(TEST_STREAM).setTraceId("abc:"); + } + }); + assertThrows( + IllegalArgumentException.class, + new ThrowingRunnable() { + @Override + public void run() throws Throwable { + StreamWriterV2.newBuilder(TEST_STREAM).setTraceId(":abc"); + } + }); + } + @Test public void testAppendSuccessAndConnectionError() throws Exception { StreamWriterV2 writer = getTestStreamWriterV2();