Skip to content

Commit b3a751f

Browse files
authored
feat(bigquery-jdbc): add Graph metadata support in getPrimaryKeys method (#14922)
1 parent 7d15467 commit b3a751f

4 files changed

Lines changed: 195 additions & 16 deletions

File tree

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

Lines changed: 105 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,9 @@ class BigQueryDatabaseMetaData implements DatabaseMetaData {
108108
private static String graphTablesSqlContent;
109109
private static final String GET_GRAPH_COLUMNS_SQL = "DatabaseMetaData_GetGraphColumns.sql";
110110
private static String graphColumnsSqlContent;
111+
private static final String GET_GRAPH_PRIMARY_KEYS_SQL =
112+
"DatabaseMetaData_GetGraphPrimaryKeys.sql";
113+
private static String graphPrimaryKeysSqlContent;
111114
static final String MEASURE_EXPRESSION_KIND = "MEASURE";
112115
static final String MEASURE_REMARKS_PREFIX = "MEASURE: ";
113116
// Declared package-private for testing.
@@ -161,8 +164,7 @@ class BigQueryDatabaseMetaData implements DatabaseMetaData {
161164
private static final Pattern VALID_SCHEMA_PATTERN = Pattern.compile("^[A-Za-z0-9_]+$");
162165
// BigQuery region qualifier for regional INFORMATION_SCHEMA queries (e.g. region-us,
163166
// region-us-central1)
164-
private static final Pattern VALID_REGION_SCOPE_PATTERN =
165-
Pattern.compile("^region-[A-Za-z0-9-]+$");
167+
private static final Pattern VALID_REGION_SCOPE_PATTERN = Pattern.compile("^region-[a-z0-9-]+$");
166168

167169
String URL;
168170
BigQueryConnection connection;
@@ -2320,25 +2322,37 @@ public ResultSet getPrimaryKeys(String catalog, String schema, String table) thr
23202322
}
23212323

23222324
final FieldList resultSchemaFields = resultSchema.getFields();
2323-
List<DatasetId> targetDatasets = getTargetDatasets(catalog, schema, false);
23242325

2326+
// A dot can only appear in a graph element name ("<graph>.<element>");
2327+
// Route accordingly — no pattern checks,
2328+
// because getPrimaryKeys takes exact names.
2329+
boolean isGraphTableName = table.contains(".");
23252330
BlockingQueue<BigQueryFieldValueListWrapper> queue =
23262331
fetchAndPopulateQueue(
23272332
() -> {
23282333
final List<FieldValueList> collectedResults =
23292334
Collections.synchronizedList(new ArrayList<>());
2330-
processTargetTablesConcurrently(
2331-
targetDatasets,
2332-
table,
2333-
false,
2334-
true,
2335-
true,
2336-
collectedResults,
2337-
resultSchemaFields,
2338-
(bqTable, results, fields) -> {
2339-
TableConstraints constraints = bqTable.getTableConstraints();
2340-
processPrimaryKey(constraints, bqTable.getTableId(), results, fields);
2341-
});
2335+
2336+
if (!isGraphTableName) {
2337+
2338+
// Only the REST path needs the dataset list (datasets.list when schema is null).
2339+
List<DatasetId> targetDatasets = getTargetDatasets(catalog, schema, false);
2340+
processTargetTablesConcurrently(
2341+
targetDatasets,
2342+
table,
2343+
false,
2344+
true,
2345+
true,
2346+
collectedResults,
2347+
resultSchemaFields,
2348+
(bqTable, results, fields) -> {
2349+
TableConstraints constraints = bqTable.getTableConstraints();
2350+
processPrimaryKey(constraints, bqTable.getTableId(), results, fields);
2351+
});
2352+
} else if (this.connection.isEnableGraphMetadata()) {
2353+
appendGraphPrimaryKeys(
2354+
catalog, schema, table, collectedResults, resultSchemaFields);
2355+
}
23422356

23432357
Comparator<FieldValueList> comparator =
23442358
defineGetPrimaryKeysComparator(resultSchemaFields);
@@ -2401,6 +2415,74 @@ private void processPrimaryKey(
24012415
}
24022416
}
24032417

2418+
private void appendGraphPrimaryKeys(
2419+
String catalog,
2420+
String schema,
2421+
String table,
2422+
List<FieldValueList> collectedResults,
2423+
FieldList resultSchemaFields)
2424+
throws SQLException {
2425+
// getPrimaryKeys takes exact identifiers, never LIKE patterns. The only "wide"
2426+
// input JDBC permits is null, which means "all".
2427+
List<String> targetCatalogs =
2428+
catalog != null ? Collections.singletonList(catalog) : getAccessibleCatalogNames();
2429+
2430+
// Scope of the INFORMATION_SCHEMA path:
2431+
// schema given -> `<project>.<schema>.INFORMATION_SCHEMA.PROPERTY_GRAPHS`
2432+
// schema null -> `<project>.region-<loc>.INFORMATION_SCHEMA.PROPERTY_GRAPHS`
2433+
//
2434+
// Unlike the REST path (datasets.list, region-agnostic), the null-schema case covers a
2435+
// single region: the connection Location, else DEFAULT_GRAPH_REGION. Graphs elsewhere
2436+
// are not returned unless Location is set or the dataset is passed as schema.
2437+
String scope;
2438+
if (schema != null) {
2439+
validateSchemaIdentifier(schema);
2440+
scope = schema;
2441+
} else {
2442+
scope = resolveGraphRegion();
2443+
validateRegionScopeIdentifier(scope);
2444+
}
2445+
2446+
for (String project : targetCatalogs) {
2447+
validateCatalogIdentifier(project);
2448+
String formattedSql = replaceSqlParameters(getGraphPrimaryKeysSqlContent(), project, scope);
2449+
2450+
try (PreparedStatement statement = this.connection.prepareStatement(formattedSql)) {
2451+
statement.setString(1, table); // exact "<graph>.<element>"
2452+
try (ResultSet rs = statement.executeQuery()) {
2453+
while (rs.next()) {
2454+
List<FieldValue> row = new ArrayList<>(6);
2455+
row.add(
2456+
createStringFieldValue(
2457+
rs.getString("TABLE_CAT"))); // 1. TABLE_CAT <- property_graph_catalog
2458+
row.add(
2459+
createStringFieldValue(
2460+
rs.getString("TABLE_SCHEM"))); // 2. TABLE_SCHEM <- property_graph_schema
2461+
row.add(
2462+
createStringFieldValue(
2463+
rs.getString("TABLE_NAME"))); // 3. TABLE_NAME <- "<graph>.<element>"
2464+
row.add(
2465+
createStringFieldValue(
2466+
rs.getString("COLUMN_NAME"))); // 4. COLUMN_NAME <- keyColumns[i]
2467+
row.add(createLongFieldValue(rs.getLong("KEY_SEQ"))); // 5. KEY_SEQ <- kc_off + 1
2468+
row.add(createNullFieldValue()); // 6. PK_NAME (always null)
2469+
collectedResults.add(FieldValueList.of(row, resultSchemaFields));
2470+
}
2471+
}
2472+
} catch (SQLException e) {
2473+
// Same "Not found" tolerance as appendGraphTableRows / appendGraphColumnRows.
2474+
if (e.getMessage() != null && e.getMessage().contains("Not found")) {
2475+
LOG.info(
2476+
"Dataset or project not found while querying PROPERTY_GRAPHS for primary keys in"
2477+
+ " project '%s'. Skipping.",
2478+
project);
2479+
} else {
2480+
throw e;
2481+
}
2482+
}
2483+
}
2484+
}
2485+
24042486
private Comparator<FieldValueList> defineGetPrimaryKeysComparator(FieldList resultSchemaFields) {
24052487
final int COLUMN_NAME_IDX = resultSchemaFields.getIndex("COLUMN_NAME");
24062488
return Comparator.comparing(
@@ -5191,13 +5273,20 @@ static synchronized String getGraphTablesSqlContent() {
51915273
return graphTablesSqlContent;
51925274
}
51935275

5194-
private static synchronized String getGraphColumnsSqlContent() {
5276+
static synchronized String getGraphColumnsSqlContent() {
51955277
if (graphColumnsSqlContent == null) {
51965278
graphColumnsSqlContent = readSqlFromFile(GET_GRAPH_COLUMNS_SQL);
51975279
}
51985280
return graphColumnsSqlContent;
51995281
}
52005282

5283+
static synchronized String getGraphPrimaryKeysSqlContent() {
5284+
if (graphPrimaryKeysSqlContent == null) {
5285+
graphPrimaryKeysSqlContent = readSqlFromFile(GET_GRAPH_PRIMARY_KEYS_SQL);
5286+
}
5287+
return graphPrimaryKeysSqlContent;
5288+
}
5289+
52015290
static String readSqlFromFile(String filename) {
52025291
try (InputStream in = BigQueryDatabaseMetaData.class.getResourceAsStream(filename)) {
52035292
if (in == null) {
Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
SELECT
18+
property_graph_catalog AS TABLE_CAT,
19+
property_graph_schema AS TABLE_SCHEM,
20+
CONCAT(property_graph_name, '.', JSON_VALUE(t, '$.name')) AS TABLE_NAME,
21+
JSON_VALUE(kc) AS COLUMN_NAME,
22+
kc_off + 1 AS KEY_SEQ,
23+
CAST(NULL AS STRING) AS PK_NAME
24+
FROM `%1$s.%2$s.INFORMATION_SCHEMA.PROPERTY_GRAPHS`,
25+
UNNEST(ARRAY_CONCAT(
26+
IFNULL(JSON_QUERY_ARRAY(property_graph_metadata_json, '$.nodeTables'), []),
27+
IFNULL(JSON_QUERY_ARRAY(property_graph_metadata_json, '$.edgeTables'), []))) AS t,
28+
UNNEST(JSON_QUERY_ARRAY(t, '$.keyColumns')) AS kc WITH OFFSET kc_off
29+
WHERE CONCAT(property_graph_name, '.', JSON_VALUE(t, '$.name')) = ?

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4009,6 +4009,9 @@ public void testResolveGraphRegion() throws SQLException {
40094009

40104010
when(bigQueryConnection.getLocation()).thenReturn("region-eu");
40114011
assertEquals("region-eu", dbMetadata.resolveGraphRegion());
4012+
assertThrows(
4013+
BigQueryJdbcException.class,
4014+
() -> BigQueryDatabaseMetaData.validateRegionScopeIdentifier("region-US"));
40124015
}
40134016

40144017
@Test
@@ -4024,6 +4027,24 @@ public void testGraphTablesSqlContent() throws SQLException {
40244027
assertThat(regionSql).contains("`my-project.region-us.INFORMATION_SCHEMA.PROPERTY_GRAPHS`");
40254028
}
40264029

4030+
@Test
4031+
public void testGraphPrimaryKeysSqlContent() throws SQLException {
4032+
String datasetSql =
4033+
dbMetadata.replaceSqlParameters(
4034+
BigQueryDatabaseMetaData.getGraphPrimaryKeysSqlContent(), "my-project", "my_dataset");
4035+
assertThat(datasetSql).contains("`my-project.my_dataset.INFORMATION_SCHEMA.PROPERTY_GRAPHS`");
4036+
4037+
String regionSql =
4038+
dbMetadata.replaceSqlParameters(
4039+
BigQueryDatabaseMetaData.getGraphPrimaryKeysSqlContent(), "my-project", "region-us");
4040+
assertThat(regionSql).contains("`my-project.region-us.INFORMATION_SCHEMA.PROPERTY_GRAPHS`");
4041+
4042+
// getPrimaryKeys takes exact names: no LIKE, one bound parameter, key columns + 1-based seq.
4043+
assertThat(regionSql).doesNotContain("LIKE");
4044+
assertThat(regionSql).contains("'$.keyColumns'");
4045+
assertThat(regionSql).contains("kc_off + 1 AS KEY_SEQ");
4046+
}
4047+
40274048
@FunctionalInterface
40284049
private interface MetadataResultSetSupplier {
40294050
ResultSet get(BigQueryDatabaseMetaData metadata) throws SQLException;

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

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -877,6 +877,46 @@ public void testDatabaseMetadataGetTables_PropertyGraph() throws SQLException {
877877
Assertions.assertTrue(graphTables.contains(graphName + ".Person"));
878878
Assertions.assertTrue(graphTables.contains(graphName + ".Knows"));
879879
Assertions.assertEquals(2, graphTables.size());
880+
881+
// 3. Verify getPrimaryKeys() for graph elements (exact names, dataset-scoped and
882+
// region-scoped)
883+
// 3a. Single-column key on the node element
884+
try (ResultSet rs = metaData.getPrimaryKeys(PROJECT_ID, DATASET, graphName + ".Person")) {
885+
Assertions.assertTrue(rs.next());
886+
Assertions.assertEquals(PROJECT_ID, rs.getString("TABLE_CAT"));
887+
Assertions.assertEquals(DATASET, rs.getString("TABLE_SCHEM"));
888+
Assertions.assertEquals(graphName + ".Person", rs.getString("TABLE_NAME"));
889+
Assertions.assertEquals("id", rs.getString("COLUMN_NAME"));
890+
Assertions.assertEquals(1, rs.getShort("KEY_SEQ"));
891+
Assertions.assertNull(rs.getString("PK_NAME"));
892+
Assertions.assertFalse(rs.next());
893+
}
894+
895+
// 3b. Composite key on the edge element — rows arrive sorted by COLUMN_NAME
896+
// (defineGetPrimaryKeysComparator), so "dst_id" precedes "id".
897+
try (ResultSet rs = metaData.getPrimaryKeys(PROJECT_ID, DATASET, graphName + ".Knows")) {
898+
Assertions.assertTrue(rs.next());
899+
Assertions.assertEquals("dst_id", rs.getString("COLUMN_NAME"));
900+
Assertions.assertEquals(2, rs.getShort("KEY_SEQ"));
901+
Assertions.assertTrue(rs.next());
902+
Assertions.assertEquals("id", rs.getString("COLUMN_NAME"));
903+
Assertions.assertEquals(1, rs.getShort("KEY_SEQ"));
904+
Assertions.assertFalse(rs.next());
905+
}
906+
907+
// 3c. schema == null -> region-scoped INFORMATION_SCHEMA path; same rows expected
908+
try (ResultSet rs = metaData.getPrimaryKeys(PROJECT_ID, null, graphName + ".Person")) {
909+
Assertions.assertTrue(rs.next());
910+
Assertions.assertEquals(DATASET, rs.getString("TABLE_SCHEM"));
911+
Assertions.assertEquals("id", rs.getString("COLUMN_NAME"));
912+
Assertions.assertFalse(rs.next());
913+
}
914+
915+
// 3d. Unknown element -> no rows (exact match, no wildcard expansion)
916+
try (ResultSet rs =
917+
metaData.getPrimaryKeys(PROJECT_ID, DATASET, graphName + ".DoesNotExist")) {
918+
Assertions.assertFalse(rs.next());
919+
}
880920
}
881921
}
882922

0 commit comments

Comments
 (0)