diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java index 00082f76999d..98f94a58cf20 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java +++ b/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 targetDatasets = getTargetDatasets(catalog, schema, false); + // A dot can only appear in a graph element name ("."); + // Route accordingly — no pattern checks, + // because getPrimaryKeys takes exact names. + boolean isGraphTableName = table.contains("."); BlockingQueue queue = fetchAndPopulateQueue( () -> { final List 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 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 comparator = defineGetPrimaryKeysComparator(resultSchemaFields); @@ -2401,6 +2415,74 @@ private void processPrimaryKey( } } + private void appendGraphPrimaryKeys( + String catalog, + String schema, + String table, + List collectedResults, + FieldList resultSchemaFields) + throws SQLException { + // getPrimaryKeys takes exact identifiers, never LIKE patterns. The only "wide" + // input JDBC permits is null, which means "all". + List targetCatalogs = + catalog != null ? Collections.singletonList(catalog) : getAccessibleCatalogNames(); + + // Scope of the INFORMATION_SCHEMA path: + // schema given -> `..INFORMATION_SCHEMA.PROPERTY_GRAPHS` + // schema null -> `.region-.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 "." + try (ResultSet rs = statement.executeQuery()) { + while (rs.next()) { + List 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 <- "." + 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 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) { diff --git a/java-bigquery-jdbc/src/main/resources/com/google/cloud/bigquery/jdbc/DatabaseMetaData_GetGraphPrimaryKeys.sql b/java-bigquery-jdbc/src/main/resources/com/google/cloud/bigquery/jdbc/DatabaseMetaData_GetGraphPrimaryKeys.sql new file mode 100644 index 000000000000..06467a5fabee --- /dev/null +++ b/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')) = ? \ No newline at end of file diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaDataTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaDataTest.java index 03fc2ddd3015..47489057b0a5 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaDataTest.java +++ b/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; diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITDatabaseMetadataTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITDatabaseMetadataTest.java index a1c62a53e547..671f0b9341f6 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITDatabaseMetadataTest.java +++ b/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()); + } } }