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

Commit f9e655b

Browse files
lqiu96lqiu96
committed
Merge branch 'main' into PR #3867 to update
2 parents e7ce819 + 77e0c35 commit f9e655b

6 files changed

Lines changed: 214 additions & 18 deletions

File tree

‎google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,9 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
131131
String sslTrustStorePassword;
132132
long maxBytesBilled;
133133
Map<String, String> labels;
134+
Integer httpConnectTimeout;
135+
Integer httpReadTimeout;
136+
String requestReason;
134137

135138
BigQueryConnection(String url) throws IOException {
136139
this.connectionUrl = url;
@@ -272,11 +275,25 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
272275
BigQueryJdbcUrlUtility.SSL_TRUST_STORE_PWD_PROPERTY_NAME,
273276
null,
274277
this.connectionClassName);
278+
this.httpConnectTimeout =
279+
BigQueryJdbcUrlUtility.parseIntProperty(
280+
url,
281+
BigQueryJdbcUrlUtility.HTTP_CONNECT_TIMEOUT_PROPERTY_NAME,
282+
null,
283+
this.connectionClassName);
284+
this.httpReadTimeout =
285+
BigQueryJdbcUrlUtility.parseIntProperty(
286+
url,
287+
BigQueryJdbcUrlUtility.HTTP_READ_TIMEOUT_PROPERTY_NAME,
288+
null,
289+
this.connectionClassName);
275290
this.httpTransportOptions =
276291
BigQueryJdbcProxyUtility.getHttpTransportOptions(
277292
proxyProperties,
278293
this.sslTrustStorePath,
279294
this.sslTrustStorePassword,
295+
this.httpConnectTimeout,
296+
this.httpReadTimeout,
280297
this.connectionClassName);
281298
this.transportChannelProvider =
282299
BigQueryJdbcProxyUtility.getTransportChannelProvider(
@@ -347,6 +364,12 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
347364
BigQueryJdbcUrlUtility.METADATA_FETCH_THREAD_COUNT_PROPERTY_NAME,
348365
BigQueryJdbcUrlUtility.DEFAULT_METADATA_FETCH_THREAD_COUNT_VALUE,
349366
this.connectionClassName);
367+
this.requestReason =
368+
BigQueryJdbcUrlUtility.parseStringProperty(
369+
url,
370+
BigQueryJdbcUrlUtility.REQUEST_REASON_PROPERTY_NAME,
371+
null,
372+
this.connectionClassName);
350373

351374
HEADER_PROVIDER = createHeaderProvider();
352375
this.bigQuery = getBigQueryConnection();
@@ -383,7 +406,12 @@ HeaderProvider createHeaderProvider() {
383406
String partnerToken = buildPartnerToken(this.connectionUrl);
384407
String headerToken =
385408
DEFAULT_JDBC_TOKEN_VALUE + "/" + getLibraryVersion(this.getClass()) + partnerToken;
386-
return FixedHeaderProvider.create("user-agent", headerToken);
409+
Map<String, String> headers = new java.util.HashMap<>();
410+
headers.put("user-agent", headerToken);
411+
if (this.requestReason != null) {
412+
headers.put("x-goog-request-reason", this.requestReason);
413+
}
414+
return FixedHeaderProvider.create(headers);
387415
}
388416

389417
protected void addOpenStatements(Statement statement) {
@@ -724,6 +752,14 @@ String getSSLTrustStorePassword() {
724752
return sslTrustStorePassword;
725753
}
726754

755+
Integer getHttpConnectTimeout() {
756+
return httpConnectTimeout;
757+
}
758+
759+
Integer getHttpReadTimeout() {
760+
return httpReadTimeout;
761+
}
762+
727763
@Override
728764
public boolean isValid(int timeout) throws SQLException {
729765
if (timeout < 0) {

‎google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcProxyUtility.java‎

Lines changed: 26 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,8 @@ static Map<String, String> parseProxyProperties(String URL, String callerClassNa
9797
(proxyHost == null && proxyPort != null) || (proxyHost != null && proxyPort == null);
9898
if (isMissingProxyHostOrPortWhenProxySet) {
9999
throw new IllegalArgumentException(
100-
"Both ProxyHost and ProxyPort parameters need to be specified. No defaulting behavior occurs.");
100+
"Both ProxyHost and ProxyPort parameters need to be specified. No defaulting behavior"
101+
+ " occurs.");
101102
}
102103
boolean isMissingProxyUidOrPwdWhenAuthSet =
103104
(proxyUid == null && proxyPwd != null) || (proxyUid != null && proxyPwd == null);
@@ -117,18 +118,35 @@ static HttpTransportOptions getHttpTransportOptions(
117118
Map<String, String> proxyProperties,
118119
String sslTrustStorePath,
119120
String sslTrustStorePassword,
121+
Integer connectTimeout,
122+
Integer readTimeout,
120123
String callerClassName) {
121124
LOG.finest("++enter++\t" + callerClassName);
122125

123-
if (!proxyProperties.containsKey(BigQueryJdbcUrlUtility.PROXY_HOST_PROPERTY_NAME)
124-
&& sslTrustStorePath == null) {
126+
boolean hasProxyOrSsl =
127+
proxyProperties.containsKey(BigQueryJdbcUrlUtility.PROXY_HOST_PROPERTY_NAME)
128+
|| sslTrustStorePath != null;
129+
boolean hasTimeoutConfig = connectTimeout != null || readTimeout != null;
130+
131+
if (!hasProxyOrSsl && !hasTimeoutConfig) {
125132
return null;
126133
}
127-
return HttpTransportOptions.newBuilder()
128-
.setHttpTransportFactory(
129-
getHttpTransportFactory(
130-
proxyProperties, sslTrustStorePath, sslTrustStorePassword, callerClassName))
131-
.build();
134+
135+
HttpTransportOptions.Builder httpTransportOptionsBuilder = HttpTransportOptions.newBuilder();
136+
if (hasProxyOrSsl) {
137+
httpTransportOptionsBuilder.setHttpTransportFactory(
138+
getHttpTransportFactory(
139+
proxyProperties, sslTrustStorePath, sslTrustStorePassword, callerClassName));
140+
}
141+
142+
if (connectTimeout != null) {
143+
httpTransportOptionsBuilder.setConnectTimeout(connectTimeout);
144+
}
145+
if (readTimeout != null) {
146+
httpTransportOptionsBuilder.setReadTimeout(readTimeout);
147+
}
148+
149+
return httpTransportOptionsBuilder.build();
132150
}
133151

134152
private static HttpTransportFactory getHttpTransportFactory(

‎google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,8 @@ final class BigQueryJdbcUrlUtility {
9898
static final String PROXY_PORT_PROPERTY_NAME = "ProxyPort";
9999
static final String PROXY_USER_ID_PROPERTY_NAME = "ProxyUid";
100100
static final String PROXY_PASSWORD_PROPERTY_NAME = "ProxyPwd";
101+
static final String HTTP_CONNECT_TIMEOUT_PROPERTY_NAME = "HttpConnectTimeout";
102+
static final String HTTP_READ_TIMEOUT_PROPERTY_NAME = "HttpReadTimeout";
101103
static final boolean DEFAULT_ENABLE_HTAPI_VALUE = false;
102104
static final boolean DEFAULT_ENABLE_SESSION_VALUE = false;
103105
static final int DEFAULT_LOG_LEVEL = 0;
@@ -160,6 +162,7 @@ final class BigQueryJdbcUrlUtility {
160162
OAUTH2_TOKEN_URI_PROPERTY_NAME,
161163
HTAPI_ENDPOINT_OVERRIDE_PROPERTY_NAME,
162164
STS_ENDPOINT_OVERRIDE_PROPERTY_NAME);
165+
static final String REQUEST_REASON_PROPERTY_NAME = "RequestReason";
163166
static final List<String> BYOID_PROPERTIES =
164167
Arrays.asList(
165168
BYOID_AUDIENCE_URI_PROPERTY_NAME,
@@ -249,7 +252,8 @@ final class BigQueryJdbcUrlUtility {
249252
BigQueryConnectionProperty.newBuilder()
250253
.setName(OAUTH_SA_IMPERSONATION_CHAIN_PROPERTY_NAME)
251254
.setDescription(
252-
"Comma separated list of service account emails in the impersonation chain.")
255+
"Comma separated list of service account emails in the impersonation"
256+
+ " chain.")
253257
.build(),
254258
BigQueryConnectionProperty.newBuilder()
255259
.setName(OAUTH_SA_IMPERSONATION_SCOPES_PROPERTY_NAME)
@@ -569,6 +573,22 @@ final class BigQueryJdbcUrlUtility {
569573
.setDescription(
570574
"The password for accessing the Java TrustStore that is specified using"
571575
+ " the property SSLTrustStore.")
576+
.build(),
577+
BigQueryConnectionProperty.newBuilder()
578+
.setName(HTTP_CONNECT_TIMEOUT_PROPERTY_NAME)
579+
.setDescription(
580+
"The timeout (in milliseconds) for establishing a connection to the"
581+
+ " server.")
582+
.build(),
583+
BigQueryConnectionProperty.newBuilder()
584+
.setName(HTTP_READ_TIMEOUT_PROPERTY_NAME)
585+
.setDescription("The timeout (in milliseconds) when reading from the server.")
586+
.build(),
587+
BigQueryConnectionProperty.newBuilder()
588+
.setName(REQUEST_REASON_PROPERTY_NAME)
589+
.setDescription(
590+
"Reason for the request, which is passed as the x-goog-request-reason"
591+
+ " header.")
572592
.build())));
573593

574594
private BigQueryJdbcUrlUtility() {}

‎google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,23 @@ public void testHeaderProviderWithInvalidPartner() throws IOException, SQLExcept
165165
}
166166
}
167167

168+
@Test
169+
public void testHeaderProviderWithRequestReason() throws IOException, SQLException {
170+
String requestReason = "Ticket123";
171+
String url =
172+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
173+
+ "OAuthType=2;ProjectId=MyBigQueryProject;"
174+
+ "OAuthAccessToken=redactedToken;OAuthClientId=redactedToken;"
175+
+ "OAuthClientSecret=redactedToken;RequestReason="
176+
+ requestReason;
177+
try (BigQueryConnection connection = new BigQueryConnection(url)) {
178+
HeaderProvider headerProvider = connection.createHeaderProvider();
179+
java.util.Map<String, String> headers = headerProvider.getHeaders();
180+
assertTrue(headers.containsKey("x-goog-request-reason"));
181+
assertEquals(requestReason, headers.get("x-goog-request-reason"));
182+
}
183+
}
184+
168185
@Test
169186
public void testWriteAPIConnectionProperties() throws SQLException {
170187
// Test without connection properties. Defaults to default values.

‎google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcProxyUtilityTest.java‎

Lines changed: 53 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,8 @@ public void testGetHttpTransportOptionsWithAuthenticatedProxy() {
138138
Map<String, String> proxyProperties =
139139
BigQueryJdbcProxyUtility.parseProxyProperties(connection_uri, null);
140140
HttpTransportOptions result =
141-
BigQueryJdbcProxyUtility.getHttpTransportOptions(proxyProperties, null, null, null);
141+
BigQueryJdbcProxyUtility.getHttpTransportOptions(
142+
proxyProperties, null, null, null, null, "TestClass");
142143
assertNotNull(result);
143144
}
144145

@@ -154,7 +155,8 @@ public void testGetHttpTransportOptionsWithNonAuthenticatedProxy() {
154155
Map<String, String> proxyProperties =
155156
BigQueryJdbcProxyUtility.parseProxyProperties(connection_uri, null);
156157
HttpTransportOptions result =
157-
BigQueryJdbcProxyUtility.getHttpTransportOptions(proxyProperties, null, null, null);
158+
BigQueryJdbcProxyUtility.getHttpTransportOptions(
159+
proxyProperties, null, null, null, null, "TestClass");
158160
assertNotNull(result);
159161
}
160162

@@ -168,7 +170,8 @@ public void testGetHttpTransportOptionsWithNoProxySettingsReturnsNull() {
168170
Map<String, String> proxyProperties =
169171
BigQueryJdbcProxyUtility.parseProxyProperties(connection_uri, null);
170172
HttpTransportOptions result =
171-
BigQueryJdbcProxyUtility.getHttpTransportOptions(proxyProperties, null, null, null);
173+
BigQueryJdbcProxyUtility.getHttpTransportOptions(
174+
proxyProperties, null, null, null, null, "TestClass");
172175
assertNull(result);
173176
}
174177

@@ -185,7 +188,7 @@ public void testGetHttpTransportOptions_withSslTrustStore_noPassword() throws Ex
185188
String trustStorePath = getTestResourcePath("test_truststore_nopass.jks");
186189
HttpTransportOptions options =
187190
BigQueryJdbcProxyUtility.getHttpTransportOptions(
188-
Collections.<String, String>emptyMap(), trustStorePath, null, "TestClass");
191+
Collections.<String, String>emptyMap(), trustStorePath, null, null, null, "TestClass");
189192
assertNotNull(options);
190193
assertNotNull(options.getHttpTransportFactory());
191194
}
@@ -195,7 +198,12 @@ public void testGetHttpTransportOptions_withSslTrustStore_withCorrectPassword()
195198
String trustStorePath = getTestResourcePath("test_truststore_withpass.jks");
196199
HttpTransportOptions options =
197200
BigQueryJdbcProxyUtility.getHttpTransportOptions(
198-
Collections.<String, String>emptyMap(), trustStorePath, "testpassword", "TestClass");
201+
Collections.<String, String>emptyMap(),
202+
trustStorePath,
203+
"testpassword",
204+
null,
205+
null,
206+
"TestClass");
199207
assertNotNull(options);
200208
assertNotNull(options.getHttpTransportFactory());
201209
}
@@ -212,6 +220,8 @@ public void testGetHttpTransportOptions_withSslTrustStore_withIncorrectPassword(
212220
Collections.<String, String>emptyMap(),
213221
trustStorePath,
214222
"wrongpassword",
223+
null,
224+
null,
215225
"TestClass"));
216226
assertThat(exception.getCause()).isInstanceOf(IOException.class);
217227
}
@@ -224,7 +234,12 @@ public void testGetHttpTransportOptions_withInvalidSslTrustStorePath() {
224234
BigQueryJdbcRuntimeException.class,
225235
() ->
226236
BigQueryJdbcProxyUtility.getHttpTransportOptions(
227-
Collections.<String, String>emptyMap(), invalidPath, null, "TestClass"));
237+
Collections.<String, String>emptyMap(),
238+
invalidPath,
239+
null,
240+
null,
241+
null,
242+
"TestClass"));
228243

229244
assertThat(exception.getCause()).isInstanceOf(FileNotFoundException.class);
230245
}
@@ -238,7 +253,7 @@ public void testGetHttpTransportOptions_withSslAndProxy() throws Exception {
238253

239254
HttpTransportOptions options =
240255
BigQueryJdbcProxyUtility.getHttpTransportOptions(
241-
proxyProperties, trustStorePath, null, "TestClass");
256+
proxyProperties, trustStorePath, null, null, null, "TestClass");
242257
assertNotNull(options);
243258
assertNotNull(options.getHttpTransportFactory());
244259
}
@@ -287,7 +302,37 @@ public void testGetTransportChannelProvider_noProxyNoSsl_returnsNull() {
287302
public void testGetHttpTransportOptions_noProxyNoSsl_returnsNull() {
288303
HttpTransportOptions options =
289304
BigQueryJdbcProxyUtility.getHttpTransportOptions(
290-
Collections.<String, String>emptyMap(), null, null, "TestClass");
305+
Collections.<String, String>emptyMap(), null, null, null, null, "TestClass");
291306
assertNull(options);
292307
}
308+
309+
@Test
310+
public void testGetHttpTransportOptions_withTimeouts_returnsOptions() {
311+
HttpTransportOptions options =
312+
BigQueryJdbcProxyUtility.getHttpTransportOptions(
313+
Collections.<String, String>emptyMap(), null, null, 10000, 20000, "TestClass");
314+
assertNotNull(options);
315+
assertThat(options.getConnectTimeout()).isEqualTo(10000);
316+
assertThat(options.getReadTimeout()).isEqualTo(20000);
317+
}
318+
319+
@Test
320+
public void testGetHttpTransportOptions_withConnectTimeoutOnly_returnsOptions() {
321+
HttpTransportOptions options =
322+
BigQueryJdbcProxyUtility.getHttpTransportOptions(
323+
Collections.<String, String>emptyMap(), null, null, 10000, null, "TestClass");
324+
assertNotNull(options);
325+
assertThat(options.getConnectTimeout()).isEqualTo(10000);
326+
// readTimeout defaults to HttpTransport's internal default, not asserted here.
327+
}
328+
329+
@Test
330+
public void testGetHttpTransportOptions_withReadTimeoutOnly_returnsOptions() {
331+
HttpTransportOptions options =
332+
BigQueryJdbcProxyUtility.getHttpTransportOptions(
333+
Collections.<String, String>emptyMap(), null, null, null, 20000, "TestClass");
334+
assertNotNull(options);
335+
assertThat(options.getReadTimeout()).isEqualTo(20000);
336+
// connectTimeout defaults to HttpTransport's internal default, not asserted here.
337+
}
293338
}

‎google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtilityTest.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -800,4 +800,64 @@ public void testParseLabelsEmpty() {
800800
Map<String, String> labels = BigQueryJdbcUrlUtility.parseLabels(connection_uri, null);
801801
assertNull(labels);
802802
}
803+
804+
@Test
805+
public void testParseHttpConnectTimeout() {
806+
String connection_uri =
807+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;OAuthType=3;ProjectId=testProject;"
808+
+ "HttpConnectTimeout=10000";
809+
810+
Integer timeout =
811+
BigQueryJdbcUrlUtility.parseIntProperty(
812+
connection_uri, BigQueryJdbcUrlUtility.HTTP_CONNECT_TIMEOUT_PROPERTY_NAME, null, null);
813+
assertEquals(Integer.valueOf(10000), timeout);
814+
}
815+
816+
@Test
817+
public void testParseHttpConnectTimeoutDefault() {
818+
String connection_uri =
819+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;OAuthType=3;ProjectId=testProject";
820+
821+
Integer timeout =
822+
BigQueryJdbcUrlUtility.parseIntProperty(
823+
connection_uri, BigQueryJdbcUrlUtility.HTTP_CONNECT_TIMEOUT_PROPERTY_NAME, null, null);
824+
assertNull(timeout);
825+
}
826+
827+
@Test
828+
public void testParseHttpReadTimeout() {
829+
String connection_uri =
830+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;OAuthType=3;ProjectId=testProject;"
831+
+ "HttpReadTimeout=20000";
832+
833+
Integer timeout =
834+
BigQueryJdbcUrlUtility.parseIntProperty(
835+
connection_uri, BigQueryJdbcUrlUtility.HTTP_READ_TIMEOUT_PROPERTY_NAME, null, null);
836+
assertEquals(Integer.valueOf(20000), timeout);
837+
}
838+
839+
@Test
840+
public void testParseHttpReadTimeoutDefault() {
841+
String connection_uri =
842+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;OAuthType=3;ProjectId=testProject";
843+
844+
Integer timeout =
845+
BigQueryJdbcUrlUtility.parseIntProperty(
846+
connection_uri, BigQueryJdbcUrlUtility.HTTP_READ_TIMEOUT_PROPERTY_NAME, null, null);
847+
assertNull(timeout);
848+
}
849+
850+
@Test
851+
public void testParseRequestReason() {
852+
String url =
853+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
854+
+ "OAuthType=3;ProjectId=testProject;RequestReason=testingRequestReason;";
855+
String requestReason =
856+
BigQueryJdbcUrlUtility.parseStringProperty(
857+
url,
858+
BigQueryJdbcUrlUtility.REQUEST_REASON_PROPERTY_NAME,
859+
null,
860+
"testParseRequestReason");
861+
assertEquals("testingRequestReason", requestReason);
862+
}
803863
}

0 commit comments

Comments
 (0)