feat: update StreamWriterV2 to support trace id (#895)

* feat: StreamWriterV2 will attach stream name in the first request in the connection and remove stream name and schema in the following ones. * feat: StreamWriterV2 will attach stream name in the first request in the connection and remove stream name and schema in the following ones * feat: StreamWriterV2 to support traceId * Enforce traceId to be the format of A:B

Author: yayi-google

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2.java

@@ -41,8 +41,6 @@
 /**
  * A BigQuery Stream Writer that can be used to write data into BigQuery Table.
  *
- * <p>TODO: Attach traceId.
- *
  * <p>TODO: Support batching.
  *
  * <p>TODO: Support schema change.
@@ -74,6 +72,11 @@ public class StreamWriterV2 implements AutoCloseable {
  */
  private final long maxInflightBytes;
 
+ /*
+ * TraceId for debugging purpose.
+ */
+ private final String traceId;
+
  /*
  * Tracks current inflight requests in the stream.
  */
@@ -143,6 +146,7 @@ private StreamWriterV2(Builder builder) throws IOException {
  this.writerSchema = builder.writerSchema;
  this.maxInflightRequests = builder.maxInflightRequest;
  this.maxInflightBytes = builder.maxInflightBytes;
+ this.traceId = builder.traceId;
  this.waitingRequestQueue = new LinkedList<AppendRequestAndResponse>();
  this.inflightRequestQueue = new LinkedList<AppendRequestAndResponse>();
  if (builder.client == null) {
@@ -433,6 +437,9 @@ private AppendRowsRequest prepareRequestBasedOnPosition(
  requestBuilder.getProtoRowsBuilder().setWriterSchema(this.writerSchema);
  }
  requestBuilder.setWriteStream(this.streamName);
+ if (this.traceId != null) {
+ requestBuilder.setTraceId(this.traceId);
+ }
  } else {
  requestBuilder.clearWriteStream();
  requestBuilder.getProtoRowsBuilder().clearWriterSchema();
@@ -539,6 +546,8 @@ public static final class Builder {
  private CredentialsProvider credentialsProvider =
  BigQueryWriteSettings.defaultCredentialsProviderBuilder().build();
 
+ private String traceId = null;
+
  private Builder(String streamName) {
  this.streamName = Preconditions.checkNotNull(streamName);
  this.client = null;
@@ -591,6 +600,20 @@ public Builder setCredentialsProvider(CredentialsProvider credentialsProvider) {
  return this;
  }
 
+ /**
+ * 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;
+ }
+
  /** Builds the {@code StreamWriterV2}. */
  public StreamWriterV2 build() throws IOException {
  return new StreamWriterV2(this);

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2Test.java

@@ -54,6 +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 = "DATAFLOW:job_id";
  private FakeScheduledExecutorService fakeExecutor;
  private FakeBigQueryWrite testBigQueryWrite;
  private static MockServiceHelper serviceHelper;
@@ -84,7 +85,7 @@ public void tearDown() throws Exception {
  }
 
  private StreamWriterV2 getTestStreamWriterV2() throws IOException {
- return StreamWriterV2.newBuilder(TEST_STREAM, client).build();
+ return StreamWriterV2.newBuilder(TEST_STREAM, client).setTraceId(TEST_TRACE_ID).build();
  }
 
  private ProtoSchema createProtoSchema() {
@@ -184,10 +185,12 @@ private void verifyAppendRequests(long appendCount) {
  // First request received by server should have schema and stream name.
  assertTrue(serverRequest.getProtoRows().hasWriterSchema());
  assertEquals(serverRequest.getWriteStream(), TEST_STREAM);
+ assertEquals(serverRequest.getTraceId(), TEST_TRACE_ID);
  } else {
  // Following request should not have schema and stream name.
  assertFalse(serverRequest.getProtoRows().hasWriterSchema());
  assertEquals(serverRequest.getWriteStream(), "");
+ assertEquals(serverRequest.getTraceId(), "");
  }
  }
  }
@@ -209,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++) {
@@ -269,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();