2121
2222import com .google .api .client .util .PemReader ;
2323import com .google .api .client .util .SecurityUtils ;
24+ import com .google .auth .http .HttpTransportFactory ;
2425import com .google .auth .oauth2 .AccessToken ;
2526import com .google .auth .oauth2 .ClientId ;
2627import com .google .auth .oauth2 .ExternalAccountCredentials ;
5152import java .net .URI ;
5253import java .net .URISyntaxException ;
5354import java .net .URL ;
55+ import java .nio .charset .StandardCharsets ;
5456import java .nio .file .Files ;
5557import java .nio .file .Paths ;
5658import java .security .GeneralSecurityException ;
@@ -262,6 +264,7 @@ static GoogleCredentials getCredentials(
262264 Map <String , String > authProperties ,
263265 Map <String , String > overrideProperties ,
264266 Boolean reqGoogleDriveScopeBool ,
267+ HttpTransportFactory httpTransportFactory ,
265268 String callerClassName ) {
266269 LOG .finer ("++enter++\t " + callerClassName );
267270
@@ -272,21 +275,26 @@ static GoogleCredentials getCredentials(
272275 switch (authType ) {
273276 case GOOGLE_SERVICE_ACCOUNT :
274277 credentials =
275- getGoogleServiceAccountCredentials (authProperties , overrideProperties , callerClassName );
278+ getGoogleServiceAccountCredentials (
279+ authProperties , overrideProperties , httpTransportFactory , callerClassName );
276280 break ;
277281 case GOOGLE_USER_ACCOUNT :
278282 credentials =
279- getGoogleUserAccountCredentials (authProperties , overrideProperties , callerClassName );
283+ getGoogleUserAccountCredentials (
284+ authProperties , overrideProperties , httpTransportFactory , callerClassName );
280285 break ;
281286 case PRE_GENERATED_TOKEN :
282287 credentials =
283- getPreGeneratedTokensCredentials (authProperties , overrideProperties , callerClassName );
288+ getPreGeneratedTokensCredentials (
289+ authProperties , overrideProperties , httpTransportFactory , callerClassName );
284290 break ;
285291 case APPLICATION_DEFAULT_CREDENTIALS :
286- credentials = getApplicationDefaultCredentials (callerClassName );
292+ credentials = getApplicationDefaultCredentials (httpTransportFactory , callerClassName );
287293 break ;
288294 case EXTERNAL_ACCOUNT_AUTH :
289- credentials = getExternalAccountAuthCredentials (authProperties , callerClassName );
295+ credentials =
296+ getExternalAccountAuthCredentials (
297+ authProperties , httpTransportFactory , callerClassName );
290298 break ;
291299 default :
292300 IllegalStateException ex = new IllegalStateException (OAUTH_TYPE_ERROR_MESSAGE );
@@ -295,7 +303,7 @@ static GoogleCredentials getCredentials(
295303 }
296304
297305 return getServiceAccountImpersonatedCredentials (
298- credentials , reqGoogleDriveScopeBool , authProperties );
306+ credentials , reqGoogleDriveScopeBool , authProperties , httpTransportFactory );
299307 }
300308
301309 private static boolean isFileExists (String filename ) {
@@ -326,6 +334,7 @@ private static boolean isJson(byte[] value) {
326334 private static GoogleCredentials getGoogleServiceAccountCredentials (
327335 Map <String , String > authProperties ,
328336 Map <String , String > overrideProperties ,
337+ HttpTransportFactory httpTransportFactory ,
329338 String callerClassName ) {
330339 LOG .finer ("++enter++\t " + callerClassName );
331340
@@ -370,6 +379,10 @@ private static GoogleCredentials getGoogleServiceAccountCredentials(
370379 throw new BigQueryJdbcRuntimeException ("No valid credentials provided." );
371380 }
372381
382+ if (httpTransportFactory != null ) {
383+ builder .setHttpTransportFactory (httpTransportFactory );
384+ }
385+
373386 if (overrideProperties .containsKey (BigQueryJdbcUrlUtility .OAUTH2_TOKEN_URI_PROPERTY_NAME )) {
374387 builder .setTokenServerUri (
375388 new URI (overrideProperties .get (BigQueryJdbcUrlUtility .OAUTH2_TOKEN_URI_PROPERTY_NAME )));
@@ -391,6 +404,7 @@ static UserAuthorizer getUserAuthorizer(
391404 Map <String , String > authProperties ,
392405 Map <String , String > overrideProperties ,
393406 int port ,
407+ HttpTransportFactory httpTransportFactory ,
394408 String callerClassName )
395409 throws URISyntaxException {
396410 LOG .finer ("++enter++\t " + callerClassName );
@@ -411,6 +425,10 @@ static UserAuthorizer getUserAuthorizer(
411425 .setScopes (scopes )
412426 .setCallbackUri (URI .create ("http://localhost:" + port ));
413427
428+ if (httpTransportFactory != null ) {
429+ userAuthorizerBuilder .setHttpTransportFactory (httpTransportFactory );
430+ }
431+
414432 if (overrideProperties .containsKey (BigQueryJdbcUrlUtility .OAUTH2_TOKEN_URI_PROPERTY_NAME )) {
415433 userAuthorizerBuilder .setTokenServerUri (
416434 new URI (overrideProperties .get (BigQueryJdbcUrlUtility .OAUTH2_TOKEN_URI_PROPERTY_NAME )));
@@ -420,22 +438,32 @@ static UserAuthorizer getUserAuthorizer(
420438 }
421439
422440 static UserCredentials getCredentialsFromCode (
423- UserAuthorizer userAuthorizer , String code , String callerClassName ) throws IOException {
441+ UserAuthorizer userAuthorizer ,
442+ String code ,
443+ HttpTransportFactory httpTransportFactory ,
444+ String callerClassName )
445+ throws IOException {
424446 LOG .finer ("++enter++\t " + callerClassName );
425- return userAuthorizer .getCredentialsFromCode (code , URI .create ("" ));
447+ UserCredentials credentials = userAuthorizer .getCredentialsFromCode (code , URI .create ("" ));
448+ if (httpTransportFactory != null ) {
449+ credentials = credentials .toBuilder ().setHttpTransportFactory (httpTransportFactory ).build ();
450+ }
451+ return credentials ;
426452 }
427453
428454 private static GoogleCredentials getGoogleUserAccountCredentials (
429455 Map <String , String > authProperties ,
430456 Map <String , String > overrideProperties ,
457+ HttpTransportFactory httpTransportFactory ,
431458 String callerClassName ) {
432459 LOG .finer ("++enter++\t " + callerClassName );
433460 try {
434461 ServerSocket serverSocket = new ServerSocket (0 );
435462 serverSocket .setSoTimeout (USER_AUTH_TIMEOUT_MS );
436463 int port = serverSocket .getLocalPort ();
437464 UserAuthorizer userAuthorizer =
438- getUserAuthorizer (authProperties , overrideProperties , port , callerClassName );
465+ getUserAuthorizer (
466+ authProperties , overrideProperties , port , httpTransportFactory , callerClassName );
439467
440468 URL authURL = userAuthorizer .getAuthorizationUrl ("user" , "" , URI .create ("" ));
441469 String code ;
@@ -468,7 +496,7 @@ private static GoogleCredentials getGoogleUserAccountCredentials(
468496 throw new BigQueryJdbcRuntimeException ("User auth only supported in desktop environments" );
469497 }
470498
471- return getCredentialsFromCode (userAuthorizer , code , callerClassName );
499+ return getCredentialsFromCode (userAuthorizer , code , httpTransportFactory , callerClassName );
472500 } catch (IOException | URISyntaxException ex ) {
473501 throw new BigQueryJdbcRuntimeException (
474502 "Failed to establish connection using User Account authentication" , ex );
@@ -503,12 +531,13 @@ private static GoogleCredentials getPreGeneratedAccessTokenCredentials(
503531 static GoogleCredentials getPreGeneratedTokensCredentials (
504532 Map <String , String > authProperties ,
505533 Map <String , String > overrideProperties ,
534+ HttpTransportFactory httpTransportFactory ,
506535 String callerClassName ) {
507536 LOG .finer ("++enter++\t " + callerClassName );
508537 if (authProperties .containsKey (BigQueryJdbcUrlUtility .OAUTH_REFRESH_TOKEN_PROPERTY_NAME )) {
509538 try {
510539 return getPreGeneratedRefreshTokenCredentials (
511- authProperties , overrideProperties , callerClassName );
540+ authProperties , overrideProperties , httpTransportFactory , callerClassName );
512541 } catch (URISyntaxException ex ) {
513542 throw new BigQueryJdbcRuntimeException (
514543 "URISyntaxException during getPreGeneratedTokensCredentials" , ex );
@@ -522,6 +551,7 @@ static GoogleCredentials getPreGeneratedTokensCredentials(
522551 static UserCredentials getPreGeneratedRefreshTokenCredentials (
523552 Map <String , String > authProperties ,
524553 Map <String , String > overrideProperties ,
554+ HttpTransportFactory httpTransportFactory ,
525555 String callerClassName )
526556 throws URISyntaxException {
527557 LOG .finer ("++enter++\t " + callerClassName );
@@ -534,6 +564,10 @@ static UserCredentials getPreGeneratedRefreshTokenCredentials(
534564 .setClientSecret (
535565 authProperties .get (BigQueryJdbcUrlUtility .OAUTH_CLIENT_SECRET_PROPERTY_NAME ));
536566
567+ if (httpTransportFactory != null ) {
568+ userCredentialsBuilder .setHttpTransportFactory (httpTransportFactory );
569+ }
570+
537571 if (overrideProperties .containsKey (BigQueryJdbcUrlUtility .OAUTH2_TOKEN_URI_PROPERTY_NAME )) {
538572 userCredentialsBuilder .setTokenServerUri (
539573 new URI (overrideProperties .get (BigQueryJdbcUrlUtility .OAUTH2_TOKEN_URI_PROPERTY_NAME )));
@@ -548,10 +582,14 @@ static UserCredentials getPreGeneratedRefreshTokenCredentials(
548582 return userCredentialsBuilder .build ();
549583 }
550584
551- private static GoogleCredentials getApplicationDefaultCredentials (String callerClassName ) {
585+ private static GoogleCredentials getApplicationDefaultCredentials (
586+ HttpTransportFactory httpTransportFactory , String callerClassName ) {
552587 LOG .finer ("++enter++\t " + callerClassName );
553588 try {
554- GoogleCredentials credentials = GoogleCredentials .getApplicationDefault ();
589+ GoogleCredentials credentials =
590+ httpTransportFactory != null
591+ ? GoogleCredentials .getApplicationDefault (httpTransportFactory )
592+ : GoogleCredentials .getApplicationDefault ();
555593 String principal = "unknown" ;
556594 if (credentials instanceof ServiceAccountCredentials ) {
557595 principal = ((ServiceAccountCredentials ) credentials ).getClientEmail ();
@@ -572,7 +610,9 @@ private static GoogleCredentials getApplicationDefaultCredentials(String callerC
572610 }
573611
574612 private static GoogleCredentials getExternalAccountAuthCredentials (
575- Map <String , String > authProperties , String callerClassName ) {
613+ Map <String , String > authProperties ,
614+ HttpTransportFactory httpTransportFactory ,
615+ String callerClassName ) {
576616 LOG .finer ("++enter++\t " + callerClassName );
577617 try {
578618 JsonObject jsonObject = null ;
@@ -609,18 +649,28 @@ private static GoogleCredentials getExternalAccountAuthCredentials(
609649 }
610650 }
611651
652+ ExternalAccountCredentials credentials ;
612653 if (credentialsPath != null ) {
613- return ExternalAccountCredentials .fromStream (
614- Files .newInputStream (Paths .get (credentialsPath )));
654+ try (InputStream stream = Files .newInputStream (Paths .get (credentialsPath ))) {
655+ credentials =
656+ httpTransportFactory != null
657+ ? ExternalAccountCredentials .fromStream (stream , httpTransportFactory )
658+ : ExternalAccountCredentials .fromStream (stream );
659+ }
615660 } else if (jsonObject != null ) {
616- return ExternalAccountCredentials .fromStream (
617- new ByteArrayInputStream (jsonObject .toString ().getBytes ()));
661+ InputStream stream =
662+ new ByteArrayInputStream (jsonObject .toString ().getBytes (StandardCharsets .UTF_8 ));
663+ credentials =
664+ httpTransportFactory != null
665+ ? ExternalAccountCredentials .fromStream (stream , httpTransportFactory )
666+ : ExternalAccountCredentials .fromStream (stream );
618667 } else {
619668 IllegalArgumentException ex =
620669 new IllegalArgumentException ("Insufficient info provided for external authentication" );
621670 LOG .severe (ex .getMessage (), ex );
622671 throw ex ;
623672 }
673+ return credentials ;
624674 } catch (IOException e ) {
625675 throw new BigQueryJdbcRuntimeException (
626676 "IOException during getExternalAccountAuthCredentials" , e );
@@ -634,7 +684,8 @@ private static GoogleCredentials getExternalAccountAuthCredentials(
634684 private static GoogleCredentials getServiceAccountImpersonatedCredentials (
635685 GoogleCredentials credentials ,
636686 Boolean reqGoogleDriveScopeBool ,
637- Map <String , String > authProperties ) {
687+ Map <String , String > authProperties ,
688+ HttpTransportFactory httpTransportFactory ) {
638689
639690 String impersonationEmail =
640691 authProperties .get (BigQueryJdbcUrlUtility .OAUTH_SA_IMPERSONATION_EMAIL_PROPERTY_NAME );
@@ -684,6 +735,15 @@ private static GoogleCredentials getServiceAccountImpersonatedCredentials(
684735 throw ex ;
685736 }
686737
738+ if (httpTransportFactory != null ) {
739+ return ImpersonatedCredentials .create (
740+ credentials ,
741+ impersonationEmail ,
742+ impersonationChain ,
743+ impersonationScopes ,
744+ impersonationLifetimeInt ,
745+ httpTransportFactory );
746+ }
687747 return ImpersonatedCredentials .create (
688748 credentials ,
689749 impersonationEmail ,
0 commit comments