Lokasi ngalangkungan proxy:   [ UP ]  
[Ngawartoskeun bug]   [Panyetelan cookie]                
Skip to content
This repository was archived by the owner on Feb 24, 2026. It is now read-only.

Commit 36964a3

Browse files
GaoleMenggcf-owl-bot[bot]
GaoleMeng
and
gcf-owl-bot[bot]
authored
fix: bug fix for streamWriter & jsonStreamWriter (#2122)
* feat: add public api to stream writer to set the maximum wait time * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * modify back the readme change from owl post processor * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * fix: Reduce the timeout to 5 minutes for the requests wait time in queue. Since in write api server side we have total timeout of 2 minutes, it does not make sense to wait 15 minutes to determine whether we have met dead connection, let's reduce the timeout here * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * fix: 1.disable refresh of stream writer when the table schema is explicitly provided 2. fix location string matching for multiplexing --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent a4e1131 commit 36964a3

4 files changed

Lines changed: 160 additions & 19 deletions

File tree

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -349,15 +349,15 @@ public void run(Throwable finalStatus) {
349349

350350
/** Schedules the writing of rows at given offset. */
351351
ApiFuture<AppendRowsResponse> append(StreamWriter streamWriter, ProtoRows rows, long offset) {
352-
if (this.location != null && this.location != streamWriter.getLocation()) {
352+
if (this.location != null && !this.location.equals(streamWriter.getLocation())) {
353353
throw new StatusRuntimeException(
354354
Status.fromCode(Code.INVALID_ARGUMENT)
355355
.withDescription(
356356
"StreamWriter with location "
357357
+ streamWriter.getLocation()
358358
+ " is scheduled to use a connection with location "
359359
+ this.location));
360-
} else if (this.location == null && streamWriter.getStreamName() != this.streamName) {
360+
} else if (this.location == null && !streamWriter.getStreamName().equals(this.streamName)) {
361361
// Location is null implies this is non-multiplexed connection.
362362
throw new StatusRuntimeException(
363363
Status.fromCode(Code.INVALID_ARGUMENT)

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

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,10 @@ public class SchemaAwareStreamWriter<T> implements AutoCloseable {
5454
private TableSchema tableSchema;
5555
private ProtoSchema protoSchema;
5656

57+
// During some sitaution we want to skip stream writer refresh for updated schema. e.g. when
58+
// the user provides the table schema, we should always use that schema.
59+
private final boolean skipRefreshStreamWriter;
60+
5761
/**
5862
* Constructs the SchemaAwareStreamWriter
5963
*
@@ -87,6 +91,7 @@ private SchemaAwareStreamWriter(Builder<T> builder)
8791
this.tableSchema = builder.tableSchema;
8892
this.toProtoConverter = builder.toProtoConverter;
8993
this.ignoreUnknownFields = builder.ignoreUnknownFields;
94+
this.skipRefreshStreamWriter = builder.skipRefreshStreamWriter;
9095
}
9196

9297
/**
@@ -125,6 +130,10 @@ private Message buildMessage(T item)
125130
return this.toProtoConverter.convertToProtoMessage(
126131
this.descriptor, this.tableSchema, item, ignoreUnknownFields);
127132
} catch (Exceptions.DataHasUnknownFieldException ex) {
133+
// Directly return error when stream writer refresh is disabled.
134+
if (this.skipRefreshStreamWriter) {
135+
throw ex;
136+
}
128137
LOG.warning(
129138
"Saw unknown field "
130139
+ ex.getFieldName()
@@ -157,7 +166,7 @@ public ApiFuture<AppendRowsResponse> append(Iterable<T> items, long offset)
157166
// Handle schema updates in a Thread-safe way by locking down the operation
158167
synchronized (this) {
159168
// Create a new stream writer internally if a new updated schema is reported from backend.
160-
if (this.streamWriter.getUpdatedSchema() != null) {
169+
if (!this.skipRefreshStreamWriter && this.streamWriter.getUpdatedSchema() != null) {
161170
refreshWriter(this.streamWriter.getUpdatedSchema());
162171
}
163172

@@ -404,6 +413,8 @@ public static final class Builder<T> {
404413
private final BigQueryWriteClient client;
405414
private final TableSchema tableSchema;
406415

416+
private final boolean skipRefreshStreamWriter;
417+
407418
private final ToProtoConverter<T> toProtoConverter;
408419
private TransportChannelProvider channelProvider;
409420
private CredentialsProvider credentialsProvider;
@@ -459,11 +470,12 @@ private Builder(
459470
.build();
460471

461472
WriteStream writeStream = this.client.getWriteStream(writeStreamRequest);
462-
463473
this.tableSchema = writeStream.getTableSchema();
464474
this.location = writeStream.getLocation();
475+
this.skipRefreshStreamWriter = false;
465476
} else {
466477
this.tableSchema = tableSchema;
478+
this.skipRefreshStreamWriter = true;
467479
}
468480
this.toProtoConverter = toProtoConverter;
469481
}

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

Lines changed: 142 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,13 @@ private JsonStreamWriter.Builder getTestJsonStreamWriterBuilder(
160160
.setExecutorProvider(InstantiatingExecutorProvider.newBuilder().build());
161161
}
162162

163+
private JsonStreamWriter.Builder getTestJsonStreamWriterBuilder(String testStream) {
164+
return JsonStreamWriter.newBuilder(testStream, client)
165+
.setChannelProvider(channelProvider)
166+
.setCredentialsProvider(NoCredentialsProvider.create())
167+
.setExecutorProvider(InstantiatingExecutorProvider.newBuilder().build());
168+
}
169+
163170
@Test
164171
public void testTwoParamNewBuilder_nullSchema() {
165172
try {
@@ -658,8 +665,13 @@ public void run() throws Throwable {
658665

659666
@Test
660667
public void testSimpleSchemaUpdate() throws Exception {
661-
try (JsonStreamWriter writer =
662-
getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA).build()) {
668+
testBigQueryWrite.addResponse(
669+
WriteStream.newBuilder()
670+
.setName(TEST_STREAM)
671+
.setTableSchema(TABLE_SCHEMA)
672+
.setLocation("us")
673+
.build());
674+
try (JsonStreamWriter writer = getTestJsonStreamWriterBuilder(TEST_STREAM).build()) {
663675
testBigQueryWrite.addResponse(
664676
AppendRowsResponse.newBuilder()
665677
.setAppendResult(
@@ -722,7 +734,6 @@ public void testSimpleSchemaUpdate() throws Exception {
722734
updatedFoo.put("bar", "bbb");
723735
JSONArray updatedJsonArr = new JSONArray();
724736
updatedJsonArr.put(updatedFoo);
725-
726737
ApiFuture<AppendRowsResponse> appendFuture4 = writer.append(updatedJsonArr);
727738

728739
assertEquals(3L, appendFuture4.get().getAppendResult().getOffset().getValue());
@@ -751,6 +762,88 @@ public void testSimpleSchemaUpdate() throws Exception {
751762
}
752763
}
753764

765+
@Test
766+
public void testSimpleSchemaUpdate_skipRefreshWriterIfSchemaProvided() throws Exception {
767+
testBigQueryWrite.addResponse(
768+
WriteStream.newBuilder()
769+
.setName(TEST_STREAM)
770+
.setTableSchema(TABLE_SCHEMA)
771+
.setLocation("us")
772+
.build());
773+
try (JsonStreamWriter writer =
774+
getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA).build()) {
775+
testBigQueryWrite.addResponse(
776+
AppendRowsResponse.newBuilder()
777+
.setAppendResult(
778+
AppendRowsResponse.AppendResult.newBuilder().setOffset(Int64Value.of(0)).build())
779+
.setUpdatedSchema(UPDATED_TABLE_SCHEMA)
780+
.build());
781+
testBigQueryWrite.addResponse(createAppendResponse(1));
782+
testBigQueryWrite.addResponse(createAppendResponse(2));
783+
testBigQueryWrite.addResponse(createAppendResponse(3));
784+
// First append
785+
JSONObject foo = new JSONObject();
786+
foo.put("foo", "aaa");
787+
JSONArray jsonArr = new JSONArray();
788+
jsonArr.put(foo);
789+
790+
ApiFuture<AppendRowsResponse> appendFuture1 = writer.append(jsonArr);
791+
ApiFuture<AppendRowsResponse> appendFuture2 = writer.append(jsonArr);
792+
ApiFuture<AppendRowsResponse> appendFuture3 = writer.append(jsonArr);
793+
794+
assertEquals(0L, appendFuture1.get().getAppendResult().getOffset().getValue());
795+
assertEquals(1L, appendFuture2.get().getAppendResult().getOffset().getValue());
796+
assertEquals(
797+
1,
798+
testBigQueryWrite
799+
.getAppendRequests()
800+
.get(0)
801+
.getProtoRows()
802+
.getRows()
803+
.getSerializedRowsCount());
804+
assertEquals(
805+
testBigQueryWrite
806+
.getAppendRequests()
807+
.get(0)
808+
.getProtoRows()
809+
.getRows()
810+
.getSerializedRows(0),
811+
FooType.newBuilder().setFoo("aaa").build().toByteString());
812+
813+
assertEquals(2L, appendFuture3.get().getAppendResult().getOffset().getValue());
814+
assertEquals(
815+
1,
816+
testBigQueryWrite
817+
.getAppendRequests()
818+
.get(1)
819+
.getProtoRows()
820+
.getRows()
821+
.getSerializedRowsCount());
822+
assertEquals(
823+
testBigQueryWrite
824+
.getAppendRequests()
825+
.get(1)
826+
.getProtoRows()
827+
.getRows()
828+
.getSerializedRows(0),
829+
FooType.newBuilder().setFoo("aaa").build().toByteString());
830+
831+
// Second append with updated schema.
832+
JSONObject updatedFoo = new JSONObject();
833+
updatedFoo.put("foo", "aaa");
834+
updatedFoo.put("bar", "bbb");
835+
JSONArray updatedJsonArr = new JSONArray();
836+
updatedJsonArr.put(updatedFoo);
837+
838+
// Schema update will not happen for writer that has schema explicitly provided.
839+
assertThrows(
840+
AppendSerializationError.class,
841+
() -> {
842+
ApiFuture<AppendRowsResponse> appendFuture4 = writer.append(updatedJsonArr);
843+
});
844+
}
845+
}
846+
754847
@Test
755848
public void testWithoutIgnoreUnknownFieldsUpdateImmeidateSuccess() throws Exception {
756849
TableSchema tableSchema = TableSchema.newBuilder().addFields(0, TEST_INT).build();
@@ -764,6 +857,10 @@ public void testWithoutIgnoreUnknownFieldsUpdateImmeidateSuccess() throws Except
764857
.setType(TableFieldSchema.Type.STRING)
765858
.setMode(Mode.NULLABLE))
766859
.build();
860+
861+
// GetWriteStream is called once and got the updated schema
862+
testBigQueryWrite.addResponse(
863+
WriteStream.newBuilder().setName(TEST_STREAM).setTableSchema(tableSchema).build());
767864
// GetWriteStream is called once and the writer is fixed to accept unknown fields.
768865
testBigQueryWrite.addResponse(
769866
WriteStream.newBuilder().setName(TEST_STREAM).setTableSchema(updatedSchema).build());
@@ -772,8 +869,7 @@ public void testWithoutIgnoreUnknownFieldsUpdateImmeidateSuccess() throws Except
772869
.setAppendResult(
773870
AppendRowsResponse.AppendResult.newBuilder().setOffset(Int64Value.of(0)).build())
774871
.build());
775-
try (JsonStreamWriter writer =
776-
getTestJsonStreamWriterBuilder(TEST_STREAM, tableSchema).build()) {
872+
try (JsonStreamWriter writer = getTestJsonStreamWriterBuilder(TEST_STREAM).build()) {
777873
JSONObject foo = new JSONObject();
778874
foo.put("test_int", 10);
779875
JSONObject bar = new JSONObject();
@@ -800,15 +896,16 @@ public void testWithoutIgnoreUnknownFieldsUpdateSecondSuccess() throws Exception
800896
.setMode(Mode.NULLABLE))
801897
.build();
802898
// GetWriteStream is called once and got the updated schema
899+
testBigQueryWrite.addResponse(
900+
WriteStream.newBuilder().setName(TEST_STREAM).setTableSchema(TABLE_SCHEMA).build());
803901
testBigQueryWrite.addResponse(
804902
WriteStream.newBuilder().setName(TEST_STREAM).setTableSchema(updatedSchema).build());
805903
testBigQueryWrite.addResponse(
806904
AppendRowsResponse.newBuilder()
807905
.setAppendResult(
808906
AppendRowsResponse.AppendResult.newBuilder().setOffset(Int64Value.of(0)).build())
809907
.build());
810-
try (JsonStreamWriter writer =
811-
getTestJsonStreamWriterBuilder(TEST_STREAM, tableSchema).build()) {
908+
try (JsonStreamWriter writer = getTestJsonStreamWriterBuilder(TEST_STREAM).build()) {
812909
JSONObject foo = new JSONObject();
813910
foo.put("test_int", 10);
814911
JSONObject bar = new JSONObject();
@@ -826,15 +923,28 @@ public void testSchemaUpdateInMultiplexing_singleConnection() throws Exception {
826923
// Set min connection count to be 1 to force sharing connection.
827924
ConnectionWorkerPool.setOptions(
828925
Settings.builder().setMinConnectionsPerRegion(1).setMaxConnectionsPerRegion(1).build());
926+
// GetWriteStream is called twice and got the updated schema
927+
testBigQueryWrite.addResponse(
928+
WriteStream.newBuilder()
929+
.setName(TEST_STREAM)
930+
.setTableSchema(TABLE_SCHEMA)
931+
.setLocation("us")
932+
.build());
933+
testBigQueryWrite.addResponse(
934+
WriteStream.newBuilder()
935+
.setName(TEST_STREAM)
936+
.setTableSchema(TABLE_SCHEMA_2)
937+
.setLocation("us")
938+
.build());
829939
// The following two writers have different stream name and schema, but will share the same
830940
// connection .
831941
JsonStreamWriter writer1 =
832-
getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA)
942+
getTestJsonStreamWriterBuilder(TEST_STREAM)
833943
.setEnableConnectionPool(true)
834944
.setLocation("us")
835945
.build();
836946
JsonStreamWriter writer2 =
837-
getTestJsonStreamWriterBuilder(TEST_STREAM_2, TABLE_SCHEMA_2)
947+
getTestJsonStreamWriterBuilder(TEST_STREAM_2)
838948
.setEnableConnectionPool(true)
839949
.setLocation("us")
840950
.build();
@@ -911,14 +1021,27 @@ public void testSchemaUpdateInMultiplexing_multipleWriterForSameStreamName() thr
9111021
ConnectionWorkerPool.setOptions(
9121022
Settings.builder().setMinConnectionsPerRegion(1).setMaxConnectionsPerRegion(1).build());
9131023

1024+
// GetWriteStream is called twice and got the updated schema
1025+
testBigQueryWrite.addResponse(
1026+
WriteStream.newBuilder()
1027+
.setName(TEST_STREAM)
1028+
.setTableSchema(TABLE_SCHEMA)
1029+
.setLocation("us")
1030+
.build());
1031+
testBigQueryWrite.addResponse(
1032+
WriteStream.newBuilder()
1033+
.setName(TEST_STREAM)
1034+
.setTableSchema(TABLE_SCHEMA)
1035+
.setLocation("us")
1036+
.build());
9141037
// Create two writers writing to the same stream.
9151038
JsonStreamWriter writer1 =
916-
getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA)
1039+
getTestJsonStreamWriterBuilder(TEST_STREAM)
9171040
.setEnableConnectionPool(true)
9181041
.setLocation("us")
9191042
.build();
9201043
JsonStreamWriter writer2 =
921-
getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA)
1044+
getTestJsonStreamWriterBuilder(TEST_STREAM)
9221045
.setEnableConnectionPool(true)
9231046
.setLocation("us")
9241047
.build();
@@ -987,10 +1110,16 @@ public void testSchemaUpdateInMultiplexing_IgnoreUpdateIfTimeStampNewer() throws
9871110
// Set min connection count to be 1 to force sharing connection.
9881111
ConnectionWorkerPool.setOptions(
9891112
Settings.builder().setMinConnectionsPerRegion(1).setMaxConnectionsPerRegion(1).build());
1113+
testBigQueryWrite.addResponse(
1114+
WriteStream.newBuilder()
1115+
.setName(TEST_STREAM)
1116+
.setTableSchema(TABLE_SCHEMA)
1117+
.setLocation("us")
1118+
.build());
9901119
// The following two writers have different stream name and schema, but will share the same
991-
// connection .
1120+
// connection.
9921121
JsonStreamWriter writer1 =
993-
getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA)
1122+
getTestJsonStreamWriterBuilder(TEST_STREAM)
9941123
.setEnableConnectionPool(true)
9951124
.setLocation("us")
9961125
.build();

‎google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -785,7 +785,7 @@ public void testJsonStreamWriterSchemaUpdate()
785785
WriteStream.newBuilder().setType(WriteStream.Type.COMMITTED).build())
786786
.build());
787787
try (JsonStreamWriter jsonStreamWriter =
788-
JsonStreamWriter.newBuilder(writeStream.getName(), writeStream.getTableSchema()).build()) {
788+
JsonStreamWriter.newBuilder(writeStream.getName(), client).build()) {
789789
// write the 1st row
790790
JSONObject foo = new JSONObject();
791791
foo.put("col1", "aaa");
@@ -895,7 +895,7 @@ public void testJsonStreamWriterSchemaUpdateConcurrent()
895895

896896
// Start writing using the JsonWriter
897897
try (JsonStreamWriter jsonStreamWriter =
898-
JsonStreamWriter.newBuilder(writeStream.getName(), writeStream.getTableSchema()).build()) {
898+
JsonStreamWriter.newBuilder(writeStream.getName(), client).build()) {
899899
int numberOfThreads = 5;
900900
ExecutorService streamTaskExecutor = Executors.newFixedThreadPool(5);
901901
CountDownLatch latch = new CountDownLatch(numberOfThreads);

0 commit comments

Comments
 (0)