Skip to content

feat(bigquery-jdbc): Migrate getImportedKeys and getCrossReference to BQ API #13692

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

Open
keshavdandeva wants to merge 10 commits into
base: main
Choose a base branch
from
Open

feat(bigquery-jdbc): Migrate getImportedKeys and getCrossReference to BQ API #13692

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 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
d1a914f
feat(bigquery-jdbc): Migrate `getImportedKeys` and `getCrossReference…
keshavdandeva Jul 8, 2026
c569ed5
simplify bqTable.getTableConstraints() call
keshavdandeva Jul 8, 2026
552d629
merge: merge origin/main into jdbc/getImportedKeys and resolve conflicts
keshavdandeva Jul 8, 2026
7513163
add null safety check
keshavdandeva Jul 8, 2026
10ead1e
Merge branch 'main' into jdbc/getImportedKeys
keshavdandeva Jul 8, 2026
9dad6e1
nit: early return
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
224 changes: 190 additions & 34 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 @@ -27,13 +27,15 @@
import com.google.cloud.bigquery.BigQuery.TableListOption;
import com.google.cloud.bigquery.BigQuery.TableOption;
import com.google.cloud.bigquery.BigQueryException;
import com.google.cloud.bigquery.ColumnReference;
import com.google.cloud.bigquery.Dataset;
import com.google.cloud.bigquery.DatasetId;
import com.google.cloud.bigquery.Field;
import com.google.cloud.bigquery.Field.Mode;
import com.google.cloud.bigquery.FieldList;
import com.google.cloud.bigquery.FieldValue;
import com.google.cloud.bigquery.FieldValueList;
import com.google.cloud.bigquery.ForeignKey;
import com.google.cloud.bigquery.PrimaryKey;
import com.google.cloud.bigquery.Routine;
import com.google.cloud.bigquery.RoutineArgument;
Expand All @@ -43,7 +45,6 @@
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;
Expand Down Expand Up @@ -98,10 +99,7 @@ class BigQueryDatabaseMetaData implements DatabaseMetaData {
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_PRIMARY_KEYS_SQL = "DatabaseMetaData_GetPrimaryKeys.sql";
private static final String GET_IMPORTED_KEYS_SQL = "DatabaseMetaData_GetImportedKeys.sql";
private static final String GET_EXPORTED_KEYS_SQL = "DatabaseMetaData_GetExportedKeys.sql";
private static final String GET_CROSS_REFERENCE_SQL = "DatabaseMetaData_GetCrossReference.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 @@ -2451,10 +2449,7 @@ public ResultSet getPrimaryKeys(String catalog, String schema, String table) thr
resultSchemaFields,
ignoreAccessErrors,
(bqTable, results, fields) -> {
TableConstraints constraints = null;
if (bqTable.getDefinition() instanceof StandardTableDefinition) {
constraints = ((StandardTableDefinition) bqTable.getDefinition()).getTableConstraints();
}
TableConstraints constraints = bqTable.getTableConstraints();
processPrimaryKey(constraints, bqTable.getTableId(), results, fields);
});

Expand Down Expand Up @@ -2529,16 +2524,47 @@ private Comparator<FieldValueList> defineGetPrimaryKeysComparator(FieldList resu
@Override
public ResultSet getImportedKeys(String catalog, String schema, String table)
throws SQLException {
String sql = readSqlFromFile(GET_IMPORTED_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 getImportedKeys", 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, schema);

boolean ignoreAccessErrors = (catalog == null);
processTargetTablesConcurrently(
targetDatasets,
table,
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();
processForeignKey(fk, pkTableId, bqTable.getTableId(), results, fields);
}
});

Comparator<FieldValueList> comparator = definePkTableSortComparator(resultSchemaFields);
sortResults(collectedResults, comparator, "getImportedKeys", LOG);

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

@Override
Expand All @@ -2565,24 +2591,56 @@ public ResultSet getCrossReference(
String foreignSchema,
String foreignTable)
throws SQLException {
String sql = readSqlFromFile(GET_CROSS_REFERENCE_SQL);
Statement stmt = this.connection.createStatement();
try {
stmt.closeOnCompletion();
String formattedSql =
replaceSqlParameters(
sql,
parentCatalog,
parentSchema,
parentTable,
foreignCatalog,
foreignSchema,
foreignTable);
return stmt.executeQuery(formattedSql);
} catch (SQLException e) {
closeStatementIgnoreException(stmt);
throw new BigQueryJdbcException("Error executing getCrossReference", e);
if ((parentCatalog != null && parentCatalog.isEmpty())
|| (parentSchema != null && parentSchema.isEmpty())
|| parentTable == null
|| parentTable.isEmpty()
|| (foreignCatalog != null && foreignCatalog.isEmpty())
|| (foreignSchema != null && foreignSchema.isEmpty())
|| foreignTable == null
|| foreignTable.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(foreignCatalog, foreignSchema);

boolean ignoreAccessErrors = (foreignCatalog == null);
processTargetTablesConcurrently(
targetDatasets,
foreignTable,
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 (!equalsOrNullMatchesAll(parentCatalog, pkTableId.getProject())
|| !equalsOrNullMatchesAll(parentSchema, pkTableId.getDataset())
|| !parentTable.equals(pkTableId.getTable())) {
continue;
}
processForeignKey(fk, pkTableId, bqTable.getTableId(), results, fields);
}
});

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

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

@Override
Expand Down Expand Up @@ -5004,6 +5062,10 @@ private List<String> getAccessibleCatalogNames() throws SQLException {
return sortedCatalogs;
}

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);
Expand Down Expand Up @@ -5152,4 +5214,98 @@ private void processTargetTablesConcurrently(
taskFutures.forEach(future -> future.cancel(true));
}
}

private Schema defineForeignKeyResultSetSchema() {
return Schema.of(
Field.of("PKTABLE_CAT", StandardSQLTypeName.STRING),
Field.of("PKTABLE_SCHEM", StandardSQLTypeName.STRING),
Field.of("PKTABLE_NAME", StandardSQLTypeName.STRING),
Field.of("PKCOLUMN_NAME", StandardSQLTypeName.STRING),
Field.of("FKTABLE_CAT", StandardSQLTypeName.STRING),
Field.of("FKTABLE_SCHEM", StandardSQLTypeName.STRING),
Field.of("FKTABLE_NAME", StandardSQLTypeName.STRING),
Field.of("FKCOLUMN_NAME", StandardSQLTypeName.STRING),
Field.of("KEY_SEQ", StandardSQLTypeName.INT64),
Field.of("UPDATE_RULE", StandardSQLTypeName.INT64),
Field.of("DELETE_RULE", StandardSQLTypeName.INT64),
Field.of("FK_NAME", StandardSQLTypeName.STRING),
Field.of("PK_NAME", StandardSQLTypeName.STRING),
Field.of("DEFERRABILITY", StandardSQLTypeName.INT64));
}

private void processForeignKey(
ForeignKey fk,
TableId pkTableId,
TableId fkTableId,
List<FieldValueList> collectedResults,
FieldList resultSchemaFields) {
if (pkTableId == null) {
LOG.warning(
String.format(
"Skipping foreign key '%s' on table '%s' because its referenced table ID is null.",
fk.getName() != null ? fk.getName() : "unnamed", fkTableId.getTable()));
return;
}
List<ColumnReference> colRefs = fk.getColumnReferences();
if (colRefs != null) {
for (int i = 0; i < colRefs.size(); i++) {
ColumnReference colRef = colRefs.get(i);
List<FieldValue> row = new ArrayList<>();
row.add(createStringFieldValue(pkTableId.getProject())); // PKTABLE_CAT
row.add(createStringFieldValue(pkTableId.getDataset())); // PKTABLE_SCHEM
row.add(createStringFieldValue(pkTableId.getTable())); // PKTABLE_NAME
row.add(createStringFieldValue(colRef.getReferencedColumn())); // PKCOLUMN_NAME
row.add(createStringFieldValue(fkTableId.getProject())); // FKTABLE_CAT
row.add(createStringFieldValue(fkTableId.getDataset())); // FKTABLE_SCHEM
row.add(createStringFieldValue(fkTableId.getTable())); // FKTABLE_NAME
row.add(createStringFieldValue(colRef.getReferencingColumn())); // FKCOLUMN_NAME
row.add(createLongFieldValue((long) (i + 1))); // KEY_SEQ
row.add(createNullFieldValue()); // UPDATE_RULE
row.add(createNullFieldValue()); // DELETE_RULE
row.add(createStringFieldValue(fk.getName())); // FK_NAME
row.add(createNullFieldValue()); // PK_NAME
row.add(createNullFieldValue()); // DEFERRABILITY

collectedResults.add(FieldValueList.of(row, resultSchemaFields));
}
}
}

private Comparator<FieldValueList> definePkTableSortComparator(FieldList resultSchemaFields) {
final int PKTABLE_CAT_IDX = resultSchemaFields.getIndex("PKTABLE_CAT");
final int PKTABLE_SCHEM_IDX = resultSchemaFields.getIndex("PKTABLE_SCHEM");
final int PKTABLE_NAME_IDX = resultSchemaFields.getIndex("PKTABLE_NAME");
final int KEY_SEQ_IDX = resultSchemaFields.getIndex("KEY_SEQ");
return Comparator.comparing(
(FieldValueList fvl) -> getStringValueOrNull(fvl, PKTABLE_CAT_IDX),
Comparator.nullsFirst(String::compareTo))
.thenComparing(
(FieldValueList fvl) -> getStringValueOrNull(fvl, PKTABLE_SCHEM_IDX),
Comparator.nullsFirst(String::compareTo))
.thenComparing(
(FieldValueList fvl) -> getStringValueOrNull(fvl, PKTABLE_NAME_IDX),
Comparator.nullsFirst(String::compareTo))
.thenComparing(
(FieldValueList fvl) -> getLongValueOrNull(fvl, KEY_SEQ_IDX),
Comparator.nullsFirst(Long::compareTo));
}

private Comparator<FieldValueList> defineFkTableSortComparator(FieldList resultSchemaFields) {
final int FKTABLE_CAT_IDX = resultSchemaFields.getIndex("FKTABLE_CAT");
final int FKTABLE_SCHEM_IDX = resultSchemaFields.getIndex("FKTABLE_SCHEM");
final int FKTABLE_NAME_IDX = resultSchemaFields.getIndex("FKTABLE_NAME");
final int KEY_SEQ_IDX = resultSchemaFields.getIndex("KEY_SEQ");
return Comparator.comparing(
(FieldValueList fvl) -> getStringValueOrNull(fvl, FKTABLE_CAT_IDX),
Comparator.nullsFirst(String::compareTo))
.thenComparing(
(FieldValueList fvl) -> getStringValueOrNull(fvl, FKTABLE_SCHEM_IDX),
Comparator.nullsFirst(String::compareTo))
.thenComparing(
(FieldValueList fvl) -> getStringValueOrNull(fvl, FKTABLE_NAME_IDX),
Comparator.nullsFirst(String::compareTo))
.thenComparing(
(FieldValueList fvl) -> getLongValueOrNull(fvl, KEY_SEQ_IDX),
Comparator.nullsFirst(Long::compareTo));
}
}
72 changes: 0 additions & 72 deletions .../src/main/resources/com/google/cloud/bigquery/jdbc/DatabaseMetaData_GetCrossReference.sql
View file
Open in desktop

This file was deleted.

Loading
Loading