perf(bigquery-jdbc): optimize getExportedKeys performance using hybrid metadata lookup (#13734)

b/534326814

This PR optimizes `DatabaseMetaData.getExportedKeys()` by restoring the
highly performant `INFORMATION_SCHEMA` SQL lookup path, correcting a
performance regression introduced when the method was migrated to
sequential REST API lookups.

To ensure complete contract safety across both standard native tables
and external REST catalog tables (PCNT tables), we implement a hybrid
routing strategy:

1. **PCNT Early Return (Fast Path)**: 
If `schema` contains a dot (`.`), we bypass SQL execution entirely and
return a structured empty `ResultSet` instantly.
2. **Wildcard Fallback (REST API Path)**: 
If either `catalog` or `schema` arguments are `null` (representing
wildcard metadata searches), the driver falls back to retrieving
constraints concurrently via the Java Client REST API.
3. **Standard Introspection (SQL Path)**: 
For targeted native table queries, the driver executes the optimized
`DatabaseMetaData_GetExportedKeys.sql` query template.

Author: keshavdandeva

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

@@ -52,11 +52,17 @@
 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.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
 import java.sql.Connection;
 import java.sql.DatabaseMetaData;
+import java.sql.PreparedStatement;
 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;
@@ -94,6 +100,8 @@ class BigQueryDatabaseMetaData implements DatabaseMetaData {
   private static final String PROCEDURE_TERM = "Procedure";
   private static final int DEFAULT_PAGE_SIZE = 500;
   private static final int DEFAULT_QUEUE_CAPACITY = 5000;
+  private static final String GET_EXPORTED_KEYS_SQL = "DatabaseMetaData_GetExportedKeys.sql";
+  private static String exportedKeysSqlContent;
   // Declared package-private for testing.
   static final String GOOGLE_SQL_QUOTED_IDENTIFIER = "`";
   // Does not include SQL:2003 Keywords as per JDBC spec.
@@ -2559,40 +2567,65 @@ public ResultSet getExportedKeys(String catalog, String schema, String table)
     final Schema resultSchema = defineForeignKeyResultSetSchema();
     final FieldList resultSchemaFields = resultSchema.getFields();
 
-    final List<FieldValueList> collectedResults = Collections.synchronizedList(new ArrayList<>());
-    List<DatasetId> targetDatasets = getTargetDatasets(catalog, null);
+    // Early return for PCNT catalog schemas (containing '.') as they do not support table
+    // constraints.
+    if (schema != null && schema.contains(".")) {
+      final BlockingQueue<BigQueryFieldValueListWrapper> queue = new LinkedBlockingQueue<>(1);
+      signalEndOfData(queue, resultSchemaFields);
+      return BigQueryJsonResultSet.of(resultSchema, 0, queue, 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;
+    // Fallback Path: If catalog or schema is null, fall back to REST API metadata scan.
+    if (catalog == null || schema == null) {
+      final List<FieldValueList> collectedResults = Collections.synchronizedList(new ArrayList<>());
+      List<DatasetId> targetDatasets = getTargetDatasets(catalog, schema);
+
+      boolean ignoreAccessErrors = (catalog == null);
+      processTargetTablesConcurrently(
+          targetDatasets,
+          null,
+          collectedResults,
+          resultSchemaFields,
+          ignoreAccessErrors,
+          (bqTable, results, fields) -> {
+            TableConstraints constraints = bqTable.getTableConstraints();
+            if (constraints == null || constraints.getForeignKeys() == null) {
+              return;
             }
-            processForeignKey(fk, pkTableId, bqTable.getTableId(), results, fields);
-          }
-        });
+            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);
+      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);
+      final BlockingQueue<BigQueryFieldValueListWrapper> queue =
+          new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);
+      Future<?> fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields);
+      return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture);
+    }
+
+    String sql = getExportedKeysSqlContent();
+    String formattedSql = replaceSqlParameters(sql, catalog, schema, table);
+    PreparedStatement stmt = this.connection.prepareStatement(formattedSql);
+    if (stmt == null) {
+      throw new BigQueryJdbcException("Failed to prepare statement for getExportedKeys");
+    }
+    try {
+      stmt.closeOnCompletion();
+      return stmt.executeQuery();
+    } catch (SQLException e) {
+      closeStatementIgnoreException(stmt);
+      throw new BigQueryJdbcException("Error executing getExportedKeys", e);
+    }
   }
 
   @Override
@@ -5355,4 +5388,43 @@ private Comparator<FieldValueList> defineFkTableSortComparator(FieldList resultS
             (FieldValueList fvl) -> getLongValueOrNull(fvl, KEY_SEQ_IDX),
             Comparator.nullsFirst(Long::compareTo));
   }
+
+  private static synchronized String getExportedKeysSqlContent() {
+    if (exportedKeysSqlContent == null) {
+      exportedKeysSqlContent = readSqlFromFile(GET_EXPORTED_KEYS_SQL);
+    }
+    return exportedKeysSqlContent;
+  }
+
+  static String readSqlFromFile(String filename) {
+    try (InputStream in = BigQueryDatabaseMetaData.class.getResourceAsStream(filename)) {
+      if (in == null) {
+        throw new IllegalArgumentException("SQL file not found: " + filename);
+      }
+      ByteArrayOutputStream result = new ByteArrayOutputStream();
+      byte[] buffer = new byte[1024];
+      int length;
+      while ((length = in.read(buffer)) != -1) {
+        result.write(buffer, 0, length);
+      }
+      return result.toString(StandardCharsets.UTF_8.name());
+    } catch (IOException e) {
+      throw new RuntimeException("Failed to read SQL file: " + filename, e);
+    }
+  }
+
+  String replaceSqlParameters(String sql, String... params) throws SQLException {
+    return String.format(sql, (Object[]) params);
+  }
+
+  private void closeStatementIgnoreException(Statement stmt) {
+    if (stmt == null) {
+      return;
+    }
+    try {
+      stmt.close();
+    } catch (SQLException e) {
+      // ignore
+    }
+  }
 }

java-bigquery-jdbc/src/main/resources/com/google/cloud/bigquery/jdbc/DatabaseMetaData_GetExportedKeys.sql

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+SELECT
+  pk.table_catalog        AS PKTABLE_CAT,
+  pk.table_schema         AS PKTABLE_SCHEM,
+  pk.table_name           AS PKTABLE_NAME,
+  pk.column_name          AS PKCOLUMN_NAME,
+  fk.table_catalog        AS FKTABLE_CAT,
+  fk.table_schema         AS FKTABLE_SCHEM,
+  fk.table_name           AS FKTABLE_NAME,
+  fk.column_name          AS FKCOLUMN_NAME,
+  fk.ordinal_position     AS KEY_SEQ,
+  1                       AS UPDATE_RULE,
+  1                       AS DELETE_RULE,
+  REGEXP_EXTRACT(fk.constraint_name, r'[^.]+$') AS FK_NAME,
+  REGEXP_EXTRACT(pk.constraint_name, r'[^.]+$') AS PK_NAME,
+  7                       AS DEFERRABILITY
+FROM `%1$s.%2$s.INFORMATION_SCHEMA.KEY_COLUMN_USAGE` AS pk
+JOIN `%1$s.%2$s.INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE` AS ccu
+  ON ccu.table_catalog = pk.table_catalog
+ AND ccu.table_schema  = pk.table_schema
+ AND ccu.table_name    = pk.table_name
+ AND ccu.column_name   = pk.column_name
+JOIN `%1$s.%2$s.INFORMATION_SCHEMA.KEY_COLUMN_USAGE` AS fk
+  ON fk.constraint_catalog = ccu.constraint_catalog
+ AND fk.constraint_schema  = ccu.constraint_schema
+ AND fk.constraint_name    = ccu.constraint_name
+ AND fk.position_in_unique_constraint = pk.ordinal_position
+JOIN `%1$s.%2$s.INFORMATION_SCHEMA.TABLE_CONSTRAINTS` AS tc
+  ON tc.constraint_catalog = fk.constraint_catalog
+ AND tc.constraint_schema  = fk.constraint_schema
+ AND tc.constraint_name    = fk.constraint_name
+ AND tc.constraint_type    = 'FOREIGN KEY'
+WHERE CASE WHEN pk.table_name = '%3$s' THEN 1 ELSE 0 END = 1
+  AND pk.position_in_unique_constraint IS NULL
+ORDER BY fk.table_schema, fk.table_name, fk.ordinal_position

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

@@ -37,6 +37,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.sql.DatabaseMetaData;
+import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
@@ -3522,7 +3523,27 @@ public void testGetImportedKeys_noKeys() throws SQLException {
   }
 
   @Test
-  public void testGetExportedKeys_hasKeys() throws SQLException {
+  public void testGetExportedKeys_infoSchema() throws SQLException {
+    PreparedStatement mockStmt = mock(PreparedStatement.class);
+    ResultSet mockRs = mock(ResultSet.class);
+    when(bigQueryConnection.prepareStatement(anyString())).thenReturn(mockStmt);
+    when(mockStmt.executeQuery()).thenReturn(mockRs);
+
+    ResultSet rs = dbMetadata.getExportedKeys("test-project", "dataset_p", "ref_table");
+    assertEquals(mockRs, rs);
+    verify(mockStmt).closeOnCompletion();
+    verify(mockStmt).executeQuery();
+  }
+
+  @Test
+  public void testGetExportedKeys_pcntSchema() throws SQLException {
+    try (ResultSet rs = dbMetadata.getExportedKeys("test-project", "dataset.p", "ref_table")) {
+      assertFalse(rs.next());
+    }
+  }
+
+  @Test
+  public void testGetExportedKeys_fallback_hasKeys() throws SQLException {
     DatasetId datasetId = DatasetId.of("test-project", "dataset_p");
     TableId tableId = TableId.of("test-project", "dataset_p", "table_p");
     TableId refTableId = TableId.of("test-project", "dataset_p", "ref_table");
@@ -3544,7 +3565,7 @@ public void testGetExportedKeys_hasKeys() throws SQLException {
     mockDatasetIteration(datasetId);
     mockTableIteration(datasetId, mockTableP);
 
-    try (ResultSet rs = dbMetadata.getExportedKeys("test-project", "dataset_p", "ref_table")) {
+    try (ResultSet rs = dbMetadata.getExportedKeys("test-project", null, "ref_table")) {
       assertTrue(rs.next());
       assertEquals("test-project", rs.getString("PKTABLE_CAT"));
       assertEquals("dataset_p", rs.getString("PKTABLE_SCHEM"));
@@ -3572,15 +3593,15 @@ public void testGetExportedKeys_hasKeys() throws SQLException {
   }
 
   @Test
-  public void testGetExportedKeys_noKeys() throws SQLException {
+  public void testGetExportedKeys_fallback_noKeys() throws SQLException {
     DatasetId datasetId = DatasetId.of("test-project", "dataset_p");
     TableId tableId = TableId.of("test-project", "dataset_p", "table_p");
 
     Table mockTableP = mockTableWithConstraints(tableId, null);
     mockDatasetIteration(datasetId);
     mockTableIteration(datasetId, mockTableP);
 
-    try (ResultSet rs = dbMetadata.getExportedKeys("test-project", "dataset_p", "ref_table")) {
+    try (ResultSet rs = dbMetadata.getExportedKeys("test-project", null, "ref_table")) {
       assertFalse(rs.next());
     }
   }

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

@@ -57,6 +57,8 @@ public class ITDatabaseMetadataTest extends ITBase {
   private static final String CONSTRAINTS_TABLE_NAME = "JDBC_CONSTRAINTS_TEST_TABLE";
   private static final String CONSTRAINTS_TABLE_NAME2 = "JDBC_CONSTRAINTS_TEST_TABLE2";
   private static final String CONSTRAINTS_TABLE_NAME3 = "JDBC_CONSTRAINTS_TEST_TABLE3";
+  private static final String PCNT_SCHEMA = "bq-drivers-test-warehouse.jdbc_pcnt_test_namespace";
+  private static final String PCNT_TABLE_NAME = "PCNT_TEST_TABLE";
   private static final Pattern VERSION_PATTERN =
       Pattern.compile("^(\\d+)\\.(\\d+)(?:\\.\\d+)+\\s*.*");
   private static final String DEFAULT_CATALOG = ServiceOptions.getDefaultProjectId();