From 30e54585c29a64ee5f094610ee3100e632b388ad Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Fri, 17 Jul 2026 19:31:55 +0000 Subject: [PATCH 01/14] refactor(bigquery-jdbc): optimize and unify concurrent metadata retrieval --- .../jdbc/BigQueryDatabaseMetaData.java | 1298 ++++++----------- .../jdbc/BigQueryDatabaseMetaDataTest.java | 147 +- 2 files changed, 457 insertions(+), 988 deletions(-) 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 79f6bdb2a2f6..8557b96864c4 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 @@ -20,11 +20,9 @@ import com.google.cloud.Tuple; import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.BigQuery.DatasetListOption; -import com.google.cloud.bigquery.BigQuery.RoutineField; import com.google.cloud.bigquery.BigQuery.RoutineListOption; import com.google.cloud.bigquery.BigQuery.RoutineOption; import com.google.cloud.bigquery.BigQuery.TableField; -import com.google.cloud.bigquery.BigQuery.TableListOption; import com.google.cloud.bigquery.BigQuery.TableOption; import com.google.cloud.bigquery.BigQueryException; import com.google.cloud.bigquery.ColumnReference; @@ -52,17 +50,11 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.jdbc.BigQueryJdbcTypeMappings.ColumnTypeInfo; import com.google.cloud.bigquery.jdbc.utils.BigQueryJdbcVersionUtility; -import java.io.ByteArrayOutputStream; -import java.io.IOException; -import java.io.InputStream; -import java.nio.charset.StandardCharsets; import java.sql.Connection; import java.sql.DatabaseMetaData; -import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.RowIdLifetime; import java.sql.SQLException; -import java.sql.Statement; import java.sql.Types; import java.util.ArrayList; import java.util.Arrays; @@ -100,8 +92,6 @@ class BigQueryDatabaseMetaData implements DatabaseMetaData { private static final String PROCEDURE_TERM = "Procedure"; private static final int DEFAULT_PAGE_SIZE = 500; private static final int DEFAULT_QUEUE_CAPACITY = 5000; - private static final String GET_EXPORTED_KEYS_SQL = "DatabaseMetaData_GetExportedKeys.sql"; - private static String exportedKeysSqlContent; // Declared package-private for testing. static final String GOOGLE_SQL_QUOTED_IDENTIFIER = "`"; // Does not include SQL:2003 Keywords as per JDBC spec. @@ -773,123 +763,44 @@ public boolean dataDefinitionIgnoredInTransactions() { } @Override - public ResultSet getProcedures( - String catalog, String schemaPattern, String procedureNamePattern) { + public ResultSet getProcedures(String catalog, String schemaPattern, String procedureNamePattern) + throws SQLException { + LOG.info( + "getProcedures called for catalog: %s, schemaPattern: %s, procedureNamePattern: %s", + catalog, schemaPattern, procedureNamePattern); + if ((catalog != null && catalog.isEmpty()) || (schemaPattern != null && schemaPattern.isEmpty()) || (procedureNamePattern != null && procedureNamePattern.isEmpty())) { - LOG.warning("Returning empty ResultSet as catalog is empty or a pattern is empty."); + LOG.warning( + "Returning empty ResultSet as one or more patterns are empty or catalog is empty."); return new BigQueryJsonResultSet(); } - LOG.info( - "getProcedures called for catalog: %s, schemaPattern: %s, procedureNamePattern: %s", - catalog, schemaPattern, procedureNamePattern); - - final Pattern schemaRegex = compileSqlLikePattern(schemaPattern); - final Pattern procedureNameRegex = compileSqlLikePattern(procedureNamePattern); final Schema resultSchema = defineGetProceduresSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - final BlockingQueue queue = - new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); final List collectedResults = Collections.synchronizedList(new ArrayList<>()); - final String catalogParam = catalog; - - Runnable procedureFetcher = - () -> { - ExecutorService apiExecutor = null; - final FieldList localResultSchemaFields = resultSchemaFields; - final List>> apiFutures = new ArrayList<>(); - - try { - List datasetsToScan = - fetchMatchingDatasets(catalogParam, schemaPattern, schemaRegex); - - if (datasetsToScan.isEmpty()) { - LOG.info("Fetcher thread found no matching datasets. Finishing."); - return; - } - - apiExecutor = connection.getMetadataExecutor(); - - LOG.fine("Submitting parallel findMatchingRoutines tasks..."); - for (Dataset dataset : datasetsToScan) { - if (Thread.currentThread().isInterrupted()) { - LOG.warning("Fetcher interrupted during dataset iteration submission."); - break; - } - - final DatasetId currentDatasetId = dataset.getDatasetId(); - Callable> apiCallable = - () -> - findMatchingBigQueryObjects( - "Routine", - () -> - bigquery.listRoutines( - currentDatasetId, RoutineListOption.pageSize(DEFAULT_PAGE_SIZE)), - (name) -> - bigquery.getRoutine( - RoutineId.of( - currentDatasetId.getProject(), - currentDatasetId.getDataset(), - name)), - (rt) -> rt.getRoutineId().getRoutine(), - procedureNamePattern, - procedureNameRegex, - false); - Future> apiFuture = apiExecutor.submit(apiCallable); - apiFutures.add(apiFuture); - } - LOG.fine("Finished submitting " + apiFutures.size() + " findMatchingRoutines tasks."); - - LOG.fine("Processing results from findMatchingRoutines tasks..."); - for (Future> apiFuture : apiFutures) { - if (Thread.currentThread().isInterrupted()) { - LOG.warning("Fetcher interrupted while processing API futures."); - break; - } - try { - List routinesResult = apiFuture.get(); - if (routinesResult != null) { - for (Routine routine : routinesResult) { - if (Thread.currentThread().isInterrupted()) break; - - if ("PROCEDURE".equalsIgnoreCase(routine.getRoutineType())) { - LOG.fine("Processing procedure sequentially: " + routine.getRoutineId()); - processProcedureInfo(routine, collectedResults, localResultSchemaFields); - } else { - LOG.finer("Skipping non-procedure routine: " + routine.getRoutineId()); - } - } - } - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - LOG.warning("Fetcher thread interrupted while waiting for API future result."); - break; - } catch (CancellationException e) { - LOG.warning("A findMatchingRoutines task was cancelled."); - } - } - - Comparator comparator = - defineGetProceduresComparator(localResultSchemaFields); - sortResults(collectedResults, comparator, "getProcedures", LOG); - populateQueue(collectedResults, queue, localResultSchemaFields); + List targetDatasets = getTargetDatasets(catalog, schemaPattern); - } catch (Throwable t) { - handleFetcherException(t, queue, "getProcedures"); - } finally { - apiFutures.forEach(f -> f.cancel(true)); - finalizeFetcher(queue, localResultSchemaFields, "Procedure fetcher"); + processTargetRoutinesConcurrently( + targetDatasets, + procedureNamePattern, + true, + collectedResults, + resultSchemaFields, + (routine, results, fields) -> { + if (!"PROCEDURE".equalsIgnoreCase(routine.getRoutineType())) { + return; } - }; - - Future fetcherFuture = connection.getExecutorService().submit(procedureFetcher); - BigQueryJsonResultSet resultSet = - BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); + processProcedureInfo(routine, results, fields); + }); - LOG.info("Submitted background task for getProcedures to metadata executor"); - return resultSet; + Comparator comparator = defineGetProceduresComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getProcedures", LOG); + final BlockingQueue queue = + new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); + Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } Schema defineGetProceduresSchema() { @@ -1001,103 +912,52 @@ Comparator defineGetProceduresComparator(FieldList resultSchemaF @Override public ResultSet getProcedureColumns( - String catalog, String schemaPattern, String procedureNamePattern, String columnNamePattern) { + String catalog, String schemaPattern, String procedureNamePattern, String columnNamePattern) + throws SQLException { + LOG.info( + "getProcedureColumns called for catalog: %s, schemaPattern: %s, procedureNamePattern: %s, columnNamePattern: %s", + catalog, schemaPattern, procedureNamePattern, columnNamePattern); - if (catalog != null && catalog.isEmpty()) { - LOG.warning("Returning empty ResultSet because catalog (project) is empty."); - return new BigQueryJsonResultSet(); - } - if ((schemaPattern != null && schemaPattern.isEmpty()) + if ((catalog != null && catalog.isEmpty()) + || (schemaPattern != null && schemaPattern.isEmpty()) || (procedureNamePattern != null && procedureNamePattern.isEmpty()) || (columnNamePattern != null && columnNamePattern.isEmpty())) { - LOG.warning("Returning empty ResultSet because an explicit empty pattern was provided."); + LOG.warning( + "Returning empty ResultSet as one or more patterns are empty or catalog is empty."); return new BigQueryJsonResultSet(); } - LOG.info( - "getProcedureColumns called for catalog: %s, schemaPattern: %s, procedureNamePattern:" - + " %s, columnNamePattern: %s", - catalog, schemaPattern, procedureNamePattern, columnNamePattern); - - final Pattern schemaRegex = compileSqlLikePattern(schemaPattern); - final Pattern procedureNameRegex = compileSqlLikePattern(procedureNamePattern); - final Pattern columnNameRegex = compileSqlLikePattern(columnNamePattern); - final Schema resultSchema = defineGetProcedureColumnsSchema(); - final BlockingQueue queue = - new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); + final FieldList resultSchemaFields = resultSchema.getFields(); final List collectedResults = Collections.synchronizedList(new ArrayList<>()); - final String catalogParam = catalog; - - Runnable procedureColumnFetcher = - () -> { - ExecutorService listRoutinesExecutor = null; - ExecutorService getRoutineDetailsExecutor = null; + List targetDatasets = getTargetDatasets(catalog, schemaPattern); + Pattern columnNameRegex = compileSqlLikePattern(columnNamePattern); + processTargetRoutinesConcurrently( + targetDatasets, + procedureNamePattern, + true, + collectedResults, + resultSchemaFields, + (routine, results, fields) -> { + if (!"PROCEDURE".equalsIgnoreCase(routine.getRoutineType())) { + return; + } try { - List datasetsToScan = - fetchMatchingDatasets(catalogParam, schemaPattern, schemaRegex); - if (datasetsToScan.isEmpty() || Thread.currentThread().isInterrupted()) { - LOG.info( - "Fetcher: No matching datasets or interrupted early. Catalog: " + catalogParam); - return; - } - - listRoutinesExecutor = connection.getMetadataExecutor(); - List procedureIdsToGet = - listMatchingProcedureIdsFromDatasets( - datasetsToScan, - procedureNamePattern, - procedureNameRegex, - listRoutinesExecutor, - catalogParam, - LOG); - listRoutinesExecutor = null; - - if (procedureIdsToGet.isEmpty() || Thread.currentThread().isInterrupted()) { - LOG.info("Fetcher: No procedure IDs found or interrupted. Catalog: " + catalogParam); - return; - } - - getRoutineDetailsExecutor = connection.getMetadataExecutor(); - List fullRoutines = - fetchFullRoutineDetailsForIds( - procedureIdsToGet, - getRoutineDetailsExecutor, - LOG, - RoutineOption.fields( - RoutineField.ROUTINE_REFERENCE, - RoutineField.ARGUMENTS, - RoutineField.ROUTINE_TYPE)); - getRoutineDetailsExecutor = null; - - if (fullRoutines.isEmpty() || Thread.currentThread().isInterrupted()) { - LOG.info( - "Fetcher: No full routines fetched or interrupted. Catalog: " + catalogParam); - return; - } - processProcedureArgumentsSequentially( - fullRoutines, columnNameRegex, collectedResults, resultSchema.getFields(), LOG); - - Comparator comparator = - defineGetProcedureColumnsComparator(resultSchema.getFields()); - sortResults(collectedResults, comparator, "getProcedureColumns", LOG); - populateQueue(collectedResults, queue, resultSchema.getFields()); - - } catch (Throwable t) { - handleFetcherException(t, queue, "getProcedureColumns"); - } finally { - finalizeFetcher(queue, resultSchema.getFields(), "Procedure column fetcher"); + Collections.singletonList(routine), columnNameRegex, results, fields, LOG); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted while processing procedure arguments", e); } - }; - - Future fetcherFuture = connection.getExecutorService().submit(procedureColumnFetcher); - BigQueryJsonResultSet resultSet = - BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); + }); - LOG.info("Started background task for getProcedureColumns for catalog: " + catalog); - return resultSet; + Comparator comparator = defineGetProcedureColumnsComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getProcedureColumns", LOG); + final BlockingQueue queue = + new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); + Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } List listMatchingProcedureIdsFromDatasets( @@ -1138,6 +998,7 @@ List listMatchingProcedureIdsFromDatasets( (rt) -> rt.getRoutineId().getRoutine(), procedureNamePattern, procedureNameRegex, + logger, false); listRoutineFutures.add(listRoutinesExecutor.submit(listCallable)); } @@ -1559,7 +1420,11 @@ Comparator defineGetProcedureColumnsComparator(FieldList resultS @Override public ResultSet getTables( - String catalog, String schemaPattern, String tableNamePattern, String[] types) { + String catalog, String schemaPattern, String tableNamePattern, String[] types) + throws SQLException { + LOG.info( + "getTables called for catalog: %s, schemaPattern: %s, tableNamePattern: %s, types: %s", + catalog, schemaPattern, tableNamePattern, Arrays.toString(types)); if ((catalog != null && catalog.isEmpty()) || (schemaPattern != null && schemaPattern.isEmpty()) @@ -1568,126 +1433,37 @@ public ResultSet getTables( "Returning empty ResultSet as one or more patterns are empty or catalog is empty."); return new BigQueryJsonResultSet(); } + final Schema resultSchema = defineGetTablesSchema(); + final FieldList resultSchemaFields = resultSchema.getFields(); + final List collectedResults = Collections.synchronizedList(new ArrayList<>()); Tuple effectiveIdentifiers = determineEffectiveCatalogAndSchema(catalog, schemaPattern); String effectiveCatalog = effectiveIdentifiers.x(); String effectiveSchemaPattern = effectiveIdentifiers.y(); - LOG.info( - "getTables called for catalog: %s, schemaPattern: %s, tableNamePattern: %s, types: %s", - effectiveCatalog, effectiveSchemaPattern, tableNamePattern, Arrays.toString(types)); + List targetDatasets = getTargetDatasets(effectiveCatalog, effectiveSchemaPattern); - final Pattern schemaRegex = compileSqlLikePattern(effectiveSchemaPattern); - final Pattern tableNameRegex = compileSqlLikePattern(tableNamePattern); final Set requestedTypes = (types == null || types.length == 0) ? null : new HashSet<>(Arrays.asList(types)); - final Schema resultSchema = defineGetTablesSchema(); - final FieldList resultSchemaFields = resultSchema.getFields(); + processTargetTablesConcurrently( + targetDatasets, + tableNamePattern, + true, + false, + collectedResults, + resultSchemaFields, + (bqTable, results, fields) -> { + processTableInfo(bqTable, requestedTypes, results, fields); + }); + Comparator comparator = defineGetTablesComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getTables", LOG); final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - final List collectedResults = Collections.synchronizedList(new ArrayList<>()); - final String catalogParam = effectiveCatalog; - final String schemaParam = effectiveSchemaPattern; - - Runnable tableFetcher = - () -> { - ExecutorService apiExecutor = null; - final FieldList localResultSchemaFields = resultSchemaFields; - final List>> apiFutures = new ArrayList<>(); - - try { - List datasetsToScan = - fetchMatchingDatasets(catalogParam, schemaParam, schemaRegex); - - if (datasetsToScan.isEmpty()) { - LOG.info("Fetcher thread found no matching datasets. Returning empty resultset."); - return; - } - - apiExecutor = connection.getMetadataExecutor(); - - LOG.fine("Submitting parallel findMatchingTables tasks..."); - for (Dataset dataset : datasetsToScan) { - if (Thread.currentThread().isInterrupted()) { - LOG.warning("Table fetcher interrupted during dataset iteration."); - break; - } - - final DatasetId currentDatasetId = dataset.getDatasetId(); - Callable> apiCallable = - () -> - findMatchingBigQueryObjects( - "Table", - () -> - bigquery.listTables( - currentDatasetId, TableListOption.pageSize(DEFAULT_PAGE_SIZE)), - (name) -> - bigquery.getTable( - TableId.of( - currentDatasetId.getProject(), - currentDatasetId.getDataset(), - name), - TableOption.fields( - TableField.TABLE_REFERENCE, - TableField.TYPE, - TableField.DESCRIPTION)), - (tbl) -> tbl.getTableId().getTable(), - tableNamePattern, - tableNameRegex, - false); - Future> apiFuture = apiExecutor.submit(apiCallable); - apiFutures.add(apiFuture); - } - LOG.fine("Finished submitting " + apiFutures.size() + " findMatchingTables tasks."); - - LOG.fine("Processing results from findMatchingTables tasks..."); - for (Future> apiFuture : apiFutures) { - if (Thread.currentThread().isInterrupted()) { - LOG.warning("Table fetcher interrupted while processing API futures."); - break; - } - try { - List tablesResult = apiFuture.get(); - if (tablesResult != null) { - for (Table table : tablesResult) { - if (Thread.currentThread().isInterrupted()) break; - - LOG.fine("Processing table sequentially: " + table.getTableId()); - processTableInfo( - table, requestedTypes, collectedResults, localResultSchemaFields); - } - } - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - LOG.warning("Fetcher thread interrupted while waiting for API future result."); - break; - } catch (CancellationException e) { - LOG.warning("A findMatchingTables task was cancelled."); - } - } - - Comparator comparator = - defineGetTablesComparator(localResultSchemaFields); - sortResults(collectedResults, comparator, "getTables", LOG); - populateQueue(collectedResults, queue, localResultSchemaFields); - - } catch (Throwable t) { - handleFetcherException(t, queue, "getTables"); - } finally { - apiFutures.forEach(f -> f.cancel(true)); - finalizeFetcher(queue, localResultSchemaFields, "Table fetcher"); - } - }; - - Future fetcherFuture = connection.getExecutorService().submit(tableFetcher); - BigQueryJsonResultSet resultSet = - BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); - - LOG.info("Started background thread for getTables"); - return resultSet; + Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } Schema defineGetTablesSchema() { @@ -1878,7 +1654,11 @@ static List prepareGetTableTypesRows(Schema schema) { @Override public ResultSet getColumns( - String catalog, String schemaPattern, String tableNamePattern, String columnNamePattern) { + String catalog, String schemaPattern, String tableNamePattern, String columnNamePattern) + throws SQLException { + LOG.info( + "getColumns called for catalog: %s, schemaPattern: %s, tableNamePattern: %s, columnNamePattern: %s", + catalog, schemaPattern, tableNamePattern, columnNamePattern); if ((catalog != null && catalog.isEmpty()) || (schemaPattern != null && schemaPattern.isEmpty()) @@ -1888,116 +1668,39 @@ public ResultSet getColumns( "Returning empty ResultSet as one or more patterns are empty or catalog is empty."); return new BigQueryJsonResultSet(); } + final Schema resultSchema = defineGetColumnsSchema(); + final FieldList resultSchemaFields = resultSchema.getFields(); + final List collectedResults = Collections.synchronizedList(new ArrayList<>()); Tuple effectiveIdentifiers = determineEffectiveCatalogAndSchema(catalog, schemaPattern); String effectiveCatalog = effectiveIdentifiers.x(); String effectiveSchemaPattern = effectiveIdentifiers.y(); - LOG.info( - "getColumns called for catalog: %s, schemaPattern: %s, tableNamePattern: %s," - + " columnNamePattern: %s", - effectiveCatalog, effectiveSchemaPattern, tableNamePattern, columnNamePattern); + List targetDatasets = getTargetDatasets(effectiveCatalog, effectiveSchemaPattern); - Pattern schemaRegex = compileSqlLikePattern(effectiveSchemaPattern); - Pattern tableNameRegex = compileSqlLikePattern(tableNamePattern); - Pattern columnNameRegex = compileSqlLikePattern(columnNamePattern); + boolean hasWildcards = + columnNamePattern != null + && (columnNamePattern.contains("%") || columnNamePattern.contains("_")); + Pattern columnNameRegex = hasWildcards ? compileSqlLikePattern(columnNamePattern) : null; - final Schema resultSchema = defineGetColumnsSchema(); - final FieldList resultSchemaFields = resultSchema.getFields(); + processTargetTablesConcurrently( + targetDatasets, + tableNamePattern, + true, + false, + collectedResults, + resultSchemaFields, + (bqTable, results, fields) -> { + processTableColumns(bqTable, columnNameRegex, results, fields); + }); + + Comparator comparator = defineGetColumnsComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getColumns", LOG); final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - final List collectedResults = Collections.synchronizedList(new ArrayList<>()); - final String catalogParam = effectiveCatalog; - final String schemaParam = effectiveSchemaPattern; - - Runnable columnFetcher = - () -> { - ExecutorService columnExecutor = null; - final List> taskFutures = new ArrayList<>(); - final FieldList localResultSchemaFields = resultSchemaFields; - - try { - List datasetsToScan = - fetchMatchingDatasets(catalogParam, schemaParam, schemaRegex); - - if (datasetsToScan.isEmpty()) { - LOG.info("Fetcher thread found no matching datasets. Returning empty resultset."); - return; - } - - columnExecutor = connection.getMetadataExecutor(); - - for (Dataset dataset : datasetsToScan) { - if (Thread.currentThread().isInterrupted()) { - LOG.warning("Fetcher interrupted during dataset iteration."); - break; - } - - DatasetId datasetId = dataset.getDatasetId(); - LOG.info("Processing dataset: " + datasetId.getDataset()); - - List
tablesToScan = - findMatchingBigQueryObjects( - "Table", - () -> - bigquery.listTables( - datasetId, TableListOption.pageSize(DEFAULT_PAGE_SIZE)), - (name) -> - bigquery.getTable( - TableId.of(datasetId.getProject(), datasetId.getDataset(), name), - TableOption.fields( - TableField.TABLE_REFERENCE, TableField.TYPE, TableField.SCHEMA)), - (tbl) -> tbl.getTableId().getTable(), - tableNamePattern, - tableNameRegex, - false); - - for (Table table : tablesToScan) { - if (Thread.currentThread().isInterrupted()) { - LOG.warning( - "Fetcher interrupted during table iteration for dataset " - + datasetId.getDataset()); - break; - } - - TableId tableId = table.getTableId(); - LOG.fine("Submitting task for table: " + tableId); - final Table finalTable = table; - Future future = - columnExecutor.submit( - () -> - processTableColumns( - finalTable, - columnNameRegex, - collectedResults, - localResultSchemaFields)); - taskFutures.add(future); - } - if (Thread.currentThread().isInterrupted()) break; - } - - waitForTasksCompletion(taskFutures); - - Comparator comparator = - defineGetColumnsComparator(localResultSchemaFields); - sortResults(collectedResults, comparator, "getColumns", LOG); - populateQueue(collectedResults, queue, localResultSchemaFields); - - } catch (Throwable t) { - handleFetcherException(t, queue, "getColumns"); - } finally { - taskFutures.forEach(f -> f.cancel(true)); - finalizeFetcher(queue, localResultSchemaFields, "Column fetcher"); - } - }; - - Future fetcherFuture = connection.getExecutorService().submit(columnFetcher); - BigQueryJsonResultSet resultSet = - BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); - - LOG.info("Started background thread for getColumns"); - return resultSet; + Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } private void processTableColumns( @@ -2413,6 +2116,8 @@ Schema defineGetVersionColumnsSchema() { @Override public ResultSet getPrimaryKeys(String catalog, String schema, String table) throws SQLException { + LOG.info( + "getPrimaryKeys called for catalog: %s, schema: %s, table: %s", catalog, schema, table); if ((catalog != null && catalog.isEmpty()) || (schema != null && schema.isEmpty()) || table == null @@ -2426,11 +2131,13 @@ public ResultSet getPrimaryKeys(String catalog, String schema, String table) thr final FieldList resultSchemaFields = resultSchema.getFields(); final List collectedResults = Collections.synchronizedList(new ArrayList<>()); - List targetDatasets = getTargetDatasets(catalog, schema); + List targetDatasets = getTargetDatasets(catalog, schema, false); processTargetTablesConcurrently( targetDatasets, table, + false, + true, collectedResults, resultSchemaFields, (bqTable, results, fields) -> { @@ -2508,12 +2215,14 @@ private Comparator defineGetPrimaryKeysComparator(FieldList resu @Override public ResultSet getImportedKeys(String catalog, String schema, String table) throws SQLException { + LOG.info( + "getImportedKeys called for catalog: %s, schema: %s, table: %s", catalog, schema, table); if ((catalog != null && catalog.isEmpty()) || (schema != null && schema.isEmpty()) || table == null || table.isEmpty()) { LOG.warning( - "Returning empty ResultSet as required parameters are null/empty, or catalog/schema parameters are empty."); + "Returning empty ResultSet as one or more patterns are empty or catalog is empty."); return new BigQueryJsonResultSet(); } @@ -2521,11 +2230,13 @@ public ResultSet getImportedKeys(String catalog, String schema, String table) final FieldList resultSchemaFields = resultSchema.getFields(); final List collectedResults = Collections.synchronizedList(new ArrayList<>()); - List targetDatasets = getTargetDatasets(catalog, schema); + List targetDatasets = getTargetDatasets(catalog, schema, false); processTargetTablesConcurrently( targetDatasets, table, + false, + true, collectedResults, resultSchemaFields, (bqTable, results, fields) -> { @@ -2551,75 +2262,54 @@ public ResultSet getImportedKeys(String catalog, String schema, String table) @Override public ResultSet getExportedKeys(String catalog, String schema, String table) throws SQLException { + LOG.info( + "getExportedKeys called for catalog: %s, schema: %s, table: %s", catalog, schema, table); if ((catalog != null && catalog.isEmpty()) || (schema != null && schema.isEmpty()) || table == null || table.isEmpty()) { LOG.warning( - "Returning empty ResultSet as required parameters are null/empty, or catalog/schema parameters are empty."); + "Returning empty ResultSet as one or more patterns are empty or catalog is empty."); return new BigQueryJsonResultSet(); } final Schema resultSchema = defineForeignKeyResultSetSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - // Early return for PCNT catalog schemas (containing '.') as they do not support table - // constraints. - if (schema != null && schema.contains(".")) { - final BlockingQueue queue = new LinkedBlockingQueue<>(1); - signalEndOfData(queue, resultSchemaFields); - return BigQueryJsonResultSet.of(resultSchema, 0, queue, null); - } - - // Fallback Path: If catalog or schema is null, fall back to REST API metadata scan. - if (catalog == null || schema == null) { - final List collectedResults = Collections.synchronizedList(new ArrayList<>()); - List targetDatasets = getTargetDatasets(catalog, null); + final List collectedResults = Collections.synchronizedList(new ArrayList<>()); + List targetDatasets = getTargetDatasets(catalog, null); - processTargetTablesConcurrently( - targetDatasets, - null, - collectedResults, - resultSchemaFields, - (bqTable, results, fields) -> { - TableConstraints constraints = bqTable.getTableConstraints(); - if (constraints == null || constraints.getForeignKeys() == null) { - return; - } - for (ForeignKey fk : constraints.getForeignKeys()) { - TableId pkTableId = fk.getReferencedTable(); - if (pkTableId == null - || !equalsOrNullMatchesAll(catalog, pkTableId.getProject()) - || !equalsOrNullMatchesAll(schema, pkTableId.getDataset()) - || !table.equals(pkTableId.getTable())) { - continue; - } - processForeignKey(fk, pkTableId, bqTable.getTableId(), results, fields); + processTargetTablesConcurrently( + targetDatasets, + null, + false, + true, + collectedResults, + resultSchemaFields, + (bqTable, results, fields) -> { + TableConstraints constraints = bqTable.getTableConstraints(); + if (constraints == null || constraints.getForeignKeys() == null) { + return; + } + for (ForeignKey fk : constraints.getForeignKeys()) { + TableId pkTableId = fk.getReferencedTable(); + if (pkTableId == null + || !equalsOrNullMatchesAll(catalog, pkTableId.getProject()) + || !equalsOrNullMatchesAll(schema, pkTableId.getDataset()) + || !table.equals(pkTableId.getTable())) { + continue; } - }); - - Comparator comparator = defineFkTableSortComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getExportedKeys", LOG); + processForeignKey(fk, pkTableId, bqTable.getTableId(), results, fields); + } + }); - final BlockingQueue queue = - new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); - return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); - } + Comparator comparator = defineFkTableSortComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getExportedKeys", LOG); - String sql = getExportedKeysSqlContent(); - String formattedSql = replaceSqlParameters(sql, catalog, schema, table); - PreparedStatement stmt = this.connection.prepareStatement(formattedSql); - if (stmt == null) { - throw new BigQueryJdbcException("Failed to prepare statement for getExportedKeys"); - } - try { - stmt.closeOnCompletion(); - return stmt.executeQuery(); - } catch (SQLException e) { - closeStatementIgnoreException(stmt); - throw new BigQueryJdbcException("Error executing getExportedKeys", e); - } + final BlockingQueue queue = + new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); + Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } @Override @@ -2631,6 +2321,9 @@ public ResultSet getCrossReference( String foreignSchema, String foreignTable) throws SQLException { + LOG.info( + "getCrossReference called for parentCatalog: %s, parentSchema: %s, parentTable: %s, foreignCatalog: %s, foreignSchema: %s, foreignTable: %s", + parentCatalog, parentSchema, parentTable, foreignCatalog, foreignSchema, foreignTable); if ((parentCatalog != null && parentCatalog.isEmpty()) || (parentSchema != null && parentSchema.isEmpty()) || parentTable == null @@ -2640,7 +2333,7 @@ public ResultSet getCrossReference( || foreignTable == null || foreignTable.isEmpty()) { LOG.warning( - "Returning empty ResultSet as required parameters are null/empty, or catalog/schema parameters are empty."); + "Returning empty ResultSet as one or more patterns are empty or catalog is empty."); return new BigQueryJsonResultSet(); } @@ -2648,11 +2341,13 @@ public ResultSet getCrossReference( final FieldList resultSchemaFields = resultSchema.getFields(); final List collectedResults = Collections.synchronizedList(new ArrayList<>()); - List targetDatasets = getTargetDatasets(foreignCatalog, foreignSchema); + List targetDatasets = getTargetDatasets(foreignCatalog, foreignSchema, false); processTargetTablesConcurrently( targetDatasets, foreignTable, + false, + true, collectedResults, resultSchemaFields, (bqTable, results, fields) -> { @@ -3594,67 +3289,30 @@ public RowIdLifetime getRowIdLifetime() { @Override public ResultSet getSchemas(String catalog, String schemaPattern) throws SQLException { + LOG.info("getSchemas called for catalog: %s, schemaPattern: %s", catalog, schemaPattern); if ((catalog != null && catalog.isEmpty()) || (schemaPattern != null && schemaPattern.isEmpty())) { - LOG.warning("Returning empty ResultSet as catalog or schemaPattern is an empty string."); + LOG.warning( + "Returning empty ResultSet as one or more patterns are empty or catalog is empty."); return new BigQueryJsonResultSet(); } - LOG.info("getSchemas called for catalog: %s, schemaPattern: %s", catalog, schemaPattern); - - final Pattern schemaRegex = compileSqlLikePattern(schemaPattern); final Schema resultSchema = defineGetSchemasSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - - if (catalog != null) { - // Single-Catalog Path: completely synchronous on caller thread - final BlockingQueue queue = new LinkedBlockingQueue<>(); - List datasets = fetchMatchingDatasets(catalog, schemaPattern, schemaRegex); - List collectedResults = new ArrayList<>(); - for (Dataset dataset : datasets) { - processSchemaInfo(dataset, collectedResults, resultSchemaFields); - } - Comparator comparator = defineGetSchemasComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getSchemas", LOG); - Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); - return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); + final Pattern schemaRegex = compileSqlLikePattern(schemaPattern); + List datasets = fetchMatchingDatasets(catalog, schemaPattern, schemaRegex); + List collectedResults = new ArrayList<>(datasets.size()); + for (Dataset dataset : datasets) { + processSchemaInfo(dataset.getDatasetId(), collectedResults, resultSchemaFields); } - // Multi-Catalog Path: fan out using connection-scoped metadataExecutor + Comparator comparator = defineGetSchemasComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getSchemas", LOG); + final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - Runnable multiSchemaFetcher = - () -> { - final FieldList localResultSchemaFields = resultSchemaFields; - final List collectedResults = new ArrayList<>(); - - try { - List datasets = fetchMatchingDatasets(catalog, schemaPattern, schemaRegex); - for (Dataset dataset : datasets) { - if (Thread.currentThread().isInterrupted()) { - break; - } - processSchemaInfo(dataset, collectedResults, localResultSchemaFields); - } - - Comparator comparator = - defineGetSchemasComparator(localResultSchemaFields); - sortResults(collectedResults, comparator, "getSchemas", LOG); - populateQueue(collectedResults, queue, localResultSchemaFields); - - } catch (Throwable t) { - handleFetcherException(t, queue, "getSchemas"); - } finally { - finalizeFetcher(queue, localResultSchemaFields, "Multi-schema fetcher"); - } - }; - - Future fetcherFuture = connection.getExecutorService().submit(multiSchemaFetcher); - BigQueryJsonResultSet resultSet = - BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); - - LOG.info("Submitted background task for multi-catalog getSchemas to query executor"); - return resultSet; + Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } Schema defineGetSchemasSchema() { @@ -3671,8 +3329,7 @@ Schema defineGetSchemasSchema() { } void processSchemaInfo( - Dataset dataset, List collectedResults, FieldList resultSchemaFields) { - DatasetId datasetId = dataset.getDatasetId(); + DatasetId datasetId, List collectedResults, FieldList resultSchemaFields) { LOG.finer("Processing schema info for dataset: " + datasetId); try { String schemaName = datasetId.getDataset(); @@ -3794,135 +3451,45 @@ Schema defineGetClientInfoPropertiesSchema() { } @Override - public ResultSet getFunctions(String catalog, String schemaPattern, String functionNamePattern) { + public ResultSet getFunctions(String catalog, String schemaPattern, String functionNamePattern) + throws SQLException { + LOG.info( + "getFunctions called for catalog: %s, schemaPattern: %s, functionNamePattern: %s", + catalog, schemaPattern, functionNamePattern); + if ((catalog != null && catalog.isEmpty()) || (schemaPattern != null && schemaPattern.isEmpty()) || (functionNamePattern != null && functionNamePattern.isEmpty())) { LOG.warning( - "Returning empty ResultSet as catalog is empty or a pattern is empty for" - + " getFunctions."); + "Returning empty ResultSet as one or more patterns are empty or catalog is empty."); return new BigQueryJsonResultSet(); } - LOG.info( - "getFunctions called for catalog: %s, schemaPattern: %s, functionNamePattern: %s", - catalog, schemaPattern, functionNamePattern); - - final Pattern schemaRegex = compileSqlLikePattern(schemaPattern); - final Pattern functionNameRegex = compileSqlLikePattern(functionNamePattern); final Schema resultSchema = defineGetFunctionsSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - final BlockingQueue queue = - new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); final List collectedResults = Collections.synchronizedList(new ArrayList<>()); - final String catalogParam = catalog; - - Runnable functionFetcher = - () -> { - ExecutorService apiExecutor = null; - final FieldList localResultSchemaFields = resultSchemaFields; - final List>> apiFutures = new ArrayList<>(); - - try { - List datasetsToScan = - fetchMatchingDatasets(catalogParam, schemaPattern, schemaRegex); - - if (datasetsToScan.isEmpty()) { - LOG.info("Fetcher thread found no matching datasets. Returning empty resultset."); - return; - } - - apiExecutor = connection.getMetadataExecutor(); + List targetDatasets = getTargetDatasets(catalog, schemaPattern); - for (Dataset dataset : datasetsToScan) { - if (Thread.currentThread().isInterrupted()) { - LOG.warning("Function fetcher interrupted during dataset iteration submission."); - break; - } - - final DatasetId currentDatasetId = dataset.getDatasetId(); - - Callable> apiCallable = - () -> { - LOG.fine( - "Fetching all routines for dataset: %s, pattern: %s", - currentDatasetId.getDataset(), functionNamePattern); - return findMatchingBigQueryObjects( - "Routine", - () -> - bigquery.listRoutines( - currentDatasetId, RoutineListOption.pageSize(DEFAULT_PAGE_SIZE)), - (name) -> - bigquery.getRoutine( - RoutineId.of( - currentDatasetId.getProject(), - currentDatasetId.getDataset(), - name)), - (rt) -> rt.getRoutineId().getRoutine(), - functionNamePattern, - functionNameRegex, - false); - }; - Future> apiFuture = apiExecutor.submit(apiCallable); - apiFutures.add(apiFuture); - } - LOG.fine( - "Finished submitting " - + apiFutures.size() - + " findMatchingRoutines (for functions) tasks."); - - for (Future> apiFuture : apiFutures) { - if (Thread.currentThread().isInterrupted()) { - LOG.warning("Function fetcher interrupted while processing API futures."); - break; - } - try { - List routinesResult = apiFuture.get(); - if (routinesResult != null) { - for (Routine routine : routinesResult) { - if (Thread.currentThread().isInterrupted()) { - break; - } - String routineType = routine.getRoutineType(); - if ("SCALAR_FUNCTION".equalsIgnoreCase(routineType) - || "TABLE_FUNCTION".equalsIgnoreCase(routineType)) { - LOG.fine( - "Processing function sequentially: " - + routine.getRoutineId() - + " of type " - + routineType); - processFunctionInfo(routine, collectedResults, localResultSchemaFields); - } - } - } - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - LOG.warning( - "Function fetcher thread interrupted while waiting for API future result."); - break; - } catch (CancellationException e) { - LOG.warning( - "Cancellation in findMatchingRoutines (for functions) task: " + e.getMessage()); - } - } - Comparator comparator = - defineGetFunctionsComparator(localResultSchemaFields); - sortResults(collectedResults, comparator, "getFunctions", LOG); - populateQueue(collectedResults, queue, localResultSchemaFields); - } catch (Throwable t) { - handleFetcherException(t, queue, "getFunctions"); - } finally { - apiFutures.forEach(f -> f.cancel(true)); - finalizeFetcher(queue, localResultSchemaFields, "Function fetcher"); + processTargetRoutinesConcurrently( + targetDatasets, + functionNamePattern, + true, + collectedResults, + resultSchemaFields, + (routine, results, fields) -> { + if (!"SCALAR_FUNCTION".equalsIgnoreCase(routine.getRoutineType()) + && !"TABLE_FUNCTION".equalsIgnoreCase(routine.getRoutineType())) { + return; } - }; - - Future fetcherFuture = connection.getExecutorService().submit(functionFetcher); - BigQueryJsonResultSet resultSet = - BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); + processFunctionInfo(routine, results, fields); + }); - LOG.info("Started background thread for getFunctions"); - return resultSet; + Comparator comparator = defineGetFunctionsComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getFunctions", LOG); + final BlockingQueue queue = + new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); + Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } Schema defineGetFunctionsSchema() { @@ -4015,97 +3582,54 @@ Comparator defineGetFunctionsComparator(FieldList resultSchemaFi @Override public ResultSet getFunctionColumns( - String catalog, String schemaPattern, String functionNamePattern, String columnNamePattern) { - if (catalog != null && catalog.isEmpty()) { - LOG.warning("Returning empty ResultSet catalog (project) is empty."); - return new BigQueryJsonResultSet(); - } - if ((schemaPattern != null && schemaPattern.isEmpty()) + String catalog, String schemaPattern, String functionNamePattern, String columnNamePattern) + throws SQLException { + LOG.info( + "getFunctionColumns called for catalog: %s, schemaPattern: %s, functionNamePattern: %s, columnNamePattern: %s", + catalog, schemaPattern, functionNamePattern, columnNamePattern); + + if ((catalog != null && catalog.isEmpty()) + || (schemaPattern != null && schemaPattern.isEmpty()) || (functionNamePattern != null && functionNamePattern.isEmpty()) || (columnNamePattern != null && columnNamePattern.isEmpty())) { - LOG.warning("Returning empty ResultSet because an explicit empty pattern was provided."); + LOG.warning( + "Returning empty ResultSet as one or more patterns are empty or catalog is empty."); return new BigQueryJsonResultSet(); } - LOG.info( - "getFunctionColumns called for catalog: %s, schemaPattern: %s, functionNamePattern: %s," - + " columnNamePattern: %s", - catalog, schemaPattern, functionNamePattern, columnNamePattern); - - final Pattern schemaRegex = compileSqlLikePattern(schemaPattern); - final Pattern functionNameRegex = compileSqlLikePattern(functionNamePattern); - final Pattern columnNameRegex = compileSqlLikePattern(columnNamePattern); - final Schema resultSchema = defineGetFunctionColumnsSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - final BlockingQueue queue = - new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); final List collectedResults = Collections.synchronizedList(new ArrayList<>()); - final String catalogParam = catalog; + List targetDatasets = getTargetDatasets(catalog, schemaPattern); - Runnable functionColumnFetcher = - () -> { - ExecutorService listRoutinesExecutor = null; - ExecutorService getRoutineDetailsExecutor = null; + Pattern columnNameRegex = compileSqlLikePattern(columnNamePattern); + processTargetRoutinesConcurrently( + targetDatasets, + functionNamePattern, + true, + collectedResults, + resultSchemaFields, + (routine, results, fields) -> { + if (!"SCALAR_FUNCTION".equalsIgnoreCase(routine.getRoutineType()) + && !"TABLE_FUNCTION".equalsIgnoreCase(routine.getRoutineType())) { + return; + } try { - List datasetsToScan = - fetchMatchingDatasets(catalogParam, schemaPattern, schemaRegex); - - if (datasetsToScan.isEmpty() || Thread.currentThread().isInterrupted()) { - LOG.info( - "Fetcher: No matching datasets or interrupted early. Catalog: " + catalogParam); - return; - } - - listRoutinesExecutor = connection.getMetadataExecutor(); - List functionIdsToGet = - listMatchingFunctionIdsFromDatasets( - datasetsToScan, - functionNamePattern, - functionNameRegex, - listRoutinesExecutor, - catalogParam, - LOG); - listRoutinesExecutor = null; - - if (functionIdsToGet.isEmpty() || Thread.currentThread().isInterrupted()) { - LOG.info("Fetcher: No function IDs found or interrupted. Catalog: " + catalogParam); - return; - } - - getRoutineDetailsExecutor = connection.getMetadataExecutor(); - List fullFunctions = - fetchFullRoutineDetailsForIds(functionIdsToGet, getRoutineDetailsExecutor, LOG); - getRoutineDetailsExecutor = null; - - if (fullFunctions.isEmpty() || Thread.currentThread().isInterrupted()) { - LOG.info( - "Fetcher: No full functions fetched or interrupted. Catalog: " + catalogParam); - return; - } - processFunctionParametersSequentially( - fullFunctions, columnNameRegex, collectedResults, resultSchemaFields, LOG); - - Comparator comparator = - defineGetFunctionColumnsComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getFunctionColumns", LOG); - populateQueue(collectedResults, queue, resultSchemaFields); - - } catch (Throwable t) { - handleFetcherException(t, queue, "getFunctionColumns"); - } finally { - finalizeFetcher(queue, resultSchemaFields, "Function column fetcher"); + Collections.singletonList(routine), columnNameRegex, results, fields, LOG); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted while processing function parameters", e); } - }; - - Future fetcherFuture = connection.getExecutorService().submit(functionColumnFetcher); - BigQueryJsonResultSet resultSet = - BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); + }); - LOG.info("Started background thread for getFunctionColumns for catalog: " + catalog); - return resultSet; + Comparator comparator = defineGetFunctionColumnsComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getFunctionColumns", LOG); + final BlockingQueue queue = + new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); + Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } Schema defineGetFunctionColumnsSchema() { @@ -4212,6 +3736,7 @@ List listMatchingFunctionIdsFromDatasets( (rt) -> rt.getRoutineId().getRoutine(), functionNamePattern, functionNameRegex, + logger, false); listRoutineFutures.add(listRoutinesExecutor.submit(listCallable)); } @@ -4648,6 +4173,7 @@ List findMatchingBigQueryObjects( Function nameExtractor, String pattern, Pattern regex, + BigQueryJdbcCustomLogger logger, boolean throwOn404) throws InterruptedException { @@ -4657,29 +4183,30 @@ List findMatchingBigQueryObjects( try { Iterable objects; if (needsList) { - LOG.info( + logger.info( "Listing all %ss (pattern: %s)...", objectTypeName, pattern == null ? "" : pattern); Page firstPage = listAllOperation.get(); objects = firstPage.iterateAll(); - LOG.fine("Retrieved initial %s list, iterating & filtering if needed...", objectTypeName); + logger.fine( + "Retrieved initial %s list, iterating & filtering if needed...", objectTypeName); } else { - LOG.info("Getting specific %s: '%s'", objectTypeName, pattern); + logger.info("Getting specific %s: '%s'", objectTypeName, pattern); T specificObject = getSpecificOperation.apply(pattern); objects = (specificObject == null) ? Collections.emptyList() : Collections.singletonList(specificObject); if (specificObject == null) { - LOG.info("Specific %s not found: '%s'", objectTypeName, pattern); + logger.info("Specific %s not found: '%s'", objectTypeName, pattern); } } boolean wasListing = needsList; for (T obj : objects) { if (Thread.currentThread().isInterrupted()) { - LOG.warning("Thread interrupted during " + objectTypeName + " processing loop."); + logger.warning("Thread interrupted during " + objectTypeName + " processing loop."); throw new InterruptedException( "Interrupted during " + objectTypeName + " processing loop"); } @@ -4697,20 +4224,20 @@ List findMatchingBigQueryObjects( } catch (BigQueryException e) { if (e.getCode() == 404 && !throwOn404) { - LOG.info("%s '%s' not found (API error 404).", objectTypeName, pattern); + logger.info("%s '%s' not found (API error 404).", objectTypeName, pattern); return Collections.emptyList(); } else { - LOG.warning( + logger.warning( "BigQueryException finding %ss for pattern '%s': %s (Code: %d)", objectTypeName, pattern, e.getMessage(), e.getCode()); throw e; } } catch (InterruptedException e) { Thread.currentThread().interrupt(); - LOG.warning("Interrupted while finding " + objectTypeName + "s."); + logger.warning("Interrupted while finding " + objectTypeName + "s."); throw e; } catch (Exception e) { - LOG.severe( + logger.severe( "Unexpected exception finding %ss for pattern '%s': %s", objectTypeName, pattern, e.getMessage()); throw new RuntimeException(e); @@ -4868,7 +4395,11 @@ private static boolean isRegexMetacharacter(char c) { } boolean needsListing(String pattern) { - return pattern == null || pattern.contains("%") || pattern.contains("_"); + return pattern == null || hasWildcards(pattern); + } + + boolean hasWildcards(String pattern) { + return pattern != null && (pattern.contains("%") || pattern.contains("_")); } FieldValue createStringFieldValue(String value) { @@ -4910,20 +4441,24 @@ private Long getLongValueOrNull(FieldValueList fvl, int index) { private void waitForTasksCompletion(List> taskFutures) throws ExecutionException { LOG.info("Waiting for %d submitted tasks to complete...", taskFutures.size()); - for (Future future : taskFutures) { - try { + try { + for (Future future : taskFutures) { if (!future.isCancelled()) { future.get(); } - } catch (CancellationException e) { - LOG.warning("A table processing task was cancelled."); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - LOG.warning( - "Fetcher thread interrupted while waiting for tasks. Attempting to cancel remaining" - + " tasks."); - taskFutures.forEach(f -> f.cancel(true)); - break; + } + } catch (CancellationException e) { + LOG.warning("A table processing task was cancelled."); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOG.warning( + "Fetcher thread interrupted while waiting for tasks. Attempting to cancel remaining" + + " tasks."); + } finally { + for (Future future : taskFutures) { + if (!future.isDone()) { + future.cancel(true); + } } } LOG.info("Finished waiting for tasks."); @@ -4966,29 +4501,6 @@ private void populateQueue( } } - private void handleFetcherException( - Throwable t, BlockingQueue queue, String fetcherName) { - if (t instanceof ExecutionException && t.getCause() != null) { - t = t.getCause(); - } - if (t instanceof InterruptedException) { - Thread.currentThread().interrupt(); - LOG.warning("Fetcher interrupted in " + fetcherName + ": " + t.getMessage()); - } else { - LOG.severe("Unexpected error in " + fetcherName + ": " + t.getMessage()); - } - writeErrorToQueue(queue, t); - } - - private void finalizeFetcher( - BlockingQueue queue, FieldList schema, String fetcherName) { - if (Thread.currentThread().isInterrupted()) { - writeErrorToQueue(queue, new SQLException("Metadata fetch interrupted.")); - } - signalEndOfData(queue, schema); - LOG.info(fetcherName + " finished."); - } - private void signalEndOfData( BlockingQueue queue, FieldList resultSchemaFields) { try { @@ -5025,6 +4537,7 @@ private List fetchDatasetsForProject( (ds) -> ds.getDatasetId().getDataset(), schemaPattern, schemaRegex, + LOG, throwOn404); return datasets != null ? datasets : Collections.emptyList(); } catch (InterruptedException e) { @@ -5120,46 +4633,31 @@ private boolean equalsOrNullMatchesAll(String expected, String actual) { return expected == null || expected.equals(actual); } - private void writeErrorToQueue(BlockingQueue queue, Throwable t) { - Exception ex = (t instanceof Exception) ? (Exception) t : new Exception(t); - BigQueryFieldValueListWrapper element = BigQueryFieldValueListWrapper.ofError(ex); - if (!queue.offer(element)) { - boolean wasInterrupted = Thread.interrupted(); - try { - queue.put(element); - } catch (InterruptedException ie) { - LOG.warning("Failed to put exception to queue due to interruption."); - wasInterrupted = true; - } finally { - if (wasInterrupted) { - Thread.currentThread().interrupt(); - } - } - } - } - - 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()); + private List getTargetDatasets(String catalog, String schemaPattern, boolean isPattern) + throws SQLException { + if (schemaPattern != null && (!isPattern || !hasWildcards(schemaPattern))) { + List projects = resolveTargetProjects(catalog); + List datasetIds = new ArrayList<>(projects.size()); + for (String project : projects) { + datasetIds.add(DatasetId.of(project, schemaPattern)); } 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)); + Pattern schemaRegex = compileSqlLikePattern(schemaPattern); + List datasets = fetchMatchingDatasets(catalog, schemaPattern, schemaRegex); + List datasetIds = new ArrayList<>(datasets.size()); + for (Dataset dataset : datasets) { + datasetIds.add(dataset.getDatasetId()); } return datasetIds; } + private List getTargetDatasets(String catalog, String schemaPattern) + throws SQLException { + return getTargetDatasets(catalog, schemaPattern, true); + } + private List resolveTargetProjects(String catalog) throws SQLException { return (catalog != null) ? Collections.singletonList(catalog) : getAccessibleCatalogNames(); } @@ -5173,6 +4671,7 @@ void process(Table bqTable, List collectedResults, FieldList res private void processSingleTable( DatasetId datasetId, String tableName, + boolean requireBaseTable, List collectedResults, FieldList resultSchemaFields, TableProcessor processor) @@ -5186,28 +4685,152 @@ private void processSingleTable( LOG.info( "Table '%s' or dataset '%s' not found in project '%s' (API error 404). Skipping.", tableName, datasetId.getDataset(), datasetId.getProject()); - bqTable = null; - } else { - throw new SQLException("Error while fetching table metadata: " + e.getMessage(), e); + return; } + 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); + + if (bqTable == null || bqTable.getDefinition() == null) { + return; + } + if (requireBaseTable && bqTable.getDefinition().getType() != TableDefinition.Type.TABLE) { + return; + } + processor.process(bqTable, collectedResults, resultSchemaFields); + } + + private interface RoutineProcessor { + void process( + Routine routine, List collectedResults, FieldList resultSchemaFields) + throws SQLException; + } + + private void processSingleRoutine( + DatasetId datasetId, + String routineName, + List collectedResults, + FieldList resultSchemaFields, + RoutineProcessor processor) + throws SQLException { + Routine routine; + try { + routine = + bigquery.getRoutine( + RoutineId.of(datasetId.getProject(), datasetId.getDataset(), routineName)); + } catch (BigQueryException e) { + if (e.getCode() == 404) { + LOG.info("Routine '%s' not found. Skipping.", routineName); + return; + } + throw new SQLException("Error while fetching routine metadata: " + e.getMessage(), e); + } catch (Exception e) { + throw new SQLException("Error while fetching routine metadata: " + e.getMessage(), e); + } + + if (routine == null) { + return; + } + processor.process(routine, collectedResults, resultSchemaFields); + } + + private void processTargetRoutinesConcurrently( + List targetDatasets, + String routineNamePattern, + boolean isPattern, + List collectedResults, + FieldList resultSchemaFields, + RoutineProcessor processor) + throws SQLException { + + boolean hasWildcards = isPattern && hasWildcards(routineNamePattern); + Pattern routineRegex = compileSqlLikePattern(routineNamePattern); + boolean isLiteralName = routineNamePattern != null && !hasWildcards; + + if (targetDatasets.size() == 1 && isLiteralName) { + processSingleRoutine( + targetDatasets.get(0), + routineNamePattern, + collectedResults, + resultSchemaFields, + processor); + return; + } + + ExecutorService executor = connection.getMetadataExecutor(); + List> taskFutures = new ArrayList<>(); + + try { + List> tasks = new ArrayList<>(); + for (DatasetId datasetId : targetDatasets) { + if (isLiteralName) { + tasks.add( + () -> { + processSingleRoutine( + datasetId, routineNamePattern, collectedResults, resultSchemaFields, processor); + return null; + }); + continue; + } + + try { + Page routinesPage = + bigquery.listRoutines( + datasetId, BigQuery.RoutineListOption.pageSize(DEFAULT_PAGE_SIZE)); + if (routinesPage == null) continue; + for (Routine routine : routinesPage.iterateAll()) { + if (routine == null) continue; + String routineName = routine.getRoutineId().getRoutine(); + if (routineRegex != null && !routineRegex.matcher(routineName).matches()) continue; + tasks.add( + () -> { + processSingleRoutine( + datasetId, routineName, collectedResults, resultSchemaFields, processor); + return null; + }); + } + } catch (BigQueryException e) { + if (e.getCode() == 404) { + LOG.info("Dataset '%s' not found while listing routines. Skipping.", datasetId); + continue; + } + throw new SQLException("Error while listing routines: " + e.getMessage(), e); + } + } + + for (Callable task : tasks) { + taskFutures.add(executor.submit(task)); + } + waitForTasksCompletion(taskFutures); + } catch (ExecutionException e) { + throw new SQLException( + "Error while fetching routine metadata: " + e.getCause().getMessage(), e); } } private void processTargetTablesConcurrently( List targetDatasets, - String tableName, + String tableNamePattern, + boolean isPattern, + boolean requireBaseTable, List collectedResults, FieldList resultSchemaFields, TableProcessor processor) throws SQLException { - if (targetDatasets.size() == 1 && tableName != null) { + + boolean hasWildcards = isPattern && hasWildcards(tableNamePattern); + Pattern tableRegex = compileSqlLikePattern(tableNamePattern); + boolean isLiteralName = tableNamePattern != null && !hasWildcards; + + if (targetDatasets.size() == 1 && isLiteralName) { processSingleTable( - targetDatasets.get(0), tableName, collectedResults, resultSchemaFields, processor); + targetDatasets.get(0), + tableNamePattern, + requireBaseTable, + collectedResults, + resultSchemaFields, + processor); return; } @@ -5217,11 +4840,16 @@ private void processTargetTablesConcurrently( try { List> tasks = new ArrayList<>(); for (DatasetId datasetId : targetDatasets) { - if (tableName != null) { + if (isLiteralName) { tasks.add( () -> { processSingleTable( - datasetId, tableName, collectedResults, resultSchemaFields, processor); + datasetId, + tableNamePattern, + requireBaseTable, + collectedResults, + resultSchemaFields, + processor); return null; }); continue; @@ -5229,20 +4857,21 @@ private void processTargetTablesConcurrently( try { Page
tablesPage = - bigquery.listTables(datasetId, TableListOption.pageSize(DEFAULT_PAGE_SIZE)); - if (tablesPage == null) { - continue; - } + bigquery.listTables(datasetId, BigQuery.TableListOption.pageSize(DEFAULT_PAGE_SIZE)); + if (tablesPage == null) continue; for (Table table : tablesPage.iterateAll()) { - if (table.getDefinition() == null - || table.getDefinition().getType() != TableDefinition.Type.TABLE) { - continue; - } + if (table == null) continue; + if (requireBaseTable + && table.getDefinition() != null + && table.getDefinition().getType() != TableDefinition.Type.TABLE) continue; + String tableName = table.getTableId().getTable(); + if (tableRegex != null && !tableRegex.matcher(tableName).matches()) continue; tasks.add( () -> { processSingleTable( datasetId, - table.getTableId().getTable(), + tableName, + requireBaseTable, collectedResults, resultSchemaFields, processor); @@ -5251,9 +4880,7 @@ private void processTargetTablesConcurrently( } } catch (BigQueryException e) { if (e.getCode() == 404) { - LOG.info( - "Dataset '%s' not found in project '%s' (API error 404). Skipping.", - datasetId.getDataset(), datasetId.getProject()); + LOG.info("Dataset '%s' not found while listing tables. Skipping.", datasetId); continue; } throw new SQLException("Error while listing tables: " + e.getMessage(), e); @@ -5264,17 +4891,9 @@ private void processTargetTablesConcurrently( taskFutures.add(executor.submit(task)); } 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)); + throw new SQLException( + "Error while fetching table metadata: " + e.getCause().getMessage(), e); } } @@ -5371,43 +4990,4 @@ private Comparator defineFkTableSortComparator(FieldList resultS (FieldValueList fvl) -> getLongValueOrNull(fvl, KEY_SEQ_IDX), Comparator.nullsFirst(Long::compareTo)); } - - private static synchronized String getExportedKeysSqlContent() { - if (exportedKeysSqlContent == null) { - exportedKeysSqlContent = readSqlFromFile(GET_EXPORTED_KEYS_SQL); - } - return exportedKeysSqlContent; - } - - static String readSqlFromFile(String filename) { - try (InputStream in = BigQueryDatabaseMetaData.class.getResourceAsStream(filename)) { - if (in == null) { - throw new IllegalArgumentException("SQL file not found: " + filename); - } - ByteArrayOutputStream result = new ByteArrayOutputStream(); - byte[] buffer = new byte[1024]; - int length; - while ((length = in.read(buffer)) != -1) { - result.write(buffer, 0, length); - } - return result.toString(StandardCharsets.UTF_8.name()); - } catch (IOException e) { - throw new RuntimeException("Failed to read SQL file: " + filename, e); - } - } - - String replaceSqlParameters(String sql, String... params) throws SQLException { - return String.format(sql, (Object[]) params); - } - - private void closeStatementIgnoreException(Statement stmt) { - if (stmt == null) { - return; - } - try { - stmt.close(); - } catch (SQLException e) { - // ignore - } - } } 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 3ef43bf648fb..fba46dca9478 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 @@ -31,24 +31,19 @@ import com.google.api.gax.paging.Page; import com.google.cloud.bigquery.*; -import com.google.cloud.bigquery.BigQuery.RoutineListOption; import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.jdbc.BigQueryJdbcTypeMappings.ColumnTypeInfo; import java.io.IOException; import java.io.InputStream; import java.sql.DatabaseMetaData; -import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Statement; import java.sql.Types; import java.util.*; -import java.util.concurrent.Callable; -import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.Future; import java.util.regex.Pattern; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -731,7 +726,7 @@ public void testProcessSchemaInfo() { String schemaName = "dataset_beta"; Dataset dataset = mockBigQueryDataset(catalog, schemaName); - dbMetadata.processSchemaInfo(dataset, collectedResults, resultSchemaFields); + dbMetadata.processSchemaInfo(dataset.getDatasetId(), collectedResults, resultSchemaFields); assertEquals(1, collectedResults.size()); FieldValueList row = collectedResults.get(0); @@ -1010,7 +1005,7 @@ public void testFindMatchingBigQueryObjects_Routines_ListWithPattern() throws Ex Routine func1 = mockBigQueryRoutine(catalog, schema, "func_123", "FUNCTION", "f1"); Routine otherProc = mockBigQueryRoutine(catalog, schema, "another_proc", "PROCEDURE", "p3"); - Page page = mock(Page.class, withSettings().withoutAnnotations()); + Page page = mock(Page.class); when(page.iterateAll()).thenReturn(Arrays.asList(proc1, func1, proc2, otherProc)); when(bigqueryClient.listRoutines(eq(datasetId), any(BigQuery.RoutineListOption[].class))) .thenReturn(page); @@ -1028,6 +1023,7 @@ public void testFindMatchingBigQueryObjects_Routines_ListWithPattern() throws Ex (rt) -> rt.getRoutineId().getRoutine(), pattern, regex, + dbMetadata.LOG, false); verify(bigqueryClient, times(1)) @@ -1054,7 +1050,7 @@ public void testFindMatchingBigQueryObjects_Routines_ListNoPattern() throws Exce Routine proc1 = mockBigQueryRoutine(catalog, schema, "proc_abc", "PROCEDURE", "p1"); Routine func1 = mockBigQueryRoutine(catalog, schema, "func_123", "FUNCTION", "f1"); - Page page = mock(Page.class, withSettings().withoutAnnotations()); + Page page = mock(Page.class); when(page.iterateAll()).thenReturn(Arrays.asList(proc1, func1)); when(bigqueryClient.listRoutines(eq(datasetId), any(BigQuery.RoutineListOption[].class))) .thenReturn(page); @@ -1070,6 +1066,7 @@ public void testFindMatchingBigQueryObjects_Routines_ListNoPattern() throws Exce (rt) -> rt.getRoutineId().getRoutine(), pattern, regex, + dbMetadata.LOG, false); verify(bigqueryClient, times(1)) @@ -1105,6 +1102,7 @@ public void testFindMatchingBigQueryObjects_Routines_GetSpecific() throws Except (rt) -> rt.getRoutineId().getRoutine(), procNameExact, regex, + dbMetadata.LOG, false); verify(bigqueryClient, times(1)).getRoutine(eq(routineId)); @@ -1130,6 +1128,7 @@ private List
invokeFindMatchingObjectsWithException( (table) -> "name", pattern, dbMetadata.compileSqlLikePattern(pattern), + dbMetadata.LOG, throwOn404); } @@ -1584,96 +1583,6 @@ public void testDefineGetProcedureColumnsComparator() { assertEquals("param_a", results.get(4).get("COLUMN_NAME").getStringValue()); } - @Test - public void testListMatchingProcedureIdsFromDatasets() throws Exception { - String catalog = "test-proj"; - String schema1Name = "dataset1"; - String schema2Name = "dataset2"; - Dataset dataset1 = mockBigQueryDataset(catalog, schema1Name); - Dataset dataset2 = mockBigQueryDataset(catalog, schema2Name); - List datasetsToScan = Arrays.asList(dataset1, dataset2); - - Routine proc1_ds1 = mockBigQueryRoutine(catalog, schema1Name, "proc_a", "PROCEDURE", "desc a"); - Routine func1_ds1 = mockBigQueryRoutine(catalog, schema1Name, "func_b", "FUNCTION", "desc b"); - Routine proc2_ds2 = mockBigQueryRoutine(catalog, schema2Name, "proc_c", "PROCEDURE", "desc c"); - - Page page1 = mock(Page.class, withSettings().withoutAnnotations()); - when(page1.iterateAll()).thenReturn(Arrays.asList(proc1_ds1, func1_ds1)); - when(bigqueryClient.listRoutines(eq(dataset1.getDatasetId()), any(RoutineListOption.class))) - .thenReturn(page1); - - Page page2 = mock(Page.class, withSettings().withoutAnnotations()); - when(page2.iterateAll()).thenReturn(Collections.singletonList(proc2_ds2)); - when(bigqueryClient.listRoutines(eq(dataset2.getDatasetId()), any(RoutineListOption.class))) - .thenReturn(page2); - - ExecutorService mockExecutor = mock(ExecutorService.class); - doAnswer( - invocation -> { - Callable callable = invocation.getArgument(0); - @SuppressWarnings("unchecked") // Suppress warning for raw Future mock - Future mockedFuture = mock(Future.class); - - try { - Object result = callable.call(); - doReturn(result).when(mockedFuture).get(); - } catch (InterruptedException interruptedException) { - doThrow(interruptedException).when(mockedFuture).get(); - } catch (Exception e) { - doThrow(new ExecutionException(e)).when(mockedFuture).get(); - } - return mockedFuture; - }) - .when(mockExecutor) - .submit(any(Callable.class)); - - List resultIds = - dbMetadata.listMatchingProcedureIdsFromDatasets( - datasetsToScan, null, null, mockExecutor, catalog, dbMetadata.LOG); - - assertEquals(2, resultIds.size()); - assertTrue(resultIds.contains(proc1_ds1.getRoutineId())); - assertTrue(resultIds.contains(proc2_ds2.getRoutineId())); - assertFalse(resultIds.contains(func1_ds1.getRoutineId())); // Should not contain functions - - verify(mockExecutor, times(2)).submit(any(Callable.class)); - } - - @Test - public void testProcessProcedureArgumentsSequentially_Basic() throws InterruptedException { - String catalog = "p"; - String schemaName = "d"; - RoutineArgument arg1 = mockRoutineArgument("arg1_name", StandardSQLTypeName.STRING, "IN"); - Routine proc1 = - mockBigQueryRoutineWithArgs( - catalog, schemaName, "proc1", "PROCEDURE", "desc1", Collections.singletonList(arg1)); - Routine func1 = - mockBigQueryRoutineWithArgs( - catalog, - schemaName, - "func1", - "FUNCTION", - "desc_func", - Collections.emptyList()); // Should be skipped - Routine proc2 = - mockBigQueryRoutineWithArgs( - catalog, schemaName, "proc2", "PROCEDURE", "desc2", Collections.emptyList()); - - List fullRoutines = Arrays.asList(proc1, func1, proc2); - Pattern columnNameRegex = null; - List collectedResults = Collections.synchronizedList(new ArrayList<>()); - Schema resultSchema = dbMetadata.defineGetProcedureColumnsSchema(); - FieldList resultSchemaFields = resultSchema.getFields(); - - dbMetadata.processProcedureArgumentsSequentially( - fullRoutines, columnNameRegex, collectedResults, resultSchemaFields, dbMetadata.LOG); - - // Only proc1 has arguments, so collectedResults should contain 1 row. - assertEquals(1, collectedResults.size()); - FieldValueList row = collectedResults.get(0); - assertEquals("arg1_name", row.get("COLUMN_NAME").getStringValue()); - } - @Test public void testDefineGetTableTypesSchema() { Schema schema = BigQueryDatabaseMetaData.defineGetTableTypesSchema(); @@ -3361,20 +3270,20 @@ public void testGetSchemas_WithProjectDiscovery() throws SQLException { when(bigQueryConnection.getDiscoveredProjects()).thenReturn(Arrays.asList("discovered-1")); when(bigQueryConnection.getAdditionalProjects()).thenReturn("additional-1"); - Page pagePrimary = mock(Page.class, withSettings().withoutAnnotations()); + Page pagePrimary = mock(Page.class); Dataset dsPrimary = mockBigQueryDataset("primary-project", "dataset_p"); when(pagePrimary.iterateAll()).thenReturn(Collections.singletonList(dsPrimary)); when(bigqueryClient.listDatasets( eq("primary-project"), any(BigQuery.DatasetListOption[].class))) .thenReturn(pagePrimary); - Page pageAdditional = mock(Page.class, withSettings().withoutAnnotations()); + Page pageAdditional = mock(Page.class); Dataset dsAdditional = mockBigQueryDataset("additional-1", "dataset_a"); when(pageAdditional.iterateAll()).thenReturn(Collections.singletonList(dsAdditional)); when(bigqueryClient.listDatasets(eq("additional-1"), any(BigQuery.DatasetListOption[].class))) .thenReturn(pageAdditional); - Page pageDiscovered = mock(Page.class, withSettings().withoutAnnotations()); + Page pageDiscovered = mock(Page.class); Dataset dsDiscovered = mockBigQueryDataset("discovered-1", "dataset_d"); when(pageDiscovered.iterateAll()).thenReturn(Collections.singletonList(dsDiscovered)); when(bigqueryClient.listDatasets(eq("discovered-1"), any(BigQuery.DatasetListOption[].class))) @@ -3406,14 +3315,14 @@ public void testGetSchemas_WithoutProjectDiscovery() throws SQLException { when(bigQueryConnection.getDiscoveredProjects()).thenReturn(Arrays.asList("discovered-1")); when(bigQueryConnection.getAdditionalProjects()).thenReturn("additional-1"); - Page pagePrimary = mock(Page.class, withSettings().withoutAnnotations()); + Page pagePrimary = mock(Page.class); Dataset dsPrimary = mockBigQueryDataset("primary-project", "dataset_p"); when(pagePrimary.iterateAll()).thenReturn(Collections.singletonList(dsPrimary)); when(bigqueryClient.listDatasets( eq("primary-project"), any(BigQuery.DatasetListOption[].class))) .thenReturn(pagePrimary); - Page pageAdditional = mock(Page.class, withSettings().withoutAnnotations()); + Page pageAdditional = mock(Page.class); Dataset dsAdditional = mockBigQueryDataset("additional-1", "dataset_a"); when(pageAdditional.iterateAll()).thenReturn(Collections.singletonList(dsAdditional)); when(bigqueryClient.listDatasets(eq("additional-1"), any(BigQuery.DatasetListOption[].class))) @@ -3440,7 +3349,7 @@ public void testGetSchemas_WithoutProjectDiscovery() throws SQLException { } private void mockDatasetIteration(DatasetId datasetId) { - Page pagePrimary = mock(Page.class, withSettings().withoutAnnotations()); + Page pagePrimary = mock(Page.class); Dataset dsPrimary = mock(Dataset.class); when(dsPrimary.getDatasetId()).thenReturn(datasetId); when(pagePrimary.iterateAll()).thenReturn(Collections.singletonList(dsPrimary)); @@ -3463,7 +3372,7 @@ private Table mockTableWithConstraints(TableId tableId, TableConstraints constra } private void mockTableIteration(DatasetId datasetId, Table... tables) { - Page
pageTables = mock(Page.class, withSettings().withoutAnnotations()); + Page
pageTables = mock(Page.class); when(pageTables.iterateAll()).thenReturn(Arrays.asList(tables)); when(bigqueryClient.listTables(eq(datasetId), any(BigQuery.TableListOption[].class))) .thenReturn(pageTables); @@ -3574,27 +3483,7 @@ public void testGetImportedKeys_noKeys() throws SQLException { } @Test - public void testGetExportedKeys_infoSchema() throws SQLException { - PreparedStatement mockStmt = mock(PreparedStatement.class); - ResultSet mockRs = mock(ResultSet.class); - when(bigQueryConnection.prepareStatement(anyString())).thenReturn(mockStmt); - when(mockStmt.executeQuery()).thenReturn(mockRs); - - ResultSet rs = dbMetadata.getExportedKeys("test-project", "dataset_p", "ref_table"); - assertEquals(mockRs, rs); - verify(mockStmt).closeOnCompletion(); - verify(mockStmt).executeQuery(); - } - - @Test - public void testGetExportedKeys_pcntSchema() throws SQLException { - try (ResultSet rs = dbMetadata.getExportedKeys("test-project", "dataset.p", "ref_table")) { - assertFalse(rs.next()); - } - } - - @Test - public void testGetExportedKeys_fallback_hasKeys() throws SQLException { + public void testGetExportedKeys_hasKeys() throws SQLException { DatasetId datasetId = DatasetId.of("test-project", "dataset_p"); TableId tableId = TableId.of("test-project", "dataset_p", "table_p"); TableId refTableId = TableId.of("test-project", "dataset_p", "ref_table"); @@ -3616,7 +3505,7 @@ public void testGetExportedKeys_fallback_hasKeys() throws SQLException { mockDatasetIteration(datasetId); mockTableIteration(datasetId, mockTableP); - try (ResultSet rs = dbMetadata.getExportedKeys("test-project", null, "ref_table")) { + try (ResultSet rs = dbMetadata.getExportedKeys("test-project", "dataset_p", "ref_table")) { assertTrue(rs.next()); assertEquals("test-project", rs.getString("PKTABLE_CAT")); assertEquals("dataset_p", rs.getString("PKTABLE_SCHEM")); @@ -3644,7 +3533,7 @@ public void testGetExportedKeys_fallback_hasKeys() throws SQLException { } @Test - public void testGetExportedKeys_fallback_noKeys() throws SQLException { + public void testGetExportedKeys_noKeys() throws SQLException { DatasetId datasetId = DatasetId.of("test-project", "dataset_p"); TableId tableId = TableId.of("test-project", "dataset_p", "table_p"); @@ -3652,7 +3541,7 @@ public void testGetExportedKeys_fallback_noKeys() throws SQLException { mockDatasetIteration(datasetId); mockTableIteration(datasetId, mockTableP); - try (ResultSet rs = dbMetadata.getExportedKeys("test-project", null, "ref_table")) { + try (ResultSet rs = dbMetadata.getExportedKeys("test-project", "dataset_p", "ref_table")) { assertFalse(rs.next()); } } From 78bdd9034fca322c84bfc62c6dfb9bf03c129f71 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Mon, 20 Jul 2026 19:32:47 +0000 Subject: [PATCH 02/14] Fix tests for fallback logic in getExportedKeys --- .../cloud/bigquery/jdbc/BigQueryDatabaseMetaDataTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 fba46dca9478..b14b0a8c97d3 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 @@ -3505,7 +3505,7 @@ public void testGetExportedKeys_hasKeys() throws SQLException { mockDatasetIteration(datasetId); mockTableIteration(datasetId, mockTableP); - try (ResultSet rs = dbMetadata.getExportedKeys("test-project", "dataset_p", "ref_table")) { + try (ResultSet rs = dbMetadata.getExportedKeys("test-project", null, "ref_table")) { assertTrue(rs.next()); assertEquals("test-project", rs.getString("PKTABLE_CAT")); assertEquals("dataset_p", rs.getString("PKTABLE_SCHEM")); @@ -3541,7 +3541,7 @@ public void testGetExportedKeys_noKeys() throws SQLException { mockDatasetIteration(datasetId); mockTableIteration(datasetId, mockTableP); - try (ResultSet rs = dbMetadata.getExportedKeys("test-project", "dataset_p", "ref_table")) { + try (ResultSet rs = dbMetadata.getExportedKeys("test-project", null, "ref_table")) { assertFalse(rs.next()); } } From c7340ddbfd77d477f648c567964c2ed942432e33 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Mon, 20 Jul 2026 19:34:48 +0000 Subject: [PATCH 03/14] Merge main's getExportedKeys SQL logic --- .../jdbc/BigQueryDatabaseMetaData.java | 132 ++++++++++++++---- 1 file changed, 102 insertions(+), 30 deletions(-) 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 8557b96864c4..c752c975f5ab 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 @@ -50,11 +50,17 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.jdbc.BigQueryJdbcTypeMappings.ColumnTypeInfo; import com.google.cloud.bigquery.jdbc.utils.BigQueryJdbcVersionUtility; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; import java.sql.Connection; import java.sql.DatabaseMetaData; +import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.RowIdLifetime; import java.sql.SQLException; +import java.sql.Statement; import java.sql.Types; import java.util.ArrayList; import java.util.Arrays; @@ -92,6 +98,8 @@ class BigQueryDatabaseMetaData implements DatabaseMetaData { private static final String PROCEDURE_TERM = "Procedure"; private static final int DEFAULT_PAGE_SIZE = 500; private static final int DEFAULT_QUEUE_CAPACITY = 5000; + private static final String GET_EXPORTED_KEYS_SQL = "DatabaseMetaData_GetExportedKeys.sql"; + private static String exportedKeysSqlContent; // Declared package-private for testing. static final String GOOGLE_SQL_QUOTED_IDENTIFIER = "`"; // Does not include SQL:2003 Keywords as per JDBC spec. @@ -2276,40 +2284,65 @@ public ResultSet getExportedKeys(String catalog, String schema, String table) final Schema resultSchema = defineForeignKeyResultSetSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - final List collectedResults = Collections.synchronizedList(new ArrayList<>()); - List targetDatasets = getTargetDatasets(catalog, null); + // Early return for PCNT catalog schemas (containing '.') as they do not support table + // constraints. + if (schema != null && schema.contains(".")) { + final BlockingQueue queue = new LinkedBlockingQueue<>(1); + signalEndOfData(queue, resultSchemaFields); + return BigQueryJsonResultSet.of(resultSchema, 0, queue, null); + } - processTargetTablesConcurrently( - targetDatasets, - null, - false, - true, - collectedResults, - resultSchemaFields, - (bqTable, results, fields) -> { - TableConstraints constraints = bqTable.getTableConstraints(); - if (constraints == null || constraints.getForeignKeys() == null) { - return; - } - for (ForeignKey fk : constraints.getForeignKeys()) { - TableId pkTableId = fk.getReferencedTable(); - if (pkTableId == null - || !equalsOrNullMatchesAll(catalog, pkTableId.getProject()) - || !equalsOrNullMatchesAll(schema, pkTableId.getDataset()) - || !table.equals(pkTableId.getTable())) { - continue; + // Fallback Path: If catalog or schema is null, fall back to REST API metadata scan. + if (catalog == null || schema == null) { + final List collectedResults = Collections.synchronizedList(new ArrayList<>()); + List targetDatasets = getTargetDatasets(catalog, schema); + + processTargetTablesConcurrently( + targetDatasets, + null, + false, + true, + collectedResults, + resultSchemaFields, + (bqTable, results, fields) -> { + TableConstraints constraints = bqTable.getTableConstraints(); + if (constraints == null || constraints.getForeignKeys() == null) { + return; } - processForeignKey(fk, pkTableId, bqTable.getTableId(), results, fields); - } - }); + for (ForeignKey fk : constraints.getForeignKeys()) { + TableId pkTableId = fk.getReferencedTable(); + if (pkTableId == null + || !equalsOrNullMatchesAll(catalog, pkTableId.getProject()) + || !equalsOrNullMatchesAll(schema, pkTableId.getDataset()) + || !table.equals(pkTableId.getTable())) { + continue; + } + processForeignKey(fk, pkTableId, bqTable.getTableId(), results, fields); + } + }); - Comparator comparator = defineFkTableSortComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getExportedKeys", LOG); + Comparator comparator = defineFkTableSortComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getExportedKeys", LOG); - final BlockingQueue queue = - new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); - return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); + final BlockingQueue queue = + new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); + Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); + } + + String sql = getExportedKeysSqlContent(); + String formattedSql = replaceSqlParameters(sql, catalog, schema, table); + PreparedStatement stmt = this.connection.prepareStatement(formattedSql); + if (stmt == null) { + throw new BigQueryJdbcException("Failed to prepare statement for getExportedKeys"); + } + try { + stmt.closeOnCompletion(); + return stmt.executeQuery(); + } catch (SQLException e) { + closeStatementIgnoreException(stmt); + throw new BigQueryJdbcException("Error executing getExportedKeys", e); + } } @Override @@ -4990,4 +5023,43 @@ private Comparator defineFkTableSortComparator(FieldList resultS (FieldValueList fvl) -> getLongValueOrNull(fvl, KEY_SEQ_IDX), Comparator.nullsFirst(Long::compareTo)); } + + private static synchronized String getExportedKeysSqlContent() { + if (exportedKeysSqlContent == null) { + exportedKeysSqlContent = readSqlFromFile(GET_EXPORTED_KEYS_SQL); + } + return exportedKeysSqlContent; + } + + static String readSqlFromFile(String filename) { + try (InputStream in = BigQueryDatabaseMetaData.class.getResourceAsStream(filename)) { + if (in == null) { + throw new IllegalArgumentException("SQL file not found: " + filename); + } + ByteArrayOutputStream result = new ByteArrayOutputStream(); + byte[] buffer = new byte[1024]; + int length; + while ((length = in.read(buffer)) != -1) { + result.write(buffer, 0, length); + } + return result.toString(StandardCharsets.UTF_8.name()); + } catch (IOException e) { + throw new RuntimeException("Failed to read SQL file: " + filename, e); + } + } + + String replaceSqlParameters(String sql, String... params) throws SQLException { + return String.format(sql, (Object[]) params); + } + + private void closeStatementIgnoreException(Statement stmt) { + if (stmt == null) { + return; + } + try { + stmt.close(); + } catch (SQLException e) { + // ignore + } + } } From ecd25fcd5ba07ab0584e809a5c90957ba2e665ec Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Mon, 20 Jul 2026 19:53:01 +0000 Subject: [PATCH 04/14] modify `processTarget...Concurrently` methods to only get full object when needed --- .../jdbc/BigQueryDatabaseMetaData.java | 88 +++++++++++++------ 1 file changed, 60 insertions(+), 28 deletions(-) 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 c752c975f5ab..cd95ac73c57f 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 @@ -794,6 +794,7 @@ public ResultSet getProcedures(String catalog, String schemaPattern, String proc targetDatasets, procedureNamePattern, true, + false, collectedResults, resultSchemaFields, (routine, results, fields) -> { @@ -945,6 +946,7 @@ public ResultSet getProcedureColumns( targetDatasets, procedureNamePattern, true, + true, collectedResults, resultSchemaFields, (routine, results, fields) -> { @@ -1460,6 +1462,7 @@ public ResultSet getTables( tableNamePattern, true, false, + false, collectedResults, resultSchemaFields, (bqTable, results, fields) -> { @@ -1697,6 +1700,7 @@ public ResultSet getColumns( tableNamePattern, true, false, + true, collectedResults, resultSchemaFields, (bqTable, results, fields) -> { @@ -2146,6 +2150,7 @@ public ResultSet getPrimaryKeys(String catalog, String schema, String table) thr table, false, true, + true, collectedResults, resultSchemaFields, (bqTable, results, fields) -> { @@ -2245,6 +2250,7 @@ public ResultSet getImportedKeys(String catalog, String schema, String table) table, false, true, + true, collectedResults, resultSchemaFields, (bqTable, results, fields) -> { @@ -2302,6 +2308,7 @@ public ResultSet getExportedKeys(String catalog, String schema, String table) null, false, true, + true, collectedResults, resultSchemaFields, (bqTable, results, fields) -> { @@ -2381,6 +2388,7 @@ public ResultSet getCrossReference( foreignTable, false, true, + true, collectedResults, resultSchemaFields, (bqTable, results, fields) -> { @@ -3507,6 +3515,7 @@ public ResultSet getFunctions(String catalog, String schemaPattern, String funct targetDatasets, functionNamePattern, true, + false, collectedResults, resultSchemaFields, (routine, results, fields) -> { @@ -3641,6 +3650,7 @@ public ResultSet getFunctionColumns( targetDatasets, functionNamePattern, true, + true, collectedResults, resultSchemaFields, (routine, results, fields) -> { @@ -4704,25 +4714,28 @@ void process(Table bqTable, List collectedResults, FieldList res private void processSingleTable( DatasetId datasetId, String tableName, + Table preFetchedTable, boolean requireBaseTable, List collectedResults, FieldList resultSchemaFields, TableProcessor processor) throws SQLException { - Table bqTable; - try { - bqTable = - bigquery.getTable(TableId.of(datasetId.getProject(), datasetId.getDataset(), tableName)); - } catch (BigQueryException e) { - if (e.getCode() == 404) { - LOG.info( - "Table '%s' or dataset '%s' not found in project '%s' (API error 404). Skipping.", - tableName, datasetId.getDataset(), datasetId.getProject()); - return; + Table bqTable = preFetchedTable; + if (bqTable == null) { + try { + bqTable = + bigquery.getTable(TableId.of(datasetId.getProject(), datasetId.getDataset(), tableName)); + } catch (BigQueryException e) { + if (e.getCode() == 404) { + LOG.info( + "Table '%s' or dataset '%s' not found in project '%s' (API error 404). Skipping.", + tableName, datasetId.getDataset(), datasetId.getProject()); + return; + } + 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); } - 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) { @@ -4743,23 +4756,26 @@ void process( private void processSingleRoutine( DatasetId datasetId, String routineName, + Routine preFetchedRoutine, List collectedResults, FieldList resultSchemaFields, RoutineProcessor processor) throws SQLException { - Routine routine; - try { - routine = - bigquery.getRoutine( - RoutineId.of(datasetId.getProject(), datasetId.getDataset(), routineName)); - } catch (BigQueryException e) { - if (e.getCode() == 404) { - LOG.info("Routine '%s' not found. Skipping.", routineName); - return; + Routine routine = preFetchedRoutine; + if (routine == null) { + try { + routine = + bigquery.getRoutine( + RoutineId.of(datasetId.getProject(), datasetId.getDataset(), routineName)); + } catch (BigQueryException e) { + if (e.getCode() == 404) { + LOG.info("Routine '%s' not found. Skipping.", routineName); + return; + } + throw new SQLException("Error while fetching routine metadata: " + e.getMessage(), e); + } catch (Exception e) { + throw new SQLException("Error while fetching routine metadata: " + e.getMessage(), e); } - throw new SQLException("Error while fetching routine metadata: " + e.getMessage(), e); - } catch (Exception e) { - throw new SQLException("Error while fetching routine metadata: " + e.getMessage(), e); } if (routine == null) { @@ -4772,6 +4788,7 @@ private void processTargetRoutinesConcurrently( List targetDatasets, String routineNamePattern, boolean isPattern, + boolean requiresFullRoutine, List collectedResults, FieldList resultSchemaFields, RoutineProcessor processor) @@ -4785,6 +4802,7 @@ private void processTargetRoutinesConcurrently( processSingleRoutine( targetDatasets.get(0), routineNamePattern, + null, collectedResults, resultSchemaFields, processor); @@ -4801,7 +4819,12 @@ private void processTargetRoutinesConcurrently( tasks.add( () -> { processSingleRoutine( - datasetId, routineNamePattern, collectedResults, resultSchemaFields, processor); + datasetId, + routineNamePattern, + null, + collectedResults, + resultSchemaFields, + processor); return null; }); continue; @@ -4819,7 +4842,12 @@ private void processTargetRoutinesConcurrently( tasks.add( () -> { processSingleRoutine( - datasetId, routineName, collectedResults, resultSchemaFields, processor); + datasetId, + requiresFullRoutine ? routineName : null, + requiresFullRoutine ? null : routine, + collectedResults, + resultSchemaFields, + processor); return null; }); } @@ -4847,6 +4875,7 @@ private void processTargetTablesConcurrently( String tableNamePattern, boolean isPattern, boolean requireBaseTable, + boolean requiresFullTable, List collectedResults, FieldList resultSchemaFields, TableProcessor processor) @@ -4860,6 +4889,7 @@ private void processTargetTablesConcurrently( processSingleTable( targetDatasets.get(0), tableNamePattern, + null, // always fetch for literal matches requireBaseTable, collectedResults, resultSchemaFields, @@ -4879,6 +4909,7 @@ private void processTargetTablesConcurrently( processSingleTable( datasetId, tableNamePattern, + null, requireBaseTable, collectedResults, resultSchemaFields, @@ -4903,7 +4934,8 @@ private void processTargetTablesConcurrently( () -> { processSingleTable( datasetId, - tableName, + requiresFullTable ? tableName : null, + requiresFullTable ? null : table, requireBaseTable, collectedResults, resultSchemaFields, From 6040ce8dccf0474211c4734897c711c42c820ea3 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Tue, 21 Jul 2026 00:07:16 +0000 Subject: [PATCH 05/14] revert to async resultset --- .../jdbc/BigQueryDatabaseMetaData.java | 585 ++++++++++-------- 1 file changed, 342 insertions(+), 243 deletions(-) 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 cd95ac73c57f..ff659b80d054 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 @@ -787,28 +787,36 @@ public ResultSet getProcedures(String catalog, String schemaPattern, String proc final Schema resultSchema = defineGetProceduresSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - final List collectedResults = Collections.synchronizedList(new ArrayList<>()); List targetDatasets = getTargetDatasets(catalog, schemaPattern); - processTargetRoutinesConcurrently( - targetDatasets, - procedureNamePattern, - true, - false, - collectedResults, - resultSchemaFields, - (routine, results, fields) -> { - if (!"PROCEDURE".equalsIgnoreCase(routine.getRoutineType())) { - return; - } - processProcedureInfo(routine, results, fields); - }); - - Comparator comparator = defineGetProceduresComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getProcedures", LOG); final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + Future fetcherFuture = + fetchAndPopulateQueueAsync( + () -> { + final List collectedResults = + Collections.synchronizedList(new ArrayList<>()); + processTargetRoutinesConcurrently( + targetDatasets, + procedureNamePattern, + true, + false, + collectedResults, + resultSchemaFields, + (routine, results, fields) -> { + if (!"PROCEDURE".equalsIgnoreCase(routine.getRoutineType())) { + return; + } + processProcedureInfo(routine, results, fields); + }); + + Comparator comparator = + defineGetProceduresComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getProcedures", LOG); + return collectedResults; + }, + queue, + resultSchemaFields); return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } @@ -938,35 +946,48 @@ public ResultSet getProcedureColumns( final Schema resultSchema = defineGetProcedureColumnsSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - final List collectedResults = Collections.synchronizedList(new ArrayList<>()); List targetDatasets = getTargetDatasets(catalog, schemaPattern); Pattern columnNameRegex = compileSqlLikePattern(columnNamePattern); - processTargetRoutinesConcurrently( - targetDatasets, - procedureNamePattern, - true, - true, - collectedResults, - resultSchemaFields, - (routine, results, fields) -> { - if (!"PROCEDURE".equalsIgnoreCase(routine.getRoutineType())) { - return; - } - try { - processProcedureArgumentsSequentially( - Collections.singletonList(routine), columnNameRegex, results, fields, LOG); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException("Interrupted while processing procedure arguments", e); - } - }); - - Comparator comparator = defineGetProcedureColumnsComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getProcedureColumns", LOG); final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + Future fetcherFuture = + fetchAndPopulateQueueAsync( + () -> { + final List collectedResults = + Collections.synchronizedList(new ArrayList<>()); + processTargetRoutinesConcurrently( + targetDatasets, + procedureNamePattern, + true, + true, + collectedResults, + resultSchemaFields, + (routine, results, fields) -> { + if (!"PROCEDURE".equalsIgnoreCase(routine.getRoutineType())) { + return; + } + try { + processProcedureArgumentsSequentially( + Collections.singletonList(routine), + columnNameRegex, + results, + fields, + LOG); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException( + "Interrupted while processing procedure arguments", e); + } + }); + + Comparator comparator = + defineGetProcedureColumnsComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getProcedureColumns", LOG); + return collectedResults; + }, + queue, + resultSchemaFields); return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } @@ -1445,8 +1466,6 @@ public ResultSet getTables( } final Schema resultSchema = defineGetTablesSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - final List collectedResults = Collections.synchronizedList(new ArrayList<>()); - Tuple effectiveIdentifiers = determineEffectiveCatalogAndSchema(catalog, schemaPattern); String effectiveCatalog = effectiveIdentifiers.x(); @@ -1457,23 +1476,32 @@ public ResultSet getTables( final Set requestedTypes = (types == null || types.length == 0) ? null : new HashSet<>(Arrays.asList(types)); - processTargetTablesConcurrently( - targetDatasets, - tableNamePattern, - true, - false, - false, - collectedResults, - resultSchemaFields, - (bqTable, results, fields) -> { - processTableInfo(bqTable, requestedTypes, results, fields); - }); - - Comparator comparator = defineGetTablesComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getTables", LOG); final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + Future fetcherFuture = + fetchAndPopulateQueueAsync( + () -> { + final List collectedResults = + Collections.synchronizedList(new ArrayList<>()); + processTargetTablesConcurrently( + targetDatasets, + tableNamePattern, + true, + false, + false, + collectedResults, + resultSchemaFields, + (bqTable, results, fields) -> { + processTableInfo(bqTable, requestedTypes, results, fields); + }); + + Comparator comparator = defineGetTablesComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getTables", LOG); + return collectedResults; + }, + queue, + resultSchemaFields); + return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } @@ -1607,7 +1635,7 @@ public ResultSet getCatalogs() throws SQLException { final BlockingQueue queue = new LinkedBlockingQueue<>(catalogRows.isEmpty() ? 1 : catalogRows.size() + 1); - Future fetcherFuture = populateQueueAsync(catalogRows, queue, schemaFields); + Future fetcherFuture = fetchAndPopulateQueueAsync(() -> catalogRows, queue, schemaFields); return BigQueryJsonResultSet.of(catalogsSchema, catalogRows.size(), queue, null, fetcherFuture); } @@ -1638,7 +1666,7 @@ public ResultSet getTableTypes() { new LinkedBlockingQueue<>(tableTypeRows.size() + 1); Future fetcherFuture = - populateQueueAsync(tableTypeRows, queue, tableTypesSchema.getFields()); + fetchAndPopulateQueueAsync(() -> tableTypeRows, queue, tableTypesSchema.getFields()); return BigQueryJsonResultSet.of( tableTypesSchema, tableTypeRows.size(), queue, null, fetcherFuture); @@ -1681,8 +1709,6 @@ public ResultSet getColumns( } final Schema resultSchema = defineGetColumnsSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - final List collectedResults = Collections.synchronizedList(new ArrayList<>()); - Tuple effectiveIdentifiers = determineEffectiveCatalogAndSchema(catalog, schemaPattern); String effectiveCatalog = effectiveIdentifiers.x(); @@ -1695,23 +1721,33 @@ public ResultSet getColumns( && (columnNamePattern.contains("%") || columnNamePattern.contains("_")); Pattern columnNameRegex = hasWildcards ? compileSqlLikePattern(columnNamePattern) : null; - processTargetTablesConcurrently( - targetDatasets, - tableNamePattern, - true, - false, - true, - collectedResults, - resultSchemaFields, - (bqTable, results, fields) -> { - processTableColumns(bqTable, columnNameRegex, results, fields); - }); - - Comparator comparator = defineGetColumnsComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getColumns", LOG); final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + Future fetcherFuture = + fetchAndPopulateQueueAsync( + () -> { + final List collectedResults = + Collections.synchronizedList(new ArrayList<>()); + processTargetTablesConcurrently( + targetDatasets, + tableNamePattern, + true, + false, + true, + collectedResults, + resultSchemaFields, + (bqTable, results, fields) -> { + processTableColumns(bqTable, columnNameRegex, results, fields); + }); + + Comparator comparator = + defineGetColumnsComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getColumns", LOG); + return collectedResults; + }, + queue, + resultSchemaFields); + return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } @@ -2141,29 +2177,36 @@ public ResultSet getPrimaryKeys(String catalog, String schema, String table) thr final Schema resultSchema = defineGetPrimaryKeysSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - - final List collectedResults = Collections.synchronizedList(new ArrayList<>()); 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); - }); - - Comparator comparator = defineGetPrimaryKeysComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getPrimaryKeys", LOG); - final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + Future fetcherFuture = + fetchAndPopulateQueueAsync( + () -> { + 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); + }); + + Comparator comparator = + defineGetPrimaryKeysComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getPrimaryKeys", LOG); + return collectedResults; + }, + queue, + resultSchemaFields); + return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } @@ -2241,35 +2284,41 @@ public ResultSet getImportedKeys(String catalog, String schema, String table) final Schema resultSchema = defineForeignKeyResultSetSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - - final List collectedResults = Collections.synchronizedList(new ArrayList<>()); List targetDatasets = getTargetDatasets(catalog, schema, false); - processTargetTablesConcurrently( - targetDatasets, - table, - false, - true, - true, - collectedResults, - resultSchemaFields, - (bqTable, results, fields) -> { - TableConstraints constraints = bqTable.getTableConstraints(); - if (constraints == null || constraints.getForeignKeys() == null) { - return; - } - for (ForeignKey fk : constraints.getForeignKeys()) { - TableId pkTableId = fk.getReferencedTable(); - processForeignKey(fk, pkTableId, bqTable.getTableId(), results, fields); - } - }); - - Comparator comparator = definePkTableSortComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getImportedKeys", LOG); - final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + Future fetcherFuture = + fetchAndPopulateQueueAsync( + () -> { + final List collectedResults = + Collections.synchronizedList(new ArrayList<>()); + processTargetTablesConcurrently( + targetDatasets, + table, + false, + true, + true, + collectedResults, + resultSchemaFields, + (bqTable, results, fields) -> { + TableConstraints constraints = bqTable.getTableConstraints(); + if (constraints == null || constraints.getForeignKeys() == null) { + return; + } + for (ForeignKey fk : constraints.getForeignKeys()) { + TableId pkTableId = fk.getReferencedTable(); + processForeignKey(fk, pkTableId, bqTable.getTableId(), results, fields); + } + }); + + Comparator comparator = + definePkTableSortComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getImportedKeys", LOG); + return collectedResults; + }, + queue, + resultSchemaFields); return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } @@ -2300,40 +2349,46 @@ public ResultSet getExportedKeys(String catalog, String schema, String table) // Fallback Path: If catalog or schema is null, fall back to REST API metadata scan. if (catalog == null || schema == null) { - final List collectedResults = Collections.synchronizedList(new ArrayList<>()); List targetDatasets = getTargetDatasets(catalog, schema); - - processTargetTablesConcurrently( - targetDatasets, - null, - false, - true, - true, - collectedResults, - resultSchemaFields, - (bqTable, results, fields) -> { - TableConstraints constraints = bqTable.getTableConstraints(); - if (constraints == null || constraints.getForeignKeys() == null) { - return; - } - for (ForeignKey fk : constraints.getForeignKeys()) { - TableId pkTableId = fk.getReferencedTable(); - if (pkTableId == null - || !equalsOrNullMatchesAll(catalog, pkTableId.getProject()) - || !equalsOrNullMatchesAll(schema, pkTableId.getDataset()) - || !table.equals(pkTableId.getTable())) { - continue; - } - processForeignKey(fk, pkTableId, bqTable.getTableId(), results, fields); - } - }); - - Comparator comparator = defineFkTableSortComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getExportedKeys", LOG); - final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + Future fetcherFuture = + fetchAndPopulateQueueAsync( + () -> { + final List collectedResults = + Collections.synchronizedList(new ArrayList<>()); + processTargetTablesConcurrently( + targetDatasets, + null, + false, + true, + true, + collectedResults, + resultSchemaFields, + (bqTable, results, fields) -> { + TableConstraints constraints = bqTable.getTableConstraints(); + if (constraints == null || constraints.getForeignKeys() == null) { + return; + } + for (ForeignKey fk : constraints.getForeignKeys()) { + TableId pkTableId = fk.getReferencedTable(); + if (pkTableId == null + || !equalsOrNullMatchesAll(catalog, pkTableId.getProject()) + || !equalsOrNullMatchesAll(schema, pkTableId.getDataset()) + || !table.equals(pkTableId.getTable())) { + continue; + } + processForeignKey(fk, pkTableId, bqTable.getTableId(), results, fields); + } + }); + + Comparator comparator = + defineFkTableSortComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getExportedKeys", LOG); + return collectedResults; + }, + queue, + resultSchemaFields); return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } @@ -2379,40 +2434,46 @@ public ResultSet getCrossReference( final Schema resultSchema = defineForeignKeyResultSetSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - - final List collectedResults = Collections.synchronizedList(new ArrayList<>()); List targetDatasets = getTargetDatasets(foreignCatalog, foreignSchema, false); - processTargetTablesConcurrently( - targetDatasets, - foreignTable, - false, - true, - true, - collectedResults, - resultSchemaFields, - (bqTable, results, fields) -> { - TableConstraints constraints = bqTable.getTableConstraints(); - if (constraints == null || constraints.getForeignKeys() == null) { - return; - } - for (ForeignKey fk : constraints.getForeignKeys()) { - TableId pkTableId = fk.getReferencedTable(); - if (!equalsOrNullMatchesAll(parentCatalog, pkTableId.getProject()) - || !equalsOrNullMatchesAll(parentSchema, pkTableId.getDataset()) - || !parentTable.equals(pkTableId.getTable())) { - continue; - } - processForeignKey(fk, pkTableId, bqTable.getTableId(), results, fields); - } - }); - - Comparator comparator = defineFkTableSortComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getCrossReference", LOG); - final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + Future fetcherFuture = + fetchAndPopulateQueueAsync( + () -> { + final List collectedResults = + Collections.synchronizedList(new ArrayList<>()); + processTargetTablesConcurrently( + targetDatasets, + foreignTable, + false, + true, + true, + collectedResults, + resultSchemaFields, + (bqTable, results, fields) -> { + TableConstraints constraints = bqTable.getTableConstraints(); + if (constraints == null || constraints.getForeignKeys() == null) { + return; + } + for (ForeignKey fk : constraints.getForeignKeys()) { + TableId pkTableId = fk.getReferencedTable(); + if (!equalsOrNullMatchesAll(parentCatalog, pkTableId.getProject()) + || !equalsOrNullMatchesAll(parentSchema, pkTableId.getDataset()) + || !parentTable.equals(pkTableId.getTable())) { + continue; + } + processForeignKey(fk, pkTableId, bqTable.getTableId(), results, fields); + } + }); + + Comparator comparator = + defineFkTableSortComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getCrossReference", LOG); + return collectedResults; + }, + queue, + resultSchemaFields); return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } @@ -2429,7 +2490,7 @@ public ResultSet getTypeInfo() { final BlockingQueue queue = new LinkedBlockingQueue<>(typeInfoRows.size() + 1); - Future fetcherFuture = populateQueueAsync(typeInfoRows, queue, schemaFields); + Future fetcherFuture = fetchAndPopulateQueueAsync(() -> typeInfoRows, queue, schemaFields); return BigQueryJsonResultSet.of( typeInfoSchema, typeInfoRows.size(), queue, null, fetcherFuture); } @@ -3341,18 +3402,23 @@ public ResultSet getSchemas(String catalog, String schemaPattern) throws SQLExce final Schema resultSchema = defineGetSchemasSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); final Pattern schemaRegex = compileSqlLikePattern(schemaPattern); - List datasets = fetchMatchingDatasets(catalog, schemaPattern, schemaRegex); - List collectedResults = new ArrayList<>(datasets.size()); - for (Dataset dataset : datasets) { - processSchemaInfo(dataset.getDatasetId(), collectedResults, resultSchemaFields); - } - - Comparator comparator = defineGetSchemasComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getSchemas", LOG); - final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + Future fetcherFuture = + fetchAndPopulateQueueAsync( + () -> { + List datasets = fetchMatchingDatasets(catalog, schemaPattern, schemaRegex); + List collectedResults = new ArrayList<>(datasets.size()); + for (Dataset dataset : datasets) { + processSchemaInfo(dataset.getDatasetId(), collectedResults, resultSchemaFields); + } + Comparator comparator = + defineGetSchemasComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getSchemas", LOG); + return collectedResults; + }, + queue, + resultSchemaFields); return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } @@ -3508,29 +3574,37 @@ public ResultSet getFunctions(String catalog, String schemaPattern, String funct final Schema resultSchema = defineGetFunctionsSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - final List collectedResults = Collections.synchronizedList(new ArrayList<>()); List targetDatasets = getTargetDatasets(catalog, schemaPattern); - processTargetRoutinesConcurrently( - targetDatasets, - functionNamePattern, - true, - false, - collectedResults, - resultSchemaFields, - (routine, results, fields) -> { - if (!"SCALAR_FUNCTION".equalsIgnoreCase(routine.getRoutineType()) - && !"TABLE_FUNCTION".equalsIgnoreCase(routine.getRoutineType())) { - return; - } - processFunctionInfo(routine, results, fields); - }); - - Comparator comparator = defineGetFunctionsComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getFunctions", LOG); final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + Future fetcherFuture = + fetchAndPopulateQueueAsync( + () -> { + final List collectedResults = + Collections.synchronizedList(new ArrayList<>()); + processTargetRoutinesConcurrently( + targetDatasets, + functionNamePattern, + true, + false, + collectedResults, + resultSchemaFields, + (routine, results, fields) -> { + if (!"SCALAR_FUNCTION".equalsIgnoreCase(routine.getRoutineType()) + && !"TABLE_FUNCTION".equalsIgnoreCase(routine.getRoutineType())) { + return; + } + processFunctionInfo(routine, results, fields); + }); + + Comparator comparator = + defineGetFunctionsComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getFunctions", LOG); + return collectedResults; + }, + queue, + resultSchemaFields); return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } @@ -3641,37 +3715,50 @@ public ResultSet getFunctionColumns( final Schema resultSchema = defineGetFunctionColumnsSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - final List collectedResults = Collections.synchronizedList(new ArrayList<>()); List targetDatasets = getTargetDatasets(catalog, schemaPattern); Pattern columnNameRegex = compileSqlLikePattern(columnNamePattern); - processTargetRoutinesConcurrently( - targetDatasets, - functionNamePattern, - true, - true, - collectedResults, - resultSchemaFields, - (routine, results, fields) -> { - if (!"SCALAR_FUNCTION".equalsIgnoreCase(routine.getRoutineType()) - && !"TABLE_FUNCTION".equalsIgnoreCase(routine.getRoutineType())) { - return; - } - try { - processFunctionParametersSequentially( - Collections.singletonList(routine), columnNameRegex, results, fields, LOG); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException("Interrupted while processing function parameters", e); - } - }); - - Comparator comparator = defineGetFunctionColumnsComparator(resultSchemaFields); - sortResults(collectedResults, comparator, "getFunctionColumns", LOG); final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); - Future fetcherFuture = populateQueueAsync(collectedResults, queue, resultSchemaFields); + Future fetcherFuture = + fetchAndPopulateQueueAsync( + () -> { + final List collectedResults = + Collections.synchronizedList(new ArrayList<>()); + processTargetRoutinesConcurrently( + targetDatasets, + functionNamePattern, + true, + true, + collectedResults, + resultSchemaFields, + (routine, results, fields) -> { + if (!"SCALAR_FUNCTION".equalsIgnoreCase(routine.getRoutineType()) + && !"TABLE_FUNCTION".equalsIgnoreCase(routine.getRoutineType())) { + return; + } + try { + processFunctionParametersSequentially( + Collections.singletonList(routine), + columnNameRegex, + results, + fields, + LOG); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException( + "Interrupted while processing function parameters", e); + } + }); + + Comparator comparator = + defineGetFunctionColumnsComparator(resultSchemaFields); + sortResults(collectedResults, comparator, "getFunctionColumns", LOG); + return collectedResults; + }, + queue, + resultSchemaFields); return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } @@ -4507,16 +4594,27 @@ private void waitForTasksCompletion(List> taskFutures) throws Executio LOG.info("Finished waiting for tasks."); } - private Future populateQueueAsync( - List collectedResults, + private Future fetchAndPopulateQueueAsync( + Callable> task, BlockingQueue queue, FieldList resultSchemaFields) { return connection .getMetadataExecutor() .submit( () -> { - populateQueue(collectedResults, queue, resultSchemaFields); - signalEndOfData(queue, resultSchemaFields); + try { + List collectedResults = task.call(); + populateQueue(collectedResults, queue, resultSchemaFields); + } catch (Exception e) { + LOG.severe("Error during async metadata fetch: " + e.getMessage()); + try { + queue.put(BigQueryFieldValueListWrapper.ofError(e)); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + } + } finally { + signalEndOfData(queue, resultSchemaFields); + } }); } @@ -4724,7 +4822,8 @@ private void processSingleTable( if (bqTable == null) { try { bqTable = - bigquery.getTable(TableId.of(datasetId.getProject(), datasetId.getDataset(), tableName)); + bigquery.getTable( + TableId.of(datasetId.getProject(), datasetId.getDataset(), tableName)); } catch (BigQueryException e) { if (e.getCode() == 404) { LOG.info( From 32f978456319850f4b7a9443f545d0f3afd9f3a5 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Tue, 21 Jul 2026 00:22:37 +0000 Subject: [PATCH 06/14] revert to parallel execution from sequential --- .../jdbc/BigQueryDatabaseMetaData.java | 222 ++++++++++++------ 1 file changed, 146 insertions(+), 76 deletions(-) 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 ff659b80d054..8b9215d5793e 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 @@ -4912,54 +4912,89 @@ private void processTargetRoutinesConcurrently( List> taskFutures = new ArrayList<>(); try { - List> tasks = new ArrayList<>(); + List>>> datasetListTasks = new ArrayList<>(); for (DatasetId datasetId : targetDatasets) { if (isLiteralName) { - tasks.add( + datasetListTasks.add( () -> { - processSingleRoutine( - datasetId, - routineNamePattern, - null, - collectedResults, - resultSchemaFields, - processor); - return null; + return Collections.singletonList( + () -> { + processSingleRoutine( + datasetId, + routineNamePattern, + null, + collectedResults, + resultSchemaFields, + processor); + return null; + }); }); continue; } - try { - Page routinesPage = - bigquery.listRoutines( - datasetId, BigQuery.RoutineListOption.pageSize(DEFAULT_PAGE_SIZE)); - if (routinesPage == null) continue; - for (Routine routine : routinesPage.iterateAll()) { - if (routine == null) continue; - String routineName = routine.getRoutineId().getRoutine(); - if (routineRegex != null && !routineRegex.matcher(routineName).matches()) continue; - tasks.add( - () -> { - processSingleRoutine( - datasetId, - requiresFullRoutine ? routineName : null, - requiresFullRoutine ? null : routine, - collectedResults, - resultSchemaFields, - processor); - return null; - }); + datasetListTasks.add( + () -> { + List> routineTasks = new ArrayList<>(); + try { + Page routinesPage = + bigquery.listRoutines( + datasetId, BigQuery.RoutineListOption.pageSize(DEFAULT_PAGE_SIZE)); + if (routinesPage != null) { + for (Routine routine : routinesPage.iterateAll()) { + if (routine == null) continue; + String routineName = routine.getRoutineId().getRoutine(); + if (routineRegex != null && !routineRegex.matcher(routineName).matches()) + continue; + routineTasks.add( + () -> { + processSingleRoutine( + datasetId, + requiresFullRoutine ? routineName : null, + requiresFullRoutine ? null : routine, + collectedResults, + resultSchemaFields, + processor); + return null; + }); + } + } + } catch (BigQueryException e) { + if (e.getCode() == 404) { + LOG.info("Dataset '%s' not found while listing routines. Skipping.", datasetId); + } else { + throw new SQLException("Error while listing routines: " + e.getMessage(), e); + } + } + return routineTasks; + }); + } + + List>>> listFutures = new ArrayList<>(); + for (Callable>> listTask : datasetListTasks) { + listFutures.add(executor.submit(listTask)); + } + + List> detailTasks = new ArrayList<>(); + try { + for (Future>> future : listFutures) { + if (!future.isCancelled()) { + detailTasks.addAll(future.get()); } - } catch (BigQueryException e) { - if (e.getCode() == 404) { - LOG.info("Dataset '%s' not found while listing routines. Skipping.", datasetId); - continue; + } + } catch (CancellationException e) { + LOG.warning("A dataset listing task was cancelled."); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOG.warning("Fetcher thread interrupted while listing datasets."); + } finally { + for (Future>> future : listFutures) { + if (!future.isDone()) { + future.cancel(true); } - throw new SQLException("Error while listing routines: " + e.getMessage(), e); } } - for (Callable task : tasks) { + for (Callable task : detailTasks) { taskFutures.add(executor.submit(task)); } waitForTasksCompletion(taskFutures); @@ -5000,58 +5035,93 @@ private void processTargetTablesConcurrently( List> taskFutures = new ArrayList<>(); try { - List> tasks = new ArrayList<>(); + List>>> datasetListTasks = new ArrayList<>(); for (DatasetId datasetId : targetDatasets) { if (isLiteralName) { - tasks.add( + datasetListTasks.add( () -> { - processSingleTable( - datasetId, - tableNamePattern, - null, - requireBaseTable, - collectedResults, - resultSchemaFields, - processor); - return null; + return Collections.singletonList( + () -> { + processSingleTable( + datasetId, + tableNamePattern, + null, + requireBaseTable, + collectedResults, + resultSchemaFields, + processor); + return null; + }); }); continue; } - try { - Page
tablesPage = - bigquery.listTables(datasetId, BigQuery.TableListOption.pageSize(DEFAULT_PAGE_SIZE)); - if (tablesPage == null) continue; - for (Table table : tablesPage.iterateAll()) { - if (table == null) continue; - if (requireBaseTable - && table.getDefinition() != null - && table.getDefinition().getType() != TableDefinition.Type.TABLE) continue; - String tableName = table.getTableId().getTable(); - if (tableRegex != null && !tableRegex.matcher(tableName).matches()) continue; - tasks.add( - () -> { - processSingleTable( - datasetId, - requiresFullTable ? tableName : null, - requiresFullTable ? null : table, - requireBaseTable, - collectedResults, - resultSchemaFields, - processor); - return null; - }); + datasetListTasks.add( + () -> { + List> tableTasks = new ArrayList<>(); + try { + Page
tablesPage = + bigquery.listTables( + datasetId, BigQuery.TableListOption.pageSize(DEFAULT_PAGE_SIZE)); + if (tablesPage != null) { + for (Table table : tablesPage.iterateAll()) { + if (table == null) continue; + if (requireBaseTable + && table.getDefinition() != null + && table.getDefinition().getType() != TableDefinition.Type.TABLE) continue; + String tableName = table.getTableId().getTable(); + if (tableRegex != null && !tableRegex.matcher(tableName).matches()) continue; + tableTasks.add( + () -> { + processSingleTable( + datasetId, + requiresFullTable ? tableName : null, + requiresFullTable ? null : table, + requireBaseTable, + collectedResults, + resultSchemaFields, + processor); + return null; + }); + } + } + } catch (BigQueryException e) { + if (e.getCode() == 404) { + LOG.info("Dataset '%s' not found while listing tables. Skipping.", datasetId); + } else { + throw new SQLException("Error while listing tables: " + e.getMessage(), e); + } + } + return tableTasks; + }); + } + + List>>> listFutures = new ArrayList<>(); + for (Callable>> listTask : datasetListTasks) { + listFutures.add(executor.submit(listTask)); + } + + List> detailTasks = new ArrayList<>(); + try { + for (Future>> future : listFutures) { + if (!future.isCancelled()) { + detailTasks.addAll(future.get()); } - } catch (BigQueryException e) { - if (e.getCode() == 404) { - LOG.info("Dataset '%s' not found while listing tables. Skipping.", datasetId); - continue; + } + } catch (CancellationException e) { + LOG.warning("A dataset listing task was cancelled."); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOG.warning("Fetcher thread interrupted while listing datasets."); + } finally { + for (Future>> future : listFutures) { + if (!future.isDone()) { + future.cancel(true); } - throw new SQLException("Error while listing tables: " + e.getMessage(), e); } } - for (Callable task : tasks) { + for (Callable task : detailTasks) { taskFutures.add(executor.submit(task)); } waitForTasksCompletion(taskFutures); From 82c639854d8f1362e12cf9d8c68a9df81ba39cee Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Tue, 21 Jul 2026 00:26:20 +0000 Subject: [PATCH 07/14] preserve error handling --- .../jdbc/BigQueryDatabaseMetaData.java | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) 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 8b9215d5793e..fa0b6d3d8854 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 @@ -4998,9 +4998,17 @@ private void processTargetRoutinesConcurrently( taskFutures.add(executor.submit(task)); } waitForTasksCompletion(taskFutures); + if (Thread.currentThread().isInterrupted()) { + throw new SQLException("Interrupted while parallel-fetching metadata"); + } } catch (ExecutionException e) { - throw new SQLException( - "Error while fetching routine metadata: " + e.getCause().getMessage(), 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)); } } @@ -5125,9 +5133,17 @@ private void processTargetTablesConcurrently( taskFutures.add(executor.submit(task)); } waitForTasksCompletion(taskFutures); + if (Thread.currentThread().isInterrupted()) { + throw new SQLException("Interrupted while parallel-fetching metadata"); + } } catch (ExecutionException e) { - throw new SQLException( - "Error while fetching table metadata: " + e.getCause().getMessage(), 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)); } } From 00ab9693e95e013f98d1dfa691e06b5a3b7a1d0f Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Tue, 21 Jul 2026 00:39:30 +0000 Subject: [PATCH 08/14] address pr feedback --- .../jdbc/BigQueryDatabaseMetaData.java | 88 ++++++++----------- 1 file changed, 37 insertions(+), 51 deletions(-) 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 fa0b6d3d8854..7c41cc85948a 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 @@ -4569,7 +4569,8 @@ private Long getLongValueOrNull(FieldValueList fvl, int index) { } } - private void waitForTasksCompletion(List> taskFutures) throws ExecutionException { + private void waitForTasksCompletion(List> taskFutures) + throws ExecutionException, SQLException { LOG.info("Waiting for %d submitted tasks to complete...", taskFutures.size()); try { for (Future future : taskFutures) { @@ -4582,8 +4583,8 @@ private void waitForTasksCompletion(List> taskFutures) throws Executio } catch (InterruptedException e) { Thread.currentThread().interrupt(); LOG.warning( - "Fetcher thread interrupted while waiting for tasks. Attempting to cancel remaining" - + " tasks."); + "Fetcher thread interrupted while waiting for tasks. Attempting to cancel remaining tasks."); + throw new SQLException("Metadata fetch interrupted.", e); } finally { for (Future future : taskFutures) { if (!future.isDone()) { @@ -4724,9 +4725,6 @@ private List fetchMatchingDatasets( } waitForTasksCompletion(taskFutures); - if (Thread.currentThread().isInterrupted()) { - throw new SQLException("Interrupted while parallel-fetching matching datasets"); - } } catch (ExecutionException e) { Throwable cause = e.getCause(); if (cause instanceof SQLException) { @@ -4939,24 +4937,21 @@ private void processTargetRoutinesConcurrently( Page routinesPage = bigquery.listRoutines( datasetId, BigQuery.RoutineListOption.pageSize(DEFAULT_PAGE_SIZE)); - if (routinesPage != null) { - for (Routine routine : routinesPage.iterateAll()) { - if (routine == null) continue; - String routineName = routine.getRoutineId().getRoutine(); - if (routineRegex != null && !routineRegex.matcher(routineName).matches()) - continue; - routineTasks.add( - () -> { - processSingleRoutine( - datasetId, - requiresFullRoutine ? routineName : null, - requiresFullRoutine ? null : routine, - collectedResults, - resultSchemaFields, - processor); - return null; - }); - } + for (Routine routine : routinesPage.iterateAll()) { + String routineName = routine.getRoutineId().getRoutine(); + if (routineRegex != null && !routineRegex.matcher(routineName).matches()) + continue; + routineTasks.add( + () -> { + processSingleRoutine( + datasetId, + requiresFullRoutine ? routineName : null, + requiresFullRoutine ? null : routine, + collectedResults, + resultSchemaFields, + processor); + return null; + }); } } catch (BigQueryException e) { if (e.getCode() == 404) { @@ -4998,9 +4993,6 @@ private void processTargetRoutinesConcurrently( taskFutures.add(executor.submit(task)); } 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) { @@ -5071,27 +5063,24 @@ private void processTargetTablesConcurrently( Page
tablesPage = bigquery.listTables( datasetId, BigQuery.TableListOption.pageSize(DEFAULT_PAGE_SIZE)); - if (tablesPage != null) { - for (Table table : tablesPage.iterateAll()) { - if (table == null) continue; - if (requireBaseTable - && table.getDefinition() != null - && table.getDefinition().getType() != TableDefinition.Type.TABLE) continue; - String tableName = table.getTableId().getTable(); - if (tableRegex != null && !tableRegex.matcher(tableName).matches()) continue; - tableTasks.add( - () -> { - processSingleTable( - datasetId, - requiresFullTable ? tableName : null, - requiresFullTable ? null : table, - requireBaseTable, - collectedResults, - resultSchemaFields, - processor); - return null; - }); - } + for (Table table : tablesPage.iterateAll()) { + if (requireBaseTable + && table.getDefinition() != null + && table.getDefinition().getType() != TableDefinition.Type.TABLE) continue; + String tableName = table.getTableId().getTable(); + if (tableRegex != null && !tableRegex.matcher(tableName).matches()) continue; + tableTasks.add( + () -> { + processSingleTable( + datasetId, + requiresFullTable ? tableName : null, + requiresFullTable ? null : table, + requireBaseTable, + collectedResults, + resultSchemaFields, + processor); + return null; + }); } } catch (BigQueryException e) { if (e.getCode() == 404) { @@ -5133,9 +5122,6 @@ private void processTargetTablesConcurrently( taskFutures.add(executor.submit(task)); } 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) { From 8b95738c80b410809a22c6180b9abb0515b50125 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Tue, 21 Jul 2026 01:08:43 +0000 Subject: [PATCH 09/14] fix tests --- .../jdbc/BigQueryDatabaseMetaData.java | 60 +++++++------------ .../jdbc/BigQueryDatabaseMetaDataTest.java | 43 +++++++++---- 2 files changed, 52 insertions(+), 51 deletions(-) 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 7c41cc85948a..6444dc3ef258 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 @@ -4606,10 +4606,10 @@ private Future fetchAndPopulateQueueAsync( try { List collectedResults = task.call(); populateQueue(collectedResults, queue, resultSchemaFields); - } catch (Exception e) { + } catch (Throwable e) { LOG.severe("Error during async metadata fetch: " + e.getMessage()); try { - queue.put(BigQueryFieldValueListWrapper.ofError(e)); + queue.put(BigQueryFieldValueListWrapper.ofError(new SQLException(e))); } catch (InterruptedException ie) { Thread.currentThread().interrupt(); } @@ -4964,29 +4964,19 @@ private void processTargetRoutinesConcurrently( }); } - List>>> listFutures = new ArrayList<>(); - for (Callable>> listTask : datasetListTasks) { - listFutures.add(executor.submit(listTask)); - } - List> detailTasks = new ArrayList<>(); try { - for (Future>> future : listFutures) { - if (!future.isCancelled()) { - detailTasks.addAll(future.get()); - } + for (Callable>> listTask : datasetListTasks) { + detailTasks.addAll(listTask.call()); } - } catch (CancellationException e) { - LOG.warning("A dataset listing task was cancelled."); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - LOG.warning("Fetcher thread interrupted while listing datasets."); - } finally { - for (Future>> future : listFutures) { - if (!future.isDone()) { - future.cancel(true); - } + } catch (Exception e) { + if (e instanceof InterruptedException) { + Thread.currentThread().interrupt(); + } + if (e instanceof SQLException) { + throw (SQLException) e; } + throw new SQLException("Error while listing routines: " + e.getMessage(), e); } for (Callable task : detailTasks) { @@ -5093,29 +5083,19 @@ private void processTargetTablesConcurrently( }); } - List>>> listFutures = new ArrayList<>(); - for (Callable>> listTask : datasetListTasks) { - listFutures.add(executor.submit(listTask)); - } - List> detailTasks = new ArrayList<>(); try { - for (Future>> future : listFutures) { - if (!future.isCancelled()) { - detailTasks.addAll(future.get()); - } + for (Callable>> listTask : datasetListTasks) { + detailTasks.addAll(listTask.call()); } - } catch (CancellationException e) { - LOG.warning("A dataset listing task was cancelled."); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - LOG.warning("Fetcher thread interrupted while listing datasets."); - } finally { - for (Future>> future : listFutures) { - if (!future.isDone()) { - future.cancel(true); - } + } catch (Exception e) { + if (e instanceof InterruptedException) { + Thread.currentThread().interrupt(); + } + if (e instanceof SQLException) { + throw (SQLException) e; } + throw new SQLException("Error while listing tables: " + e.getMessage(), e); } for (Callable task : detailTasks) { 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 b14b0a8c97d3..861d771d53a0 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 @@ -36,6 +36,7 @@ import java.io.IOException; import java.io.InputStream; import java.sql.DatabaseMetaData; +import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; @@ -1005,7 +1006,7 @@ public void testFindMatchingBigQueryObjects_Routines_ListWithPattern() throws Ex Routine func1 = mockBigQueryRoutine(catalog, schema, "func_123", "FUNCTION", "f1"); Routine otherProc = mockBigQueryRoutine(catalog, schema, "another_proc", "PROCEDURE", "p3"); - Page page = mock(Page.class); + Page page = mock(Page.class, withSettings().withoutAnnotations()); when(page.iterateAll()).thenReturn(Arrays.asList(proc1, func1, proc2, otherProc)); when(bigqueryClient.listRoutines(eq(datasetId), any(BigQuery.RoutineListOption[].class))) .thenReturn(page); @@ -1050,7 +1051,7 @@ public void testFindMatchingBigQueryObjects_Routines_ListNoPattern() throws Exce Routine proc1 = mockBigQueryRoutine(catalog, schema, "proc_abc", "PROCEDURE", "p1"); Routine func1 = mockBigQueryRoutine(catalog, schema, "func_123", "FUNCTION", "f1"); - Page page = mock(Page.class); + Page page = mock(Page.class, withSettings().withoutAnnotations()); when(page.iterateAll()).thenReturn(Arrays.asList(proc1, func1)); when(bigqueryClient.listRoutines(eq(datasetId), any(BigQuery.RoutineListOption[].class))) .thenReturn(page); @@ -3270,20 +3271,20 @@ public void testGetSchemas_WithProjectDiscovery() throws SQLException { when(bigQueryConnection.getDiscoveredProjects()).thenReturn(Arrays.asList("discovered-1")); when(bigQueryConnection.getAdditionalProjects()).thenReturn("additional-1"); - Page pagePrimary = mock(Page.class); + Page pagePrimary = mock(Page.class, withSettings().withoutAnnotations()); Dataset dsPrimary = mockBigQueryDataset("primary-project", "dataset_p"); when(pagePrimary.iterateAll()).thenReturn(Collections.singletonList(dsPrimary)); when(bigqueryClient.listDatasets( eq("primary-project"), any(BigQuery.DatasetListOption[].class))) .thenReturn(pagePrimary); - Page pageAdditional = mock(Page.class); + Page pageAdditional = mock(Page.class, withSettings().withoutAnnotations()); Dataset dsAdditional = mockBigQueryDataset("additional-1", "dataset_a"); when(pageAdditional.iterateAll()).thenReturn(Collections.singletonList(dsAdditional)); when(bigqueryClient.listDatasets(eq("additional-1"), any(BigQuery.DatasetListOption[].class))) .thenReturn(pageAdditional); - Page pageDiscovered = mock(Page.class); + Page pageDiscovered = mock(Page.class, withSettings().withoutAnnotations()); Dataset dsDiscovered = mockBigQueryDataset("discovered-1", "dataset_d"); when(pageDiscovered.iterateAll()).thenReturn(Collections.singletonList(dsDiscovered)); when(bigqueryClient.listDatasets(eq("discovered-1"), any(BigQuery.DatasetListOption[].class))) @@ -3315,14 +3316,14 @@ public void testGetSchemas_WithoutProjectDiscovery() throws SQLException { when(bigQueryConnection.getDiscoveredProjects()).thenReturn(Arrays.asList("discovered-1")); when(bigQueryConnection.getAdditionalProjects()).thenReturn("additional-1"); - Page pagePrimary = mock(Page.class); + Page pagePrimary = mock(Page.class, withSettings().withoutAnnotations()); Dataset dsPrimary = mockBigQueryDataset("primary-project", "dataset_p"); when(pagePrimary.iterateAll()).thenReturn(Collections.singletonList(dsPrimary)); when(bigqueryClient.listDatasets( eq("primary-project"), any(BigQuery.DatasetListOption[].class))) .thenReturn(pagePrimary); - Page pageAdditional = mock(Page.class); + Page pageAdditional = mock(Page.class, withSettings().withoutAnnotations()); Dataset dsAdditional = mockBigQueryDataset("additional-1", "dataset_a"); when(pageAdditional.iterateAll()).thenReturn(Collections.singletonList(dsAdditional)); when(bigqueryClient.listDatasets(eq("additional-1"), any(BigQuery.DatasetListOption[].class))) @@ -3349,7 +3350,7 @@ public void testGetSchemas_WithoutProjectDiscovery() throws SQLException { } private void mockDatasetIteration(DatasetId datasetId) { - Page pagePrimary = mock(Page.class); + Page pagePrimary = mock(Page.class, withSettings().withoutAnnotations()); Dataset dsPrimary = mock(Dataset.class); when(dsPrimary.getDatasetId()).thenReturn(datasetId); when(pagePrimary.iterateAll()).thenReturn(Collections.singletonList(dsPrimary)); @@ -3372,7 +3373,7 @@ private Table mockTableWithConstraints(TableId tableId, TableConstraints constra } private void mockTableIteration(DatasetId datasetId, Table... tables) { - Page
pageTables = mock(Page.class); + Page
pageTables = mock(Page.class, withSettings().withoutAnnotations()); when(pageTables.iterateAll()).thenReturn(Arrays.asList(tables)); when(bigqueryClient.listTables(eq(datasetId), any(BigQuery.TableListOption[].class))) .thenReturn(pageTables); @@ -3483,7 +3484,27 @@ public void testGetImportedKeys_noKeys() throws SQLException { } @Test - public void testGetExportedKeys_hasKeys() throws SQLException { + public void testGetExportedKeys_infoSchema() throws SQLException { + PreparedStatement mockStmt = mock(PreparedStatement.class); + ResultSet mockRs = mock(ResultSet.class); + when(bigQueryConnection.prepareStatement(anyString())).thenReturn(mockStmt); + when(mockStmt.executeQuery()).thenReturn(mockRs); + + ResultSet rs = dbMetadata.getExportedKeys("test-project", "dataset_p", "ref_table"); + assertEquals(mockRs, rs); + verify(mockStmt).closeOnCompletion(); + verify(mockStmt).executeQuery(); + } + + @Test + public void testGetExportedKeys_pcntSchema() throws SQLException { + try (ResultSet rs = dbMetadata.getExportedKeys("test-project", "dataset.p", "ref_table")) { + assertFalse(rs.next()); + } + } + + @Test + public void testGetExportedKeys_fallback_hasKeys() throws SQLException { DatasetId datasetId = DatasetId.of("test-project", "dataset_p"); TableId tableId = TableId.of("test-project", "dataset_p", "table_p"); TableId refTableId = TableId.of("test-project", "dataset_p", "ref_table"); @@ -3533,7 +3554,7 @@ public void testGetExportedKeys_hasKeys() throws SQLException { } @Test - public void testGetExportedKeys_noKeys() throws SQLException { + public void testGetExportedKeys_fallback_noKeys() throws SQLException { DatasetId datasetId = DatasetId.of("test-project", "dataset_p"); TableId tableId = TableId.of("test-project", "dataset_p", "table_p"); From 9fa2509b44dd4c231daa443de3790956e9e2d7cf Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Tue, 21 Jul 2026 01:33:08 +0000 Subject: [PATCH 10/14] fix wrong conflict resolution during merge --- .../jdbc/BigQueryDatabaseMetaData.java | 342 +----------------- .../jdbc/BigQueryDatabaseMetaDataTest.java | 4 - 2 files changed, 12 insertions(+), 334 deletions(-) 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 6444dc3ef258..fb3bae4f9829 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 @@ -20,8 +20,6 @@ import com.google.cloud.Tuple; import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.BigQuery.DatasetListOption; -import com.google.cloud.bigquery.BigQuery.RoutineListOption; -import com.google.cloud.bigquery.BigQuery.RoutineOption; import com.google.cloud.bigquery.BigQuery.TableField; import com.google.cloud.bigquery.BigQuery.TableOption; import com.google.cloud.bigquery.BigQueryException; @@ -967,18 +965,7 @@ public ResultSet getProcedureColumns( if (!"PROCEDURE".equalsIgnoreCase(routine.getRoutineType())) { return; } - try { - processProcedureArgumentsSequentially( - Collections.singletonList(routine), - columnNameRegex, - results, - fields, - LOG); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException( - "Interrupted while processing procedure arguments", e); - } + processProcedureArguments(routine, columnNameRegex, results, fields); }); Comparator comparator = @@ -991,180 +978,6 @@ public ResultSet getProcedureColumns( return BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); } - List listMatchingProcedureIdsFromDatasets( - List datasetsToScan, - String procedureNamePattern, - Pattern procedureNameRegex, - ExecutorService listRoutinesExecutor, - String catalogParam, - BigQueryJdbcCustomLogger logger) - throws InterruptedException, ExecutionException { - - logger.fine( - "Listing matching procedure IDs from %d datasets for catalog '%s'.", - datasetsToScan.size(), catalogParam); - final List>> listRoutineFutures = new ArrayList<>(); - final List procedureIdsToGet = Collections.synchronizedList(new ArrayList<>()); - - for (Dataset dataset : datasetsToScan) { - if (Thread.currentThread().isInterrupted()) { - InterruptedException ex = - new InterruptedException( - "Interrupted while listing routines for catalog: " + catalogParam); - logger.severe(ex.getMessage(), ex); - throw ex; - } - final DatasetId currentDatasetId = dataset.getDatasetId(); - Callable> listCallable = - () -> - findMatchingBigQueryObjects( - "Routine", - () -> - bigquery.listRoutines( - currentDatasetId, RoutineListOption.pageSize(DEFAULT_PAGE_SIZE)), - (name) -> - bigquery.getRoutine( - RoutineId.of( - currentDatasetId.getProject(), currentDatasetId.getDataset(), name)), - (rt) -> rt.getRoutineId().getRoutine(), - procedureNamePattern, - procedureNameRegex, - logger, - false); - listRoutineFutures.add(listRoutinesExecutor.submit(listCallable)); - } - logger.fine( - "Submitted " - + listRoutineFutures.size() - + " routine list tasks for catalog: " - + catalogParam); - - for (Future> listFuture : listRoutineFutures) { - if (Thread.currentThread().isInterrupted()) { - listRoutineFutures.forEach(f -> f.cancel(true)); - InterruptedException ex = - new InterruptedException( - "Interrupted while collecting routine lists for catalog: " + catalogParam); - logger.severe(ex.getMessage(), ex); - throw ex; - } - try { - List listedRoutines = listFuture.get(); - if (listedRoutines != null) { - for (Routine listedRoutine : listedRoutines) { - if (listedRoutine != null - && "PROCEDURE".equalsIgnoreCase(listedRoutine.getRoutineType())) { - if (listedRoutine.getRoutineId() != null) { - procedureIdsToGet.add(listedRoutine.getRoutineId()); - } else { - logger.warning( - "Found a procedure type routine with a null ID during listing phase for" - + " catalog: " - + catalogParam); - } - } - } - } - } catch (CancellationException e) { - logger.warning("Routine list task cancelled for catalog: " + catalogParam); - } - } - logger.info( - "Found %d procedure IDs to fetch details for in catalog '%s'.", - procedureIdsToGet.size(), catalogParam); - return procedureIdsToGet; - } - - List fetchFullRoutineDetailsForIds( - List procedureIdsToGet, - ExecutorService getRoutineDetailsExecutor, - BigQueryJdbcCustomLogger logger, - RoutineOption... options) - throws InterruptedException, ExecutionException { - logger.fine("Fetching full details for %d procedure IDs.", procedureIdsToGet.size()); - final List> getRoutineFutures = new ArrayList<>(); - final List fullRoutines = Collections.synchronizedList(new ArrayList<>()); - - for (RoutineId procId : procedureIdsToGet) { - if (Thread.currentThread().isInterrupted()) { - InterruptedException ex = - new InterruptedException("Interrupted while submitting getRoutine tasks"); - logger.severe(ex.getMessage(), ex); - throw ex; - } - final RoutineId currentProcId = procId; - Callable getCallable = - () -> { - try { - return bigquery.getRoutine(currentProcId, options); - } catch (Exception e) { - logger.warning( - "Failed to get full details for routine " - + currentProcId - + ": " - + e.getMessage()); - return null; - } - }; - getRoutineFutures.add(getRoutineDetailsExecutor.submit(getCallable)); - } - logger.fine("Submitted " + getRoutineFutures.size() + " getRoutine detail tasks."); - - for (Future getFuture : getRoutineFutures) { - if (Thread.currentThread().isInterrupted()) { - getRoutineFutures.forEach(f -> f.cancel(true)); // Cancel remaining - InterruptedException ex = - new InterruptedException("Interrupted while collecting Routine details"); - logger.severe(ex.getMessage(), ex); - throw ex; - } - try { - Routine fullRoutine = getFuture.get(); - if (fullRoutine != null) { - fullRoutines.add(fullRoutine); - } - } catch (CancellationException e) { - logger.warning("getRoutine detail task cancelled."); - } - } - logger.info("Successfully fetched full details for %d routines.", fullRoutines.size()); - return fullRoutines; - } - - void processProcedureArgumentsSequentially( - List fullRoutines, - Pattern columnNameRegex, - List collectedResults, - FieldList resultSchemaFields, - BigQueryJdbcCustomLogger logger) - throws InterruptedException { - logger.fine("Processing argument jobs sequentially for %d routines.", fullRoutines.size()); - - for (Routine fullRoutine : fullRoutines) { - if (Thread.currentThread().isInterrupted()) { - InterruptedException ex = - new InterruptedException("Interrupted while processing argument jobs sequentially"); - logger.severe(ex.getMessage(), ex); - throw ex; - } - if (fullRoutine != null) { - if ("PROCEDURE".equalsIgnoreCase(fullRoutine.getRoutineType())) { - processProcedureArguments( - fullRoutine, columnNameRegex, collectedResults, resultSchemaFields); - } else { - logger.warning( - "Routine " - + (fullRoutine.getRoutineId() != null - ? fullRoutine.getRoutineId().toString() - : "UNKNOWN_ID") - + " fetched via getRoutine was not of type PROCEDURE (Type: " - + fullRoutine.getRoutineType() - + "). Skipping argument processing."); - } - } - } - } - Schema defineGetProcedureColumnsSchema() { List fields = new ArrayList<>(20); fields.add( @@ -3738,18 +3551,8 @@ public ResultSet getFunctionColumns( && !"TABLE_FUNCTION".equalsIgnoreCase(routine.getRoutineType())) { return; } - try { - processFunctionParametersSequentially( - Collections.singletonList(routine), - columnNameRegex, - results, - fields, - LOG); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException( - "Interrupted while processing function parameters", e); - } + processFunctionParametersAndReturnValue( + routine, columnNameRegex, results, fields); }); Comparator comparator = @@ -3829,124 +3632,6 @@ Schema defineGetFunctionColumnsSchema() { return Schema.of(fields); } - List listMatchingFunctionIdsFromDatasets( - List datasetsToScan, - String functionNamePattern, - Pattern functionNameRegex, - ExecutorService listRoutinesExecutor, - String catalogParam, - BigQueryJdbcCustomLogger logger) - throws InterruptedException, ExecutionException { - - logger.fine( - "Listing matching function IDs from %d datasets for catalog '%s'.", - datasetsToScan.size(), catalogParam); - final List>> listRoutineFutures = new ArrayList<>(); - final List functionIdsToGet = Collections.synchronizedList(new ArrayList<>()); - - for (Dataset dataset : datasetsToScan) { - if (Thread.currentThread().isInterrupted()) { - logger.warning( - "Interrupted during submission of routine (function) listing tasks for catalog: " - + catalogParam); - throw new InterruptedException("Interrupted while listing functions"); - } - final DatasetId currentDatasetId = dataset.getDatasetId(); - Callable> listCallable = - () -> - findMatchingBigQueryObjects( - "Routine", - () -> - bigquery.listRoutines( - currentDatasetId, RoutineListOption.pageSize(DEFAULT_PAGE_SIZE)), - (name) -> - bigquery.getRoutine( - RoutineId.of( - currentDatasetId.getProject(), currentDatasetId.getDataset(), name)), - (rt) -> rt.getRoutineId().getRoutine(), - functionNamePattern, - functionNameRegex, - logger, - false); - listRoutineFutures.add(listRoutinesExecutor.submit(listCallable)); - } - logger.fine( - "Submitted " - + listRoutineFutures.size() - + " routine (function) list tasks for catalog: " - + catalogParam); - - for (Future> listFuture : listRoutineFutures) { - if (Thread.currentThread().isInterrupted()) { - logger.warning( - "Interrupted while collecting routine (function) list results for catalog: " - + catalogParam); - listRoutineFutures.forEach(f -> f.cancel(true)); - throw new InterruptedException("Interrupted while collecting function lists"); - } - try { - List listedRoutines = listFuture.get(); - if (listedRoutines != null) { - for (Routine listedRoutine : listedRoutines) { - if (listedRoutine != null - && ("SCALAR_FUNCTION".equalsIgnoreCase(listedRoutine.getRoutineType()) - || "TABLE_FUNCTION".equalsIgnoreCase(listedRoutine.getRoutineType()))) { - if (listedRoutine.getRoutineId() != null) { - functionIdsToGet.add(listedRoutine.getRoutineId()); - } else { - logger.warning( - "Found a function type routine with a null ID during listing phase for catalog:" - + " " - + catalogParam); - } - } - } - } - } catch (CancellationException e) { - logger.warning("Routine (function) list task cancelled for catalog: " + catalogParam); - } - } - logger.info( - "Found %d function IDs to fetch details for in catalog '%s'.", - functionIdsToGet.size(), catalogParam); - return functionIdsToGet; - } - - void processFunctionParametersSequentially( - List fullFunctions, - Pattern columnNameRegex, - List collectedResults, - FieldList resultSchemaFields, - BigQueryJdbcCustomLogger logger) - throws InterruptedException { - logger.fine("Processing parameter jobs sequentially for %d functions.", fullFunctions.size()); - - for (Routine fullFunction : fullFunctions) { - if (Thread.currentThread().isInterrupted()) { - logger.warning("Interrupted during function parameter processing."); - throw new InterruptedException( - "Interrupted while processing function parameters sequentially"); - } - if (fullFunction != null) { - String routineType = fullFunction.getRoutineType(); - if ("SCALAR_FUNCTION".equalsIgnoreCase(routineType) - || "TABLE_FUNCTION".equalsIgnoreCase(routineType)) { - processFunctionParametersAndReturnValue( - fullFunction, columnNameRegex, collectedResults, resultSchemaFields); - } else { - logger.warning( - "Routine " - + (fullFunction.getRoutineId() != null - ? fullFunction.getRoutineId().toString() - : "UNKNOWN_ID") - + " fetched for getFunctionColumns was not of a function type (Type: " - + routineType - + "). Skipping parameter processing."); - } - } - } - } - void processFunctionParametersAndReturnValue( Routine routine, Pattern columnNameRegex, @@ -4303,7 +3988,6 @@ List findMatchingBigQueryObjects( Function nameExtractor, String pattern, Pattern regex, - BigQueryJdbcCustomLogger logger, boolean throwOn404) throws InterruptedException { @@ -4313,30 +3997,29 @@ List findMatchingBigQueryObjects( try { Iterable objects; if (needsList) { - logger.info( + LOG.info( "Listing all %ss (pattern: %s)...", objectTypeName, pattern == null ? "" : pattern); Page firstPage = listAllOperation.get(); objects = firstPage.iterateAll(); - logger.fine( - "Retrieved initial %s list, iterating & filtering if needed...", objectTypeName); + LOG.fine("Retrieved initial %s list, iterating & filtering if needed...", objectTypeName); } else { - logger.info("Getting specific %s: '%s'", objectTypeName, pattern); + LOG.info("Getting specific %s: '%s'", objectTypeName, pattern); T specificObject = getSpecificOperation.apply(pattern); objects = (specificObject == null) ? Collections.emptyList() : Collections.singletonList(specificObject); if (specificObject == null) { - logger.info("Specific %s not found: '%s'", objectTypeName, pattern); + LOG.info("Specific %s not found: '%s'", objectTypeName, pattern); } } boolean wasListing = needsList; for (T obj : objects) { if (Thread.currentThread().isInterrupted()) { - logger.warning("Thread interrupted during " + objectTypeName + " processing loop."); + LOG.warning("Thread interrupted during " + objectTypeName + " processing loop."); throw new InterruptedException( "Interrupted during " + objectTypeName + " processing loop"); } @@ -4354,20 +4037,20 @@ List findMatchingBigQueryObjects( } catch (BigQueryException e) { if (e.getCode() == 404 && !throwOn404) { - logger.info("%s '%s' not found (API error 404).", objectTypeName, pattern); + LOG.info("%s '%s' not found (API error 404).", objectTypeName, pattern); return Collections.emptyList(); } else { - logger.warning( + LOG.warning( "BigQueryException finding %ss for pattern '%s': %s (Code: %d)", objectTypeName, pattern, e.getMessage(), e.getCode()); throw e; } } catch (InterruptedException e) { Thread.currentThread().interrupt(); - logger.warning("Interrupted while finding " + objectTypeName + "s."); + LOG.warning("Interrupted while finding " + objectTypeName + "s."); throw e; } catch (Exception e) { - logger.severe( + LOG.severe( "Unexpected exception finding %ss for pattern '%s': %s", objectTypeName, pattern, e.getMessage()); throw new RuntimeException(e); @@ -4679,7 +4362,6 @@ private List fetchDatasetsForProject( (ds) -> ds.getDatasetId().getDataset(), schemaPattern, schemaRegex, - LOG, throwOn404); return datasets != null ? datasets : Collections.emptyList(); } catch (InterruptedException e) { 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 861d771d53a0..795eeada50a3 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 @@ -1024,7 +1024,6 @@ public void testFindMatchingBigQueryObjects_Routines_ListWithPattern() throws Ex (rt) -> rt.getRoutineId().getRoutine(), pattern, regex, - dbMetadata.LOG, false); verify(bigqueryClient, times(1)) @@ -1067,7 +1066,6 @@ public void testFindMatchingBigQueryObjects_Routines_ListNoPattern() throws Exce (rt) -> rt.getRoutineId().getRoutine(), pattern, regex, - dbMetadata.LOG, false); verify(bigqueryClient, times(1)) @@ -1103,7 +1101,6 @@ public void testFindMatchingBigQueryObjects_Routines_GetSpecific() throws Except (rt) -> rt.getRoutineId().getRoutine(), procNameExact, regex, - dbMetadata.LOG, false); verify(bigqueryClient, times(1)).getRoutine(eq(routineId)); @@ -1129,7 +1126,6 @@ private List
invokeFindMatchingObjectsWithException( (table) -> "name", pattern, dbMetadata.compileSqlLikePattern(pattern), - dbMetadata.LOG, throwOn404); } From 61ad2f44b689395139303900bc1a2434b76a0dff Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Tue, 21 Jul 2026 01:59:25 +0000 Subject: [PATCH 11/14] use uninterruptibles --- .../jdbc/BigQueryDatabaseMetaData.java | 225 +++++++----------- 1 file changed, 85 insertions(+), 140 deletions(-) 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 fb3bae4f9829..5f044bd35b0e 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 @@ -48,6 +48,7 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.jdbc.BigQueryJdbcTypeMappings.ColumnTypeInfo; import com.google.cloud.bigquery.jdbc.utils.BigQueryJdbcVersionUtility; +import com.google.common.util.concurrent.Uninterruptibles; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStream; @@ -4291,11 +4292,9 @@ private Future fetchAndPopulateQueueAsync( populateQueue(collectedResults, queue, resultSchemaFields); } catch (Throwable e) { LOG.severe("Error during async metadata fetch: " + e.getMessage()); - try { - queue.put(BigQueryFieldValueListWrapper.ofError(new SQLException(e))); - } catch (InterruptedException ie) { - Thread.currentThread().interrupt(); - } + BigQueryFieldValueListWrapper errorElement = + BigQueryFieldValueListWrapper.ofError(new SQLException(e)); + Uninterruptibles.putUninterruptibly(queue, errorElement); } finally { signalEndOfData(queue, resultSchemaFields); } @@ -4332,19 +4331,7 @@ private void signalEndOfData( LOG.info("Adding end signal to queue."); BigQueryFieldValueListWrapper element = BigQueryFieldValueListWrapper.ofEndOfStream(resultSchemaFields); - if (!queue.offer(element)) { - boolean wasInterrupted = Thread.interrupted(); - try { - queue.put(element); - } catch (InterruptedException e) { - LOG.warning("Interrupted while sending end signal to queue."); - wasInterrupted = true; - } finally { - if (wasInterrupted) { - Thread.currentThread().interrupt(); - } - } - } + Uninterruptibles.putUninterruptibly(queue, element); } catch (Exception e) { LOG.severe("Exception while sending end signal to queue: " + e.getMessage()); } @@ -4592,73 +4579,52 @@ private void processTargetRoutinesConcurrently( List> taskFutures = new ArrayList<>(); try { - List>>> datasetListTasks = new ArrayList<>(); + List> detailTasks = new ArrayList<>(); for (DatasetId datasetId : targetDatasets) { if (isLiteralName) { - datasetListTasks.add( + detailTasks.add( () -> { - return Collections.singletonList( - () -> { - processSingleRoutine( - datasetId, - routineNamePattern, - null, - collectedResults, - resultSchemaFields, - processor); - return null; - }); + processSingleRoutine( + datasetId, + routineNamePattern, + null, + collectedResults, + resultSchemaFields, + processor); + return null; }); continue; } - datasetListTasks.add( - () -> { - List> routineTasks = new ArrayList<>(); - try { - Page routinesPage = - bigquery.listRoutines( - datasetId, BigQuery.RoutineListOption.pageSize(DEFAULT_PAGE_SIZE)); - for (Routine routine : routinesPage.iterateAll()) { - String routineName = routine.getRoutineId().getRoutine(); - if (routineRegex != null && !routineRegex.matcher(routineName).matches()) - continue; - routineTasks.add( - () -> { - processSingleRoutine( - datasetId, - requiresFullRoutine ? routineName : null, - requiresFullRoutine ? null : routine, - collectedResults, - resultSchemaFields, - processor); - return null; - }); - } - } catch (BigQueryException e) { - if (e.getCode() == 404) { - LOG.info("Dataset '%s' not found while listing routines. Skipping.", datasetId); - } else { - throw new SQLException("Error while listing routines: " + e.getMessage(), e); - } - } - return routineTasks; - }); - } - - List> detailTasks = new ArrayList<>(); - try { - for (Callable>> listTask : datasetListTasks) { - detailTasks.addAll(listTask.call()); - } - } catch (Exception e) { - if (e instanceof InterruptedException) { - Thread.currentThread().interrupt(); - } - if (e instanceof SQLException) { - throw (SQLException) e; + try { + Page routinesPage = + bigquery.listRoutines( + datasetId, BigQuery.RoutineListOption.pageSize(DEFAULT_PAGE_SIZE)); + for (Routine routine : routinesPage.iterateAll()) { + if (Thread.currentThread().isInterrupted()) { + throw new SQLException("Interrupted while listing routines."); + } + String routineName = routine.getRoutineId().getRoutine(); + if (routineRegex != null && !routineRegex.matcher(routineName).matches()) continue; + detailTasks.add( + () -> { + processSingleRoutine( + datasetId, + requiresFullRoutine ? routineName : null, + requiresFullRoutine ? null : routine, + collectedResults, + resultSchemaFields, + processor); + return null; + }); + } + } catch (BigQueryException e) { + if (e.getCode() == 404) { + LOG.info("Dataset '%s' not found while listing routines. Skipping.", datasetId); + } else { + throw new SQLException("Error while listing routines: " + e.getMessage(), e); + } } - throw new SQLException("Error while listing routines: " + e.getMessage(), e); } for (Callable task : detailTasks) { @@ -4707,77 +4673,56 @@ private void processTargetTablesConcurrently( List> taskFutures = new ArrayList<>(); try { - List>>> datasetListTasks = new ArrayList<>(); + List> detailTasks = new ArrayList<>(); for (DatasetId datasetId : targetDatasets) { if (isLiteralName) { - datasetListTasks.add( + detailTasks.add( () -> { - return Collections.singletonList( - () -> { - processSingleTable( - datasetId, - tableNamePattern, - null, - requireBaseTable, - collectedResults, - resultSchemaFields, - processor); - return null; - }); + processSingleTable( + datasetId, + tableNamePattern, + null, + requireBaseTable, + collectedResults, + resultSchemaFields, + processor); + return null; }); continue; } - datasetListTasks.add( - () -> { - List> tableTasks = new ArrayList<>(); - try { - Page
tablesPage = - bigquery.listTables( - datasetId, BigQuery.TableListOption.pageSize(DEFAULT_PAGE_SIZE)); - for (Table table : tablesPage.iterateAll()) { - if (requireBaseTable - && table.getDefinition() != null - && table.getDefinition().getType() != TableDefinition.Type.TABLE) continue; - String tableName = table.getTableId().getTable(); - if (tableRegex != null && !tableRegex.matcher(tableName).matches()) continue; - tableTasks.add( - () -> { - processSingleTable( - datasetId, - requiresFullTable ? tableName : null, - requiresFullTable ? null : table, - requireBaseTable, - collectedResults, - resultSchemaFields, - processor); - return null; - }); - } - } catch (BigQueryException e) { - if (e.getCode() == 404) { - LOG.info("Dataset '%s' not found while listing tables. Skipping.", datasetId); - } else { - throw new SQLException("Error while listing tables: " + e.getMessage(), e); - } - } - return tableTasks; - }); - } - - List> detailTasks = new ArrayList<>(); - try { - for (Callable>> listTask : datasetListTasks) { - detailTasks.addAll(listTask.call()); - } - } catch (Exception e) { - if (e instanceof InterruptedException) { - Thread.currentThread().interrupt(); - } - if (e instanceof SQLException) { - throw (SQLException) e; + try { + Page
tablesPage = + bigquery.listTables(datasetId, BigQuery.TableListOption.pageSize(DEFAULT_PAGE_SIZE)); + for (Table table : tablesPage.iterateAll()) { + if (Thread.currentThread().isInterrupted()) { + throw new SQLException("Interrupted while listing tables."); + } + if (requireBaseTable + && table.getDefinition() != null + && table.getDefinition().getType() != TableDefinition.Type.TABLE) continue; + String tableName = table.getTableId().getTable(); + if (tableRegex != null && !tableRegex.matcher(tableName).matches()) continue; + detailTasks.add( + () -> { + processSingleTable( + datasetId, + requiresFullTable ? tableName : null, + requiresFullTable ? null : table, + requireBaseTable, + collectedResults, + resultSchemaFields, + processor); + return null; + }); + } + } catch (BigQueryException e) { + if (e.getCode() == 404) { + LOG.info("Dataset '%s' not found while listing tables. Skipping.", datasetId); + } else { + throw new SQLException("Error while listing tables: " + e.getMessage(), e); + } } - throw new SQLException("Error while listing tables: " + e.getMessage(), e); } for (Callable task : detailTasks) { From 3be489d87b3fe98542cdeb8f3a44cffc6dbb055a Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Tue, 21 Jul 2026 02:19:19 +0000 Subject: [PATCH 12/14] fix getExportedKeys --- .../google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 5f044bd35b0e..51913cafc87a 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 @@ -2163,7 +2163,7 @@ public ResultSet getExportedKeys(String catalog, String schema, String table) // Fallback Path: If catalog or schema is null, fall back to REST API metadata scan. if (catalog == null || schema == null) { - List targetDatasets = getTargetDatasets(catalog, schema); + List targetDatasets = getTargetDatasets(catalog, schema, false); final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); Future fetcherFuture = From b473c0869108451a46bc0fa8d90bec92dfb155e9 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Tue, 21 Jul 2026 02:27:04 +0000 Subject: [PATCH 13/14] finalise logic and flow --- .../jdbc/BigQueryDatabaseMetaData.java | 200 +++++++++++------- 1 file changed, 124 insertions(+), 76 deletions(-) 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 51913cafc87a..0606a8664e8e 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 @@ -4284,7 +4284,7 @@ private Future fetchAndPopulateQueueAsync( BlockingQueue queue, FieldList resultSchemaFields) { return connection - .getMetadataExecutor() + .getExecutorService() .submit( () -> { try { @@ -4575,13 +4575,14 @@ private void processTargetRoutinesConcurrently( return; } - ExecutorService executor = connection.getMetadataExecutor(); - List> taskFutures = new ArrayList<>(); + ExecutorService apiExecutor = connection.getMetadataExecutor(); + List> activeFutures = new ArrayList<>(); try { List> detailTasks = new ArrayList<>(); - for (DatasetId datasetId : targetDatasets) { - if (isLiteralName) { + + if (isLiteralName) { + for (DatasetId datasetId : targetDatasets) { detailTasks.add( () -> { processSingleRoutine( @@ -4593,52 +4594,74 @@ private void processTargetRoutinesConcurrently( processor); return null; }); - continue; + } + } else { + List>>> listFutures = new ArrayList<>(); + for (DatasetId datasetId : targetDatasets) { + Future>> future = + apiExecutor.submit( + () -> { + List> tasks = new ArrayList<>(); + try { + Page routinesPage = + bigquery.listRoutines( + datasetId, BigQuery.RoutineListOption.pageSize(DEFAULT_PAGE_SIZE)); + for (Routine routine : routinesPage.iterateAll()) { + if (Thread.currentThread().isInterrupted()) { + throw new InterruptedException("Interrupted while listing routines."); + } + String routineName = routine.getRoutineId().getRoutine(); + if (routineRegex != null && !routineRegex.matcher(routineName).matches()) { + continue; + } + tasks.add( + () -> { + processSingleRoutine( + datasetId, + requiresFullRoutine ? routineName : null, + requiresFullRoutine ? null : routine, + collectedResults, + resultSchemaFields, + processor); + return null; + }); + } + } catch (BigQueryException e) { + if (e.getCode() == 404) { + LOG.info( + "Dataset '%s' not found while listing routines. Skipping.", datasetId); + } else { + throw new SQLException( + "Error while listing routines: " + e.getMessage(), e); + } + } + return tasks; + }); + listFutures.add(future); + activeFutures.add(future); } - try { - Page routinesPage = - bigquery.listRoutines( - datasetId, BigQuery.RoutineListOption.pageSize(DEFAULT_PAGE_SIZE)); - for (Routine routine : routinesPage.iterateAll()) { - if (Thread.currentThread().isInterrupted()) { - throw new SQLException("Interrupted while listing routines."); - } - String routineName = routine.getRoutineId().getRoutine(); - if (routineRegex != null && !routineRegex.matcher(routineName).matches()) continue; - detailTasks.add( - () -> { - processSingleRoutine( - datasetId, - requiresFullRoutine ? routineName : null, - requiresFullRoutine ? null : routine, - collectedResults, - resultSchemaFields, - processor); - return null; - }); - } - } catch (BigQueryException e) { - if (e.getCode() == 404) { - LOG.info("Dataset '%s' not found while listing routines. Skipping.", datasetId); - } else { - throw new SQLException("Error while listing routines: " + e.getMessage(), e); - } + for (Future>> future : listFutures) { + detailTasks.addAll(future.get()); } + activeFutures.clear(); // Clear so we don't hold references unnecessarily } for (Callable task : detailTasks) { - taskFutures.add(executor.submit(task)); + activeFutures.add(apiExecutor.submit(task)); } - waitForTasksCompletion(taskFutures); + waitForTasksCompletion(activeFutures); } catch (ExecutionException e) { Throwable cause = e.getCause(); if (cause instanceof SQLException) { throw (SQLException) cause; } throw new SQLException("Error while fetching metadata", e); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new SQLException("Interrupted while processing routines.", e); } finally { - taskFutures.forEach(future -> future.cancel(true)); + activeFutures.forEach(future -> future.cancel(true)); } } @@ -4669,13 +4692,14 @@ private void processTargetTablesConcurrently( return; } - ExecutorService executor = connection.getMetadataExecutor(); - List> taskFutures = new ArrayList<>(); + ExecutorService apiExecutor = connection.getMetadataExecutor(); + List> activeFutures = new ArrayList<>(); try { List> detailTasks = new ArrayList<>(); - for (DatasetId datasetId : targetDatasets) { - if (isLiteralName) { + + if (isLiteralName) { + for (DatasetId datasetId : targetDatasets) { detailTasks.add( () -> { processSingleTable( @@ -4688,55 +4712,79 @@ private void processTargetTablesConcurrently( processor); return null; }); - continue; + } + } else { + List>>> listFutures = new ArrayList<>(); + for (DatasetId datasetId : targetDatasets) { + Future>> future = + apiExecutor.submit( + () -> { + List> tasks = new ArrayList<>(); + try { + Page
tablesPage = + bigquery.listTables( + datasetId, BigQuery.TableListOption.pageSize(DEFAULT_PAGE_SIZE)); + for (Table table : tablesPage.iterateAll()) { + if (Thread.currentThread().isInterrupted()) { + throw new InterruptedException("Interrupted while listing tables."); + } + if (requireBaseTable + && table.getDefinition() != null + && table.getDefinition().getType() != TableDefinition.Type.TABLE) { + continue; + } + String tableName = table.getTableId().getTable(); + if (tableRegex != null && !tableRegex.matcher(tableName).matches()) { + continue; + } + tasks.add( + () -> { + processSingleTable( + datasetId, + requiresFullTable ? tableName : null, + requiresFullTable ? null : table, + requireBaseTable, + collectedResults, + resultSchemaFields, + processor); + return null; + }); + } + } catch (BigQueryException e) { + if (e.getCode() == 404) { + LOG.info( + "Dataset '%s' not found while listing tables. Skipping.", datasetId); + } else { + throw new SQLException("Error while listing tables: " + e.getMessage(), e); + } + } + return tasks; + }); + listFutures.add(future); + activeFutures.add(future); } - try { - Page
tablesPage = - bigquery.listTables(datasetId, BigQuery.TableListOption.pageSize(DEFAULT_PAGE_SIZE)); - for (Table table : tablesPage.iterateAll()) { - if (Thread.currentThread().isInterrupted()) { - throw new SQLException("Interrupted while listing tables."); - } - if (requireBaseTable - && table.getDefinition() != null - && table.getDefinition().getType() != TableDefinition.Type.TABLE) continue; - String tableName = table.getTableId().getTable(); - if (tableRegex != null && !tableRegex.matcher(tableName).matches()) continue; - detailTasks.add( - () -> { - processSingleTable( - datasetId, - requiresFullTable ? tableName : null, - requiresFullTable ? null : table, - requireBaseTable, - collectedResults, - resultSchemaFields, - processor); - return null; - }); - } - } catch (BigQueryException e) { - if (e.getCode() == 404) { - LOG.info("Dataset '%s' not found while listing tables. Skipping.", datasetId); - } else { - throw new SQLException("Error while listing tables: " + e.getMessage(), e); - } + for (Future>> future : listFutures) { + detailTasks.addAll(future.get()); } + activeFutures.clear(); // Clear so we don't hold references unnecessarily } for (Callable task : detailTasks) { - taskFutures.add(executor.submit(task)); + activeFutures.add(apiExecutor.submit(task)); } - waitForTasksCompletion(taskFutures); + waitForTasksCompletion(activeFutures); } catch (ExecutionException e) { Throwable cause = e.getCause(); if (cause instanceof SQLException) { throw (SQLException) cause; } throw new SQLException("Error while fetching metadata", e); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new SQLException("Interrupted while processing tables.", e); } finally { - taskFutures.forEach(future -> future.cancel(true)); + activeFutures.forEach(future -> future.cancel(true)); } } From 9de9652461df945027fefdf792e64516939b90a8 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Tue, 21 Jul 2026 13:50:35 +0000 Subject: [PATCH 14/14] address nit pr feedback --- .../google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) 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 0606a8664e8e..daec78870df2 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 @@ -1530,10 +1530,7 @@ public ResultSet getColumns( List targetDatasets = getTargetDatasets(effectiveCatalog, effectiveSchemaPattern); - boolean hasWildcards = - columnNamePattern != null - && (columnNamePattern.contains("%") || columnNamePattern.contains("_")); - Pattern columnNameRegex = hasWildcards ? compileSqlLikePattern(columnNamePattern) : null; + Pattern columnNameRegex = compileSqlLikePattern(columnNamePattern); final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);