diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java index 1d7d89e3f1f1..e5cf2df9dbd5 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java @@ -27,6 +27,9 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch; import com.google.cloud.bigquery.storage.v1.ArrowSchema; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; import java.io.IOException; import java.math.BigDecimal; import java.sql.Date; @@ -236,7 +239,10 @@ public boolean next() throws SQLException { || this.currentBatchRowIndex == (this.vectorSchemaRoot.getRowCount() - 1)) { /* Start of iteration or we have exhausted the current batch */ // Advance the cursor. Potentially blocking operation. - BigQueryArrowBatchWrapper batchWrapper = this.buffer.take(); + BigQueryArrowBatchWrapper batchWrapper; + try (Scope scope = Context.current().with(Span.wrap(originalSpanContext)).makeCurrent()) { + batchWrapper = this.buffer.take(); + } if (batchWrapper.getException() != null) { throw new BigQueryJdbcRuntimeException(batchWrapper.getException()); } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java index 4ff4acad6b2a..9b20354ad014 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java @@ -27,6 +27,8 @@ import com.google.cloud.bigquery.exception.BigQueryConversionException; import com.google.cloud.bigquery.exception.BigQueryJdbcCoercionException; import com.google.cloud.bigquery.exception.BigQueryJdbcCoercionNotFoundException; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanContext; import java.io.InputStream; import java.io.Reader; import java.io.StringReader; @@ -58,6 +60,7 @@ public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet protected boolean isClosed = false; protected boolean wasNull = false; protected final BigQueryTypeCoercer bigQueryTypeCoercer = BigQueryTypeCoercionUtility.INSTANCE; + protected final SpanContext originalSpanContext; protected BigQueryBaseResultSet( BigQuery bigQuery, BigQueryStatement statement, Schema schema, boolean isNested) { @@ -66,6 +69,7 @@ protected BigQueryBaseResultSet( this.schema = schema; this.schemaFieldList = schema != null ? schema.getFields() : null; this.isNested = isNested; + this.originalSpanContext = Span.current().getSpanContext(); } public QueryStatistics getQueryStatistics() { diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java index f08e29215c1a..87352cf8cbed 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java @@ -42,6 +42,12 @@ import com.google.cloud.bigquery.TableDefinition; import com.google.cloud.bigquery.TableId; import com.google.cloud.bigquery.exception.BigQueryJdbcException; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; import java.io.BufferedReader; import java.io.IOException; import java.io.InputStream; @@ -1698,8 +1704,16 @@ 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 { + return withTracing( + "BigQueryDatabaseMetaData.getTables", + () -> getTablesImpl(catalog, schemaPattern, tableNamePattern, types)); + } + private ResultSet getTablesImpl( + String catalog, String schemaPattern, String tableNamePattern, String[] types) + throws SQLException { Tuple effectiveIdentifiers = determineEffectiveCatalogAndSchema(catalog, schemaPattern); String effectiveCatalog = effectiveIdentifiers.x(); @@ -1731,141 +1745,159 @@ public ResultSet getTables( final String catalogParam = effectiveCatalog; final String schemaParam = effectiveSchemaPattern; + Tracer tracer = BigQueryJdbcOpenTelemetry.getSafeTracer(this.connection); + SpanContext parentSpanContext = Span.current().getSpanContext(); Runnable tableFetcher = () -> { - ExecutorService apiExecutor = null; - ExecutorService tableProcessorExecutor = null; - final FieldList localResultSchemaFields = resultSchemaFields; - final List>> apiFutures = new ArrayList<>(); - final List> processingFutures = new ArrayList<>(); + Span backgroundSpan = + tracer + .spanBuilder("BigQueryDatabaseMetaData.getTables.background") + .setNoParent() + .addLink(parentSpanContext) + .startSpan(); + + try (Scope backgroundScope = backgroundSpan.makeCurrent()) { + ExecutorService apiExecutor = null; + ExecutorService tableProcessorExecutor = null; + final FieldList localResultSchemaFields = resultSchemaFields; + final List>> apiFutures = new ArrayList<>(); + final List> processingFutures = new ArrayList<>(); - try { - List datasetsToScan = - findMatchingBigQueryObjects( - "Dataset", - () -> - bigquery.listDatasets( - catalogParam, DatasetListOption.pageSize(DEFAULT_PAGE_SIZE)), - (name) -> bigquery.getDataset(DatasetId.of(catalogParam, name)), - (ds) -> ds.getDatasetId().getDataset(), - schemaParam, - schemaRegex, - LOG); + try { + List datasetsToScan = + findMatchingBigQueryObjects( + "Dataset", + () -> + bigquery.listDatasets( + catalogParam, DatasetListOption.pageSize(DEFAULT_PAGE_SIZE)), + (name) -> bigquery.getDataset(DatasetId.of(catalogParam, name)), + (ds) -> ds.getDatasetId().getDataset(), + schemaParam, + schemaRegex, + LOG); - if (datasetsToScan.isEmpty()) { - LOG.info("Fetcher thread found no matching datasets. Returning empty resultset."); - return; - } + if (datasetsToScan.isEmpty()) { + LOG.info("Fetcher thread found no matching datasets. Returning empty resultset."); + return; + } - apiExecutor = Executors.newFixedThreadPool(API_EXECUTOR_POOL_SIZE); - tableProcessorExecutor = Executors.newFixedThreadPool(this.metadataFetchThreadCount); + apiExecutor = Executors.newFixedThreadPool(API_EXECUTOR_POOL_SIZE); + tableProcessorExecutor = Executors.newFixedThreadPool(this.metadataFetchThreadCount); - LOG.fine("Submitting parallel findMatchingTables tasks..."); - for (Dataset dataset : datasetsToScan) { - if (Thread.currentThread().isInterrupted()) { - LOG.warning("Table fetcher interrupted during dataset iteration."); - break; + 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)), + (tbl) -> tbl.getTableId().getTable(), + tableNamePattern, + tableNameRegex, + LOG); + + Callable> wrappedApiCallable = Context.current().wrap(apiCallable); + Future> apiFuture = apiExecutor.submit(wrappedApiCallable); + apiFutures.add(apiFuture); } + LOG.fine("Finished submitting " + apiFutures.size() + " findMatchingTables tasks."); + apiExecutor.shutdown(); - final DatasetId currentDatasetId = dataset.getDatasetId(); - Callable> apiCallable = - () -> - findMatchingBigQueryObjects( - "Table", + 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; + + final Table currentTable = table; + Runnable processRunnable = () -> - bigquery.listTables( - currentDatasetId, TableListOption.pageSize(DEFAULT_PAGE_SIZE)), - (name) -> - bigquery.getTable( - TableId.of( - currentDatasetId.getProject(), - currentDatasetId.getDataset(), - name)), - (tbl) -> tbl.getTableId().getTable(), - tableNamePattern, - tableNameRegex, - LOG); - Future> apiFuture = apiExecutor.submit(apiCallable); - apiFutures.add(apiFuture); - } - LOG.fine("Finished submitting " + apiFutures.size() + " findMatchingTables tasks."); - apiExecutor.shutdown(); + processTableInfo( + currentTable, + requestedTypes, + collectedResults, + localResultSchemaFields); + Runnable wrappedProcessRunnable = Context.current().wrap(processRunnable); + Future processFuture = + tableProcessorExecutor.submit(wrappedProcessRunnable); + processingFutures.add(processFuture); + } + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOG.warning("Fetcher thread interrupted while waiting for API future result."); + break; + } catch (ExecutionException e) { + LOG.warning( + "Error executing findMatchingTables task: " + + e.getMessage() + + ". Cause: " + + e.getCause()); + } catch (CancellationException e) { + LOG.warning("A findMatchingTables task was cancelled."); + } + } + + LOG.fine( + "Finished submitting " + processingFutures.size() + " processTableInfo 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; + LOG.warning( + "Fetcher interrupted before waiting for processing tasks; cancelling remaining."); + processingFutures.forEach(f -> f.cancel(true)); + } else { + LOG.fine("Waiting for processTableInfo tasks to complete..."); + waitForTasksCompletion(processingFutures); + LOG.fine("All processTableInfo tasks completed."); } - try { - List
tablesResult = apiFuture.get(); - if (tablesResult != null) { - for (Table table : tablesResult) { - if (Thread.currentThread().isInterrupted()) break; - final Table currentTable = table; - Future processFuture = - tableProcessorExecutor.submit( - () -> - processTableInfo( - currentTable, - requestedTypes, - collectedResults, - localResultSchemaFields)); - processingFutures.add(processFuture); - } - } - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - LOG.warning("Fetcher thread interrupted while waiting for API future result."); - break; - } catch (ExecutionException e) { - LOG.warning( - "Error executing findMatchingTables task: " - + e.getMessage() - + ". Cause: " - + e.getCause()); - } catch (CancellationException e) { - LOG.warning("A findMatchingTables task was cancelled."); + if (!Thread.currentThread().isInterrupted()) { + Comparator comparator = + defineGetTablesComparator(localResultSchemaFields); + sortResults(collectedResults, comparator, "getTables", LOG); } - } - LOG.fine( - "Finished submitting " + processingFutures.size() + " processTableInfo tasks."); + if (!Thread.currentThread().isInterrupted()) { + populateQueue(collectedResults, queue, localResultSchemaFields); + } - if (Thread.currentThread().isInterrupted()) { - LOG.warning( - "Fetcher interrupted before waiting for processing tasks; cancelling remaining."); + } catch (Throwable t) { + LOG.severe("Unexpected error in table fetcher runnable: " + t.getMessage()); + apiFutures.forEach(f -> f.cancel(true)); processingFutures.forEach(f -> f.cancel(true)); - } else { - LOG.fine("Waiting for processTableInfo tasks to complete..."); - waitForTasksCompletion(processingFutures); - LOG.fine("All processTableInfo tasks completed."); - } - - if (!Thread.currentThread().isInterrupted()) { - Comparator comparator = - defineGetTablesComparator(localResultSchemaFields); - sortResults(collectedResults, comparator, "getTables", LOG); + } finally { + signalEndOfData(queue, localResultSchemaFields); + shutdownExecutor(apiExecutor); + shutdownExecutor(tableProcessorExecutor); + LOG.info("Table fetcher thread finished."); } - - if (!Thread.currentThread().isInterrupted()) { - populateQueue(collectedResults, queue, localResultSchemaFields); - } - - } catch (Throwable t) { - LOG.severe("Unexpected error in table fetcher runnable: " + t.getMessage()); - apiFutures.forEach(f -> f.cancel(true)); - processingFutures.forEach(f -> f.cancel(true)); } finally { - signalEndOfData(queue, localResultSchemaFields); - shutdownExecutor(apiExecutor); - shutdownExecutor(tableProcessorExecutor); - LOG.info("Table fetcher thread finished."); + backgroundSpan.end(); } }; - Thread fetcherThread = new Thread(tableFetcher, "getTables-fetcher-" + effectiveCatalog); + Runnable wrappedTableFetcher = Context.current().wrap(tableFetcher); + Thread fetcherThread = new Thread(wrappedTableFetcher, "getTables-fetcher-" + effectiveCatalog); BigQueryJsonResultSet resultSet = BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread}); @@ -1985,16 +2017,19 @@ Comparator defineGetTablesComparator(FieldList resultSchemaField } @Override - public ResultSet getSchemas() { + public ResultSet getSchemas() throws SQLException { LOG.info("getSchemas() called"); return getSchemas(null, null); } @Override - public ResultSet getCatalogs() { - LOG.info("getCatalogs() called"); + public ResultSet getCatalogs() throws SQLException { + return withTracing("BigQueryDatabaseMetaData.getCatalogs", () -> getCatalogsImpl()); + } + private ResultSet getCatalogsImpl() throws SQLException { + LOG.info("getCatalogs() called"); final List accessibleCatalogs = getAccessibleCatalogNames(); final Schema catalogsSchema = defineGetCatalogsSchema(); final FieldList schemaFields = catalogsSchema.getFields(); @@ -2063,8 +2098,16 @@ 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 { + return withTracing( + "BigQueryDatabaseMetaData.getColumns", + () -> getColumnsImpl(catalog, schemaPattern, tableNamePattern, columnNamePattern)); + } + private ResultSet getColumnsImpl( + String catalog, String schemaPattern, String tableNamePattern, String columnNamePattern) + throws SQLException { Tuple effectiveIdentifiers = determineEffectiveCatalogAndSchema(catalog, schemaPattern); String effectiveCatalog = effectiveIdentifiers.x(); @@ -2096,102 +2139,119 @@ public ResultSet getColumns( final String catalogParam = effectiveCatalog; final String schemaParam = effectiveSchemaPattern; + Tracer tracer = BigQueryJdbcOpenTelemetry.getSafeTracer(this.connection); + SpanContext parentSpanContext = Span.current().getSpanContext(); Runnable columnFetcher = () -> { - ExecutorService columnExecutor = null; - final List> taskFutures = new ArrayList<>(); - final FieldList localResultSchemaFields = resultSchemaFields; - - try { - List datasetsToScan = - findMatchingBigQueryObjects( - "Dataset", - () -> - bigquery.listDatasets( - catalogParam, DatasetListOption.pageSize(DEFAULT_PAGE_SIZE)), - (name) -> bigquery.getDataset(DatasetId.of(catalogParam, name)), - (ds) -> ds.getDatasetId().getDataset(), - schemaParam, - schemaRegex, - LOG); - - if (datasetsToScan.isEmpty()) { - LOG.info("Fetcher thread found no matching datasets. Returning empty resultset."); - return; - } - - columnExecutor = Executors.newFixedThreadPool(this.metadataFetchThreadCount); - - 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()); + Span backgroundSpan = + tracer + .spanBuilder("BigQueryDatabaseMetaData.getColumns.background") + .setNoParent() + .addLink(parentSpanContext) + .startSpan(); + + try (Scope scope = backgroundSpan.makeCurrent()) { + ExecutorService columnExecutor = null; + final List> taskFutures = new ArrayList<>(); + final FieldList localResultSchemaFields = resultSchemaFields; - List
tablesToScan = + try { + List datasetsToScan = findMatchingBigQueryObjects( - "Table", + "Dataset", () -> - bigquery.listTables( - datasetId, TableListOption.pageSize(DEFAULT_PAGE_SIZE)), - (name) -> - bigquery.getTable( - TableId.of(datasetId.getProject(), datasetId.getDataset(), name)), - (tbl) -> tbl.getTableId().getTable(), - tableNamePattern, - tableNameRegex, + bigquery.listDatasets( + catalogParam, DatasetListOption.pageSize(DEFAULT_PAGE_SIZE)), + (name) -> bigquery.getDataset(DatasetId.of(catalogParam, name)), + (ds) -> ds.getDatasetId().getDataset(), + schemaParam, + schemaRegex, LOG); - for (Table table : tablesToScan) { + if (datasetsToScan.isEmpty()) { + LOG.info("Fetcher thread found no matching datasets. Returning empty resultset."); + return; + } + + columnExecutor = Executors.newFixedThreadPool(this.metadataFetchThreadCount); + + for (Dataset dataset : datasetsToScan) { if (Thread.currentThread().isInterrupted()) { - LOG.warning( - "Fetcher interrupted during table iteration for dataset " - + datasetId.getDataset()); + LOG.warning("Fetcher interrupted during dataset iteration."); break; } - TableId tableId = table.getTableId(); - LOG.fine("Submitting task for table: " + tableId); - final Table finalTable = table; - Future future = - columnExecutor.submit( + DatasetId datasetId = dataset.getDatasetId(); + LOG.info("Processing dataset: " + datasetId.getDataset()); + + List
tablesToScan = + findMatchingBigQueryObjects( + "Table", () -> - processTableColumns( - finalTable, - columnNameRegex, - collectedResults, - localResultSchemaFields)); - taskFutures.add(future); + bigquery.listTables( + datasetId, TableListOption.pageSize(DEFAULT_PAGE_SIZE)), + (name) -> + bigquery.getTable( + TableId.of(datasetId.getProject(), datasetId.getDataset(), name)), + (tbl) -> tbl.getTableId().getTable(), + tableNamePattern, + tableNameRegex, + LOG); + + 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; + + Runnable columnTask = + () -> + processTableColumns( + finalTable, + columnNameRegex, + collectedResults, + localResultSchemaFields); + Runnable wrappedColumnTask = Context.current().wrap(columnTask); + Future future = columnExecutor.submit(wrappedColumnTask); + taskFutures.add(future); + } + if (Thread.currentThread().isInterrupted()) break; } - if (Thread.currentThread().isInterrupted()) break; - } - waitForTasksCompletion(taskFutures); + waitForTasksCompletion(taskFutures); - if (!Thread.currentThread().isInterrupted()) { - Comparator comparator = - defineGetColumnsComparator(localResultSchemaFields); - sortResults(collectedResults, comparator, "getColumns", LOG); - } + if (!Thread.currentThread().isInterrupted()) { + Comparator comparator = + defineGetColumnsComparator(localResultSchemaFields); + sortResults(collectedResults, comparator, "getColumns", LOG); + } - if (!Thread.currentThread().isInterrupted()) { - populateQueue(collectedResults, queue, localResultSchemaFields); - } + if (!Thread.currentThread().isInterrupted()) { + populateQueue(collectedResults, queue, localResultSchemaFields); + } - } catch (Throwable t) { - LOG.severe("Unexpected error in column fetcher runnable: " + t.getMessage()); - taskFutures.forEach(f -> f.cancel(true)); + } catch (Throwable t) { + LOG.severe("Unexpected error in column fetcher runnable: " + t.getMessage()); + taskFutures.forEach(f -> f.cancel(true)); + } finally { + signalEndOfData(queue, localResultSchemaFields); + shutdownExecutor(columnExecutor); + LOG.info("Column fetcher thread finished."); + } } finally { - signalEndOfData(queue, localResultSchemaFields); - shutdownExecutor(columnExecutor); - LOG.info("Column fetcher thread finished."); + backgroundSpan.end(); } }; - Thread fetcherThread = new Thread(columnFetcher, "getColumns-fetcher-" + effectiveCatalog); + Runnable wrappedColumnFetcher = Context.current().wrap(columnFetcher); + Thread fetcherThread = + new Thread(wrappedColumnFetcher, "getColumns-fetcher-" + effectiveCatalog); BigQueryJsonResultSet resultSet = BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread}); @@ -3616,7 +3676,12 @@ public RowIdLifetime getRowIdLifetime() { } @Override - public ResultSet getSchemas(String catalog, String schemaPattern) { + public ResultSet getSchemas(String catalog, String schemaPattern) throws SQLException { + return withTracing( + "BigQueryDatabaseMetaData.getSchemas", () -> getSchemasImpl(catalog, schemaPattern)); + } + + private ResultSet getSchemasImpl(String catalog, String schemaPattern) throws SQLException { if ((catalog != null && catalog.isEmpty()) || (schemaPattern != null && schemaPattern.isEmpty())) { LOG.warning("Returning empty ResultSet as catalog or schemaPattern is an empty string."); @@ -3634,84 +3699,98 @@ public ResultSet getSchemas(String catalog, String schemaPattern) { final List collectedResults = Collections.synchronizedList(new ArrayList<>()); final String catalogParam = catalog; + Tracer tracer = BigQueryJdbcOpenTelemetry.getSafeTracer(this.connection); + SpanContext parentSpanContext = Span.current().getSpanContext(); Runnable schemaFetcher = () -> { - final FieldList localResultSchemaFields = resultSchemaFields; - List projectsToScanList = new ArrayList<>(); - - if (catalogParam != null) { - projectsToScanList.add(catalogParam); - } else { - projectsToScanList.addAll(getAccessibleCatalogNames()); - } - - if (projectsToScanList.isEmpty()) { - LOG.info( - "No valid projects to scan (primary, specified, or additional). Returning empty" - + " resultset."); - return; - } - - try { - for (String currentProjectToScan : projectsToScanList) { - if (Thread.currentThread().isInterrupted()) { - LOG.warning( - "Schema fetcher interrupted during project iteration for project: " - + currentProjectToScan); - break; - } - LOG.info("Fetching schemas for project: " + currentProjectToScan); - List datasetsInProject = - findMatchingBigQueryObjects( - "Dataset", - () -> - bigquery.listDatasets( - currentProjectToScan, - BigQuery.DatasetListOption.pageSize(DEFAULT_PAGE_SIZE)), - (name) -> bigquery.getDataset(DatasetId.of(currentProjectToScan, name)), - (ds) -> ds.getDatasetId().getDataset(), - schemaPattern, - schemaRegex, - LOG); + Span backgroundSpan = + tracer + .spanBuilder("BigQueryDatabaseMetaData.getSchemas.background") + .setNoParent() + .addLink(parentSpanContext) + .startSpan(); + + try (Scope backgroundScope = backgroundSpan.makeCurrent()) { + final FieldList localResultSchemaFields = resultSchemaFields; + List projectsToScanList = new ArrayList<>(); + + if (catalogParam != null) { + projectsToScanList.add(catalogParam); + } else { + projectsToScanList.addAll(getAccessibleCatalogNames()); + } - if (datasetsInProject.isEmpty() || Thread.currentThread().isInterrupted()) { - LOG.info( - "Fetcher thread found no matching datasets in project: " - + currentProjectToScan); - continue; - } + if (projectsToScanList.isEmpty()) { + LOG.info( + "No valid projects to scan (primary, specified, or additional). Returning empty" + + " resultset."); + return; + } - LOG.fine("Processing found datasets for project: " + currentProjectToScan); - for (Dataset dataset : datasetsInProject) { + try { + for (String currentProjectToScan : projectsToScanList) { if (Thread.currentThread().isInterrupted()) { LOG.warning( - "Schema fetcher interrupted during dataset iteration for project: " + "Schema fetcher interrupted during project iteration for project: " + currentProjectToScan); break; } - processSchemaInfo(dataset, collectedResults, localResultSchemaFields); + LOG.info("Fetching schemas for project: " + currentProjectToScan); + List datasetsInProject = + findMatchingBigQueryObjects( + "Dataset", + () -> + bigquery.listDatasets( + currentProjectToScan, + BigQuery.DatasetListOption.pageSize(DEFAULT_PAGE_SIZE)), + (name) -> bigquery.getDataset(DatasetId.of(currentProjectToScan, name)), + (ds) -> ds.getDatasetId().getDataset(), + schemaPattern, + schemaRegex, + LOG); + + if (datasetsInProject.isEmpty() || Thread.currentThread().isInterrupted()) { + LOG.info( + "Fetcher thread found no matching datasets in project: " + + currentProjectToScan); + continue; + } + + LOG.fine("Processing found datasets for project: " + currentProjectToScan); + for (Dataset dataset : datasetsInProject) { + if (Thread.currentThread().isInterrupted()) { + LOG.warning( + "Schema fetcher interrupted during dataset iteration for project: " + + currentProjectToScan); + break; + } + processSchemaInfo(dataset, collectedResults, localResultSchemaFields); + } } - } - if (!Thread.currentThread().isInterrupted()) { - Comparator comparator = - defineGetSchemasComparator(localResultSchemaFields); - sortResults(collectedResults, comparator, "getSchemas", LOG); - } + if (!Thread.currentThread().isInterrupted()) { + Comparator comparator = + defineGetSchemasComparator(localResultSchemaFields); + sortResults(collectedResults, comparator, "getSchemas", LOG); + } - if (!Thread.currentThread().isInterrupted()) { - populateQueue(collectedResults, queue, localResultSchemaFields); - } + if (!Thread.currentThread().isInterrupted()) { + populateQueue(collectedResults, queue, localResultSchemaFields); + } - } catch (Throwable t) { - LOG.severe("Unexpected error in schema fetcher runnable: " + t.getMessage()); + } catch (Throwable t) { + LOG.severe("Unexpected error in schema fetcher runnable: " + t.getMessage()); + } finally { + signalEndOfData(queue, localResultSchemaFields); + LOG.info("Schema fetcher thread finished."); + } } finally { - signalEndOfData(queue, localResultSchemaFields); - LOG.info("Schema fetcher thread finished."); + backgroundSpan.end(); } }; - Thread fetcherThread = new Thread(schemaFetcher, "getSchemas-fetcher-" + catalog); + Runnable wrappedFetcher = Context.current().wrap(schemaFetcher); + Thread fetcherThread = new Thread(wrappedFetcher, "getSchemas-fetcher-" + catalog); BigQueryJsonResultSet resultSet = BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread}); @@ -5291,4 +5370,23 @@ private void loadDriverVersionProperties() { throw new IllegalStateException(errorMessage, e); } } + + private interface TracedMetadataOperation { + T run() throws SQLException; + } + + private T withTracing(String spanName, TracedMetadataOperation operation) + throws SQLException { + Tracer tracer = BigQueryJdbcOpenTelemetry.getSafeTracer(this.connection); + Span span = tracer.spanBuilder(spanName).startSpan(); + try (Scope scope = span.makeCurrent()) { + return operation.run(); + } catch (Exception ex) { + span.recordException(ex); + span.setStatus(StatusCode.ERROR, ex.getMessage()); + throw ex; + } finally { + span.end(); + } + } } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java index 181e15629c5b..592569bde70c 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java @@ -49,4 +49,12 @@ public static OpenTelemetry getOpenTelemetry( public static Tracer getTracer(OpenTelemetry openTelemetry) { return openTelemetry.getTracer(INSTRUMENTATION_SCOPE_NAME); } + + /** Gets a Tracer for the JDBC driver, fallback to noop if connection has no tracer. */ + public static Tracer getSafeTracer(BigQueryConnection connection) { + if (connection != null && connection.getTracer() != null) { + return connection.getTracer(); + } + return getOpenTelemetry(false, false, null).getTracer(INSTRUMENTATION_SCOPE_NAME); + } } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java index da2ade028e9f..fe613120b91f 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java @@ -25,6 +25,9 @@ import com.google.cloud.bigquery.FieldValue.Attribute; import com.google.cloud.bigquery.Schema; import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; import java.sql.ResultSet; import java.sql.SQLException; import java.util.concurrent.BlockingQueue; @@ -158,7 +161,9 @@ public boolean next() throws SQLException { } try { // Advance the cursor,Potentially blocking operation - this.cursor = this.buffer.take(); + try (Scope scope = Context.current().with(Span.wrap(originalSpanContext)).makeCurrent()) { + this.cursor = this.buffer.take(); + } if (this.cursor.getException() != null) { throw new BigQueryJdbcRuntimeException(this.cursor.getException()); } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 4740e55aa33c..a846540b722b 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -60,6 +60,7 @@ import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanBuilder; +import io.opentelemetry.api.trace.SpanContext; import io.opentelemetry.api.trace.StatusCode; import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.context.Context; @@ -848,88 +849,85 @@ Thread populateArrowBufferedQueue( Runnable arrowStreamProcessor = Context.current() .wrap( - () -> { - long rowsRead = 0; - int retryCount = 0; - try { - // Use the first stream to perform reading. - String streamName = readSession.getStreams(0).getName(); - - while (true) { - try { - ReadRowsRequest readRowsRequest = - ReadRowsRequest.newBuilder() - .setReadStream(streamName) - .setOffset(rowsRead) - .build(); - - // Process each block of rows as they arrive and decode using our simple row - // reader. - com.google.api.gax.rpc.ServerStream stream = - bqReadClient.readRowsCallable().call(readRowsRequest); - for (ReadRowsResponse response : stream) { - if (Thread.currentThread().isInterrupted() - || queryTaskExecutor.isShutdown()) { - break; - } - - ArrowRecordBatch currentBatch = response.getArrowRecordBatch(); - Uninterruptibles.putUninterruptibly( - arrowBatchWrapperBlockingQueue, - BigQueryArrowBatchWrapper.of(currentBatch)); - rowsRead += response.getRowCount(); - } - break; - } catch (com.google.api.gax.rpc.ApiException e) { - if (e.getStatusCode().getCode() - == com.google.api.gax.rpc.StatusCode.Code.NOT_FOUND) { - LOG.warning("Read session expired or not found: %s", e.getMessage()); - enqueueError(arrowBatchWrapperBlockingQueue, e); - break; - } - if (retryCount >= MAX_RETRY_COUNT) { - LOG.log( - Level.SEVERE, - "\n" - + Thread.currentThread().getName() - + " Interrupted @ arrowStreamProcessor, max retries exceeded", - e); - enqueueError(arrowBatchWrapperBlockingQueue, e); - break; - } - retryCount++; - LOG.warning( - "Connection interrupted during arrow stream read, retrying. attempt: %d", - retryCount); - Thread.sleep(RETRY_DELAY_MS); - } - } - - } catch (InterruptedException e) { - LOG.log( - Level.WARNING, - "\n" - + Thread.currentThread().getName() - + " Interrupted @ arrowStreamProcessor", - e); - enqueueError(arrowBatchWrapperBlockingQueue, e); - Thread.currentThread().interrupt(); - } catch (Exception e) { - LOG.log( - Level.WARNING, - "\n" + Thread.currentThread().getName() + " Error @ arrowStreamProcessor", - e); - enqueueError(arrowBatchWrapperBlockingQueue, e); - } finally { // logic needed for graceful shutdown - enqueueEndOfStream(arrowBatchWrapperBlockingQueue); - } - }); + () -> + processArrowStream(readSession, arrowBatchWrapperBlockingQueue, bqReadClient)); Thread populateBufferWorker = JDBC_THREAD_FACTORY.newThread(arrowStreamProcessor); populateBufferWorker.start(); return populateBufferWorker; } + private void processArrowStream( + ReadSession readSession, + BlockingQueue arrowBatchWrapperBlockingQueue, + BigQueryReadClient bqReadClient) { + long rowsRead = 0; + int retryCount = 0; + try { + // Use the first stream to perform reading. + String streamName = readSession.getStreams(0).getName(); + + while (true) { + try { + ReadRowsRequest readRowsRequest = + ReadRowsRequest.newBuilder().setReadStream(streamName).setOffset(rowsRead).build(); + + // Process each block of rows as they arrive and decode using our simple row + // reader. + com.google.api.gax.rpc.ServerStream stream = + bqReadClient.readRowsCallable().call(readRowsRequest); + for (ReadRowsResponse response : stream) { + if (Thread.currentThread().isInterrupted() || queryTaskExecutor.isShutdown()) { + break; + } + + ArrowRecordBatch currentBatch = response.getArrowRecordBatch(); + Uninterruptibles.putUninterruptibly( + arrowBatchWrapperBlockingQueue, BigQueryArrowBatchWrapper.of(currentBatch)); + rowsRead += response.getRowCount(); + } + break; + } catch (com.google.api.gax.rpc.ApiException e) { + if (e.getStatusCode().getCode() == com.google.api.gax.rpc.StatusCode.Code.NOT_FOUND) { + LOG.warning("Read session expired or not found: %s", e.getMessage()); + enqueueError(arrowBatchWrapperBlockingQueue, e); + break; + } + if (retryCount >= MAX_RETRY_COUNT) { + LOG.log( + Level.SEVERE, + "\n" + + Thread.currentThread().getName() + + " Interrupted @ arrowStreamProcessor, max retries exceeded", + e); + enqueueError(arrowBatchWrapperBlockingQueue, e); + break; + } + retryCount++; + LOG.warning( + "Connection interrupted during arrow stream read, retrying. attempt: %d", retryCount); + Thread.sleep(RETRY_DELAY_MS); + } + } + + } catch (InterruptedException e) { + LOG.log( + Level.WARNING, + "\n" + Thread.currentThread().getName() + " Interrupted @ arrowStreamProcessor", + e); + enqueueError(arrowBatchWrapperBlockingQueue, e); + Thread.currentThread().interrupt(); + } catch (Exception e) { + LOG.log( + Level.WARNING, + "\n" + Thread.currentThread().getName() + " Error @ arrowStreamProcessor", + e); + enqueueError(arrowBatchWrapperBlockingQueue, e); + } finally { // logic needed for graceful shutdown + enqueueEndOfStream(arrowBatchWrapperBlockingQueue); + } + } + /** Executes SQL query using either fast query path or read API */ void processQueryResponse(String query, TableResult results) throws SQLException { LOG.finest( @@ -1544,7 +1542,7 @@ private interface TracedOperation { } private T withTracing(String spanName, TracedOperation operation) throws SQLException { - Tracer tracer = getSafeTracer(); + Tracer tracer = BigQueryJdbcOpenTelemetry.getSafeTracer(this.connection); Span span = tracer.spanBuilder(spanName).startSpan(); try (Scope scope = span.makeCurrent()) { return operation.run(span); @@ -1563,7 +1561,8 @@ private void fetchNextPages( BlockingQueue> rpcResponseQueue, BlockingQueue bigQueryFieldValueListWrapperBlockingQueue, TableResult result) { - Tracer tracer = getSafeTracer(); + Tracer tracer = BigQueryJdbcOpenTelemetry.getSafeTracer(this.connection); + SpanContext parentSpanContext = Span.current().getSpanContext(); String currentPageToken = firstPageToken; TableResult currentResults = result; TableId destinationTable = null; @@ -1579,6 +1578,9 @@ private void fetchNextPages( } SpanBuilder spanBuilder = tracer.spanBuilder("BigQueryStatement.pagination"); + if (parentSpanContext.isValid()) { + spanBuilder.addLink(parentSpanContext); + } Span paginationSpan = spanBuilder.startSpan(); try (Scope scope = paginationSpan.makeCurrent()) { paginationSpan.setAttribute("db.pagination.page_token", currentPageToken); @@ -1672,12 +1674,4 @@ private void parseAndPopulateRpcData( enqueueBufferEndOfStream(bigQueryFieldValueListWrapperBlockingQueue); } } - - private Tracer getSafeTracer() { - if (connection != null && connection.getTracer() != null) { - return connection.getTracer(); - } - return BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, false, null) - .getTracer(BigQueryJdbcOpenTelemetry.INSTRUMENTATION_SCOPE_NAME); - } } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaDataTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaDataTest.java index 4d108ee54f8b..81fe6d4d13cb 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaDataTest.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaDataTest.java @@ -2917,7 +2917,7 @@ public void testPrepareGetCatalogsRows() { } @Test - public void testGetSchemas_NoArgs_DelegatesCorrectly() { + public void testGetSchemas_NoArgs_DelegatesCorrectly() throws Exception { BigQueryDatabaseMetaData spiedDbMetadata = spy(dbMetadata); ResultSet mockResultSet = mock(ResultSet.class); doReturn(mockResultSet).when(spiedDbMetadata).getSchemas(null, null);