Skip to content

feat(bigquery-jdbc): migrate getPrimaryKeys to use BQ API #13691

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 8, 2026
Merged

feat(bigquery-jdbc): migrate getPrimaryKeys to use BQ API #13691

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
3af7ff2
feat(bigquery-jdbc): migrate getPrimaryKeys to use BQ API
keshavdandeva Jul 7, 2026
be0fd22
add sqlexception handling and proper cancellation
keshavdandeva Jul 7, 2026
f7af9a2
re-use helper methods and preserve exception details
keshavdandeva Jul 7, 2026
621e53a
lint
keshavdandeva Jul 7, 2026
767d271
address pr feedback
keshavdandeva Jul 8, 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
220 changes: 211 additions & 9 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 @@ -34,6 +34,7 @@
import com.google.cloud.bigquery.FieldList;
import com.google.cloud.bigquery.FieldValue;
import com.google.cloud.bigquery.FieldValueList;
import com.google.cloud.bigquery.PrimaryKey;
import com.google.cloud.bigquery.Routine;
import com.google.cloud.bigquery.RoutineArgument;
import com.google.cloud.bigquery.RoutineId;
Expand All @@ -42,7 +43,9 @@
import com.google.cloud.bigquery.StandardSQLField;
import com.google.cloud.bigquery.StandardSQLTableType;
import com.google.cloud.bigquery.StandardSQLTypeName;
import com.google.cloud.bigquery.StandardTableDefinition;
import com.google.cloud.bigquery.Table;
import com.google.cloud.bigquery.TableConstraints;
import com.google.cloud.bigquery.TableDefinition;
import com.google.cloud.bigquery.TableId;
import com.google.cloud.bigquery.exception.BigQueryJdbcException;
Expand Down Expand Up @@ -2425,16 +2428,102 @@ private void closeStatementIgnoreException(Statement statement) {

@Override
public ResultSet getPrimaryKeys(String catalog, String schema, String table) throws SQLException {
String sql = readSqlFromFile(GET_PRIMARY_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 getPrimaryKeys", e);
if ((catalog != null && catalog.isEmpty())
|| (schema != null && schema.isEmpty())
|| table == null
|| table.isEmpty()) {
LOG.warning(
"Returning empty ResultSet as one or more patterns are empty or catalog is empty.");
return new BigQueryJsonResultSet();
}

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

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

boolean ignoreAccessErrors = (catalog == null);
processTargetTablesConcurrently(
targetDatasets,
table,
collectedResults,
resultSchemaFields,
ignoreAccessErrors,
(bqTable, results, fields) -> {
TableConstraints constraints = null;
if (bqTable.getDefinition() instanceof StandardTableDefinition) {
constraints = ((StandardTableDefinition) bqTable.getDefinition()).getTableConstraints();
}
processPrimaryKey(constraints, bqTable.getTableId(), results, fields);
});

Comparator<FieldValueList> comparator = defineGetPrimaryKeysComparator(resultSchemaFields);
sortResults(collectedResults, comparator, "getPrimaryKeys", LOG);

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

private Schema defineGetPrimaryKeysSchema() {
List<Field> fields = new ArrayList<>(6);
fields.add(
Field.newBuilder("TABLE_CAT", StandardSQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build());
fields.add(
Field.newBuilder("TABLE_SCHEM", StandardSQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build());
fields.add(
Field.newBuilder("TABLE_NAME", StandardSQLTypeName.STRING)
.setMode(Field.Mode.REQUIRED)
.build());
fields.add(
Field.newBuilder("COLUMN_NAME", StandardSQLTypeName.STRING)
.setMode(Field.Mode.REQUIRED)
.build());
fields.add(
Field.newBuilder("KEY_SEQ", StandardSQLTypeName.INT64)
.setMode(Field.Mode.REQUIRED)
.build());
fields.add(
Field.newBuilder("PK_NAME", StandardSQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build());
return Schema.of(fields);
}

private void processPrimaryKey(
TableConstraints constraints,
TableId tableId,
List<FieldValueList> collectedResults,
FieldList resultSchemaFields) {
if (constraints == null || constraints.getPrimaryKey() == null) {
return;
}
PrimaryKey pk = constraints.getPrimaryKey();
List<String> pkColumns = pk.getColumns();
for (int i = 0; i < pkColumns.size(); i++) {
List<FieldValue> row = new ArrayList<>(6);
row.add(createStringFieldValue(tableId.getProject())); // 1. TABLE_CAT
row.add(createStringFieldValue(tableId.getDataset())); // 2. TABLE_SCHEM
row.add(createStringFieldValue(tableId.getTable())); // 3. TABLE_NAME
row.add(createStringFieldValue(pkColumns.get(i))); // 4. COLUMN_NAME
row.add(createLongFieldValue((long) (i + 1))); // 5. KEY_SEQ
row.add(createNullFieldValue()); // 6. PK_NAME
collectedResults.add(FieldValueList.of(row, resultSchemaFields));
}
}

private Comparator<FieldValueList> defineGetPrimaryKeysComparator(FieldList resultSchemaFields) {
final int COLUMN_NAME_IDX = resultSchemaFields.getIndex("COLUMN_NAME");
return Comparator.comparing(
(FieldValueList fvl) -> getStringValueOrNull(fvl, COLUMN_NAME_IDX),
Comparator.nullsFirst(String::compareTo));
}

@Override
Expand Down Expand Up @@ -4950,4 +5039,117 @@ private void writeErrorToQueue(BlockingQueue<BigQueryFieldValueListWrapper> queu
}
}
}

private List<DatasetId> getTargetDatasets(String catalog, String schema) throws SQLException {
if (schema == null) {
List<Dataset> datasets = fetchMatchingDatasets(catalog, null, null);
List<DatasetId> datasetIds = new ArrayList<>(datasets.size());
for (Dataset dataset : datasets) {
datasetIds.add(dataset.getDatasetId());
}
return datasetIds;
}

if (schema.isEmpty()) {
return Collections.emptyList();
}

List<String> projects = resolveTargetProjects(catalog);
List<DatasetId> datasetIds = new ArrayList<>(projects.size());
for (String project : projects) {
datasetIds.add(DatasetId.of(project, schema));

@logachev logachev Jul 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What happens if specific project doesn't have that dataset? e.g. have 3 projects, but only 2 have dataset named foo

@keshavdandeva keshavdandeva Jul 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch! I used the same logic here to throw exception (instead of logging/swallowing when there is an error with one of the catalogs) but it shouldn't be that way.
I've updated processSingleTable to catch BigQueryException. During catalog-wide searches (catalog == null), if a dataset/table is missing (404) or inaccessible (403), we now log it at INFO level and skip it instead of failing the entire metadata query. If a specific catalog was queried, these errors will still correctly propagate

}
return datasetIds;
}

private List<String> resolveTargetProjects(String catalog) throws SQLException {
return (catalog != null) ? Collections.singletonList(catalog) : getAccessibleCatalogNames();
}

@FunctionalInterface
private interface TableProcessor {
void process(Table bqTable, List<FieldValueList> collectedResults, FieldList resultSchemaFields)
throws SQLException;
}

private void processSingleTable(
DatasetId datasetId,
String tableName,
List<FieldValueList> collectedResults,
FieldList resultSchemaFields,
boolean ignoreAccessErrors,
TableProcessor processor)
throws SQLException {
Table bqTable;
try {
bqTable =
bigquery.getTable(TableId.of(datasetId.getProject(), datasetId.getDataset(), tableName));
} catch (BigQueryException e) {
if (ignoreAccessErrors && (e.getCode() == 404 || e.getCode() == 403)) {
LOG.info(
"Table '%s' or dataset '%s' not found/accessible in project '%s' (API error %d). Skipping.",
tableName, datasetId.getDataset(), datasetId.getProject(), e.getCode());
bqTable = null;
} else {
throw new SQLException("Error while fetching table metadata: " + e.getMessage(), e);
}
} catch (Exception e) {
throw new SQLException("Error while fetching table metadata: " + e.getMessage(), e);
}
if (bqTable != null && bqTable.getDefinition() != null) {
processor.process(bqTable, collectedResults, resultSchemaFields);
}
}

private void processTargetTablesConcurrently(
List<DatasetId> targetDatasets,
String tableName,
List<FieldValueList> collectedResults,
FieldList resultSchemaFields,
boolean ignoreAccessErrors,
TableProcessor processor)
throws SQLException {
if (targetDatasets.size() == 1) {
processSingleTable(
targetDatasets.get(0),
tableName,
collectedResults,
resultSchemaFields,
ignoreAccessErrors,
processor);
return;
}

ExecutorService executor = connection.getMetadataExecutor();
List<Future<?>> taskFutures = new ArrayList<>();

try {
for (DatasetId datasetId : targetDatasets) {
taskFutures.add(
executor.submit(
() -> {
processSingleTable(
datasetId,
tableName,
collectedResults,
resultSchemaFields,
ignoreAccessErrors,
processor);
return null;
}));
}
waitForTasksCompletion(taskFutures);
if (Thread.currentThread().isInterrupted()) {
throw new SQLException("Interrupted while parallel-fetching metadata");
}
} catch (ExecutionException e) {
Throwable cause = e.getCause();
if (cause instanceof SQLException) {
throw (SQLException) cause;
}
throw new SQLException("Error while fetching metadata", e);
} finally {
taskFutures.forEach(future -> future.cancel(true));
}
Comment thread
keshavdandeva marked this conversation as resolved.
}
}
30 changes: 0 additions & 30 deletions ...dbc/src/main/resources/com/google/cloud/bigquery/jdbc/DatabaseMetaData_GetPrimaryKeys.sql
View file
Open in desktop

This file was deleted.

49 changes: 0 additions & 49 deletions ...query-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaDataTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,25 +159,6 @@ public void testBigqueryDatabaseMetaDataGetters() throws SQLException {
assertEquals("Project", dbMetadata.getCatalogTerm());
}

@Test
public void testReadSqlFromFile() throws SQLException {
BigQueryDatabaseMetaData dbMetadata = new BigQueryDatabaseMetaData(bigQueryConnection);

String primaryKeysQuery =
BigQueryDatabaseMetaData.readSqlFromFile("DatabaseMetaData_GetPrimaryKeys.sql");
assertTrue(primaryKeysQuery.contains("pk$"));

try {
when(bigQueryConnection.prepareStatement(primaryKeysQuery)).thenCallRealMethod();
String sql =
dbMetadata.replaceSqlParameters(
primaryKeysQuery, "project_name", "dataset_name", "table_name");
assertTrue(sql.contains("project_name.dataset_name.INFORMATION_SCHEMA.KEY_COLUMN_USAGE"));
} catch (SQLException e) {
throw new RuntimeException(e);
}
}

@Test
public void testNeedsListing() {
assertTrue(dbMetadata.needsListing(null), "Null pattern should require listing");
Expand Down Expand Up @@ -3251,36 +3232,6 @@ public void testWrapperMethods() throws SQLException {
assertThat((Throwable) e).hasMessageThat().contains("Cannot unwrap to java.sql.Connection");
}

@Test
public void testMetadataMethodsDoNotInterfere() throws SQLException {
Statement mockStatement1 = mock(Statement.class);
Statement mockStatement2 = mock(Statement.class);
ResultSet mockResultSet1 = mock(ResultSet.class);
ResultSet mockResultSet2 = mock(ResultSet.class);

when(bigQueryConnection.createStatement())
.thenReturn(mockStatement1)
.thenReturn(mockStatement2);

when(mockStatement1.executeQuery(any())).thenReturn(mockResultSet1);
when(mockStatement2.executeQuery(any())).thenReturn(mockResultSet2);

// Call first metadata method
ResultSet rs1 = dbMetadata.getPrimaryKeys("cat", "schema", "table");
assertSame(mockResultSet1, rs1);

// Call second metadata method
ResultSet rs2 = dbMetadata.getImportedKeys("cat", "schema", "table");
assertSame(mockResultSet2, rs2);

// Verify closeOnCompletion was called on both statements
verify(mockStatement1).closeOnCompletion();
verify(mockStatement2).closeOnCompletion();

// Verify connection.createStatement() was called twice
verify(bigQueryConnection, times(2)).createStatement();
}

@ParameterizedTest
@EnumSource(StandardSQLTypeName.class)
public void testMetadataAndResultSetMetadataTypeMappingConsistency(StandardSQLTypeName type) {
Expand Down
Loading