diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java index 1146f4d46d95..8a2fc25ddcf3 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java @@ -24,6 +24,7 @@ import com.google.api.gax.rpc.HeaderProvider; import com.google.api.gax.rpc.TransportChannelProvider; import com.google.auth.Credentials; +import com.google.auth.http.HttpTransportFactory; import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.BigQueryException; import com.google.cloud.bigquery.BigQueryOptions; @@ -265,11 +266,34 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { String.valueOf(ds.getRequestGoogleDriveScope()), BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME); + Map proxyProperties = + BigQueryJdbcProxyUtility.parseProxyProperties(ds, this.connectionClassName); + + this.sslTrustStorePath = ds.getSSLTrustStorePath(); + this.sslTrustStorePassword = ds.getSSLTrustStorePassword(); + this.httpConnectTimeout = ds.getHttpConnectTimeout(); + this.httpReadTimeout = ds.getHttpReadTimeout(); + + this.httpTransportOptions = + BigQueryJdbcProxyUtility.getHttpTransportOptions( + proxyProperties, + this.sslTrustStorePath, + this.sslTrustStorePassword, + this.httpConnectTimeout, + this.httpReadTimeout, + this.connectionClassName); + + HttpTransportFactory httpTransportFactory = + this.httpTransportOptions != null + ? this.httpTransportOptions.getHttpTransportFactory() + : null; + this.credentials = BigQueryJdbcOAuthUtility.getCredentials( authProperties, overrideProperties, this.reqGoogleDriveScope, + httpTransportFactory, this.connectionClassName); String defaultDatasetString = ds.getDefaultDataset(); if (defaultDatasetString == null || defaultDatasetString.trim().isEmpty()) { @@ -302,22 +326,6 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { this.destinationDataset = ds.getDestinationDataset(); this.destinationDatasetExpirationTime = ds.getDestinationDatasetExpirationTime(); this.kmsKeyName = ds.getKmsKeyName(); - Map proxyProperties = - BigQueryJdbcProxyUtility.parseProxyProperties(ds, this.connectionClassName); - - this.sslTrustStorePath = ds.getSSLTrustStorePath(); - this.sslTrustStorePassword = ds.getSSLTrustStorePassword(); - this.httpConnectTimeout = ds.getHttpConnectTimeout(); - this.httpReadTimeout = ds.getHttpReadTimeout(); - - this.httpTransportOptions = - BigQueryJdbcProxyUtility.getHttpTransportOptions( - proxyProperties, - this.sslTrustStorePath, - this.sslTrustStorePassword, - this.httpConnectTimeout, - this.httpReadTimeout, - this.connectionClassName); this.transportChannelProvider = BigQueryJdbcProxyUtility.getTransportChannelProvider( proxyProperties, diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtility.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtility.java index 0b9166d4cda2..d2236a7ad90b 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtility.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtility.java @@ -21,6 +21,7 @@ import com.google.api.client.util.PemReader; import com.google.api.client.util.SecurityUtils; +import com.google.auth.http.HttpTransportFactory; import com.google.auth.oauth2.AccessToken; import com.google.auth.oauth2.ClientId; import com.google.auth.oauth2.ExternalAccountCredentials; @@ -51,6 +52,7 @@ import java.net.URI; import java.net.URISyntaxException; import java.net.URL; +import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Paths; import java.security.GeneralSecurityException; @@ -262,6 +264,7 @@ static GoogleCredentials getCredentials( Map authProperties, Map overrideProperties, Boolean reqGoogleDriveScopeBool, + HttpTransportFactory httpTransportFactory, String callerClassName) { LOG.finer("++enter++\t" + callerClassName); @@ -272,21 +275,26 @@ static GoogleCredentials getCredentials( switch (authType) { case GOOGLE_SERVICE_ACCOUNT: credentials = - getGoogleServiceAccountCredentials(authProperties, overrideProperties, callerClassName); + getGoogleServiceAccountCredentials( + authProperties, overrideProperties, httpTransportFactory, callerClassName); break; case GOOGLE_USER_ACCOUNT: credentials = - getGoogleUserAccountCredentials(authProperties, overrideProperties, callerClassName); + getGoogleUserAccountCredentials( + authProperties, overrideProperties, httpTransportFactory, callerClassName); break; case PRE_GENERATED_TOKEN: credentials = - getPreGeneratedTokensCredentials(authProperties, overrideProperties, callerClassName); + getPreGeneratedTokensCredentials( + authProperties, overrideProperties, httpTransportFactory, callerClassName); break; case APPLICATION_DEFAULT_CREDENTIALS: - credentials = getApplicationDefaultCredentials(callerClassName); + credentials = getApplicationDefaultCredentials(httpTransportFactory, callerClassName); break; case EXTERNAL_ACCOUNT_AUTH: - credentials = getExternalAccountAuthCredentials(authProperties, callerClassName); + credentials = + getExternalAccountAuthCredentials( + authProperties, httpTransportFactory, callerClassName); break; default: IllegalStateException ex = new IllegalStateException(OAUTH_TYPE_ERROR_MESSAGE); @@ -295,7 +303,7 @@ static GoogleCredentials getCredentials( } return getServiceAccountImpersonatedCredentials( - credentials, reqGoogleDriveScopeBool, authProperties); + credentials, reqGoogleDriveScopeBool, authProperties, httpTransportFactory); } private static boolean isFileExists(String filename) { @@ -326,6 +334,7 @@ private static boolean isJson(byte[] value) { private static GoogleCredentials getGoogleServiceAccountCredentials( Map authProperties, Map overrideProperties, + HttpTransportFactory httpTransportFactory, String callerClassName) { LOG.finer("++enter++\t" + callerClassName); @@ -370,6 +379,10 @@ private static GoogleCredentials getGoogleServiceAccountCredentials( throw new BigQueryJdbcRuntimeException("No valid credentials provided."); } + if (httpTransportFactory != null) { + builder.setHttpTransportFactory(httpTransportFactory); + } + if (overrideProperties.containsKey(BigQueryJdbcUrlUtility.OAUTH2_TOKEN_URI_PROPERTY_NAME)) { builder.setTokenServerUri( new URI(overrideProperties.get(BigQueryJdbcUrlUtility.OAUTH2_TOKEN_URI_PROPERTY_NAME))); @@ -391,6 +404,7 @@ static UserAuthorizer getUserAuthorizer( Map authProperties, Map overrideProperties, int port, + HttpTransportFactory httpTransportFactory, String callerClassName) throws URISyntaxException { LOG.finer("++enter++\t" + callerClassName); @@ -411,6 +425,10 @@ static UserAuthorizer getUserAuthorizer( .setScopes(scopes) .setCallbackUri(URI.create("http://localhost:" + port)); + if (httpTransportFactory != null) { + userAuthorizerBuilder.setHttpTransportFactory(httpTransportFactory); + } + if (overrideProperties.containsKey(BigQueryJdbcUrlUtility.OAUTH2_TOKEN_URI_PROPERTY_NAME)) { userAuthorizerBuilder.setTokenServerUri( new URI(overrideProperties.get(BigQueryJdbcUrlUtility.OAUTH2_TOKEN_URI_PROPERTY_NAME))); @@ -420,14 +438,23 @@ static UserAuthorizer getUserAuthorizer( } static UserCredentials getCredentialsFromCode( - UserAuthorizer userAuthorizer, String code, String callerClassName) throws IOException { + UserAuthorizer userAuthorizer, + String code, + HttpTransportFactory httpTransportFactory, + String callerClassName) + throws IOException { LOG.finer("++enter++\t" + callerClassName); - return userAuthorizer.getCredentialsFromCode(code, URI.create("")); + UserCredentials credentials = userAuthorizer.getCredentialsFromCode(code, URI.create("")); + if (httpTransportFactory != null) { + credentials = credentials.toBuilder().setHttpTransportFactory(httpTransportFactory).build(); + } + return credentials; } private static GoogleCredentials getGoogleUserAccountCredentials( Map authProperties, Map overrideProperties, + HttpTransportFactory httpTransportFactory, String callerClassName) { LOG.finer("++enter++\t" + callerClassName); try { @@ -435,7 +462,8 @@ private static GoogleCredentials getGoogleUserAccountCredentials( serverSocket.setSoTimeout(USER_AUTH_TIMEOUT_MS); int port = serverSocket.getLocalPort(); UserAuthorizer userAuthorizer = - getUserAuthorizer(authProperties, overrideProperties, port, callerClassName); + getUserAuthorizer( + authProperties, overrideProperties, port, httpTransportFactory, callerClassName); URL authURL = userAuthorizer.getAuthorizationUrl("user", "", URI.create("")); String code; @@ -468,7 +496,7 @@ private static GoogleCredentials getGoogleUserAccountCredentials( throw new BigQueryJdbcRuntimeException("User auth only supported in desktop environments"); } - return getCredentialsFromCode(userAuthorizer, code, callerClassName); + return getCredentialsFromCode(userAuthorizer, code, httpTransportFactory, callerClassName); } catch (IOException | URISyntaxException ex) { throw new BigQueryJdbcRuntimeException( "Failed to establish connection using User Account authentication", ex); @@ -503,12 +531,13 @@ private static GoogleCredentials getPreGeneratedAccessTokenCredentials( static GoogleCredentials getPreGeneratedTokensCredentials( Map authProperties, Map overrideProperties, + HttpTransportFactory httpTransportFactory, String callerClassName) { LOG.finer("++enter++\t" + callerClassName); if (authProperties.containsKey(BigQueryJdbcUrlUtility.OAUTH_REFRESH_TOKEN_PROPERTY_NAME)) { try { return getPreGeneratedRefreshTokenCredentials( - authProperties, overrideProperties, callerClassName); + authProperties, overrideProperties, httpTransportFactory, callerClassName); } catch (URISyntaxException ex) { throw new BigQueryJdbcRuntimeException( "URISyntaxException during getPreGeneratedTokensCredentials", ex); @@ -522,6 +551,7 @@ static GoogleCredentials getPreGeneratedTokensCredentials( static UserCredentials getPreGeneratedRefreshTokenCredentials( Map authProperties, Map overrideProperties, + HttpTransportFactory httpTransportFactory, String callerClassName) throws URISyntaxException { LOG.finer("++enter++\t" + callerClassName); @@ -534,6 +564,10 @@ static UserCredentials getPreGeneratedRefreshTokenCredentials( .setClientSecret( authProperties.get(BigQueryJdbcUrlUtility.OAUTH_CLIENT_SECRET_PROPERTY_NAME)); + if (httpTransportFactory != null) { + userCredentialsBuilder.setHttpTransportFactory(httpTransportFactory); + } + if (overrideProperties.containsKey(BigQueryJdbcUrlUtility.OAUTH2_TOKEN_URI_PROPERTY_NAME)) { userCredentialsBuilder.setTokenServerUri( new URI(overrideProperties.get(BigQueryJdbcUrlUtility.OAUTH2_TOKEN_URI_PROPERTY_NAME))); @@ -548,10 +582,14 @@ static UserCredentials getPreGeneratedRefreshTokenCredentials( return userCredentialsBuilder.build(); } - private static GoogleCredentials getApplicationDefaultCredentials(String callerClassName) { + private static GoogleCredentials getApplicationDefaultCredentials( + HttpTransportFactory httpTransportFactory, String callerClassName) { LOG.finer("++enter++\t" + callerClassName); try { - GoogleCredentials credentials = GoogleCredentials.getApplicationDefault(); + GoogleCredentials credentials = + httpTransportFactory != null + ? GoogleCredentials.getApplicationDefault(httpTransportFactory) + : GoogleCredentials.getApplicationDefault(); String principal = "unknown"; if (credentials instanceof ServiceAccountCredentials) { principal = ((ServiceAccountCredentials) credentials).getClientEmail(); @@ -572,7 +610,9 @@ private static GoogleCredentials getApplicationDefaultCredentials(String callerC } private static GoogleCredentials getExternalAccountAuthCredentials( - Map authProperties, String callerClassName) { + Map authProperties, + HttpTransportFactory httpTransportFactory, + String callerClassName) { LOG.finer("++enter++\t" + callerClassName); try { JsonObject jsonObject = null; @@ -609,18 +649,28 @@ private static GoogleCredentials getExternalAccountAuthCredentials( } } + ExternalAccountCredentials credentials; if (credentialsPath != null) { - return ExternalAccountCredentials.fromStream( - Files.newInputStream(Paths.get(credentialsPath))); + try (InputStream stream = Files.newInputStream(Paths.get(credentialsPath))) { + credentials = + httpTransportFactory != null + ? ExternalAccountCredentials.fromStream(stream, httpTransportFactory) + : ExternalAccountCredentials.fromStream(stream); + } } else if (jsonObject != null) { - return ExternalAccountCredentials.fromStream( - new ByteArrayInputStream(jsonObject.toString().getBytes())); + InputStream stream = + new ByteArrayInputStream(jsonObject.toString().getBytes(StandardCharsets.UTF_8)); + credentials = + httpTransportFactory != null + ? ExternalAccountCredentials.fromStream(stream, httpTransportFactory) + : ExternalAccountCredentials.fromStream(stream); } else { IllegalArgumentException ex = new IllegalArgumentException("Insufficient info provided for external authentication"); LOG.severe(ex.getMessage(), ex); throw ex; } + return credentials; } catch (IOException e) { throw new BigQueryJdbcRuntimeException( "IOException during getExternalAccountAuthCredentials", e); @@ -634,7 +684,8 @@ private static GoogleCredentials getExternalAccountAuthCredentials( private static GoogleCredentials getServiceAccountImpersonatedCredentials( GoogleCredentials credentials, Boolean reqGoogleDriveScopeBool, - Map authProperties) { + Map authProperties, + HttpTransportFactory httpTransportFactory) { String impersonationEmail = authProperties.get(BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_EMAIL_PROPERTY_NAME); @@ -684,6 +735,15 @@ private static GoogleCredentials getServiceAccountImpersonatedCredentials( throw ex; } + if (httpTransportFactory != null) { + return ImpersonatedCredentials.create( + credentials, + impersonationEmail, + impersonationChain, + impersonationScopes, + impersonationLifetimeInt, + httpTransportFactory); + } return ImpersonatedCredentials.create( credentials, impersonationEmail, diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtilityTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtilityTest.java index f785173c00c1..1a086782e341 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtilityTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtilityTest.java @@ -22,8 +22,10 @@ import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; +import com.google.auth.http.HttpTransportFactory; import com.google.auth.oauth2.GoogleCredentials; import com.google.auth.oauth2.ImpersonatedCredentials; +import com.google.auth.oauth2.ServiceAccountCredentials; import com.google.auth.oauth2.UserAuthorizer; import com.google.auth.oauth2.UserCredentials; import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; @@ -108,7 +110,8 @@ public void testInvalidTokenUriForAuthType0() { DataSource.fromUrl(connectionString).getOverrideProperties(); try { - BigQueryJdbcOAuthUtility.getCredentials(oauthProperties, overrideProperties, false, null); + BigQueryJdbcOAuthUtility.getCredentials( + oauthProperties, overrideProperties, false, null, null); Assertions.fail(); } catch (BigQueryJdbcRuntimeException e) { assertThat(e.getMessage()).contains("Validation failure"); @@ -164,7 +167,8 @@ public void testGetCredentialsForPreGeneratedToken() { null); GoogleCredentials credentials = - BigQueryJdbcOAuthUtility.getCredentials(authProperties, Collections.EMPTY_MAP, false, null); + BigQueryJdbcOAuthUtility.getCredentials( + authProperties, Collections.EMPTY_MAP, false, null, null); assertThat(credentials).isNotNull(); } @@ -184,7 +188,8 @@ public void testGetCredentialsForPreGeneratedTokenTPC() throws IOException { Map overrideProperties = new HashMap<>(stringStringMap); GoogleCredentials credentials = - BigQueryJdbcOAuthUtility.getCredentials(authProperties, overrideProperties, false, null); + BigQueryJdbcOAuthUtility.getCredentials( + authProperties, overrideProperties, false, null, null); assertThat(credentials.getUniverseDomain()).isEqualTo("testDomain"); } @@ -199,7 +204,7 @@ public void testGetCredentialsForApplicationDefault() { null); GoogleCredentials credentials = - BigQueryJdbcOAuthUtility.getCredentials(authProperties, null, false, null); + BigQueryJdbcOAuthUtility.getCredentials(authProperties, null, false, null, null); assertThat(credentials).isNotNull(); } @@ -243,7 +248,7 @@ public void testGenerateUserAuthURL() { UserAuthorizer userAuthorizer = BigQueryJdbcOAuthUtility.getUserAuthorizer( - authProperties, new HashMap(), USER_AUTH_PORT, null); + authProperties, new HashMap(), USER_AUTH_PORT, null, null); String userId = "test_user"; String state = "test_state"; @@ -276,7 +281,7 @@ public void testGenerateUserAuthURLOverrideOauthEndpoint() { UserAuthorizer userAuthorizer = BigQueryJdbcOAuthUtility.getUserAuthorizer( - authProperties, overrideProperties, USER_AUTH_PORT, null); + authProperties, overrideProperties, USER_AUTH_PORT, null, null); assertThat(overrideTokenSeverURI).isEqualTo(userAuthorizer.toBuilder().getTokenServerUri()); } catch (URISyntaxException e) { @@ -317,7 +322,7 @@ public void testParseOverridePropsForRefreshTokenAuth() { UserCredentials userCredentials = BigQueryJdbcOAuthUtility.getPreGeneratedRefreshTokenCredentials( - authProperties, overrideProperties, null); + authProperties, overrideProperties, null, null); assertThat(userCredentials.toBuilder().getTokenServerUri()) .isEqualTo(URI.create("https://oauth2-private.p.googleapis.com/token")); @@ -427,7 +432,7 @@ public void testGetServiceAccountImpersonatedCredentialsForADC() throws Exceptio GoogleCredentials credentials = BigQueryJdbcOAuthUtility.getCredentials( - authProperties, java.util.Collections.EMPTY_MAP, false, null); + authProperties, java.util.Collections.EMPTY_MAP, false, null, null); assertThat(credentials).isInstanceOf(ImpersonatedCredentials.class); assertThat(((ImpersonatedCredentials) credentials).getSourceCredentials()) @@ -445,7 +450,8 @@ public void testGetServiceAccountImpersonatedCredentials() { .toString()), ""); GoogleCredentials credentials = - BigQueryJdbcOAuthUtility.getCredentials(authProperties, Collections.EMPTY_MAP, false, null); + BigQueryJdbcOAuthUtility.getCredentials( + authProperties, Collections.EMPTY_MAP, false, null, null); assertThat(credentials).isInstanceOf(ImpersonatedCredentials.class); } @@ -489,4 +495,74 @@ public void testPrivateKeyFromP12Bytes_wrong_password() { assertTrue(false); } } + + @Test + public void testGetCredentialsPropagatesHttpTransportFactory() { + Map authProperties = + BigQueryJdbcOAuthUtility.parseOAuthProperties( + DataSource.fromUrl( + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" + + "ProjectId=MyBigQueryProject;OAuthType=0;" + + "OAuthServiceAcctEmail=dummytest@dummytest.iam.gserviceaccount.com;" + + "OAuthPvtKey=" + + fake_pkcs8_key + + ";"), + null); + + HttpTransportFactory dummyFactory = () -> null; + + GoogleCredentials credentials = + BigQueryJdbcOAuthUtility.getCredentials( + authProperties, Collections.emptyMap(), false, dummyFactory, null); + + assertThat(credentials).isInstanceOf(ServiceAccountCredentials.class); + assertThat(((ServiceAccountCredentials) credentials).toBuilder().getHttpTransportFactory()) + .isEqualTo(dummyFactory); + } + + @Test + public void testGetImpersonatedCredentialsPropagatesHttpTransportFactory() { + Map authProperties = + BigQueryJdbcOAuthUtility.parseOAuthProperties( + DataSource.fromUrl( + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" + + "ProjectId=MyBigQueryProject;OAuthType=0;" + + "OAuthServiceAcctEmail=dummytest@dummytest.iam.gserviceaccount.com;" + + "OAuthPvtKey=" + + fake_pkcs8_key + + ";" + + "ServiceAccountImpersonationEmail=impersonated@email.com;"), + null); + + HttpTransportFactory dummyFactory = () -> null; + + GoogleCredentials credentials = + BigQueryJdbcOAuthUtility.getCredentials( + authProperties, Collections.emptyMap(), false, dummyFactory, null); + + assertThat(credentials).isInstanceOf(ImpersonatedCredentials.class); + assertThat(((ImpersonatedCredentials) credentials).toBuilder().getHttpTransportFactory()) + .isEqualTo(dummyFactory); + } + + @Test + public void testGetPreGeneratedRefreshTokenCredentialsPropagatesHttpTransportFactory() { + Map authProperties = + BigQueryJdbcOAuthUtility.parseOAuthProperties( + DataSource.fromUrl( + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" + + "ProjectId=MyBigQueryProject;OAuthType=2;" + + "OAuthRefreshToken=dummy_refresh_token;OAuthClientId=dummy_client_id;OAuthClientSecret=dummy_client_secret;"), + null); + + HttpTransportFactory dummyFactory = () -> null; + + GoogleCredentials credentials = + BigQueryJdbcOAuthUtility.getCredentials( + authProperties, Collections.emptyMap(), false, dummyFactory, null); + + assertThat(credentials).isInstanceOf(UserCredentials.class); + assertThat(((UserCredentials) credentials).toBuilder().getHttpTransportFactory()) + .isEqualTo(dummyFactory); + } }