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

b/532245342

Migrates the `getPrimaryKeys` metadata methods in
`BigQueryDatabaseMetaData` away from executing SQL queries on
`INFORMATION_SCHEMA` tables, replacing them with native BigQuery Java
Client SDK calls.

This PR also introduces the shared thread-safe concurrent dataset
scanning infrastructure that will be utilized by subsequent metadata
migrations.

### **Key Changes:**
1. **`getPrimaryKeys`**: Rewritten to fetch metadata using
`bigquery.getTable(...)` and parsing table constraints in-memory.
2. **Concurrency Infrastructure**:
- Introduced `processTargetTablesConcurrently(...)` and
`processSingleTable(...)` to handle concurrent metadata retrieval across
target datasets using the metadata executor thread pool.
- Added a short-circuit optimization for single target datasets to
execute synchronously on the calling thread.
- Extracted project list resolution into `resolveTargetProjects(...)` to
support wildcard schema matching.
3. **Cleanup**:
- Deleted the obsolete `DatabaseMetaData_GetPrimaryKeys.sql` resource
file.
- Removed two obsolete SQL-parsing unit tests (`testReadSqlFromFile` and
`testMetadataMethodsDoNotInterfere`) from
`BigQueryDatabaseMetaDataTest.java` as they are no longer applicable.

Author: keshavdandeva

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

@@ -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;
@@ -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;
@@ -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
@@ -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));
+    }
+    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));
+    }
+  }
 }

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

@@ -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");
@@ -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) {