Skip to content

Commit cd57169

Browse files
logachevlogachev
authored
feat(bqjdbc): Bypass dry-run job for read-only tokens. (#12961)
1 parent 856febf commit cd57169

9 files changed

Lines changed: 146 additions & 60 deletions

File tree

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
143143
String partnerToken;
144144
DatabaseMetaData databaseMetaData;
145145
Boolean reqGoogleDriveScope;
146+
private boolean isReadOnlyTokenUsed = false;
146147

147148
BigQueryConnection(String url) throws IOException {
148149
this(url, DataSource.fromUrl(url));
@@ -172,6 +173,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
172173
this.jobTimeoutInSeconds = ds.getJobTimeout();
173174
this.authProperties =
174175
BigQueryJdbcOAuthUtility.parseOAuthProperties(ds, this.connectionClassName);
176+
this.isReadOnlyTokenUsed = checkIsReadOnlyTokenUsed(this.authProperties);
175177
this.catalog = ds.getProjectId();
176178
this.universeDomain = ds.getUniverseDomain();
177179

@@ -1193,4 +1195,18 @@ public CallableStatement prepareCall(
11931195
}
11941196
return prepareCall(sql);
11951197
}
1198+
1199+
public boolean isReadOnlyTokenUsed() {
1200+
return this.isReadOnlyTokenUsed;
1201+
}
1202+
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);
1209+
}
1210+
return false;
1211+
}
11961212
}

‎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: 19 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -333,6 +333,13 @@ private boolean executeImpl(String sql) throws SQLException {
333333

334334
StatementType getStatementType(QueryJobConfiguration queryJobConfiguration) throws SQLException {
335335
LOG.finest("++enter++");
336+
// BQ Read-only tokens are not recommended to use, they have a lot of known flaws.
337+
// We're supporting them in a limited capacity, for pure SELECT statements.
338+
if (this.connection.isReadOnlyTokenUsed()) {
339+
LOG.warning(
340+
"Read-only token detected, skipping dry run and assuming StatementType is SELECT.");
341+
return StatementType.SELECT;
342+
}
336343
QueryJobConfiguration dryRunJobConfiguration =
337344
queryJobConfiguration.toBuilder().setDryRun(true).build();
338345
Job job;
@@ -574,31 +581,21 @@ ExecuteResult executeJob(QueryJobConfiguration jobConfiguration)
574581
// so we need to explicitly set it;
575582
// Do not set custom JobId here or it will disable jobless queries.
576583
JobId jobId = JobId.newBuilder().setLocation(connection.getLocation()).build();
577-
if (connection.getUseStatelessQueryMode()) {
578-
Object result = bigQuery.queryWithTimeout(jobConfiguration, jobId, null);
579-
if (result instanceof TableResult) {
580-
TableResult tableResult = (TableResult) result;
581-
if (tableResult.getJobId() != null) {
582-
return new ExecuteResult(tableResult, bigQuery.getJob(tableResult.getJobId()));
583-
}
584-
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()));
585589
}
590+
return new ExecuteResult((TableResult) result, null);
591+
}
586592

587-
if (result instanceof Job) {
588-
job = (Job) result;
589-
} else {
590-
throw new BigQueryJdbcException("Unexpected result type from queryWithTimeout");
591-
}
593+
if (result instanceof Job) {
594+
job = (Job) result;
592595
} else {
593-
// Update jobId with custom JobId if jobless query is disabled.
594-
jobId = jobId.toBuilder().setJob(generateJobId()).build();
595-
JobInfo jobInfo = JobInfo.newBuilder(jobConfiguration).setJobId(jobId).build();
596-
job = bigQuery.create(jobInfo);
596+
throw new BigQueryJdbcException("Unexpected result type from queryWithTimeout");
597597
}
598598

599-
if (job == null) {
600-
throw new BigQueryJdbcException("Failed to create BQ Job.");
601-
}
602599
synchronized (cancelLock) {
603600
if (isCanceled) {
604601
job.cancel();
@@ -608,12 +605,12 @@ ExecuteResult executeJob(QueryJobConfiguration jobConfiguration)
608605
jobIds.add(jobId);
609606
}
610607
LOG.info("Query submitted with Job ID: " + job.getJobId().getJob());
611-
TableResult result =
608+
TableResult tableResult =
612609
job.getQueryResults(QueryResultsOption.pageSize(querySettings.getMaxResultPerPage()));
613610
synchronized (cancelLock) {
614611
jobIds.remove(jobId);
615612
}
616-
return new ExecuteResult(result, job);
613+
return new ExecuteResult(tableResult, job);
617614
}
618615

619616
/**

‎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: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@
3232
import java.util.Properties;
3333
import org.junit.jupiter.api.BeforeEach;
3434
import org.junit.jupiter.api.Test;
35+
import org.junit.jupiter.params.ParameterizedTest;
36+
import org.junit.jupiter.params.provider.CsvSource;
3537

3638
public class BigQueryConnectionTest {
3739

@@ -437,4 +439,21 @@ public void testWithDriveScopeDefault() throws Exception {
437439
assertFalse(connection.reqGoogleDriveScope);
438440
}
439441
}
442+
443+
@ParameterizedTest
444+
@CsvSource({"1, true", "0, false", "true, true", "false, false"})
445+
public void testIsReadOnlyTokenProvided(String readonlyProp, boolean expectedIsReadOnly)
446+
throws Exception {
447+
String url =
448+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
449+
+ "OAuthType=2;ProjectId=MyBigQueryProject;"
450+
+ "OAuthAccessToken=redacted;"
451+
+ "OAuthAccessTokenReadonly="
452+
+ readonlyProp
453+
+ ";";
454+
455+
try (BigQueryConnection connection = new BigQueryConnection(url)) {
456+
assertEquals(expectedIsReadOnly, connection.isReadOnlyTokenUsed());
457+
}
458+
}
440459
}

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

Lines changed: 31 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,8 @@
6868
import org.junit.jupiter.api.BeforeEach;
6969
import org.junit.jupiter.api.Disabled;
7070
import org.junit.jupiter.api.Test;
71+
import org.junit.jupiter.params.ParameterizedTest;
72+
import org.junit.jupiter.params.provider.ValueSource;
7173
import org.mockito.ArgumentCaptor;
7274
import org.mockito.Mockito;
7375

@@ -212,7 +214,7 @@ public void testExecSlowQueryPath() throws SQLException, InterruptedException {
212214
.build();
213215
Job job = getJobMock(tableResult, queryJobConfiguration, StatementType.SELECT);
214216

215-
doReturn(job).when(bigquery).create(any(JobInfo.class));
217+
doReturn(job).when(bigquery).queryWithTimeout(any(), any(), any());
216218

217219
doReturn(jobIdWrapper)
218220
.when(bigQueryStatementSpy)
@@ -297,14 +299,15 @@ public void setQueryTimeoutTest() throws Exception {
297299
QueryJobConfiguration.newBuilder(query).setJobTimeoutMs(10000L).build();
298300

299301
Job job = getJobMock(result, jobConfiguration, StatementType.SELECT);
300-
doReturn(job).when(bigquery).create(any(JobInfo.class));
302+
doReturn(job).when(bigquery).queryWithTimeout(any(), any(), any());
301303

302304
doReturn(jsonResultSet).when(bigQueryStatementSpy).processJsonResultSet(result);
303-
ArgumentCaptor<JobInfo> captor = ArgumentCaptor.forClass(JobInfo.class);
305+
ArgumentCaptor<QueryJobConfiguration> captor =
306+
ArgumentCaptor.forClass(QueryJobConfiguration.class);
304307

305308
bigQueryStatementSpy.runQuery(query, jobConfiguration);
306-
verify(bigquery).create(captor.capture());
307-
QueryJobConfiguration jobConfig = captor.getValue().getConfiguration();
309+
verify(bigquery).queryWithTimeout(captor.capture(), any(), any());
310+
QueryJobConfiguration jobConfig = captor.getValue();
308311
assertEquals(3000L, jobConfig.getJobTimeoutMs().longValue());
309312
}
310313

@@ -393,23 +396,16 @@ public void testJoblessQuery() throws SQLException, InterruptedException {
393396
TableResult tableResultJobfulMock = mock(TableResult.class);
394397
QueryJobConfiguration jobConf = QueryJobConfiguration.newBuilder("SELECT 1").build();
395398
Job jobMock = getJobMock(tableResultJobfulMock, jobConf, StatementType.SELECT);
396-
ArgumentCaptor<JobInfo> jobfulCaptor = ArgumentCaptor.forClass(JobInfo.class);
397-
doReturn(jobMock).when(bigquery).create(jobfulCaptor.capture());
399+
doReturn(jobMock)
400+
.when(bigquery)
401+
.queryWithTimeout(any(QueryJobConfiguration.class), any(), any());
398402
doReturn(mock(BigQueryJsonResultSet.class))
399403
.when(jobfulStatementSpy)
400404
.processJsonResultSet(tableResultJobfulMock);
401405

402406
jobfulStatementSpy.executeQuery("SELECT 1");
403407

404-
verify(bigquery).create(any(JobInfo.class));
405-
assertTrue(
406-
jobfulCaptor.getAllValues().stream()
407-
.noneMatch(
408-
jobInfo ->
409-
Boolean.TRUE.equals(
410-
((QueryJobConfiguration) jobInfo.getConfiguration()).dryRun())));
411-
verify(bigquery, Mockito.never())
412-
.queryWithTimeout(any(QueryJobConfiguration.class), any(), any());
408+
verify(bigquery).queryWithTimeout(any(QueryJobConfiguration.class), any(), any());
413409
}
414410

415411
@Test
@@ -422,7 +418,7 @@ public void testCloseCancelsJob() throws SQLException, InterruptedException {
422418
QueryJobConfiguration.newBuilder(query).setPriority(Priority.BATCH).build();
423419
Job job = getJobMock(tableResult, queryJobConfiguration, StatementType.SELECT);
424420

425-
doReturn(job).when(bigquery).create(any(JobInfo.class));
421+
doReturn(job).when(bigquery).queryWithTimeout(any(), any(), any());
426422
doReturn(false).when(bigQueryStatementSpy).useReadAPI(eq(tableResult));
427423
doReturn(mock(JobId.class)).when(tableResult).getJobId();
428424
Mockito.when(job.getQueryResults(any(QueryResultsOption.class)))
@@ -480,4 +476,22 @@ public void testCancelWithJoblessQuery() throws SQLException, InterruptedExcepti
480476
// And no backend cancellation was attempted
481477
verify(bigquery, Mockito.never()).cancel(any(JobId.class));
482478
}
479+
480+
@ParameterizedTest
481+
@ValueSource(booleans = {true, false})
482+
public void testGetStatementType(boolean isReadOnlyTokenUsed) throws Exception {
483+
doReturn(isReadOnlyTokenUsed).when(bigQueryConnection).isReadOnlyTokenUsed();
484+
485+
Job dryRunJobMock = getJobMock(null, null, StatementType.SELECT);
486+
doReturn(dryRunJobMock).when(bigquery).create(any(JobInfo.class));
487+
488+
BigQueryStatement statementSpy = Mockito.spy(bigQueryStatement);
489+
QueryJobConfiguration queryJobConfiguration = QueryJobConfiguration.newBuilder(query).build();
490+
491+
StatementType type = statementSpy.getStatementType(queryJobConfiguration);
492+
493+
assertThat(type).isEqualTo(StatementType.SELECT);
494+
verify(bigquery, isReadOnlyTokenUsed ? Mockito.never() : Mockito.times(1))
495+
.create(any(JobInfo.class));
496+
}
483497
}

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

Lines changed: 30 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -21,23 +21,29 @@
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();
@@ -283,26 +289,31 @@ public void testValidExternalAccountAuthenticationRawJson() throws SQLException
283289
connection.close();
284290
}
285291

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");
292+
@ParameterizedTest
293+
@CsvSource({
294+
"https://www.googleapis.com/auth/bigquery.readonly, true",
295+
"https://www.googleapis.com/auth/bigquery, false"
296+
})
297+
public void testValidPreGeneratedAccessTokenAuthentication(String scope, boolean isReadOnly)
298+
throws Exception {
299+
final JsonObject authJson = getAuthJson();
300+
InputStream stream =
301+
new ByteArrayInputStream(authJson.toString().getBytes(StandardCharsets.UTF_8));
302+
GoogleCredentials credentials =
303+
GoogleCredentials.fromStream(stream).createScoped(Arrays.asList(scope));
304+
credentials.refresh();
305+
String accessToken = credentials.getAccessToken().getTokenValue();
306+
307+
String connectionUri =
308+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;ProjectId="
309+
+ authJson.get("project_id").getAsString()
310+
+ ";OAuthType=2"
311+
+ ";OAuthAccessToken="
312+
+ accessToken
313+
+ ";OAuthAccessTokenReadonly="
314+
+ isReadOnly;
303315

304-
assertEquals(50, resultSetRowCount(resultSet));
305-
connection.close();
316+
validateConnection(connectionUri);
306317
}
307318

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

0 commit comments

Comments
 (0)