Skip to content

Commit b1279ab

Browse files
logachevlogachev
authored
fix(bqjdbc): avoid reusing statement in DatabaseMetaData (#13224)
1 parent 1c6fe13 commit b1279ab

3 files changed

Lines changed: 91 additions & 17 deletions

File tree

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

Lines changed: 27 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,6 @@ class BigQueryDatabaseMetaData implements DatabaseMetaData {
139139

140140
String URL;
141141
BigQueryConnection connection;
142-
Statement statement = null;
143142
private final BigQuery bigquery;
144143
private final int metadataFetchThreadCount;
145144
private static final AtomicReference<String> parsedDriverVersion = new AtomicReference<>(null);
@@ -2639,16 +2638,27 @@ Schema defineGetVersionColumnsSchema() {
26392638
return Schema.of(fields);
26402639
}
26412640

2641+
private void closeStatementIgnoreException(Statement statement) {
2642+
if (statement == null) {
2643+
return;
2644+
}
2645+
try {
2646+
statement.close();
2647+
} catch (SQLException e) {
2648+
// pass
2649+
}
2650+
}
2651+
26422652
@Override
26432653
public ResultSet getPrimaryKeys(String catalog, String schema, String table) throws SQLException {
26442654
String sql = readSqlFromFile(GET_PRIMARY_KEYS_SQL);
2655+
Statement stmt = this.connection.createStatement();
26452656
try {
2646-
if (this.statement == null) {
2647-
this.statement = this.connection.createStatement();
2648-
}
2657+
stmt.closeOnCompletion();
26492658
String formattedSql = replaceSqlParameters(sql, catalog, schema, table);
2650-
return this.statement.executeQuery(formattedSql);
2659+
return stmt.executeQuery(formattedSql);
26512660
} catch (SQLException e) {
2661+
closeStatementIgnoreException(stmt);
26522662
throw new BigQueryJdbcException("Error executing getPrimaryKeys", e);
26532663
}
26542664
}
@@ -2657,13 +2667,13 @@ public ResultSet getPrimaryKeys(String catalog, String schema, String table) thr
26572667
public ResultSet getImportedKeys(String catalog, String schema, String table)
26582668
throws SQLException {
26592669
String sql = readSqlFromFile(GET_IMPORTED_KEYS_SQL);
2670+
Statement stmt = this.connection.createStatement();
26602671
try {
2661-
if (this.statement == null) {
2662-
this.statement = this.connection.createStatement();
2663-
}
2672+
stmt.closeOnCompletion();
26642673
String formattedSql = replaceSqlParameters(sql, catalog, schema, table);
2665-
return this.statement.executeQuery(formattedSql);
2674+
return stmt.executeQuery(formattedSql);
26662675
} catch (SQLException e) {
2676+
closeStatementIgnoreException(stmt);
26672677
throw new BigQueryJdbcException("Error executing getImportedKeys", e);
26682678
}
26692679
}
@@ -2672,13 +2682,13 @@ public ResultSet getImportedKeys(String catalog, String schema, String table)
26722682
public ResultSet getExportedKeys(String catalog, String schema, String table)
26732683
throws SQLException {
26742684
String sql = readSqlFromFile(GET_EXPORTED_KEYS_SQL);
2685+
Statement stmt = this.connection.createStatement();
26752686
try {
2676-
if (this.statement == null) {
2677-
this.statement = this.connection.createStatement();
2678-
}
2687+
stmt.closeOnCompletion();
26792688
String formattedSql = replaceSqlParameters(sql, catalog, schema, table);
2680-
return this.statement.executeQuery(formattedSql);
2689+
return stmt.executeQuery(formattedSql);
26812690
} catch (SQLException e) {
2691+
closeStatementIgnoreException(stmt);
26822692
throw new BigQueryJdbcException("Error executing getExportedKeys", e);
26832693
}
26842694
}
@@ -2693,10 +2703,9 @@ public ResultSet getCrossReference(
26932703
String foreignTable)
26942704
throws SQLException {
26952705
String sql = readSqlFromFile(GET_CROSS_REFERENCE_SQL);
2706+
Statement stmt = this.connection.createStatement();
26962707
try {
2697-
if (this.statement == null) {
2698-
this.statement = this.connection.createStatement();
2699-
}
2708+
stmt.closeOnCompletion();
27002709
String formattedSql =
27012710
replaceSqlParameters(
27022711
sql,
@@ -2706,8 +2715,9 @@ public ResultSet getCrossReference(
27062715
foreignCatalog,
27072716
foreignSchema,
27082717
foreignTable);
2709-
return this.statement.executeQuery(formattedSql);
2718+
return stmt.executeQuery(formattedSql);
27102719
} catch (SQLException e) {
2720+
closeStatementIgnoreException(stmt);
27112721
throw new BigQueryJdbcException("Error executing getCrossReference", e);
27122722
}
27132723
}

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3206,4 +3206,34 @@ public void testSupportsResultSetConcurrency() throws SQLException {
32063206
public void testGetSQLStateType() throws SQLException {
32073207
assertEquals(DatabaseMetaData.sqlStateSQL, dbMetadata.getSQLStateType());
32083208
}
3209+
3210+
@Test
3211+
public void testMetadataMethodsDoNotInterfere() throws SQLException {
3212+
Statement mockStatement1 = mock(Statement.class);
3213+
Statement mockStatement2 = mock(Statement.class);
3214+
ResultSet mockResultSet1 = mock(ResultSet.class);
3215+
ResultSet mockResultSet2 = mock(ResultSet.class);
3216+
3217+
when(bigQueryConnection.createStatement())
3218+
.thenReturn(mockStatement1)
3219+
.thenReturn(mockStatement2);
3220+
3221+
when(mockStatement1.executeQuery(any())).thenReturn(mockResultSet1);
3222+
when(mockStatement2.executeQuery(any())).thenReturn(mockResultSet2);
3223+
3224+
// Call first metadata method
3225+
ResultSet rs1 = dbMetadata.getPrimaryKeys("cat", "schema", "table");
3226+
assertSame(mockResultSet1, rs1);
3227+
3228+
// Call second metadata method
3229+
ResultSet rs2 = dbMetadata.getImportedKeys("cat", "schema", "table");
3230+
assertSame(mockResultSet2, rs2);
3231+
3232+
// Verify closeOnCompletion was called on both statements
3233+
verify(mockStatement1).closeOnCompletion();
3234+
verify(mockStatement2).closeOnCompletion();
3235+
3236+
// Verify connection.createStatement() was called twice
3237+
verify(bigQueryConnection, times(2)).createStatement();
3238+
}
32093239
}

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -341,6 +341,40 @@ public void testTableConstraints() throws SQLException {
341341
connection.close();
342342
}
343343

344+
@Test
345+
public void testMetadataResultSetsDoNotInterfere() throws SQLException {
346+
try (Connection connection =
347+
DriverManager.getConnection(String.format(connectionUrl, DEFAULT_CATALOG))) {
348+
DatabaseMetaData metaData = connection.getMetaData();
349+
350+
// Get primary keys for table 1
351+
ResultSet primaryKeys1 =
352+
metaData.getPrimaryKeys(PROJECT_ID, CONSTRAINTS_DATASET, CONSTRAINTS_TABLE_NAME);
353+
354+
// Get imported keys for table 1, BEFORE fully consuming primaryKeys1
355+
ResultSet importedKeys =
356+
metaData.getImportedKeys(PROJECT_ID, CONSTRAINTS_DATASET, CONSTRAINTS_TABLE_NAME);
357+
358+
// Now try to read from primaryKeys1
359+
Assertions.assertTrue(primaryKeys1.next());
360+
Assertions.assertEquals("id", primaryKeys1.getString("COLUMN_NAME"));
361+
362+
// Read from importedKeys
363+
Assertions.assertTrue(importedKeys.next());
364+
Assertions.assertEquals(CONSTRAINTS_TABLE_NAME2, importedKeys.getString("PKTABLE_NAME"));
365+
366+
// Read more from primaryKeys1 (should be finished now)
367+
Assertions.assertFalse(primaryKeys1.next());
368+
369+
// Read more from importedKeys
370+
Assertions.assertTrue(importedKeys.next());
371+
Assertions.assertEquals(CONSTRAINTS_TABLE_NAME2, importedKeys.getString("PKTABLE_NAME"));
372+
373+
primaryKeys1.close();
374+
importedKeys.close();
375+
}
376+
}
377+
344378
@Test
345379
public void testDatabaseMetadataGetCatalogs() throws SQLException {
346380

0 commit comments

Comments
 (0)