Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: StreamWriterV2 will handle schema/streamName attachment #877

Merged
merged 3 commits into from Feb 25, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
Expand Up @@ -296,6 +296,7 @@ public void close() {
* It takes requests from waiting queue and sends them to server.
*/
private void appendLoop() {
boolean isFirstRequestInConnection = true;
Deque<AppendRequestAndResponse> localQueue = new LinkedList<AppendRequestAndResponse>();
while (!waitingQueueDrained()) {
this.lock.lock();
Expand All @@ -322,7 +323,11 @@ private void appendLoop() {

// TODO: Add reconnection here.
while (!localQueue.isEmpty()) {
this.streamConnection.send(localQueue.pollFirst().message);
AppendRowsRequest preparedRequest =
prepareRequestBasedOnPosition(
localQueue.pollFirst().message, isFirstRequestInConnection);
this.streamConnection.send(preparedRequest);
isFirstRequestInConnection = false;
yayi-google marked this conversation as resolved.
Show resolved Hide resolved
}
}

Expand Down Expand Up @@ -371,6 +376,18 @@ private void waitForDoneCallback() {
}
}

private AppendRowsRequest prepareRequestBasedOnPosition(
AppendRowsRequest original, boolean isFirstRequest) {
AppendRowsRequest.Builder requestBuilder = original.toBuilder();
if (isFirstRequest) {
requestBuilder.setWriteStream(this.streamName);
} else {
requestBuilder.clearWriteStream();
requestBuilder.getProtoRowsBuilder().clearWriterSchema();
}
return requestBuilder.build();
}

private void cleanupInflightRequests() {
Throwable finalStatus;
Deque<AppendRequestAndResponse> localQueue = new LinkedList<AppendRequestAndResponse>();
Expand Down
Expand Up @@ -16,6 +16,7 @@
package com.google.cloud.bigquery.storage.v1beta2;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;

Expand Down Expand Up @@ -192,6 +193,19 @@ public void testAppendSuccess() throws Exception {
}
assertEquals(appendCount, testBigQueryWrite.getAppendRequests().size());

for (int i = 0; i < appendCount; i++) {
AppendRowsRequest serverRequest = testBigQueryWrite.getAppendRequests().get(i);
if (i == 0) {
// First request received by server should have schema and stream name.
assertTrue(serverRequest.getProtoRows().hasWriterSchema());
assertEquals(serverRequest.getWriteStream(), TEST_STREAM);
} else {
// Following request should not have schema and stream name.
assertFalse(serverRequest.getProtoRows().hasWriterSchema());
assertEquals(serverRequest.getWriteStream(), "");
}
}

writer.close();
}

Expand Down