Skip to content

Commit d38d9f0

Browse files
keshavdandevakeshavdandeva
authored
fix(bigquery-jdbc): resolve statementType via getJob fallback to avoid post-execution dry run for DDL (#14265)
1 parent 700e70a commit d38d9f0

3 files changed

Lines changed: 64 additions & 0 deletions

File tree

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -617,12 +617,23 @@ private void saveSessionIdIfPresent(TableResult tableResult) {
617617
}
618618

619619
private StatementType getStatementType(ExecuteResult executeResult) {
620+
// Fast path: Read statementType directly from TableResult
620621
if (executeResult.tableResult.getStatementType() != null) {
621622
return executeResult.tableResult.getStatementType();
622623
}
624+
// Jobful path: Read statementType from Job statistics when executed via JobCreationMode=1 or
625+
// jobs.insert
623626
if (executeResult.job != null && executeResult.job.getStatistics() instanceof QueryStatistics) {
624627
return ((QueryStatistics) executeResult.job.getStatistics()).getStatementType();
625628
}
629+
// Fallback path: Lazily fetch completed Job metadata to resolve statementType without dry
630+
// runs if omitted in TableResult
631+
if (executeResult.tableResult.getJobId() != null && this.bigQuery != null) {
632+
Job job = this.bigQuery.getJob(executeResult.tableResult.getJobId());
633+
if (job != null && job.getStatistics() instanceof QueryStatistics) {
634+
return ((QueryStatistics) job.getStatistics()).getStatementType();
635+
}
636+
}
626637
return null;
627638
}
628639

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1136,4 +1136,27 @@ public void testStatelessDmlExecutionUsesTableResultWithoutDryRunOrGetJob() thro
11361136
verify(bigquery, Mockito.never()).create(any(JobInfo.class));
11371137
verify(bigquery, Mockito.never()).getJob(any(JobId.class));
11381138
}
1139+
1140+
@Test
1141+
public void testNullStatementTypeInTableResultFallsBackToGetJobWithoutDryRun() throws Exception {
1142+
TableResult tableResultMock = mock(TableResult.class);
1143+
doReturn(this.jobId).when(tableResultMock).getJobId();
1144+
doReturn(null).when(tableResultMock).getStatementType();
1145+
doReturn(0L).when(tableResultMock).getTotalRows();
1146+
doReturn(Schema.of()).when(tableResultMock).getSchema();
1147+
1148+
doReturn(tableResultMock)
1149+
.when(bigquery)
1150+
.queryWithTimeout(any(QueryJobConfiguration.class), any(), any());
1151+
1152+
Job jobMock = getJobMock(null, null, StatementType.CREATE_TABLE);
1153+
doReturn(jobMock).when(bigquery).getJob(eq(this.jobId));
1154+
1155+
boolean result = bigQueryStatement.execute("CREATE TABLE dataset.my_table (x INT64)");
1156+
1157+
assertFalse(result);
1158+
assertNull(bigQueryStatement.getResultSet());
1159+
verify(bigquery, Mockito.times(1)).getJob(eq(this.jobId));
1160+
verify(bigquery, Mockito.never()).create(any(JobInfo.class));
1161+
}
11391162
}

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -460,4 +460,34 @@ public void testSessionPersistenceAcrossQueries() throws SQLException {
460460
}
461461
}
462462
}
463+
464+
@Test
465+
public void testNonIdempotentCreateTableAndDrop() throws SQLException {
466+
String tempTableName = "NON_IDEMPOTENT_TABLE_" + Math.abs(random.nextInt(100000));
467+
String createTableQuery =
468+
String.format("CREATE TABLE %s.%s (`id` INT64, `name` STRING);", DATASET, tempTableName);
469+
String dropTableQuery = String.format("DROP TABLE %s.%s;", DATASET, tempTableName);
470+
471+
try (Connection connection = DriverManager.getConnection(ITBase.connectionUrl);
472+
Statement statement = connection.createStatement()) {
473+
boolean hasResultSet = statement.execute(createTableQuery);
474+
assertFalse(hasResultSet);
475+
assertEquals(0, statement.getUpdateCount());
476+
477+
try (ResultSet rs =
478+
statement.executeQuery(
479+
String.format("SELECT count(*) FROM %s.%s;", DATASET, tempTableName))) {
480+
assertTrue(rs.next());
481+
assertEquals(0L, rs.getLong(1));
482+
assertFalse(rs.next());
483+
}
484+
} finally {
485+
try (Connection connection = DriverManager.getConnection(ITBase.connectionUrl);
486+
Statement statement = connection.createStatement()) {
487+
statement.execute(dropTableQuery);
488+
} catch (SQLException ignored) {
489+
// Ignore cleanup exception if table was not created
490+
}
491+
}
492+
}
463493
}

0 commit comments

Comments
 (0)