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 752100f4888c..210c379e8373 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 @@ -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 collectedResults = Collections.synchronizedList(new ArrayList<>()); + List 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 comparator = defineGetPrimaryKeysComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getPrimaryKeys", LOG); + + final BlockingQueue queue = + new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); + populateQueue(collectedResults, queue, resultSchemaFields); + signalEndOfData(queue, resultSchemaFields); + return BigQueryJsonResultSet.of(resultSchema, -1, queue, null); + } + + private Schema defineGetPrimaryKeysSchema() { + List 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 collectedResults, + FieldList resultSchemaFields) { + if (constraints == null || constraints.getPrimaryKey() == null) { + return; } + PrimaryKey pk = constraints.getPrimaryKey(); + List pkColumns = pk.getColumns(); + for (int i = 0; i < pkColumns.size(); i++) { + List 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 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 queu } } } + + private List getTargetDatasets(String catalog, String schema) throws SQLException { + if (schema == null) { + List datasets = fetchMatchingDatasets(catalog, null, null); + List datasetIds = new ArrayList<>(datasets.size()); + for (Dataset dataset : datasets) { + datasetIds.add(dataset.getDatasetId()); + } + return datasetIds; + } + + if (schema.isEmpty()) { + return Collections.emptyList(); + } + + List projects = resolveTargetProjects(catalog); + List datasetIds = new ArrayList<>(projects.size()); + for (String project : projects) { + datasetIds.add(DatasetId.of(project, schema)); + } + return datasetIds; + } + + private List resolveTargetProjects(String catalog) throws SQLException { + return (catalog != null) ? Collections.singletonList(catalog) : getAccessibleCatalogNames(); + } + + @FunctionalInterface + private interface TableProcessor { + void process(Table bqTable, List collectedResults, FieldList resultSchemaFields) + throws SQLException; + } + + private void processSingleTable( + DatasetId datasetId, + String tableName, + List 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 targetDatasets, + String tableName, + List 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> 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)); + } + } } diff --git a/java-bigquery-jdbc/src/main/resources/com/google/cloud/bigquery/jdbc/DatabaseMetaData_GetPrimaryKeys.sql b/java-bigquery-jdbc/src/main/resources/com/google/cloud/bigquery/jdbc/DatabaseMetaData_GetPrimaryKeys.sql deleted file mode 100644 index 282910fb9721..000000000000 --- a/java-bigquery-jdbc/src/main/resources/com/google/cloud/bigquery/jdbc/DatabaseMetaData_GetPrimaryKeys.sql +++ /dev/null @@ -1,30 +0,0 @@ -/* - * Copyright 2024 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 table_catalog AS TABLE_CAT, - table_schema AS TABLE_SCHEM, - table_name AS TABLE_NAME, - column_name AS COLUMN_NAME, - ordinal_position AS KEY_SEQ, - constraint_name AS PK_NAME -FROM - %s.%s.INFORMATION_SCHEMA.KEY_COLUMN_USAGE -WHERE - table_name = '%s' - AND CONTAINS_SUBSTR(constraint_name - , 'pk$') -ORDER BY - COLUMN_NAME; 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 2b72ed01b8a3..f99ba0bbe5dd 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 @@ -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) {