Skip to content

Commit 50aa549

Browse files
logachevlogachev
committed
fix(bigquery-jdbc): ensure largeResults are handled in PreparedStatement
1 parent 7b5efb5 commit 50aa549

4 files changed

Lines changed: 123 additions & 49 deletions

File tree

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,4 +284,8 @@ StandardSQLTypeName getSqlType(String name) {
284284
}
285285
return null;
286286
}
287+
288+
int getParametersArraySize() {
289+
return this.parametersArraySize;
290+
}
287291
}

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

Lines changed: 6 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,7 @@
6464
class BigQueryPreparedStatement extends BigQueryStatement implements PreparedStatement {
6565
private final BigQueryJdbcCustomLogger LOG = new BigQueryJdbcCustomLogger(this.toString());
6666
private static final char POSITIONAL_PARAMETER_CHAR = '?';
67-
// Making this protected so BigQueryCallableStatement subclass can access the parameters.
68-
protected final BigQueryParameterHandler parameterHandler;
67+
// parameterHandler is inherited from BigQueryStatement
6968
protected int parameterCount = 0;
7069
protected String currentQuery;
7170
private Queue<ArrayList<BigQueryJdbcParameter>> batchParameters = new LinkedList<>();
@@ -91,48 +90,25 @@ private int getParameterCount(String query) {
9190
@Override
9291
public ResultSet executeQuery() throws SQLException {
9392
logQueryExecutionStart(this.currentQuery);
94-
try {
95-
QueryJobConfiguration.Builder jobConfiguration = getJobConfig(this.currentQuery);
96-
jobConfiguration.setParameterMode("POSITIONAL");
97-
jobConfiguration = this.parameterHandler.configureParameters(jobConfiguration);
98-
runQuery(this.currentQuery, jobConfiguration.build());
99-
} catch (InterruptedException ex) {
100-
throw new BigQueryJdbcRuntimeException("Interrupted during executeQuery", ex);
101-
}
102-
return getCurrentResultSet();
93+
return super.executeQuery(this.currentQuery);
10394
}
10495

10596
@Override
10697
public long executeLargeUpdate() throws SQLException {
10798
logQueryExecutionStart(this.currentQuery);
108-
try {
109-
QueryJobConfiguration.Builder jobConfiguration = getJobConfig(this.currentQuery);
110-
jobConfiguration.setParameterMode("POSITIONAL");
111-
jobConfiguration = this.parameterHandler.configureParameters(jobConfiguration);
112-
runQuery(this.currentQuery, jobConfiguration.build());
113-
} catch (InterruptedException ex) {
114-
throw new BigQueryJdbcRuntimeException("Interrupted during executeLargeUpdate", ex);
115-
}
116-
return this.currentUpdateCount;
99+
return super.executeLargeUpdate(this.currentQuery);
117100
}
118101

119102
@Override
120103
public int executeUpdate() throws SQLException {
121-
return checkUpdateCount(executeLargeUpdate());
104+
logQueryExecutionStart(this.currentQuery);
105+
return super.executeUpdate(this.currentQuery);
122106
}
123107

124108
@Override
125109
public boolean execute() throws SQLException {
126110
logQueryExecutionStart(this.currentQuery);
127-
try {
128-
QueryJobConfiguration.Builder jobConfiguration = getJobConfig(this.currentQuery);
129-
jobConfiguration.setParameterMode("POSITIONAL");
130-
jobConfiguration = this.parameterHandler.configureParameters(jobConfiguration);
131-
runQuery(this.currentQuery, jobConfiguration.build());
132-
} catch (InterruptedException ex) {
133-
throw new BigQueryJdbcRuntimeException("Interrupted during execute", ex);
134-
}
135-
return getCurrentResultSet() != null;
111+
return super.execute(this.currentQuery);
136112
}
137113

138114
@Override

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

Lines changed: 33 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ public class BigQueryStatement extends BigQueryNoOpsStatement {
107107
protected int currentJobIdIndex = -1;
108108
protected List<String> batchQueries = new ArrayList<>();
109109
protected BigQueryConnection connection;
110+
protected BigQueryParameterHandler parameterHandler = null;
110111
protected String connectionId;
111112
protected int maxFieldSize = 0;
112113
protected int maxRows = 0;
@@ -244,9 +245,10 @@ public ResultSet executeQuery(String sql) throws SQLException {
244245
private ResultSet executeQueryImpl(String sql) throws SQLException {
245246
logQueryExecutionStart(sql);
246247
try {
247-
QueryJobConfiguration jobConfiguration =
248-
setDestinationDatasetAndTableInJobConfig(getJobConfig(sql).build());
249-
runQuery(sql, jobConfiguration);
248+
QueryJobConfiguration.Builder jobConfiguration = getJobConfig(sql);
249+
jobConfiguration = applyParametersIfPresent(jobConfiguration);
250+
jobConfiguration = setDestinationDatasetAndTableInJobConfig(jobConfiguration);
251+
runQuery(sql, jobConfiguration.build());
250252
} catch (InterruptedException ex) {
251253
throw new BigQueryJdbcException("Interrupted during executeQuery", ex);
252254
}
@@ -268,6 +270,7 @@ private long executeLargeUpdateImpl(String sql) throws SQLException {
268270
logQueryExecutionStart(sql);
269271
try {
270272
QueryJobConfiguration.Builder jobConfiguration = getJobConfig(sql);
273+
jobConfiguration = applyParametersIfPresent(jobConfiguration);
271274
runQuery(sql, jobConfiguration.build());
272275
} catch (InterruptedException ex) {
273276
throw new BigQueryJdbcRuntimeException("Interrupted during executeLargeUpdate", ex);
@@ -303,12 +306,13 @@ public boolean execute(String sql) throws SQLException {
303306
private boolean executeImpl(String sql) throws SQLException {
304307
logQueryExecutionStart(sql);
305308
try {
306-
QueryJobConfiguration jobConfiguration = getJobConfig(sql).build();
307-
// If Large Results are enabled, ensure query type is SELECT
308-
if (isLargeResultsEnabled() && getQueryType(jobConfiguration, null) == SqlType.SELECT) {
309+
QueryJobConfiguration.Builder jobConfiguration = getJobConfig(sql);
310+
jobConfiguration = applyParametersIfPresent(jobConfiguration);
311+
if (isLargeResultsEnabled()
312+
&& getQueryType(jobConfiguration.build(), null) == SqlType.SELECT) {
309313
jobConfiguration = setDestinationDatasetAndTableInJobConfig(jobConfiguration);
310314
}
311-
runQuery(sql, jobConfiguration);
315+
runQuery(sql, jobConfiguration.build());
312316
} catch (InterruptedException ex) {
313317
throw new BigQueryJdbcRuntimeException("Interrupted during execute", ex);
314318
}
@@ -625,35 +629,43 @@ void runQuery(String query, QueryJobConfiguration jobConfiguration)
625629
}
626630
}
627631

628-
private boolean isLargeResultsEnabled() {
632+
protected QueryJobConfiguration.Builder applyParametersIfPresent(
633+
QueryJobConfiguration.Builder jobConfigurationBuilder) throws SQLException {
634+
if (this.parameterHandler != null && this.parameterHandler.getParametersArraySize() > 0) {
635+
jobConfigurationBuilder.setParameterMode("POSITIONAL");
636+
jobConfigurationBuilder = this.parameterHandler.configureParameters(jobConfigurationBuilder);
637+
}
638+
return jobConfigurationBuilder;
639+
}
640+
641+
boolean isLargeResultsEnabled() {
629642
String destinationTable = this.querySettings.getDestinationTable();
630643
String destinationDataset = this.querySettings.getDestinationDataset();
631644
return destinationDataset != null || destinationTable != null;
632645
}
633646

634-
private QueryJobConfiguration setDestinationDatasetAndTableInJobConfig(
635-
QueryJobConfiguration jobConfiguration) {
647+
QueryJobConfiguration.Builder setDestinationDatasetAndTableInJobConfig(
648+
QueryJobConfiguration.Builder jobConfigurationBuilder) {
636649
String destinationTable = this.querySettings.getDestinationTable();
637650
String destinationDataset = this.querySettings.getDestinationDataset();
638651
if (destinationDataset != null || destinationTable != null) {
639652
if (destinationDataset != null) {
640653
checkIfDatasetExistElseCreate(destinationDataset);
641654
}
642-
if (jobConfiguration.useLegacySql() && destinationDataset == null) {
655+
if (getUseLegacySql() && destinationDataset == null) {
643656
checkIfDatasetExistElseCreate(DEFAULT_DATASET_NAME);
644657
destinationDataset = DEFAULT_DATASET_NAME;
645658
}
646659
if (destinationTable == null) {
647660
destinationTable = getDefaultDestinationTable();
648661
}
649-
return jobConfiguration.toBuilder()
662+
return jobConfigurationBuilder
650663
.setAllowLargeResults(this.querySettings.getAllowLargeResults())
651664
.setDestinationTable(TableId.of(destinationDataset, destinationTable))
652665
.setCreateDisposition(JobInfo.CreateDisposition.CREATE_IF_NEEDED)
653-
.setWriteDisposition(JobInfo.WriteDisposition.WRITE_TRUNCATE)
654-
.build();
666+
.setWriteDisposition(JobInfo.WriteDisposition.WRITE_TRUNCATE);
655667
}
656-
return jobConfiguration;
668+
return jobConfigurationBuilder;
657669
}
658670

659671
Job getNextJob() {
@@ -1349,14 +1361,16 @@ QueryJobConfiguration.Builder getJobConfig(String query) {
13491361
if (this.querySettings.getQueryProperties() != null) {
13501362
queryConfigBuilder.setConnectionProperties(this.querySettings.getQueryProperties());
13511363
}
1352-
boolean useLegacy =
1353-
QueryDialectType.BIG_QUERY.equals(
1354-
QueryDialectType.valueOf(this.querySettings.getQueryDialect()));
1355-
queryConfigBuilder.setUseLegacySql(useLegacy);
1364+
queryConfigBuilder.setUseLegacySql(getUseLegacySql());
13561365

13571366
return queryConfigBuilder;
13581367
}
13591368

1369+
private boolean getUseLegacySql() {
1370+
return QueryDialectType.BIG_QUERY.equals(
1371+
QueryDialectType.valueOf(this.querySettings.getQueryDialect()));
1372+
}
1373+
13601374
private void checkIfDatasetExistElseCreate(String datasetName) {
13611375
Dataset dataset = bigQuery.getDataset(DatasetId.of(datasetName));
13621376
if (dataset == null) {

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

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -689,4 +689,84 @@ public void testWrapperMethods() throws SQLException {
689689
BigQueryJdbcException.class, () -> bigQueryStatement.unwrap(java.sql.Connection.class));
690690
assertTrue(e.getMessage().contains("Cannot unwrap to java.sql.Connection"));
691691
}
692+
693+
@Test
694+
public void testPreparedStatementExecuteQueryWithLargeResults() throws Exception {
695+
// Setup connection mocks to return large results settings
696+
doReturn(true).when(bigQueryConnection).isAllowLargeResults();
697+
doReturn("test_dataset").when(bigQueryConnection).getDestinationDataset();
698+
doReturn("test_table").when(bigQueryConnection).getDestinationTable();
699+
700+
com.google.cloud.bigquery.Dataset dataset = mock(com.google.cloud.bigquery.Dataset.class);
701+
doReturn(dataset).when(bigquery).getDataset(any(com.google.cloud.bigquery.DatasetId.class));
702+
703+
// Create PreparedStatement
704+
BigQueryPreparedStatement preparedStatement =
705+
new BigQueryPreparedStatement(bigQueryConnection, query);
706+
BigQueryPreparedStatement preparedStatementSpy = Mockito.spy(preparedStatement);
707+
708+
TableResult result = Mockito.mock(TableResult.class);
709+
BigQueryJsonResultSet jsonResultSet = mock(BigQueryJsonResultSet.class);
710+
QueryJobConfiguration jobConfiguration = QueryJobConfiguration.newBuilder(query).build();
711+
Job job = getJobMock(result, jobConfiguration, StatementType.SELECT);
712+
713+
doReturn(job).when(bigquery).queryWithTimeout(any(), any(), any());
714+
doReturn(jsonResultSet).when(preparedStatementSpy).processJsonResultSet(result);
715+
716+
Job dryRunJob = getJobMock(null, jobConfiguration, StatementType.SELECT);
717+
doReturn(dryRunJob).when(bigquery).create(any(JobInfo.class));
718+
719+
// Act
720+
preparedStatementSpy.executeQuery();
721+
722+
// Assert
723+
ArgumentCaptor<QueryJobConfiguration> captor =
724+
ArgumentCaptor.forClass(QueryJobConfiguration.class);
725+
verify(bigquery).queryWithTimeout(captor.capture(), any(), any());
726+
QueryJobConfiguration capturedConfig = captor.getValue();
727+
728+
assertThat(capturedConfig.getDestinationTable())
729+
.isEqualTo(TableId.of("test_dataset", "test_table"));
730+
assertThat(capturedConfig.allowLargeResults()).isTrue();
731+
}
732+
733+
@Test
734+
public void testPreparedStatementExecuteWithLargeResults() throws Exception {
735+
// Setup connection mocks to return large results settings
736+
doReturn(true).when(bigQueryConnection).isAllowLargeResults();
737+
doReturn("test_dataset").when(bigQueryConnection).getDestinationDataset();
738+
doReturn("test_table").when(bigQueryConnection).getDestinationTable();
739+
740+
com.google.cloud.bigquery.Dataset dataset = mock(com.google.cloud.bigquery.Dataset.class);
741+
doReturn(dataset).when(bigquery).getDataset(any(com.google.cloud.bigquery.DatasetId.class));
742+
743+
// Create PreparedStatement
744+
BigQueryPreparedStatement preparedStatement =
745+
new BigQueryPreparedStatement(bigQueryConnection, query);
746+
BigQueryPreparedStatement preparedStatementSpy = Mockito.spy(preparedStatement);
747+
748+
TableResult result = Mockito.mock(TableResult.class);
749+
BigQueryJsonResultSet jsonResultSet = mock(BigQueryJsonResultSet.class);
750+
QueryJobConfiguration jobConfiguration = QueryJobConfiguration.newBuilder(query).build();
751+
Job job = getJobMock(result, jobConfiguration, StatementType.SELECT);
752+
753+
doReturn(job).when(bigquery).queryWithTimeout(any(), any(), any());
754+
doReturn(jsonResultSet).when(preparedStatementSpy).processJsonResultSet(result);
755+
756+
Job dryRunJob = getJobMock(null, jobConfiguration, StatementType.SELECT);
757+
doReturn(dryRunJob).when(bigquery).create(any(JobInfo.class));
758+
759+
// Act
760+
preparedStatementSpy.execute();
761+
762+
// Assert
763+
ArgumentCaptor<QueryJobConfiguration> captor =
764+
ArgumentCaptor.forClass(QueryJobConfiguration.class);
765+
verify(bigquery).queryWithTimeout(captor.capture(), any(), any());
766+
QueryJobConfiguration capturedConfig = captor.getValue();
767+
768+
assertThat(capturedConfig.getDestinationTable())
769+
.isEqualTo(TableId.of("test_dataset", "test_table"));
770+
assertThat(capturedConfig.allowLargeResults()).isTrue();
771+
}
692772
}

0 commit comments

Comments
 (0)