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 273d07b75602..752100f4888c 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,12 @@ 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.Dataset; import com.google.cloud.bigquery.DatasetId; @@ -1056,7 +1060,14 @@ public ResultSet getProcedureColumns( getRoutineDetailsExecutor = connection.getMetadataExecutor(); List fullRoutines = - fetchFullRoutineDetailsForIds(procedureIdsToGet, getRoutineDetailsExecutor, LOG); + fetchFullRoutineDetailsForIds( + procedureIdsToGet, + getRoutineDetailsExecutor, + LOG, + RoutineOption.fields( + RoutineField.ROUTINE_REFERENCE, + RoutineField.ARGUMENTS, + RoutineField.ROUTINE_TYPE)); getRoutineDetailsExecutor = null; if (fullRoutines.isEmpty() || Thread.currentThread().isInterrupted()) { @@ -1174,7 +1185,8 @@ List listMatchingProcedureIdsFromDatasets( List fetchFullRoutineDetailsForIds( List procedureIdsToGet, ExecutorService getRoutineDetailsExecutor, - BigQueryJdbcCustomLogger logger) + BigQueryJdbcCustomLogger logger, + RoutineOption... options) throws InterruptedException, ExecutionException { logger.fine("Fetching full details for %d procedure IDs.", procedureIdsToGet.size()); final List> getRoutineFutures = new ArrayList<>(); @@ -1191,7 +1203,7 @@ List fetchFullRoutineDetailsForIds( Callable getCallable = () -> { try { - return bigquery.getRoutine(currentProcId); + return bigquery.getRoutine(currentProcId, options); } catch (Exception e) { logger.warning( "Failed to get full details for routine " @@ -1616,7 +1628,11 @@ public ResultSet getTables( TableId.of( currentDatasetId.getProject(), currentDatasetId.getDataset(), - name)), + name), + TableOption.fields( + TableField.TABLE_REFERENCE, + TableField.TYPE, + TableField.DESCRIPTION)), (tbl) -> tbl.getTableId().getTable(), tableNamePattern, tableNameRegex, @@ -1930,7 +1946,9 @@ public ResultSet getColumns( datasetId, TableListOption.pageSize(DEFAULT_PAGE_SIZE)), (name) -> bigquery.getTable( - TableId.of(datasetId.getProject(), datasetId.getDataset(), name)), + TableId.of(datasetId.getProject(), datasetId.getDataset(), name), + TableOption.fields( + TableField.TABLE_REFERENCE, TableField.TYPE, TableField.SCHEMA)), (tbl) -> tbl.getTableId().getTable(), tableNamePattern, tableNameRegex, @@ -1999,7 +2017,10 @@ private void processTableColumns( "Schema not included in table object for " + tableId + ", fetching full table details..."); - Table fullTable = bigquery.getTable(tableId); + Table fullTable = + bigquery.getTable( + tableId, + TableOption.fields(TableField.TABLE_REFERENCE, TableField.TYPE, TableField.SCHEMA)); if (fullTable != null) { definition = fullTable.getDefinition(); tableSchema = (definition != null) ? definition.getSchema() : null; @@ -3423,49 +3444,11 @@ public ResultSet getSchemas(String catalog, String schemaPattern) throws SQLExce Runnable multiSchemaFetcher = () -> { final FieldList localResultSchemaFields = resultSchemaFields; - final List>> apiFutures = new ArrayList<>(); - final List collectedDatasets = new ArrayList<>(); final List collectedResults = new ArrayList<>(); try { - List projectsToScanList = getAccessibleCatalogNames(); - - if (projectsToScanList.isEmpty()) { - LOG.info( - "No valid projects to scan (primary, specified, or additional). Returning empty" - + " resultset."); - return; - } - - ExecutorService apiExecutor = connection.getMetadataExecutor(); - - LOG.fine("Submitting parallel fetchMatchingDatasets tasks..."); - for (String currentProjectToScan : projectsToScanList) { - if (Thread.currentThread().isInterrupted()) { - LOG.warning("Fetcher interrupted during project iteration submission."); - break; - } - - Callable> apiCallable = - () -> fetchMatchingDatasets(currentProjectToScan, schemaPattern, schemaRegex); - Future> apiFuture = apiExecutor.submit(apiCallable); - apiFutures.add(apiFuture); - } - LOG.fine("Finished submitting " + apiFutures.size() + " fetchMatchingDatasets tasks."); - - LOG.fine("Processing results from fetchMatchingDatasets tasks..."); - for (Future> apiFuture : apiFutures) { - if (Thread.currentThread().isInterrupted()) { - LOG.warning("Fetcher interrupted while processing API futures."); - break; - } - List datasetsResult = apiFuture.get(); - if (datasetsResult != null) { - collectedDatasets.addAll(datasetsResult); - } - } - - for (Dataset dataset : collectedDatasets) { + List datasets = fetchMatchingDatasets(catalog, schemaPattern, schemaRegex); + for (Dataset dataset : datasets) { if (Thread.currentThread().isInterrupted()) { break; } @@ -3480,7 +3463,6 @@ public ResultSet getSchemas(String catalog, String schemaPattern) throws SQLExce } catch (Throwable t) { handleFetcherException(t, queue, "getSchemas"); } finally { - apiFutures.forEach(f -> f.cancel(true)); finalizeFetcher(queue, localResultSchemaFields, "Multi-schema fetcher"); } }; @@ -4834,35 +4816,72 @@ private void signalEndOfData( } } + private List fetchDatasetsForProject( + String project, String schemaPattern, Pattern schemaRegex) throws SQLException { + try { + List datasets = + findMatchingBigQueryObjects( + "Dataset", + () -> bigquery.listDatasets(project, DatasetListOption.pageSize(DEFAULT_PAGE_SIZE)), + (name) -> bigquery.getDataset(DatasetId.of(project, name)), + (ds) -> ds.getDatasetId().getDataset(), + schemaPattern, + schemaRegex, + LOG); + return datasets != null ? datasets : Collections.emptyList(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new SQLException( + "Interrupted while fetching matching datasets for project " + project, e); + } catch (Exception e) { + throw new SQLException("Failed to fetch matching datasets for project " + project, e); + } + } + private List fetchMatchingDatasets( String catalog, String schemaPattern, Pattern schemaRegex) throws SQLException { List projects = (catalog != null) ? Collections.singletonList(catalog) : getAccessibleCatalogNames(); - List allDatasets = new ArrayList<>(); - for (String project : projects) { + + if (projects.isEmpty()) { + return Collections.emptyList(); + } + + // Single project path + if (projects.size() == 1) { + return fetchDatasetsForProject(projects.get(0), schemaPattern, schemaRegex); + } + + // Multi-project path + final List allDatasets = Collections.synchronizedList(new ArrayList<>()); + final List> taskFutures = new ArrayList<>(); + ExecutorService executor = connection.getMetadataExecutor(); + + try { + for (String project : projects) { + Callable task = + () -> { + List datasets = fetchDatasetsForProject(project, schemaPattern, schemaRegex); + allDatasets.addAll(datasets); + return null; + }; + taskFutures.add(executor.submit(task)); + } + + waitForTasksCompletion(taskFutures); if (Thread.currentThread().isInterrupted()) { - throw new SQLException("Interrupted while fetching matching datasets"); + throw new SQLException("Interrupted while parallel-fetching matching datasets"); } - try { - List datasets = - findMatchingBigQueryObjects( - "Dataset", - () -> bigquery.listDatasets(project, DatasetListOption.pageSize(DEFAULT_PAGE_SIZE)), - (name) -> bigquery.getDataset(DatasetId.of(project, name)), - (ds) -> ds.getDatasetId().getDataset(), - schemaPattern, - schemaRegex, - LOG); - if (datasets != null) { - allDatasets.addAll(datasets); - } - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new SQLException("Interrupted while fetching matching datasets", e); - } catch (Exception e) { - throw new SQLException("Failed to fetch matching datasets for project " + project, e); + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof SQLException) { + throw (SQLException) cause; } + throw new SQLException("Error parallel-fetching matching datasets", e); + } finally { + taskFutures.forEach(f -> f.cancel(true)); } + return allDatasets; }