Skip to content

Commit f0d82a6

Browse files
keshavdandevakeshavdandeva
committed
fix(bigquery-jdbc): propagate connection proxy settings to auth library
1 parent 94315ce commit f0d82a6

3 files changed

Lines changed: 158 additions & 29 deletions

File tree

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

Lines changed: 24 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import com.google.api.gax.rpc.HeaderProvider;
2525
import com.google.api.gax.rpc.TransportChannelProvider;
2626
import com.google.auth.Credentials;
27+
import com.google.auth.http.HttpTransportFactory;
2728
import com.google.cloud.bigquery.BigQuery;
2829
import com.google.cloud.bigquery.BigQueryException;
2930
import com.google.cloud.bigquery.BigQueryOptions;
@@ -265,11 +266,34 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
265266
String.valueOf(ds.getRequestGoogleDriveScope()),
266267
BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME);
267268

269+
Map<String, String> proxyProperties =
270+
BigQueryJdbcProxyUtility.parseProxyProperties(ds, this.connectionClassName);
271+
272+
this.sslTrustStorePath = ds.getSSLTrustStorePath();
273+
this.sslTrustStorePassword = ds.getSSLTrustStorePassword();
274+
this.httpConnectTimeout = ds.getHttpConnectTimeout();
275+
this.httpReadTimeout = ds.getHttpReadTimeout();
276+
277+
this.httpTransportOptions =
278+
BigQueryJdbcProxyUtility.getHttpTransportOptions(
279+
proxyProperties,
280+
this.sslTrustStorePath,
281+
this.sslTrustStorePassword,
282+
this.httpConnectTimeout,
283+
this.httpReadTimeout,
284+
this.connectionClassName);
285+
286+
HttpTransportFactory httpTransportFactory =
287+
this.httpTransportOptions != null
288+
? this.httpTransportOptions.getHttpTransportFactory()
289+
: null;
290+
268291
this.credentials =
269292
BigQueryJdbcOAuthUtility.getCredentials(
270293
authProperties,
271294
overrideProperties,
272295
this.reqGoogleDriveScope,
296+
httpTransportFactory,
273297
this.connectionClassName);
274298
String defaultDatasetString = ds.getDefaultDataset();
275299
if (defaultDatasetString == null || defaultDatasetString.trim().isEmpty()) {
@@ -302,22 +326,6 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
302326
this.destinationDataset = ds.getDestinationDataset();
303327
this.destinationDatasetExpirationTime = ds.getDestinationDatasetExpirationTime();
304328
this.kmsKeyName = ds.getKmsKeyName();
305-
Map<String, String> proxyProperties =
306-
BigQueryJdbcProxyUtility.parseProxyProperties(ds, this.connectionClassName);
307-
308-
this.sslTrustStorePath = ds.getSSLTrustStorePath();
309-
this.sslTrustStorePassword = ds.getSSLTrustStorePassword();
310-
this.httpConnectTimeout = ds.getHttpConnectTimeout();
311-
this.httpReadTimeout = ds.getHttpReadTimeout();
312-
313-
this.httpTransportOptions =
314-
BigQueryJdbcProxyUtility.getHttpTransportOptions(
315-
proxyProperties,
316-
this.sslTrustStorePath,
317-
this.sslTrustStorePassword,
318-
this.httpConnectTimeout,
319-
this.httpReadTimeout,
320-
this.connectionClassName);
321329
this.transportChannelProvider =
322330
BigQueryJdbcProxyUtility.getTransportChannelProvider(
323331
proxyProperties,

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

Lines changed: 83 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import com.google.api.client.util.PemReader;
2323
import com.google.api.client.util.SecurityUtils;
24+
import com.google.auth.http.HttpTransportFactory;
2425
import com.google.auth.oauth2.AccessToken;
2526
import com.google.auth.oauth2.ClientId;
2627
import com.google.auth.oauth2.ExternalAccountCredentials;
@@ -263,6 +264,16 @@ static GoogleCredentials getCredentials(
263264
Map<String, String> overrideProperties,
264265
Boolean reqGoogleDriveScopeBool,
265266
String callerClassName) {
267+
return getCredentials(
268+
authProperties, overrideProperties, reqGoogleDriveScopeBool, null, callerClassName);
269+
}
270+
271+
static GoogleCredentials getCredentials(
272+
Map<String, String> authProperties,
273+
Map<String, String> overrideProperties,
274+
Boolean reqGoogleDriveScopeBool,
275+
HttpTransportFactory httpTransportFactory,
276+
String callerClassName) {
266277
LOG.finer("++enter++\t" + callerClassName);
267278

268279
AuthType authType =
@@ -272,11 +283,13 @@ static GoogleCredentials getCredentials(
272283
switch (authType) {
273284
case GOOGLE_SERVICE_ACCOUNT:
274285
credentials =
275-
getGoogleServiceAccountCredentials(authProperties, overrideProperties, callerClassName);
286+
getGoogleServiceAccountCredentials(
287+
authProperties, overrideProperties, httpTransportFactory, callerClassName);
276288
break;
277289
case GOOGLE_USER_ACCOUNT:
278290
credentials =
279-
getGoogleUserAccountCredentials(authProperties, overrideProperties, callerClassName);
291+
getGoogleUserAccountCredentials(
292+
authProperties, overrideProperties, httpTransportFactory, callerClassName);
280293
break;
281294
case PRE_GENERATED_TOKEN:
282295
credentials =
@@ -286,7 +299,9 @@ static GoogleCredentials getCredentials(
286299
credentials = getApplicationDefaultCredentials(callerClassName);
287300
break;
288301
case EXTERNAL_ACCOUNT_AUTH:
289-
credentials = getExternalAccountAuthCredentials(authProperties, callerClassName);
302+
credentials =
303+
getExternalAccountAuthCredentials(
304+
authProperties, httpTransportFactory, callerClassName);
290305
break;
291306
default:
292307
IllegalStateException ex = new IllegalStateException(OAUTH_TYPE_ERROR_MESSAGE);
@@ -295,7 +310,7 @@ static GoogleCredentials getCredentials(
295310
}
296311

297312
return getServiceAccountImpersonatedCredentials(
298-
credentials, reqGoogleDriveScopeBool, authProperties);
313+
credentials, reqGoogleDriveScopeBool, authProperties, httpTransportFactory);
299314
}
300315

301316
private static boolean isFileExists(String filename) {
@@ -326,6 +341,7 @@ private static boolean isJson(byte[] value) {
326341
private static GoogleCredentials getGoogleServiceAccountCredentials(
327342
Map<String, String> authProperties,
328343
Map<String, String> overrideProperties,
344+
HttpTransportFactory httpTransportFactory,
329345
String callerClassName) {
330346
LOG.finer("++enter++\t" + callerClassName);
331347

@@ -370,6 +386,10 @@ private static GoogleCredentials getGoogleServiceAccountCredentials(
370386
throw new BigQueryJdbcRuntimeException("No valid credentials provided.");
371387
}
372388

389+
if (httpTransportFactory != null) {
390+
builder.setHttpTransportFactory(httpTransportFactory);
391+
}
392+
373393
if (overrideProperties.containsKey(BigQueryJdbcUrlUtility.OAUTH2_TOKEN_URI_PROPERTY_NAME)) {
374394
builder.setTokenServerUri(
375395
new URI(overrideProperties.get(BigQueryJdbcUrlUtility.OAUTH2_TOKEN_URI_PROPERTY_NAME)));
@@ -393,6 +413,16 @@ static UserAuthorizer getUserAuthorizer(
393413
int port,
394414
String callerClassName)
395415
throws URISyntaxException {
416+
return getUserAuthorizer(authProperties, overrideProperties, port, null, callerClassName);
417+
}
418+
419+
static UserAuthorizer getUserAuthorizer(
420+
Map<String, String> authProperties,
421+
Map<String, String> overrideProperties,
422+
int port,
423+
HttpTransportFactory httpTransportFactory,
424+
String callerClassName)
425+
throws URISyntaxException {
396426
LOG.finer("++enter++\t" + callerClassName);
397427

398428
List<String> scopes = new ArrayList<>();
@@ -411,6 +441,10 @@ static UserAuthorizer getUserAuthorizer(
411441
.setScopes(scopes)
412442
.setCallbackUri(URI.create("http://localhost:" + port));
413443

444+
if (httpTransportFactory != null) {
445+
userAuthorizerBuilder.setHttpTransportFactory(httpTransportFactory);
446+
}
447+
414448
if (overrideProperties.containsKey(BigQueryJdbcUrlUtility.OAUTH2_TOKEN_URI_PROPERTY_NAME)) {
415449
userAuthorizerBuilder.setTokenServerUri(
416450
new URI(overrideProperties.get(BigQueryJdbcUrlUtility.OAUTH2_TOKEN_URI_PROPERTY_NAME)));
@@ -421,21 +455,38 @@ static UserAuthorizer getUserAuthorizer(
421455

422456
static UserCredentials getCredentialsFromCode(
423457
UserAuthorizer userAuthorizer, String code, String callerClassName) throws IOException {
458+
return getCredentialsFromCode(userAuthorizer, code, null, callerClassName);
459+
}
460+
461+
static UserCredentials getCredentialsFromCode(
462+
UserAuthorizer userAuthorizer,
463+
String code,
464+
HttpTransportFactory httpTransportFactory,
465+
String callerClassName)
466+
throws IOException {
424467
LOG.finer("++enter++\t" + callerClassName);
425-
return userAuthorizer.getCredentialsFromCode(code, URI.create(""));
468+
UserCredentials credentials = userAuthorizer.getCredentialsFromCode(code, URI.create(""));
469+
if (httpTransportFactory != null) {
470+
credentials =
471+
(UserCredentials)
472+
credentials.toBuilder().setHttpTransportFactory(httpTransportFactory).build();
473+
}
474+
return credentials;
426475
}
427476

428477
private static GoogleCredentials getGoogleUserAccountCredentials(
429478
Map<String, String> authProperties,
430479
Map<String, String> overrideProperties,
480+
HttpTransportFactory httpTransportFactory,
431481
String callerClassName) {
432482
LOG.finer("++enter++\t" + callerClassName);
433483
try {
434484
ServerSocket serverSocket = new ServerSocket(0);
435485
serverSocket.setSoTimeout(USER_AUTH_TIMEOUT_MS);
436486
int port = serverSocket.getLocalPort();
437487
UserAuthorizer userAuthorizer =
438-
getUserAuthorizer(authProperties, overrideProperties, port, callerClassName);
488+
getUserAuthorizer(
489+
authProperties, overrideProperties, port, httpTransportFactory, callerClassName);
439490

440491
URL authURL = userAuthorizer.getAuthorizationUrl("user", "", URI.create(""));
441492
String code;
@@ -468,7 +519,7 @@ private static GoogleCredentials getGoogleUserAccountCredentials(
468519
throw new BigQueryJdbcRuntimeException("User auth only supported in desktop environments");
469520
}
470521

471-
return getCredentialsFromCode(userAuthorizer, code, callerClassName);
522+
return getCredentialsFromCode(userAuthorizer, code, httpTransportFactory, callerClassName);
472523
} catch (IOException | URISyntaxException ex) {
473524
throw new BigQueryJdbcRuntimeException(
474525
"Failed to establish connection using User Account authentication", ex);
@@ -572,7 +623,9 @@ private static GoogleCredentials getApplicationDefaultCredentials(String callerC
572623
}
573624

574625
private static GoogleCredentials getExternalAccountAuthCredentials(
575-
Map<String, String> authProperties, String callerClassName) {
626+
Map<String, String> authProperties,
627+
HttpTransportFactory httpTransportFactory,
628+
String callerClassName) {
576629
LOG.finer("++enter++\t" + callerClassName);
577630
try {
578631
JsonObject jsonObject = null;
@@ -609,18 +662,25 @@ private static GoogleCredentials getExternalAccountAuthCredentials(
609662
}
610663
}
611664

665+
ExternalAccountCredentials credentials;
612666
if (credentialsPath != null) {
613-
return ExternalAccountCredentials.fromStream(
614-
Files.newInputStream(Paths.get(credentialsPath)));
667+
credentials =
668+
(ExternalAccountCredentials)
669+
ExternalAccountCredentials.fromStream(
670+
Files.newInputStream(Paths.get(credentialsPath)), httpTransportFactory);
615671
} else if (jsonObject != null) {
616-
return ExternalAccountCredentials.fromStream(
617-
new ByteArrayInputStream(jsonObject.toString().getBytes()));
672+
credentials =
673+
(ExternalAccountCredentials)
674+
ExternalAccountCredentials.fromStream(
675+
new ByteArrayInputStream(jsonObject.toString().getBytes()),
676+
httpTransportFactory);
618677
} else {
619678
IllegalArgumentException ex =
620679
new IllegalArgumentException("Insufficient info provided for external authentication");
621680
LOG.severe(ex.getMessage(), ex);
622681
throw ex;
623682
}
683+
return credentials;
624684
} catch (IOException e) {
625685
throw new BigQueryJdbcRuntimeException(
626686
"IOException during getExternalAccountAuthCredentials", e);
@@ -634,7 +694,8 @@ private static GoogleCredentials getExternalAccountAuthCredentials(
634694
private static GoogleCredentials getServiceAccountImpersonatedCredentials(
635695
GoogleCredentials credentials,
636696
Boolean reqGoogleDriveScopeBool,
637-
Map<String, String> authProperties) {
697+
Map<String, String> authProperties,
698+
HttpTransportFactory httpTransportFactory) {
638699

639700
String impersonationEmail =
640701
authProperties.get(BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_EMAIL_PROPERTY_NAME);
@@ -684,6 +745,15 @@ private static GoogleCredentials getServiceAccountImpersonatedCredentials(
684745
throw ex;
685746
}
686747

748+
if (httpTransportFactory != null) {
749+
return ImpersonatedCredentials.create(
750+
credentials,
751+
impersonationEmail,
752+
impersonationChain,
753+
impersonationScopes,
754+
impersonationLifetimeInt,
755+
httpTransportFactory);
756+
}
687757
return ImpersonatedCredentials.create(
688758
credentials,
689759
impersonationEmail,

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

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,10 @@
2222
import static org.junit.jupiter.api.Assertions.assertNull;
2323
import static org.junit.jupiter.api.Assertions.assertTrue;
2424

25+
import com.google.auth.http.HttpTransportFactory;
2526
import com.google.auth.oauth2.GoogleCredentials;
2627
import com.google.auth.oauth2.ImpersonatedCredentials;
28+
import com.google.auth.oauth2.ServiceAccountCredentials;
2729
import com.google.auth.oauth2.UserAuthorizer;
2830
import com.google.auth.oauth2.UserCredentials;
2931
import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
@@ -489,4 +491,53 @@ public void testPrivateKeyFromP12Bytes_wrong_password() {
489491
assertTrue(false);
490492
}
491493
}
494+
495+
@Test
496+
public void testGetCredentialsPropagatesHttpTransportFactory() {
497+
Map<String, String> authProperties =
498+
BigQueryJdbcOAuthUtility.parseOAuthProperties(
499+
DataSource.fromUrl(
500+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
501+
+ "ProjectId=MyBigQueryProject;OAuthType=0;"
502+
+ "OAuthServiceAcctEmail=dummytest@dummytest.iam.gserviceaccount.com;"
503+
+ "OAuthPvtKey="
504+
+ fake_pkcs8_key
505+
+ ";"),
506+
null);
507+
508+
HttpTransportFactory dummyFactory = () -> null;
509+
510+
GoogleCredentials credentials =
511+
BigQueryJdbcOAuthUtility.getCredentials(
512+
authProperties, Collections.emptyMap(), false, dummyFactory, null);
513+
514+
assertThat(credentials).isInstanceOf(ServiceAccountCredentials.class);
515+
assertThat(((ServiceAccountCredentials) credentials).toBuilder().getHttpTransportFactory())
516+
.isEqualTo(dummyFactory);
517+
}
518+
519+
@Test
520+
public void testGetImpersonatedCredentialsPropagatesHttpTransportFactory() {
521+
Map<String, String> authProperties =
522+
BigQueryJdbcOAuthUtility.parseOAuthProperties(
523+
DataSource.fromUrl(
524+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
525+
+ "ProjectId=MyBigQueryProject;OAuthType=0;"
526+
+ "OAuthServiceAcctEmail=dummytest@dummytest.iam.gserviceaccount.com;"
527+
+ "OAuthPvtKey="
528+
+ fake_pkcs8_key
529+
+ ";"
530+
+ "ServiceAccountImpersonationEmail=impersonated@email.com;"),
531+
null);
532+
533+
HttpTransportFactory dummyFactory = () -> null;
534+
535+
GoogleCredentials credentials =
536+
BigQueryJdbcOAuthUtility.getCredentials(
537+
authProperties, Collections.emptyMap(), false, dummyFactory, null);
538+
539+
assertThat(credentials).isInstanceOf(ImpersonatedCredentials.class);
540+
assertThat(((ImpersonatedCredentials) credentials).toBuilder().getHttpTransportFactory())
541+
.isEqualTo(dummyFactory);
542+
}
492543
}

0 commit comments

Comments
 (0)