feat(bigquery-jdbc): add Graph metadata support in getPrimaryKeys method

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

@@ -108,6 +108,9 @@ class BigQueryDatabaseMetaData implements DatabaseMetaData {
   private static String graphTablesSqlContent;
   private static final String GET_GRAPH_COLUMNS_SQL = "DatabaseMetaData_GetGraphColumns.sql";
   private static String graphColumnsSqlContent;
+  private static final String GET_GRAPH_PRIMARY_KEYS_SQL =
+      "DatabaseMetaData_GetGraphPrimaryKeys.sql";
+  private static String graphPrimaryKeysSqlContent;
   static final String MEASURE_EXPRESSION_KIND = "MEASURE";
   static final String MEASURE_REMARKS_PREFIX = "MEASURE: ";
   // Declared package-private for testing.
@@ -161,8 +164,7 @@ class BigQueryDatabaseMetaData implements DatabaseMetaData {
   private static final Pattern VALID_SCHEMA_PATTERN = Pattern.compile("^[A-Za-z0-9_]+$");
   // BigQuery region qualifier for regional INFORMATION_SCHEMA queries (e.g. region-us,
   // region-us-central1)
-  private static final Pattern VALID_REGION_SCOPE_PATTERN =
-      Pattern.compile("^region-[A-Za-z0-9-]+$");
+  private static final Pattern VALID_REGION_SCOPE_PATTERN = Pattern.compile("^region-[a-z0-9-]+$");
 
   String URL;
   BigQueryConnection connection;
@@ -2320,25 +2322,37 @@ public ResultSet getPrimaryKeys(String catalog, String schema, String table) thr
     }
 
     final FieldList resultSchemaFields = resultSchema.getFields();
-    List<DatasetId> targetDatasets = getTargetDatasets(catalog, schema, false);
 
+    // A dot can only appear in a graph element name ("<graph>.<element>");
+    // Route accordingly — no pattern checks,
+    // because getPrimaryKeys takes exact names.
+    boolean isGraphTableName = table.contains(".");
     BlockingQueue<BigQueryFieldValueListWrapper> queue =
         fetchAndPopulateQueue(
             () -> {
               final List<FieldValueList> collectedResults =
                   Collections.synchronizedList(new ArrayList<>());
-              processTargetTablesConcurrently(
-                  targetDatasets,
-                  table,
-                  false,
-                  true,
-                  true,
-                  collectedResults,
-                  resultSchemaFields,
-                  (bqTable, results, fields) -> {
-                    TableConstraints constraints = bqTable.getTableConstraints();
-                    processPrimaryKey(constraints, bqTable.getTableId(), results, fields);
-                  });
+
+              if (!isGraphTableName) {
+
+                // Only the REST path needs the dataset list (datasets.list when schema is null).
+                List<DatasetId> targetDatasets = getTargetDatasets(catalog, schema, false);
+                processTargetTablesConcurrently(
+                    targetDatasets,
+                    table,
+                    false,
+                    true,
+                    true,
+                    collectedResults,
+                    resultSchemaFields,
+                    (bqTable, results, fields) -> {
+                      TableConstraints constraints = bqTable.getTableConstraints();
+                      processPrimaryKey(constraints, bqTable.getTableId(), results, fields);
+                    });
+              } else if (this.connection.isEnableGraphMetadata()) {
+                appendGraphPrimaryKeys(
+                    catalog, schema, table, collectedResults, resultSchemaFields);
+              }
 
               Comparator<FieldValueList> comparator =
                   defineGetPrimaryKeysComparator(resultSchemaFields);
@@ -2401,6 +2415,74 @@ private void processPrimaryKey(
     }
   }
 
+  private void appendGraphPrimaryKeys(
+      String catalog,
+      String schema,
+      String table,
+      List<FieldValueList> collectedResults,
+      FieldList resultSchemaFields)
+      throws SQLException {
+    // getPrimaryKeys takes exact identifiers, never LIKE patterns. The only "wide"
+    // input JDBC permits is null, which means "all".
+    List<String> targetCatalogs =
+        catalog != null ? Collections.singletonList(catalog) : getAccessibleCatalogNames();
+
+    // Scope of the INFORMATION_SCHEMA path:
+    //   schema given -> `<project>.<schema>.INFORMATION_SCHEMA.PROPERTY_GRAPHS`
+    //   schema null  -> `<project>.region-<loc>.INFORMATION_SCHEMA.PROPERTY_GRAPHS`
+    //
+    //  Unlike the REST path (datasets.list, region-agnostic), the null-schema case covers a
+    // single region: the connection Location, else DEFAULT_GRAPH_REGION. Graphs elsewhere
+    // are not returned unless Location is set or the dataset is passed as schema.
+    String scope;
+    if (schema != null) {
+      validateSchemaIdentifier(schema);
+      scope = schema;
+    } else {
+      scope = resolveGraphRegion();
+      validateRegionScopeIdentifier(scope);
+    }
+
+    for (String project : targetCatalogs) {
+      validateCatalogIdentifier(project);
+      String formattedSql = replaceSqlParameters(getGraphPrimaryKeysSqlContent(), project, scope);
+
+      try (PreparedStatement statement = this.connection.prepareStatement(formattedSql)) {
+        statement.setString(1, table); // exact "<graph>.<element>"
+        try (ResultSet rs = statement.executeQuery()) {
+          while (rs.next()) {
+            List<FieldValue> row = new ArrayList<>(6);
+            row.add(
+                createStringFieldValue(
+                    rs.getString("TABLE_CAT"))); // 1. TABLE_CAT   <- property_graph_catalog
+            row.add(
+                createStringFieldValue(
+                    rs.getString("TABLE_SCHEM"))); // 2. TABLE_SCHEM <- property_graph_schema
+            row.add(
+                createStringFieldValue(
+                    rs.getString("TABLE_NAME"))); // 3. TABLE_NAME  <- "<graph>.<element>"
+            row.add(
+                createStringFieldValue(
+                    rs.getString("COLUMN_NAME"))); // 4. COLUMN_NAME <- keyColumns[i]
+            row.add(createLongFieldValue(rs.getLong("KEY_SEQ"))); // 5. KEY_SEQ     <- kc_off + 1
+            row.add(createNullFieldValue()); // 6. PK_NAME     (always null)
+            collectedResults.add(FieldValueList.of(row, resultSchemaFields));
+          }
+        }
+      } catch (SQLException e) {
+        // Same "Not found" tolerance as appendGraphTableRows / appendGraphColumnRows.
+        if (e.getMessage() != null && e.getMessage().contains("Not found")) {
+          LOG.info(
+              "Dataset or project not found while querying PROPERTY_GRAPHS for primary keys in"
+                  + " project '%s'. Skipping.",
+              project);
+        } else {
+          throw e;
+        }
+      }
+    }
+  }
+
   private Comparator<FieldValueList> defineGetPrimaryKeysComparator(FieldList resultSchemaFields) {
     final int COLUMN_NAME_IDX = resultSchemaFields.getIndex("COLUMN_NAME");
     return Comparator.comparing(
@@ -5191,13 +5273,20 @@ static synchronized String getGraphTablesSqlContent() {
     return graphTablesSqlContent;
   }
 
-  private static synchronized String getGraphColumnsSqlContent() {
+  static synchronized String getGraphColumnsSqlContent() {
     if (graphColumnsSqlContent == null) {
       graphColumnsSqlContent = readSqlFromFile(GET_GRAPH_COLUMNS_SQL);
     }
     return graphColumnsSqlContent;
   }
 
+  static synchronized String getGraphPrimaryKeysSqlContent() {
+    if (graphPrimaryKeysSqlContent == null) {
+      graphPrimaryKeysSqlContent = readSqlFromFile(GET_GRAPH_PRIMARY_KEYS_SQL);
+    }
+    return graphPrimaryKeysSqlContent;
+  }
+
   static String readSqlFromFile(String filename) {
     try (InputStream in = BigQueryDatabaseMetaData.class.getResourceAsStream(filename)) {
       if (in == null) {

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

@@ -0,0 +1,29 @@
+/*
+ * 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
+  property_graph_catalog AS TABLE_CAT,
+  property_graph_schema AS TABLE_SCHEM,
+  CONCAT(property_graph_name, '.', JSON_VALUE(t, '$.name')) AS TABLE_NAME,
+  JSON_VALUE(kc) AS COLUMN_NAME,
+  kc_off + 1 AS KEY_SEQ,
+  CAST(NULL AS STRING) AS PK_NAME
+FROM `%1$s.%2$s.INFORMATION_SCHEMA.PROPERTY_GRAPHS`,
+UNNEST(ARRAY_CONCAT(
+  IFNULL(JSON_QUERY_ARRAY(property_graph_metadata_json, '$.nodeTables'), []),
+  IFNULL(JSON_QUERY_ARRAY(property_graph_metadata_json, '$.edgeTables'), []))) AS t,
+UNNEST(JSON_QUERY_ARRAY(t, '$.keyColumns')) AS kc WITH OFFSET kc_off
+WHERE CONCAT(property_graph_name, '.', JSON_VALUE(t, '$.name')) = ?

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

@@ -4009,6 +4009,9 @@ public void testResolveGraphRegion() throws SQLException {
 
     when(bigQueryConnection.getLocation()).thenReturn("region-eu");
     assertEquals("region-eu", dbMetadata.resolveGraphRegion());
+    assertThrows(
+        BigQueryJdbcException.class,
+        () -> BigQueryDatabaseMetaData.validateRegionScopeIdentifier("region-US"));
   }
 
   @Test
@@ -4024,6 +4027,24 @@ public void testGraphTablesSqlContent() throws SQLException {
     assertThat(regionSql).contains("`my-project.region-us.INFORMATION_SCHEMA.PROPERTY_GRAPHS`");
   }
 
+  @Test
+  public void testGraphPrimaryKeysSqlContent() throws SQLException {
+    String datasetSql =
+        dbMetadata.replaceSqlParameters(
+            BigQueryDatabaseMetaData.getGraphPrimaryKeysSqlContent(), "my-project", "my_dataset");
+    assertThat(datasetSql).contains("`my-project.my_dataset.INFORMATION_SCHEMA.PROPERTY_GRAPHS`");
+
+    String regionSql =
+        dbMetadata.replaceSqlParameters(
+            BigQueryDatabaseMetaData.getGraphPrimaryKeysSqlContent(), "my-project", "region-us");
+    assertThat(regionSql).contains("`my-project.region-us.INFORMATION_SCHEMA.PROPERTY_GRAPHS`");
+
+    // getPrimaryKeys takes exact names: no LIKE, one bound parameter, key columns + 1-based seq.
+    assertThat(regionSql).doesNotContain("LIKE");
+    assertThat(regionSql).contains("'$.keyColumns'");
+    assertThat(regionSql).contains("kc_off + 1 AS KEY_SEQ");
+  }
+
   @FunctionalInterface
   private interface MetadataResultSetSupplier {
     ResultSet get(BigQueryDatabaseMetaData metadata) throws SQLException;

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

@@ -877,6 +877,46 @@ public void testDatabaseMetadataGetTables_PropertyGraph() throws SQLException {
       Assertions.assertTrue(graphTables.contains(graphName + ".Person"));
       Assertions.assertTrue(graphTables.contains(graphName + ".Knows"));
       Assertions.assertEquals(2, graphTables.size());
+
+      // 3. Verify getPrimaryKeys() for graph elements (exact names, dataset-scoped and
+      // region-scoped)
+      // 3a. Single-column key on the node element
+      try (ResultSet rs = metaData.getPrimaryKeys(PROJECT_ID, DATASET, graphName + ".Person")) {
+        Assertions.assertTrue(rs.next());
+        Assertions.assertEquals(PROJECT_ID, rs.getString("TABLE_CAT"));
+        Assertions.assertEquals(DATASET, rs.getString("TABLE_SCHEM"));
+        Assertions.assertEquals(graphName + ".Person", rs.getString("TABLE_NAME"));
+        Assertions.assertEquals("id", rs.getString("COLUMN_NAME"));
+        Assertions.assertEquals(1, rs.getShort("KEY_SEQ"));
+        Assertions.assertNull(rs.getString("PK_NAME"));
+        Assertions.assertFalse(rs.next());
+      }
+
+      // 3b. Composite key on the edge element — rows arrive sorted by COLUMN_NAME
+      //     (defineGetPrimaryKeysComparator), so "dst_id" precedes "id".
+      try (ResultSet rs = metaData.getPrimaryKeys(PROJECT_ID, DATASET, graphName + ".Knows")) {
+        Assertions.assertTrue(rs.next());
+        Assertions.assertEquals("dst_id", rs.getString("COLUMN_NAME"));
+        Assertions.assertEquals(2, rs.getShort("KEY_SEQ"));
+        Assertions.assertTrue(rs.next());
+        Assertions.assertEquals("id", rs.getString("COLUMN_NAME"));
+        Assertions.assertEquals(1, rs.getShort("KEY_SEQ"));
+        Assertions.assertFalse(rs.next());
+      }
+
+      // 3c. schema == null -> region-scoped INFORMATION_SCHEMA path; same rows expected
+      try (ResultSet rs = metaData.getPrimaryKeys(PROJECT_ID, null, graphName + ".Person")) {
+        Assertions.assertTrue(rs.next());
+        Assertions.assertEquals(DATASET, rs.getString("TABLE_SCHEM"));
+        Assertions.assertEquals("id", rs.getString("COLUMN_NAME"));
+        Assertions.assertFalse(rs.next());
+      }
+
+      // 3d. Unknown element -> no rows (exact match, no wildcard expansion)
+      try (ResultSet rs =
+          metaData.getPrimaryKeys(PROJECT_ID, DATASET, graphName + ".DoesNotExist")) {
+        Assertions.assertFalse(rs.next());
+      }
     }
   }