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..0c4175ab77af 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; @@ -263,6 +265,16 @@ static GoogleCredentials getCredentials( Map overrideProperties, Boolean reqGoogleDriveScopeBool, String callerClassName) { + return getCredentials( + authProperties, overrideProperties, reqGoogleDriveScopeBool, null, callerClassName); + } + + static GoogleCredentials getCredentials( + Map authProperties, + Map overrideProperties, + Boolean reqGoogleDriveScopeBool, + HttpTransportFactory httpTransportFactory, + String callerClassName) { LOG.finer("++enter++\t" + callerClassName); AuthType authType = @@ -272,21 +284,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 +312,7 @@ static GoogleCredentials getCredentials( } return getServiceAccountImpersonatedCredentials( - credentials, reqGoogleDriveScopeBool, authProperties); + credentials, reqGoogleDriveScopeBool, authProperties, httpTransportFactory); } private static boolean isFileExists(String filename) { @@ -326,6 +343,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 +388,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))); @@ -393,6 +415,16 @@ static UserAuthorizer getUserAuthorizer( int port, String callerClassName) throws URISyntaxException { + return getUserAuthorizer(authProperties, overrideProperties, port, null, callerClassName); + } + + static UserAuthorizer getUserAuthorizer( + Map authProperties, + Map overrideProperties, + int port, + HttpTransportFactory httpTransportFactory, + String callerClassName) + throws URISyntaxException { LOG.finer("++enter++\t" + callerClassName); List scopes = new ArrayList<>(); @@ -411,6 +443,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))); @@ -421,13 +457,27 @@ static UserAuthorizer getUserAuthorizer( static UserCredentials getCredentialsFromCode( UserAuthorizer userAuthorizer, String code, String callerClassName) throws IOException { + return getCredentialsFromCode(userAuthorizer, code, null, callerClassName); + } + + static UserCredentials getCredentialsFromCode( + 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 +485,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 +519,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 +554,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); @@ -524,6 +576,16 @@ static UserCredentials getPreGeneratedRefreshTokenCredentials( Map overrideProperties, String callerClassName) throws URISyntaxException { + return getPreGeneratedRefreshTokenCredentials( + authProperties, overrideProperties, null, callerClassName); + } + + static UserCredentials getPreGeneratedRefreshTokenCredentials( + Map authProperties, + Map overrideProperties, + HttpTransportFactory httpTransportFactory, + String callerClassName) + throws URISyntaxException { LOG.finer("++enter++\t" + callerClassName); UserCredentials.Builder userCredentialsBuilder = @@ -534,6 +596,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 +614,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 +642,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 +681,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 +716,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 +767,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..77d861a73e6c 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; @@ -489,4 +491,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); + } }