Skip to content

feat(bigquery-jdbc): migrate getExportedKeys to BQ API #13711

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
keshavdandeva merged 5 commits into from
Jul 9, 2026
+176 −144
Merged

feat(bigquery-jdbc): migrate getExportedKeys to BQ API #13711

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
039a3a6
feat(bigquery-jdbc): migrate `getExportedKeys` to BQ API
keshavdandeva Jul 8, 2026
7c889f7
resolve deadlock risk via async queue population
keshavdandeva Jul 8, 2026
0c39a62
address pr feedback
keshavdandeva Jul 8, 2026
ccef49e
nit: early returns
keshavdandeva Jul 8, 2026
19c350d
address pr feedback
keshavdandeva Jul 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
191 changes: 118 additions & 73 deletions ...-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,23 +52,18 @@
import com.google.cloud.bigquery.exception.BigQueryJdbcException;
import com.google.cloud.bigquery.jdbc.BigQueryJdbcTypeMappings.ColumnTypeInfo;
import com.google.cloud.bigquery.jdbc.utils.BigQueryJdbcVersionUtility;
import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.RowIdLifetime;
import java.sql.SQLException;
import java.sql.Statement;
import java.sql.Types;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashSet;
import java.util.List;
import java.util.Scanner;
import java.util.Set;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
Expand All @@ -89,17 +84,14 @@
*
* @see BigQueryStatement
*/
// TODO(neenu): test and verify after post MVP implementation.
class BigQueryDatabaseMetaData implements DatabaseMetaData {
final BigQueryJdbcCustomLogger LOG = new BigQueryJdbcCustomLogger(this.toString());
private static final String DATABASE_PRODUCT_NAME = "Google BigQuery";
private static final String DATABASE_PRODUCT_VERSION = "2.0";
private static final String DRIVER_NAME = "GoogleJDBCDriverForGoogleBigQuery";

private static final String SCHEMA_TERM = "Dataset";
private static final String CATALOG_TERM = "Project";
private static final String PROCEDURE_TERM = "Procedure";
private static final String GET_EXPORTED_KEYS_SQL = "DatabaseMetaData_GetExportedKeys.sql";
private static final int DEFAULT_PAGE_SIZE = 500;
private static final int DEFAULT_QUEUE_CAPACITY = 5000;
// Declared package-private for testing.
Expand Down Expand Up @@ -1820,11 +1812,9 @@ public ResultSet getCatalogs() throws SQLException {
final BlockingQueue<BigQueryFieldValueListWrapper> queue =
new LinkedBlockingQueue<>(catalogRows.isEmpty() ? 1 : catalogRows.size() + 1);

populateQueue(catalogRows, queue, schemaFields);
signalEndOfData(queue, schemaFields);
Future<?> fetcherFuture = populateQueueAsync(catalogRows, queue, schemaFields);

return BigQueryJsonResultSet.of(
catalogsSchema, catalogRows.size(), queue, null, new Future<?>[0]);
return BigQueryJsonResultSet.of(catalogsSchema, catalogRows.size(), queue, null, fetcherFuture);
}

Schema defineGetCatalogsSchema() {
Expand Down Expand Up @@ -1852,11 +1842,11 @@ public ResultSet getTableTypes() {
BlockingQueue<BigQueryFieldValueListWrapper> queue =
new LinkedBlockingQueue<>(tableTypeRows.size() + 1);

populateQueue(tableTypeRows, queue, tableTypesSchema.getFields());
signalEndOfData(queue, tableTypesSchema.getFields());
Future<?> fetcherFuture =
populateQueueAsync(tableTypeRows, queue, tableTypesSchema.getFields());

return BigQueryJsonResultSet.of(
tableTypesSchema, tableTypeRows.size(), queue, null, new Future<?>[0]);
tableTypesSchema, tableTypeRows.size(), queue, null, fetcherFuture);
}

static Schema defineGetTableTypesSchema() {
Expand Down Expand Up @@ -2413,17 +2403,6 @@ Schema defineGetVersionColumnsSchema() {
return Schema.of(fields);
}

private void closeStatementIgnoreException(Statement statement) {
if (statement == null) {
return;
}
try {
statement.close();
} catch (SQLException e) {
// pass
}
}

@Override
public ResultSet getPrimaryKeys(String catalog, String schema, String table) throws SQLException {
if ((catalog != null && catalog.isEmpty())
Expand Down Expand Up @@ -2458,9 +2437,8 @@ public ResultSet getPrimaryKeys(String catalog, String schema, String table) thr

final BlockingQueue<BigQueryFieldValueListWrapper> queue =
new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);
populateQueue(collectedResults, queue, resultSchemaFields);
signalEndOfData(queue, resultSchemaFields);
return BigQueryJsonResultSet.of(resultSchema, -1, queue, null);
Future<?> fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields);
return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture);
}

private Schema defineGetPrimaryKeysSchema() {
Expand Down Expand Up @@ -2561,24 +2539,59 @@ public ResultSet getImportedKeys(String catalog, String schema, String table)

final BlockingQueue<BigQueryFieldValueListWrapper> queue =
new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);
populateQueue(collectedResults, queue, resultSchemaFields);
signalEndOfData(queue, resultSchemaFields);
return BigQueryJsonResultSet.of(resultSchema, -1, queue, null);
Future<?> fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields);
return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture);
}

@Override
public ResultSet getExportedKeys(String catalog, String schema, String table)
throws SQLException {
String sql = readSqlFromFile(GET_EXPORTED_KEYS_SQL);
Statement stmt = this.connection.createStatement();
try {
stmt.closeOnCompletion();
String formattedSql = replaceSqlParameters(sql, catalog, schema, table);
return stmt.executeQuery(formattedSql);
} catch (SQLException e) {
closeStatementIgnoreException(stmt);
throw new BigQueryJdbcException("Error executing getExportedKeys", e);
if ((catalog != null && catalog.isEmpty())
|| (schema != null && schema.isEmpty())
|| table == null
|| table.isEmpty()) {
LOG.warning(
"Returning empty ResultSet as required parameters are null/empty, or catalog/schema parameters are empty.");
return new BigQueryJsonResultSet();
}

final Schema resultSchema = defineForeignKeyResultSetSchema();
final FieldList resultSchemaFields = resultSchema.getFields();

final List<FieldValueList> collectedResults = Collections.synchronizedList(new ArrayList<>());
List<DatasetId> targetDatasets = getTargetDatasets(catalog, null);

boolean ignoreAccessErrors = (catalog == null);
processTargetTablesConcurrently(
targetDatasets,
null,
collectedResults,
resultSchemaFields,
ignoreAccessErrors,
(bqTable, results, fields) -> {
TableConstraints constraints = bqTable.getTableConstraints();
if (constraints == null || constraints.getForeignKeys() == null) {
return;
}
for (ForeignKey fk : constraints.getForeignKeys()) {
TableId pkTableId = fk.getReferencedTable();
if (pkTableId == null
|| !equalsOrNullMatchesAll(catalog, pkTableId.getProject())
|| !equalsOrNullMatchesAll(schema, pkTableId.getDataset())
|| !table.equals(pkTableId.getTable())) {
continue;
}
processForeignKey(fk, pkTableId, bqTable.getTableId(), results, fields);
}
});

Comparator<FieldValueList> comparator = defineFkTableSortComparator(resultSchemaFields);
sortResults(collectedResults, comparator, "getExportedKeys", LOG);

final BlockingQueue<BigQueryFieldValueListWrapper> queue =
new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);
Future<?> fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields);
return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture);
}

@Override
Expand Down Expand Up @@ -2635,9 +2648,8 @@ && equalsOrNullMatchesAll(parentSchema, pkTableId.getDataset())

final BlockingQueue<BigQueryFieldValueListWrapper> queue =
new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);
populateQueue(collectedResults, queue, resultSchemaFields);
signalEndOfData(queue, resultSchemaFields);
return BigQueryJsonResultSet.of(resultSchema, -1, queue, null);
Future<?> fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields);
return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture);
}

@Override
Expand All @@ -2653,10 +2665,9 @@ public ResultSet getTypeInfo() {
final BlockingQueue<BigQueryFieldValueListWrapper> queue =
new LinkedBlockingQueue<>(typeInfoRows.size() + 1);

populateQueue(typeInfoRows, queue, schemaFields);
signalEndOfData(queue, schemaFields);
Future<?> fetcherFuture = populateQueueAsync(typeInfoRows, queue, schemaFields);
return BigQueryJsonResultSet.of(
typeInfoSchema, typeInfoRows.size(), queue, null, new Future<?>[0]);
typeInfoSchema, typeInfoRows.size(), queue, null, fetcherFuture);
}

Schema defineGetTypeInfoSchema() {
Expand Down Expand Up @@ -3577,9 +3588,8 @@ public ResultSet getSchemas(String catalog, String schemaPattern) throws SQLExce
}
Comparator<FieldValueList> comparator = defineGetSchemasComparator(resultSchemaFields);
sortResults(collectedResults, comparator, "getSchemas", LOG);
populateQueue(collectedResults, queue, resultSchemaFields);
signalEndOfData(queue, resultSchemaFields);
return BigQueryJsonResultSet.of(resultSchema, -1, queue, null);
Future<?> fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields);
return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture);
}

// Multi-Catalog Path: fan out using connection-scoped metadataExecutor
Expand Down Expand Up @@ -4891,6 +4901,19 @@ private void waitForTasksCompletion(List<Future<?>> taskFutures) throws Executio
LOG.info("Finished waiting for tasks.");
}

private Future<?> populateQueueAsync(
List<FieldValueList> collectedResults,
BlockingQueue<BigQueryFieldValueListWrapper> queue,
FieldList resultSchemaFields) {
return connection
.getMetadataExecutor()
.submit(
() -> {
populateQueue(collectedResults, queue, resultSchemaFields);
signalEndOfData(queue, resultSchemaFields);
});
}

private void populateQueue(
List<FieldValueList> collectedResults,
BlockingQueue<BigQueryFieldValueListWrapper> queue,
Expand Down Expand Up @@ -5063,24 +5086,6 @@ private boolean equalsOrNullMatchesAll(String expected, String actual) {
return expected == null || expected.equals(actual);
}

static String readSqlFromFile(String filename) {
InputStream in;
in = BigQueryDatabaseMetaData.class.getResourceAsStream(filename);
BufferedReader reader = new BufferedReader(new InputStreamReader(in));
StringBuilder builder = new StringBuilder();
try (Scanner scanner = new Scanner(reader)) {
while (scanner.hasNextLine()) {
String line = scanner.nextLine();
builder.append(line).append("\n");
}
}
return builder.toString();
}

String replaceSqlParameters(String sql, String... params) throws SQLException {
return String.format(sql, (Object[]) params);
}

private void writeErrorToQueue(BlockingQueue<BigQueryFieldValueListWrapper> queue, Throwable t) {
Exception ex = (t instanceof Exception) ? (Exception) t : new Exception(t);
BigQueryFieldValueListWrapper element = BigQueryFieldValueListWrapper.ofError(ex);
Expand Down Expand Up @@ -5168,7 +5173,7 @@ private void processTargetTablesConcurrently(
boolean ignoreAccessErrors,
TableProcessor processor)
throws SQLException {
if (targetDatasets.size() == 1) {
if (targetDatasets.size() == 1 && tableName != null) {
processSingleTable(
targetDatasets.get(0),
tableName,
Expand All @@ -5183,19 +5188,59 @@ private void processTargetTablesConcurrently(
List<Future<?>> taskFutures = new ArrayList<>();

try {
List<Callable<Void>> tasks = new ArrayList<>();
for (DatasetId datasetId : targetDatasets) {
taskFutures.add(
executor.submit(
if (tableName != null) {
Comment thread
keshavdandeva marked this conversation as resolved.
tasks.add(
() -> {
processSingleTable(
datasetId,
tableName,
collectedResults,
resultSchemaFields,
ignoreAccessErrors,
processor);
return null;
});
continue;
}

try {
Page<Table> tablesPage =
bigquery.listTables(datasetId, TableListOption.pageSize(DEFAULT_PAGE_SIZE));
if (tablesPage == null) {
continue;
}
for (Table table : tablesPage.iterateAll()) {
if (table.getDefinition() == null
|| table.getDefinition().getType() != TableDefinition.Type.TABLE) {
continue;
}
tasks.add(
() -> {
processSingleTable(
datasetId,
tableName,
table.getTableId().getTable(),
collectedResults,
resultSchemaFields,
ignoreAccessErrors,
processor);
return null;
}));
});
}
} catch (BigQueryException e) {
if (ignoreAccessErrors && (e.getCode() == 404 || e.getCode() == 403)) {
LOG.info(
"Dataset '%s' not found/accessible in project '%s' (API error %d). Skipping.",
datasetId.getDataset(), datasetId.getProject(), e.getCode());
continue;
}
throw new SQLException("Error while listing tables: " + e.getMessage(), e);
}
}

for (Callable<Void> task : tasks) {
taskFutures.add(executor.submit(task));
}
waitForTasksCompletion(taskFutures);
Comment thread
keshavdandeva marked this conversation as resolved.
if (Thread.currentThread().isInterrupted()) {
Expand Down
Loading
Loading