@@ -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 ) {
0 commit comments