Skip to content
This repository was archived by the owner on Feb 24, 2026. It is now read-only.

Commit 9c1ed55

Browse files
yirutangyirutang
authored
fix: It seems GAPIC is not generated expected headers for bidi streaming client lib, apply a temp fix to unblock customers (#1017)
* repro * . * . * . * . * . * . * . * . * format * .
1 parent 827764e commit 9c1ed55

File tree

3 files changed

+160
-0
lines changed

3 files changed

+160
-0
lines changed

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import com.google.api.core.ApiFuture;
1919
import com.google.api.core.SettableApiFuture;
2020
import com.google.api.gax.core.CredentialsProvider;
21+
import com.google.api.gax.rpc.FixedHeaderProvider;
2122
import com.google.api.gax.rpc.TransportChannelProvider;
2223
import com.google.cloud.bigquery.storage.v1beta2.AppendRowsRequest.ProtoData;
2324
import com.google.cloud.bigquery.storage.v1beta2.StreamConnection.DoneCallback;
@@ -160,6 +161,10 @@ private StreamWriterV2(Builder builder) throws IOException {
160161
.setCredentialsProvider(builder.credentialsProvider)
161162
.setTransportChannelProvider(builder.channelProvider)
162163
.setEndpoint(builder.endpoint)
164+
// (b/185842996): Temporily fix this by explicitly providing the header.
165+
.setHeaderProvider(
166+
FixedHeaderProvider.create(
167+
"x-goog-request-params", "write_stream=" + this.streamName))
163168
.build();
164169
this.client = BigQueryWriteClient.create(stubSettings);
165170
this.ownsBigQueryWriteClient = true;

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,15 +49,18 @@ public class ITBigQueryWriteManualClientTest {
4949
private static final Logger LOG =
5050
Logger.getLogger(ITBigQueryWriteManualClientTest.class.getName());
5151
private static final String DATASET = RemoteBigQueryHelper.generateDatasetName();
52+
private static final String DATASET_EU = RemoteBigQueryHelper.generateDatasetName();
5253
private static final String TABLE = "testtable";
5354
private static final String TABLE2 = "complicatedtable";
5455
private static final String DESCRIPTION = "BigQuery Write Java manual client test dataset";
5556

5657
private static BigQueryWriteClient client;
5758
private static TableInfo tableInfo;
5859
private static TableInfo tableInfo2;
60+
private static TableInfo tableInfoEU;
5961
private static String tableId;
6062
private static String tableId2;
63+
private static String tableIdEU;
6164
private static BigQuery bigquery;
6265

6366
@BeforeClass
@@ -110,6 +113,25 @@ public static void beforeClass() throws IOException {
110113
String.format(
111114
"projects/%s/datasets/%s/tables/%s",
112115
ServiceOptions.getDefaultProjectId(), DATASET, TABLE2);
116+
DatasetInfo datasetInfoEU =
117+
DatasetInfo.newBuilder(/* datasetId = */ DATASET_EU)
118+
.setLocation("EU")
119+
.setDescription(DESCRIPTION)
120+
.build();
121+
bigquery.create(datasetInfoEU);
122+
tableInfoEU =
123+
TableInfo.newBuilder(
124+
TableId.of(DATASET_EU, TABLE),
125+
StandardTableDefinition.of(
126+
Schema.of(
127+
com.google.cloud.bigquery.Field.newBuilder("foo", LegacySQLTypeName.STRING)
128+
.build())))
129+
.build();
130+
tableIdEU =
131+
String.format(
132+
"projects/%s/datasets/%s/tables/%s",
133+
ServiceOptions.getDefaultProjectId(), DATASET_EU, TABLE);
134+
bigquery.create(tableInfoEU);
113135
}
114136

115137
@AfterClass
@@ -206,6 +228,54 @@ public void testBatchWriteWithCommittedStream()
206228
}
207229
}
208230

231+
ProtoRows CreateProtoRows(String[] messages) {
232+
ProtoRows.Builder rows = ProtoRows.newBuilder();
233+
for (String message : messages) {
234+
FooType foo = FooType.newBuilder().setFoo(message).build();
235+
rows.addSerializedRows(foo.toByteString());
236+
}
237+
return rows.build();
238+
}
239+
240+
@Test
241+
public void testBatchWriteWithCommittedStreamEU()
242+
throws IOException, InterruptedException, ExecutionException {
243+
WriteStream writeStream =
244+
client.createWriteStream(
245+
CreateWriteStreamRequest.newBuilder()
246+
.setParent(tableIdEU)
247+
.setWriteStream(
248+
WriteStream.newBuilder().setType(WriteStream.Type.COMMITTED).build())
249+
.build());
250+
StreamWriterV2 streamWriter =
251+
StreamWriterV2.newBuilder(writeStream.getName())
252+
.setWriterSchema(ProtoSchemaConverter.convert(FooType.getDescriptor()))
253+
.build();
254+
LOG.info("Sending one message");
255+
256+
ApiFuture<AppendRowsResponse> response =
257+
streamWriter.append(CreateProtoRows(new String[] {"aaa"}), 0);
258+
assertEquals(0, response.get().getAppendResult().getOffset().getValue());
259+
260+
LOG.info("Sending two more messages");
261+
ApiFuture<AppendRowsResponse> response1 =
262+
streamWriter.append(CreateProtoRows(new String[] {"bbb", "ccc"}), 1);
263+
ApiFuture<AppendRowsResponse> response2 =
264+
streamWriter.append(CreateProtoRows(new String[] {"ddd"}), 3);
265+
assertEquals(1, response1.get().getAppendResult().getOffset().getValue());
266+
assertEquals(3, response2.get().getAppendResult().getOffset().getValue());
267+
268+
TableResult result =
269+
bigquery.listTableData(
270+
tableInfoEU.getTableId(), BigQuery.TableDataListOption.startIndex(0L));
271+
Iterator<FieldValueList> iter = result.getValues().iterator();
272+
assertEquals("aaa", iter.next().get(0).getStringValue());
273+
assertEquals("bbb", iter.next().get(0).getStringValue());
274+
assertEquals("ccc", iter.next().get(0).getStringValue());
275+
assertEquals("ddd", iter.next().get(0).getStringValue());
276+
assertEquals(false, iter.hasNext());
277+
}
278+
209279
@Test
210280
public void testJsonStreamWriterCommittedStream()
211281
throws IOException, InterruptedException, ExecutionException,

‎google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/stub/ResourceHeaderTest.java‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,12 @@
2525
import com.google.cloud.bigquery.storage.v1beta2.BigQueryReadClient;
2626
import com.google.cloud.bigquery.storage.v1beta2.BigQueryReadGrpc.BigQueryReadImplBase;
2727
import com.google.cloud.bigquery.storage.v1beta2.BigQueryReadSettings;
28+
import com.google.cloud.bigquery.storage.v1beta2.BigQueryWriteClient;
29+
import com.google.cloud.bigquery.storage.v1beta2.BigQueryWriteSettings;
2830
import com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest;
2931
import com.google.cloud.bigquery.storage.v1beta2.ReadSession;
3032
import com.google.cloud.bigquery.storage.v1beta2.SplitReadStreamRequest;
33+
import com.google.cloud.bigquery.storage.v1beta2.WriteStream;
3134
import java.util.regex.Pattern;
3235
import org.junit.After;
3336
import org.junit.AfterClass;
@@ -43,6 +46,9 @@ public class ResourceHeaderTest {
4346
private static final String TEST_TABLE_REFERENCE =
4447
"projects/project/datasets/dataset/tables/table";
4548

49+
private static final String WRITE_STREAM_NAME =
50+
"projects/project/datasets/dataset/tables/table/streams/stream";
51+
4652
private static final String TEST_STREAM_NAME = "streamName";
4753

4854
private static final String NAME = "resource-header-test:123";
@@ -52,6 +58,20 @@ public class ResourceHeaderTest {
5258
private static final Pattern READ_SESSION_NAME_PATTERN =
5359
Pattern.compile(
5460
".*" + "read_session\\.table=projects/project/datasets/dataset/tables/table" + ".*");
61+
62+
private static final Pattern PARENT_PATTERN =
63+
Pattern.compile(".*" + "parent=projects/project/datasets/dataset/tables/table" + ".*");
64+
65+
private static final Pattern NAME_PATTERN =
66+
Pattern.compile(
67+
".*" + "name=projects/project/datasets/dataset/tables/table/streams/stream" + ".*");
68+
69+
private static final Pattern WRITE_STREAM_PATTERN =
70+
Pattern.compile(
71+
".*"
72+
+ "write_stream=projects/project/datasets/dataset/tables/table/streams/stream"
73+
+ ".*");
74+
5575
private static final Pattern READ_STREAM_PATTERN =
5676
Pattern.compile(".*" + "read_stream=streamName" + ".*");
5777
private static final Pattern STREAM_NAME_PATTERN =
@@ -64,7 +84,9 @@ public class ResourceHeaderTest {
6484
private static InProcessServer<?> server;
6585

6686
private LocalChannelProvider channelProvider;
87+
private LocalChannelProvider channelProvider2;
6788
private BigQueryReadClient client;
89+
private BigQueryWriteClient writeClient;
6890

6991
@BeforeClass
7092
public static void setUpClass() throws Exception {
@@ -81,6 +103,12 @@ public void setUp() throws Exception {
81103
.setHeaderProvider(FixedHeaderProvider.create(TEST_HEADER_NAME, TEST_HEADER_VALUE))
82104
.setTransportChannelProvider(channelProvider);
83105
client = BigQueryReadClient.create(settingsBuilder.build());
106+
channelProvider2 = LocalChannelProvider.create(NAME);
107+
BigQueryWriteSettings.Builder writeSettingsBuilder =
108+
BigQueryWriteSettings.newBuilder()
109+
.setCredentialsProvider(NoCredentialsProvider.create())
110+
.setTransportChannelProvider(channelProvider2);
111+
writeClient = BigQueryWriteClient.create(writeSettingsBuilder.build());
84112
}
85113

86114
@After
@@ -129,6 +157,63 @@ public void splitReadStreamTest() {
129157
verifyHeaderSent(STREAM_NAME_PATTERN);
130158
}
131159

160+
@Test
161+
public void createWriteStreamTest() {
162+
try {
163+
writeClient.createWriteStream(
164+
"projects/project/datasets/dataset/tables/table",
165+
WriteStream.newBuilder().setType(WriteStream.Type.BUFFERED).build());
166+
} catch (UnimplementedException e) {
167+
// Ignore the error: none of the methods are actually implemented.
168+
}
169+
boolean headerSent = channelProvider2.isHeaderSent(HEADER_NAME, PARENT_PATTERN);
170+
assertWithMessage("Generated header was sent").that(headerSent).isTrue();
171+
}
172+
173+
@Test
174+
public void getWriteStreamTest() {
175+
try {
176+
writeClient.getWriteStream(WRITE_STREAM_NAME);
177+
} catch (UnimplementedException e) {
178+
// Ignore the error: none of the methods are actually implemented.
179+
}
180+
boolean headerSent = channelProvider2.isHeaderSent(HEADER_NAME, NAME_PATTERN);
181+
assertWithMessage("Generated header was sent").that(headerSent).isTrue();
182+
}
183+
184+
// Following tests will work after b/185842996 is fixed.
185+
// @Test
186+
// public void appendRowsTest() {
187+
// try {
188+
// AppendRowsRequest req =
189+
// AppendRowsRequest.newBuilder().setWriteStream(WRITE_STREAM_NAME).build();
190+
// BidiStream<AppendRowsRequest, AppendRowsResponse> bidiStream =
191+
// writeClient.appendRowsCallable().call();
192+
// bidiStream.send(req);
193+
// } catch (UnimplementedException e) {
194+
// // Ignore the error: none of the methods are actually implemented.
195+
// }
196+
// boolean headerSent = channelProvider2.isHeaderSent(HEADER_NAME, WRITE_STREAM_PATTERN);
197+
// assertWithMessage("Generated header was sent").that(headerSent).isTrue();
198+
// }
199+
//
200+
// @Test
201+
// public void appendRowsManualTest() {
202+
// try {
203+
// StreamWriterV2 streamWriter =
204+
// StreamWriterV2.newBuilder(WRITE_STREAM_NAME, writeClient)
205+
// .setWriterSchema(ProtoSchema.newBuilder().build())
206+
// .build();
207+
// streamWriter.append(ProtoRows.newBuilder().build(), 1);
208+
// } catch (UnimplementedException e) {
209+
// // Ignore the error: none of the methods are actually implemented.
210+
// } catch (IOException e) {
211+
// // Ignore the error: none of the methods are actually implemented.
212+
// }
213+
// boolean headerSent = channelProvider2.isHeaderSent(HEADER_NAME, WRITE_STREAM_PATTERN);
214+
// assertWithMessage("Generated header was sent").that(headerSent).isTrue();
215+
// }
216+
132217
private void verifyHeaderSent(Pattern... patterns) {
133218
for (Pattern pattern : patterns) {
134219
boolean headerSent = channelProvider.isHeaderSent(HEADER_NAME, pattern);

0 commit comments

Comments
 (0)