Skip to content
Open
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -265,11 +266,34 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
String.valueOf(ds.getRequestGoogleDriveScope()),
BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME);

Map<String, String> 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()) {
Expand Down Expand Up @@ -302,22 +326,6 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
this.destinationDataset = ds.getDestinationDataset();
this.destinationDatasetExpirationTime = ds.getDestinationDatasetExpirationTime();
this.kmsKeyName = ds.getKmsKeyName();
Map<String, String> 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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -262,6 +264,7 @@ static GoogleCredentials getCredentials(
Map<String, String> authProperties,
Map<String, String> overrideProperties,
Boolean reqGoogleDriveScopeBool,
HttpTransportFactory httpTransportFactory,
String callerClassName) {
LOG.finer("++enter++\t" + callerClassName);

Expand All @@ -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);
Expand All @@ -295,7 +303,7 @@ static GoogleCredentials getCredentials(
}

return getServiceAccountImpersonatedCredentials(
credentials, reqGoogleDriveScopeBool, authProperties);
credentials, reqGoogleDriveScopeBool, authProperties, httpTransportFactory);
}

private static boolean isFileExists(String filename) {
Expand Down Expand Up @@ -326,6 +334,7 @@ private static boolean isJson(byte[] value) {
private static GoogleCredentials getGoogleServiceAccountCredentials(
Map<String, String> authProperties,
Map<String, String> overrideProperties,
HttpTransportFactory httpTransportFactory,
String callerClassName) {
LOG.finer("++enter++\t" + callerClassName);

Expand Down Expand Up @@ -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)));
Expand All @@ -391,6 +404,7 @@ static UserAuthorizer getUserAuthorizer(
Map<String, String> authProperties,
Map<String, String> overrideProperties,
int port,
HttpTransportFactory httpTransportFactory,
String callerClassName)
throws URISyntaxException {
LOG.finer("++enter++\t" + callerClassName);
Expand All @@ -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)));
Expand All @@ -420,22 +438,32 @@ 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<String, String> authProperties,
Map<String, String> overrideProperties,
HttpTransportFactory httpTransportFactory,
String callerClassName) {
LOG.finer("++enter++\t" + callerClassName);
try {
ServerSocket serverSocket = new ServerSocket(0);
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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -503,12 +531,13 @@ private static GoogleCredentials getPreGeneratedAccessTokenCredentials(
static GoogleCredentials getPreGeneratedTokensCredentials(
Map<String, String> authProperties,
Map<String, String> 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);
Expand All @@ -522,6 +551,7 @@ static GoogleCredentials getPreGeneratedTokensCredentials(
static UserCredentials getPreGeneratedRefreshTokenCredentials(
Map<String, String> authProperties,
Map<String, String> overrideProperties,
HttpTransportFactory httpTransportFactory,
String callerClassName)
throws URISyntaxException {
LOG.finer("++enter++\t" + callerClassName);
Expand All @@ -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)));
Expand All @@ -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();
Expand All @@ -572,7 +610,9 @@ private static GoogleCredentials getApplicationDefaultCredentials(String callerC
}

private static GoogleCredentials getExternalAccountAuthCredentials(
Map<String, String> authProperties, String callerClassName) {
Map<String, String> authProperties,
HttpTransportFactory httpTransportFactory,
String callerClassName) {
LOG.finer("++enter++\t" + callerClassName);
try {
JsonObject jsonObject = null;
Expand Down Expand Up @@ -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);
Comment thread
keshavdandeva marked this conversation as resolved.
} else {
IllegalArgumentException ex =
new IllegalArgumentException("Insufficient info provided for external authentication");
LOG.severe(ex.getMessage(), ex);
throw ex;
}
return credentials;
Comment thread
keshavdandeva marked this conversation as resolved.
Comment thread
keshavdandeva marked this conversation as resolved.
} catch (IOException e) {
throw new BigQueryJdbcRuntimeException(
"IOException during getExternalAccountAuthCredentials", e);
Expand All @@ -634,7 +684,8 @@ private static GoogleCredentials getExternalAccountAuthCredentials(
private static GoogleCredentials getServiceAccountImpersonatedCredentials(
GoogleCredentials credentials,
Boolean reqGoogleDriveScopeBool,
Map<String, String> authProperties) {
Map<String, String> authProperties,
HttpTransportFactory httpTransportFactory) {

String impersonationEmail =
authProperties.get(BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_EMAIL_PROPERTY_NAME);
Expand Down Expand Up @@ -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,
Expand Down
Loading
Loading