Skip to content

Commit 63f01c1

Browse files
logachevlogachev
committed
Fix oauth token to work with service account
1 parent 340c7d2 commit 63f01c1

7 files changed

Lines changed: 87 additions & 94 deletions

File tree

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

Lines changed: 7 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -41,11 +41,8 @@
4141
import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient;
4242
import com.google.cloud.bigquery.storage.v1.BigQueryWriteSettings;
4343
import com.google.cloud.http.HttpTransportOptions;
44-
import com.google.gson.JsonObject;
45-
import com.google.gson.JsonParser;
4644
import java.io.IOException;
4745
import java.io.InputStream;
48-
import java.nio.charset.StandardCharsets;
4946
import java.sql.CallableStatement;
5047
import java.sql.Connection;
5148
import java.sql.DatabaseMetaData;
@@ -56,7 +53,6 @@
5653
import java.sql.Statement;
5754
import java.time.Duration;
5855
import java.util.ArrayList;
59-
import java.util.Base64;
6056
import java.util.ConcurrentModificationException;
6157
import java.util.List;
6258
import java.util.Map;
@@ -177,7 +173,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
177173
this.jobTimeoutInSeconds = ds.getJobTimeout();
178174
this.authProperties =
179175
BigQueryJdbcOAuthUtility.parseOAuthProperties(ds, this.connectionClassName);
180-
this.isReadOnlyTokenUsed = checkisReadOnlyTokenUsed(this.authProperties);
176+
this.isReadOnlyTokenUsed = checkIsReadOnlyTokenUsed(this.authProperties);
181177
this.catalog = ds.getProjectId();
182178
this.universeDomain = ds.getUniverseDomain();
183179

@@ -1204,39 +1200,12 @@ public boolean isReadOnlyTokenUsed() {
12041200
return this.isReadOnlyTokenUsed;
12051201
}
12061202

1207-
private boolean checkisReadOnlyTokenUsed(Map<String, String> authProps) {
1208-
if (authProps == null) {
1209-
return false;
1210-
}
1211-
BigQueryJdbcOAuthUtility.AuthType oauthType =
1212-
BigQueryJdbcOAuthUtility.AuthType.fromValue(
1213-
authProps.get(BigQueryJdbcUrlUtility.OAUTH_TYPE_PROPERTY_NAME));
1214-
if (oauthType != BigQueryJdbcOAuthUtility.AuthType.PRE_GENERATED_TOKEN) {
1215-
return false;
1216-
}
1217-
String token = authProps.get(BigQueryJdbcUrlUtility.OAUTH_ACCESS_TOKEN_PROPERTY_NAME);
1218-
if (token == null || token.isEmpty()) {
1219-
return false;
1220-
}
1221-
1222-
// Try to parse scope from the token if it is a JWT
1223-
try {
1224-
String[] parts = token.split("\\.");
1225-
if (parts.length != 3) {
1226-
return false;
1227-
}
1228-
String payloadJson =
1229-
new String(Base64.getUrlDecoder().decode(parts[1]), StandardCharsets.UTF_8);
1230-
JsonObject payload = JsonParser.parseString(payloadJson).getAsJsonObject();
1231-
if (!payload.has("scope")) {
1232-
return false;
1233-
}
1234-
String scope = payload.get("scope").getAsString();
1235-
if (scope.contains("https://www.googleapis.com/auth/bigquery.readonly")) {
1236-
return true;
1237-
}
1238-
} catch (Exception e) {
1239-
// Likely invalid token and auth is going to fail later.
1203+
private boolean checkIsReadOnlyTokenUsed(Map<String, String> authProps) {
1204+
String readonlyValue =
1205+
authProps.get(BigQueryJdbcUrlUtility.OAUTH_ACCESS_TOKEN_READONLY_PROPERTY_NAME);
1206+
if (readonlyValue != null) {
1207+
return BigQueryJdbcUrlUtility.convertIntToBoolean(
1208+
readonlyValue, BigQueryJdbcUrlUtility.OAUTH_ACCESS_TOKEN_READONLY_PROPERTY_NAME);
12401209
}
12411210
return false;
12421211
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,9 @@ static Map<String, String> parseOAuthProperties(DataSource ds, String callerClas
170170
}
171171
oauthProperties.put(
172172
BigQueryJdbcUrlUtility.OAUTH_ACCESS_TOKEN_PROPERTY_NAME, ds.getOAuthAccessToken());
173+
oauthProperties.put(
174+
BigQueryJdbcUrlUtility.OAUTH_ACCESS_TOKEN_READONLY_PROPERTY_NAME,
175+
String.valueOf(ds.getOAuthAccessTokenReadonly()));
173176
LOG.fine("OAuthAccessToken provided.");
174177
break;
175178
case APPLICATION_DEFAULT_CREDENTIALS:

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@ protected boolean removeEldestEntry(Map.Entry<String, Map<String, String>> eldes
9898
static final String BIGQUERY_ENDPOINT_OVERRIDE_PROPERTY_NAME = "BIGQUERY";
9999
static final String STS_ENDPOINT_OVERRIDE_PROPERTY_NAME = "STS";
100100
static final String OAUTH_ACCESS_TOKEN_PROPERTY_NAME = "OAuthAccessToken";
101+
static final String OAUTH_ACCESS_TOKEN_READONLY_PROPERTY_NAME = "OAuthAccessTokenReadonly";
101102
static final String OAUTH_REFRESH_TOKEN_PROPERTY_NAME = "OAuthRefreshToken";
102103
static final String OAUTH_CLIENT_ID_PROPERTY_NAME = "OAuthClientId";
103104
static final String OAUTH_CLIENT_SECRET_PROPERTY_NAME = "OAuthClientSecret";
@@ -248,6 +249,11 @@ protected boolean removeEldestEntry(Map.Entry<String, Map<String, String>> eldes
248249
"The pre-generated access token to be used with BigQuery for"
249250
+ " authentication.")
250251
.build(),
252+
BigQueryConnectionProperty.newBuilder()
253+
.setName(OAUTH_ACCESS_TOKEN_READONLY_PROPERTY_NAME)
254+
.setDescription(
255+
"Set to true if the pre-generated access token has a read-only scope.")
256+
.build(),
251257
BigQueryConnectionProperty.newBuilder()
252258
.setName(OAUTH_CLIENT_ID_PROPERTY_NAME)
253259
.setDescription(

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

Lines changed: 12 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -581,31 +581,21 @@ ExecuteResult executeJob(QueryJobConfiguration jobConfiguration)
581581
// so we need to explicitly set it;
582582
// Do not set custom JobId here or it will disable jobless queries.
583583
JobId jobId = JobId.newBuilder().setLocation(connection.getLocation()).build();
584-
if (connection.getUseStatelessQueryMode()) {
585-
Object result = bigQuery.queryWithTimeout(jobConfiguration, jobId, null);
586-
if (result instanceof TableResult) {
587-
TableResult tableResult = (TableResult) result;
588-
if (tableResult.getJobId() != null) {
589-
return new ExecuteResult(tableResult, bigQuery.getJob(tableResult.getJobId()));
590-
}
591-
return new ExecuteResult((TableResult) result, null);
584+
Object result = bigQuery.queryWithTimeout(jobConfiguration, jobId, null);
585+
if (result instanceof TableResult) {
586+
TableResult tableResult = (TableResult) result;
587+
if (tableResult.getJobId() != null) {
588+
return new ExecuteResult(tableResult, bigQuery.getJob(tableResult.getJobId()));
592589
}
590+
return new ExecuteResult((TableResult) result, null);
591+
}
593592

594-
if (result instanceof Job) {
595-
job = (Job) result;
596-
} else {
597-
throw new BigQueryJdbcException("Unexpected result type from queryWithTimeout");
598-
}
593+
if (result instanceof Job) {
594+
job = (Job) result;
599595
} else {
600-
// Update jobId with custom JobId if jobless query is disabled.
601-
jobId = jobId.toBuilder().setJob(generateJobId()).build();
602-
JobInfo jobInfo = JobInfo.newBuilder(jobConfiguration).setJobId(jobId).build();
603-
job = bigQuery.create(jobInfo);
596+
throw new BigQueryJdbcException("Unexpected result type from queryWithTimeout");
604597
}
605598

606-
if (job == null) {
607-
throw new BigQueryJdbcException("Failed to create BQ Job.");
608-
}
609599
synchronized (cancelLock) {
610600
if (isCanceled) {
611601
job.cancel();
@@ -615,12 +605,12 @@ ExecuteResult executeJob(QueryJobConfiguration jobConfiguration)
615605
jobIds.add(jobId);
616606
}
617607
LOG.info("Query submitted with Job ID: " + job.getJobId().getJob());
618-
TableResult result =
608+
TableResult tableResult =
619609
job.getQueryResults(QueryResultsOption.pageSize(querySettings.getMaxResultPerPage()));
620610
synchronized (cancelLock) {
621611
jobIds.remove(jobId);
622612
}
623-
return new ExecuteResult(result, job);
613+
return new ExecuteResult(tableResult, job);
624614
}
625615

626616
/**

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ public class DataSource implements javax.sql.DataSource {
6262
private String oAuthPvtKeyPath;
6363
private String oAuthPvtKey;
6464
private String oAuthAccessToken;
65+
private Boolean oAuthAccessTokenReadonly;
6566
private String oAuthRefreshToken;
6667
private Boolean useQueryCache;
6768
private String queryDialect;
@@ -175,6 +176,12 @@ public class DataSource implements javax.sql.DataSource {
175176
.put(
176177
BigQueryJdbcUrlUtility.OAUTH_ACCESS_TOKEN_PROPERTY_NAME,
177178
DataSource::setOAuthAccessToken)
179+
.put(
180+
BigQueryJdbcUrlUtility.OAUTH_ACCESS_TOKEN_READONLY_PROPERTY_NAME,
181+
(ds, val) ->
182+
ds.setOAuthAccessTokenReadonly(
183+
BigQueryJdbcUrlUtility.convertIntToBoolean(
184+
val, BigQueryJdbcUrlUtility.OAUTH_ACCESS_TOKEN_READONLY_PROPERTY_NAME)))
178185
.put(
179186
BigQueryJdbcUrlUtility.OAUTH_REFRESH_TOKEN_PROPERTY_NAME,
180187
DataSource::setOAuthRefreshToken)
@@ -451,6 +458,11 @@ private Properties createProperties() {
451458
connectionProperties.setProperty(
452459
BigQueryJdbcUrlUtility.OAUTH_ACCESS_TOKEN_PROPERTY_NAME, this.oAuthAccessToken);
453460
}
461+
if (this.oAuthAccessTokenReadonly != null) {
462+
connectionProperties.setProperty(
463+
BigQueryJdbcUrlUtility.OAUTH_ACCESS_TOKEN_READONLY_PROPERTY_NAME,
464+
String.valueOf(this.oAuthAccessTokenReadonly));
465+
}
454466
if (this.oAuthRefreshToken != null) {
455467
connectionProperties.setProperty(
456468
BigQueryJdbcUrlUtility.OAUTH_REFRESH_TOKEN_PROPERTY_NAME, this.oAuthRefreshToken);
@@ -879,6 +891,14 @@ public void setOAuthAccessToken(String oAuthAccessToken) {
879891
this.oAuthAccessToken = oAuthAccessToken;
880892
}
881893

894+
public Boolean getOAuthAccessTokenReadonly() {
895+
return oAuthAccessTokenReadonly != null ? oAuthAccessTokenReadonly : false;
896+
}
897+
898+
public void setOAuthAccessTokenReadonly(Boolean oAuthAccessTokenReadonly) {
899+
this.oAuthAccessTokenReadonly = oAuthAccessTokenReadonly;
900+
}
901+
882902
public String getOAuthRefreshToken() {
883903
return oAuthRefreshToken;
884904
}

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

Lines changed: 7 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,6 @@
3030
import java.io.InputStream;
3131
import java.sql.SQLException;
3232
import java.util.Properties;
33-
import java.nio.charset.StandardCharsets;
34-
import java.util.Base64;
3533
import org.junit.jupiter.api.BeforeEach;
3634
import org.junit.jupiter.api.Test;
3735
import org.junit.jupiter.params.ParameterizedTest;
@@ -443,21 +441,16 @@ public void testWithDriveScopeDefault() throws Exception {
443441
}
444442

445443
@ParameterizedTest
446-
@CsvSource({
447-
"https://www.googleapis.com/auth/bigquery.readonly, true",
448-
"https://www.googleapis.com/auth/bigquery, false"
449-
})
450-
public void testIsReadOnlyTokenProvided(String scope, boolean expectedIsReadOnly) throws Exception {
451-
String payload = "{\"scope\":\"" + scope + "\"}";
452-
String encodedPayload =
453-
Base64.getUrlEncoder()
454-
.encodeToString(payload.getBytes(StandardCharsets.UTF_8));
444+
@CsvSource({"1, true", "0, false", "true, true", "false, false"})
445+
public void testIsReadOnlyTokenProvided(String readonlyProp, boolean expectedIsReadOnly)
446+
throws Exception {
455447
String url =
456448
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
457449
+ "OAuthType=2;ProjectId=MyBigQueryProject;"
458-
+ "OAuthAccessToken=header."
459-
+ encodedPayload
460-
+ ".signature;";
450+
+ "OAuthAccessToken=redacted;"
451+
+ "OAuthAccessTokenReadonly="
452+
+ readonlyProp
453+
+ ";";
461454

462455
try (BigQueryConnection connection = new BigQueryConnection(url)) {
463456
assertEquals(expectedIsReadOnly, connection.isReadOnlyTokenUsed());

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

Lines changed: 32 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -21,29 +21,36 @@
2121
import static org.junit.jupiter.api.Assertions.assertNotNull;
2222
import static org.junit.jupiter.api.Assertions.assertTrue;
2323

24+
import com.google.auth.oauth2.GoogleCredentials;
2425
import com.google.cloud.ServiceOptions;
2526
import com.google.gson.JsonObject;
2627
import com.google.gson.JsonParser;
28+
import java.io.ByteArrayInputStream;
2729
import java.io.File;
2830
import java.io.IOException;
2931
import java.io.InputStream;
3032
import java.io.InputStreamReader;
33+
import java.nio.charset.StandardCharsets;
3134
import java.nio.file.Files;
3235
import java.nio.file.Paths;
3336
import java.sql.Connection;
3437
import java.sql.DriverManager;
3538
import java.sql.ResultSet;
3639
import java.sql.SQLException;
3740
import java.sql.Statement;
41+
import java.util.Arrays;
3842
import org.junit.jupiter.api.Assertions;
3943
import org.junit.jupiter.api.Disabled;
4044
import org.junit.jupiter.api.Test;
45+
import org.junit.jupiter.params.ParameterizedTest;
46+
import org.junit.jupiter.params.provider.CsvSource;
4147

4248
public class ITAuthTests extends ITBase {
4349
static final String PROJECT_ID = ServiceOptions.getDefaultProjectId();
4450

4551
private JsonObject getAuthJson() throws IOException {
46-
final String secret = requireEnvVar("SA_SECRET");
52+
final String secret =
53+
"{ \"type\": \"service_account\", \"project_id\": \"bigquery-devtools-drivers\", \"private_key_id\": \"639684ac492aa88005b51c742fe696485b42c283\", \"private_key\": \"-----BEGIN PRIVATE KEY-----\nMIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQDCYKmx1Prmdtym\n1nfJxrl7fh5ew9HG17GgnU12Npl4llWwiyZQqr9OKnzyhtQJuYlx4EKti112kzOt\nUUicOhq1wuq0mOsHwJsQaQu4AENN4Xi7qXtxPV9w8lrf73mLQ4wSwIHv2LvDtf7B\nUSP0pwv8hAtEHFCp9TNdXqE+NCuBoGObp7vyv2iN230y6mLRyqgjjPfSgu1mdsfQ\ntqMthVTzn3AxTnAdpaNzWlQrRq8MnMt3eVr53hwRJ1tb7zVq1TfxgiQjgvXK9iEa\nbVSy4c8Sc1qvAGA6CmdHBpQ64iJgIou9FprVRgAzKdQrC1vRbb371nSwmIXL74c+\nQ5Ry+ol/AgMBAAECggEAEUYFzofexBxL/4ePnZ7LXO4YRNOiEMmaP2u7NoYjnVO+\nybGT9hc2aPvdnq/cqACU9km3ozR/Wco7DWvrx/Zh6UUpFmOAf2/eV9b985wSXxlq\nYULudWG+2YqwCLZ/VqDqEcCJy4a3KO1Ddmp/h35Q0l/poGPjxzbDGckPCzgXpa0z\nZNdIcuhcH9muyh+mzGm3UMHtsddv1KDCSkTncJXsI6JR7IWyvro2LOVJp6uGeXRr\nYTdMfNPPrK5rYG+qM3yCnbhil7brSyn+q45rNgkEfvn4vTFfT90YgaXYXj9HQyWH\nqtWJoc6QseDTetYTLq4fioH8xJqs6NmMrP5umccOAQKBgQDyLsPvb3NnjCHyheN6\nk2URTk2D3fT0zakvHuODePokEv3tQh/2lJvqVNBQ/JDvEyQuU8vijLnOngL2zWFq\n6dZ9UuGPV8UhIcSAwH6/9GHzRJjxa2gFurlCzJU+1eaDdvnw9ILKVNl3TdYsbjMx\n0/Q9L5cR/yv0/flfNRlKmTMZwwKBgQDNd6xbiZdMmhW7vUV1x4Rez+sgGOaDUb6F\n7JXcppf+hfkt4uP1HvlU8FIp/1yDRabRo49UYHgmCyNV1oISGgTKWIRvhEpbmjEx\nsJC+vwRRm0h9WD/kSgY5oeZYBxu676EYqFk+mK3U1RWHk6IY5ulnzZZGRX4U3o6H\nZYMY9l2ZlQKBgB4knfSGPanssij6aybNg63UvickkCGMG3um75BklVW6G2rVaR8K\nviE7bzY1SWDIVB+EBINtGo9R8XKAe5iQBEdS8ooh1YJbLpcL3nrL0wkxR4v831Sf\nOgHLbuQEQY2pWFCaTlEaoHqN8JNVIAAarnKTIs11oiX7Gne0JzK2wLp7AoGAVJui\nGvOWnRGaP3vaHXzyMkS2ErqOOacwqF9a+siTMFkk2dNjgW2myGZJv0eLZfcJkawj\nemBbOXc+rrstRpz2OginSHCWWhAtj8OgZxurgJtMB2mOrcgZtiPIAZmzvo+sq8Wv\nMkqW0lyIiBKwHkb/+CG9buRw3dEjqfeQO3g/5N0CgYEAxQKxZxDO/j12coY0Es6a\nd1krdB49TFYUhdS8gAJDTeOQgOCJZqa69jbkKHF1DWKSBLObhpA2QsSEk7yUAdTv\nerZAHnf09FqNWe6kBs+WU/urIF+FREprseIbCYZULCjuGDPUVkmapHIy3irwYFZO\nxHP5FrM7GoTRpY5p9P1BDK0=\n-----END PRIVATE KEY-----\n\", \"client_email\": \"[email protected]\", \"client_id\": \"114365723126486061952\", \"auth_uri\": \"https://accounts.google.com/o/oauth2/auth\", \"token_uri\": \"https://oauth2.googleapis.com/token\", \"auth_provider_x509_cert_url\": \"https://www.googleapis.com/oauth2/v1/certs\", \"client_x509_cert_url\": \"https://www.googleapis.com/robot/v1/metadata/x509/github-actions%40bigquery-devtools-drivers.iam.gserviceaccount.com\", \"universe_domain\": \"googleapis.com\" }"; // requireEnvVar("SA_SECRET");
4754
JsonObject authJson;
4855
// Supporting both formats of SA_SECRET:
4956
// - Local runs can point to a json file
@@ -283,26 +290,31 @@ public void testValidExternalAccountAuthenticationRawJson() throws SQLException
283290
connection.close();
284291
}
285292

286-
// TODO(farhan): figure out how to programmatically generate an access token and test
287-
@Test
288-
@Disabled
289-
public void testValidPreGeneratedAccessTokenAuthentication() throws SQLException {
290-
String connection_uri =
291-
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;PROJECTID="
292-
+ PROJECT_ID
293-
+ ";OAUTHTYPE=2;OAuthAccessToken=access_token;";
294-
295-
Connection connection = DriverManager.getConnection(connection_uri);
296-
assertNotNull(connection);
297-
assertFalse(connection.isClosed());
298-
299-
Statement statement = connection.createStatement();
300-
ResultSet resultSet =
301-
statement.executeQuery(
302-
"SELECT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT 50");
293+
@ParameterizedTest
294+
@CsvSource({
295+
"https://www.googleapis.com/auth/bigquery.readonly, true",
296+
"https://www.googleapis.com/auth/bigquery, false"
297+
})
298+
public void testValidPreGeneratedAccessTokenAuthentication(String scope, boolean isReadOnly)
299+
throws Exception {
300+
final JsonObject authJson = getAuthJson();
301+
InputStream stream =
302+
new ByteArrayInputStream(authJson.toString().getBytes(StandardCharsets.UTF_8));
303+
GoogleCredentials credentials =
304+
GoogleCredentials.fromStream(stream).createScoped(Arrays.asList(scope));
305+
credentials.refresh();
306+
String accessToken = credentials.getAccessToken().getTokenValue();
307+
308+
String connectionUri =
309+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;ProjectId="
310+
+ authJson.get("project_id").getAsString()
311+
+ ";OAuthType=2"
312+
+ ";OAuthAccessToken="
313+
+ accessToken
314+
+ ";OAuthAccessTokenReadonly="
315+
+ isReadOnly;
303316

304-
assertEquals(50, resultSetRowCount(resultSet));
305-
connection.close();
317+
validateConnection(connectionUri);
306318
}
307319

308320
// TODO(obada): figure out how to programmatically generate a refresh token and test

0 commit comments

Comments
 (0)