diff --git a/java-bigquery-jdbc/pom.xml b/java-bigquery-jdbc/pom.xml index ed141e41ef35..024507f1285e 100644 --- a/java-bigquery-jdbc/pom.xml +++ b/java-bigquery-jdbc/pom.xml @@ -83,6 +83,10 @@ org.apache.httpcomponents.*:* io.grpc:* + + io.opentelemetry.contrib:opentelemetry-gcp-auth-extension + io.opentelemetry:opentelemetry-exporter-otlp + @@ -167,6 +171,17 @@ io.opentelemetry com.google.bqjdbc.shaded.io.opentelemetry + + + io.opentelemetry.api.* + io.opentelemetry.context.* + io.perfmark @@ -217,6 +232,11 @@ google-cloud-bigquerystorage 3.30.0-SNAPSHOT + + com.google.cloud + google-cloud-logging + 3.32.0 + com.google.http-client google-http-client-apache-v5 @@ -304,6 +324,16 @@ io.grpc grpc-netty-shaded + + io.grpc + grpc-opentelemetry + + + io.opentelemetry + opentelemetry-api + + + org.apache.arrow @@ -322,6 +352,42 @@ httpcore5 + + + io.opentelemetry + opentelemetry-api + + + io.opentelemetry + opentelemetry-context + + + + io.opentelemetry + opentelemetry-sdk + + + io.opentelemetry + opentelemetry-sdk-trace + + + io.opentelemetry + opentelemetry-exporter-otlp + + + io.opentelemetry + opentelemetry-sdk-extension-autoconfigure + + + io.opentelemetry + opentelemetry-sdk-extension-autoconfigure-spi + + + io.opentelemetry.contrib + opentelemetry-gcp-auth-extension + 1.56.0-alpha + + @@ -371,6 +437,22 @@ junit-platform-suite-engine test + + io.opentelemetry + opentelemetry-sdk-testing + test + + + io.opentelemetry + opentelemetry-sdk-logs + test + + + com.google.cloud + google-cloud-trace + 2.92.0 + test + diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java index 64269c7f74be..27351e5e8a3e 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java @@ -28,6 +28,7 @@ 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.context.Scope; import java.io.IOException; import java.math.BigDecimal; import java.sql.Date; @@ -263,29 +264,31 @@ 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(); - if (batchWrapper.getException() != null) { - throw new BigQueryJdbcRuntimeException(batchWrapper.getException()); - } - if (batchWrapper.isLast()) { - /* Marks the end of the records */ - if (this.vectorSchemaRoot != null) { - // IMP: To avoid memory leak: clear vectorSchemaRoot as it still holds - // the last batch - this.vectorSchemaRoot.clear(); + try (Scope scope = makeOriginalContextCurrent()) { + BigQueryArrowBatchWrapper batchWrapper = this.buffer.take(); + if (batchWrapper.getException() != null) { + throw new BigQueryJdbcRuntimeException(batchWrapper.getException()); + } + if (batchWrapper.isLast()) { + /* Marks the end of the records */ + if (this.vectorSchemaRoot != null) { + // IMP: To avoid memory leak: clear vectorSchemaRoot as it still holds + // the last batch + this.vectorSchemaRoot.clear(); + } + this.hasReachedEnd = true; + this.rowCount++; + return false; } - this.hasReachedEnd = true; + // Valid batch, process it + ArrowRecordBatch arrowBatch = batchWrapper.getCurrentArrowBatch(); + // Populates vectorSchemaRoot + this.arrowDeserializer.deserializeArrowBatch(arrowBatch); + // Pointing to the first row in this fresh batch + this.currentBatchRowIndex = 0; this.rowCount++; - return false; + return true; } - // Valid batch, process it - ArrowRecordBatch arrowBatch = batchWrapper.getCurrentArrowBatch(); - // Populates vectorSchemaRoot - this.arrowDeserializer.deserializeArrowBatch(arrowBatch); - // Pointing to the first row in this fresh batch - this.currentBatchRowIndex = 0; - this.rowCount++; - return true; } // There are rows left in the current batch. else if (this.currentBatchRowIndex < this.vectorSchemaRoot.getRowCount()) { diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java index 5731372e6774..c650658ce4a4 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java @@ -30,6 +30,10 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcCoercionException; import com.google.cloud.bigquery.exception.BigQueryJdbcCoercionNotFoundException; import com.google.cloud.bigquery.exception.BigQueryJdbcException; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; import java.io.InputStream; import java.io.Reader; import java.io.StringReader; @@ -67,6 +71,7 @@ public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet private SQLWarning warnings; private boolean warningsLoaded = false; protected final BigQueryTypeCoercer bigQueryTypeCoercer = BigQueryTypeCoercionUtility.INSTANCE; + protected final SpanContext originalSpanContext; protected BigQueryBaseResultSet( BigQuery bigQuery, BigQueryStatement statement, Schema schema, boolean isNested) { @@ -80,6 +85,7 @@ protected BigQueryBaseResultSet( this.schema = schema; this.schemaFieldList = schema != null ? schema.getFields() : null; this.isNested = isNested; + this.originalSpanContext = Span.current().getSpanContext(); this.job = job; if (job != null) { this.jobId = job.getJobId(); @@ -89,6 +95,10 @@ protected BigQueryBaseResultSet( this.getClass(), statement != null ? statement.connectionId : null); } + protected Scope makeOriginalContextCurrent() { + return Context.current().with(Span.wrap(this.originalSpanContext)).makeCurrent(); + } + public QueryStatistics getQueryStatistics() { if (queryStatistics != null) { return queryStatistics; diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java index 586a5c329405..422c64fad973 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java @@ -41,7 +41,15 @@ import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient; import com.google.cloud.bigquery.storage.v1.BigQueryWriteSettings; import com.google.cloud.http.HttpTransportOptions; +import com.google.cloud.logging.Logging; +import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableSortedSet; +import io.grpc.ManagedChannelBuilder; +import io.grpc.opentelemetry.GrpcOpenTelemetry; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.baggage.Baggage; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Context; import java.io.IOException; import java.io.InputStream; import java.sql.CallableStatement; @@ -152,6 +160,8 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { int transactionIsolation; List sqlWarnings; String catalog; + String gcpTelemetryCredentials; + String gcpTelemetryProjectId; int holdability; long retryTimeoutInSeconds; Duration retryTimeoutDuration; @@ -206,6 +216,14 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { Long connectionPoolSize; Long listenerPoolSize; String partnerToken; + boolean enableGcpTraceExporter; + boolean enableGcpLogExporter; + OpenTelemetry customOpenTelemetry; + boolean useGlobalOpenTelemetry; + private OpenTelemetry openTelemetry; + private Context otelContext; + Tracer tracer = + OpenTelemetry.noop().getTracer(BigQueryJdbcOpenTelemetry.INSTRUMENTATION_SCOPE_NAME); DatabaseMetaData databaseMetaData; Boolean reqGoogleDriveScope; private final Properties clientInfo = new Properties(); @@ -219,6 +237,11 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { BigQueryConnection(String url, DataSource ds) throws IOException { this.connectionId = UUID.randomUUID().toString(); + Baggage baggage = + Baggage.builder() + .put(BigQueryJdbcOpenTelemetry.CONNECTION_ID_BAGGAGE_KEY, this.connectionId) + .build(); + this.otelContext = Context.current().with(baggage); try (BigQueryJdbcMdc.MdcCloseable mdc = BigQueryJdbcMdc.registerInstance(this.connectionId)) { this.connectionUrl = url; if (LOG.isLoggable(java.util.logging.Level.CONFIG)) { @@ -241,6 +264,8 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { this.labels = ds.getLabels() != null ? ds.getLabels() : new java.util.HashMap<>(); this.maxBytesBilled = ds.getMaximumBytesBilled(); + this.gcpTelemetryCredentials = ds.getGcpTelemetryCredentials(); + this.gcpTelemetryProjectId = ds.getGcpTelemetryProjectId(); this.retryTimeoutInSeconds = ds.getTimeout(); this.retryTimeoutDuration = Duration.ofMillis(retryTimeoutInSeconds * 1000L); this.retryInitialDelayInSeconds = ds.getRetryInitialDelay(); @@ -346,6 +371,11 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { this.partnerToken = ds.getPartnerToken(); this.headerProvider = createHeaderProvider(); + this.enableGcpTraceExporter = ds.getEnableGcpTraceExporter(); + this.enableGcpLogExporter = ds.getEnableGcpLogExporter(); + this.customOpenTelemetry = ds.getCustomOpenTelemetry(); + this.useGlobalOpenTelemetry = ds.getUseGlobalOpenTelemetry(); + this.openTelemetry = getOpenTelemetryInstance(); this.bigQuery = getBigQueryConnection(); this.metadataExecutor = BigQueryJdbcMdc.newFixedThreadPool(metadataFetchThreadCount); this.queryExecutor = BigQueryJdbcMdc.newCachedThreadPool(); @@ -419,7 +449,8 @@ String getConnectionUrl() { return connectionUrl; } - String getConnectionId() { + @VisibleForTesting + public String getConnectionId() { return this.connectionId; } @@ -1034,6 +1065,7 @@ private void closeImpl() throws SQLException { } finally { BigQueryJdbcMdc.clear(); BigQueryJdbcRootLogger.closeConnectionHandler(this.connectionId); + BigQueryJdbcOpenTelemetry.unregisterConnection(this.connectionId); } if (exceptionToThrow != null) { throw exceptionToThrow; @@ -1107,6 +1139,65 @@ void removeStatement(Statement statement) { this.openStatements.remove(statement); } + private OpenTelemetry getOpenTelemetryInstance() { + + String effectiveProjectId = + (this.gcpTelemetryProjectId != null) ? this.gcpTelemetryProjectId : this.catalog; + String effectiveCredentials = resolveEffectiveCredentials(); + + validateTraceConfiguration(this.enableGcpTraceExporter, effectiveCredentials); + + OpenTelemetry openTelemetry = + BigQueryJdbcOpenTelemetry.getOpenTelemetry( + this.useGlobalOpenTelemetry, + this.enableGcpTraceExporter, + this.enableGcpLogExporter, + this.customOpenTelemetry, + effectiveCredentials, + effectiveProjectId); + + boolean hasExternalOtel = this.customOpenTelemetry != null || this.useGlobalOpenTelemetry; + Logging localLoggingClient = null; + if (this.enableGcpLogExporter && !hasExternalOtel) { + localLoggingClient = + BigQueryJdbcOpenTelemetry.createLoggingClient( + true, null, effectiveCredentials, effectiveProjectId, this.credentials); + } + + if (this.enableGcpLogExporter || hasExternalOtel) { + BigQueryJdbcOpenTelemetry.registerConnection( + this.connectionId, + openTelemetry, + localLoggingClient, + this.enableGcpLogExporter && !hasExternalOtel); + } + + return openTelemetry; + } + + private String resolveEffectiveCredentials() { + String creds = this.gcpTelemetryCredentials; + String authTypeStr = this.authProperties.get(BigQueryJdbcUrlUtility.OAUTH_TYPE_PROPERTY_NAME); + if (creds == null + && BigQueryJdbcOAuthUtility.AuthType.GOOGLE_SERVICE_ACCOUNT.name().equals(authTypeStr)) { + return this.authProperties.get(BigQueryJdbcUrlUtility.OAUTH_PVT_KEY_PROPERTY_NAME); + } + return creds; + } + + private void validateTraceConfiguration(boolean isTraceEnabled, String effectiveCredentials) { + if (isTraceEnabled && effectiveCredentials == null) { + String authTypeStr = this.authProperties.get(BigQueryJdbcUrlUtility.OAUTH_TYPE_PROPERTY_NAME); + if (!BigQueryJdbcOAuthUtility.AuthType.GOOGLE_SERVICE_ACCOUNT.name().equals(authTypeStr) + && !BigQueryJdbcOAuthUtility.AuthType.APPLICATION_DEFAULT_CREDENTIALS + .name() + .equals(authTypeStr)) { + throw new BigQueryJdbcRuntimeException( + "Exporting traces to Google Cloud is only supported when using Application Default Credentials (ADC) or Service Account authentication."); + } + } + } + private BigQuery getBigQueryConnection() { BigQueryOptions.Builder bigQueryOptions = BigQueryOptions.newBuilder(); if (this.retryTimeoutInSeconds > 0L @@ -1143,6 +1234,15 @@ private BigQuery getBigQueryConnection() { if (this.httpTransportOptions != null) { bigQueryOptions.setTransportOptions(this.httpTransportOptions); } + if (this.enableGcpTraceExporter + || this.customOpenTelemetry != null + || this.useGlobalOpenTelemetry) { + Tracer sdkTracer = this.openTelemetry.getTracer(BigQueryJdbcOpenTelemetry.BIGQUERY_NAMESPACE); + bigQueryOptions.setOpenTelemetryTracer(sdkTracer); + bigQueryOptions.setEnableOpenTelemetryTracing(true); + this.tracer = + this.openTelemetry.getTracer(BigQueryJdbcOpenTelemetry.INSTRUMENTATION_SCOPE_NAME); + } BigQueryOptions options = bigQueryOptions.setHeaderProvider(this.headerProvider).build(); options.setDefaultJobCreationMode( @@ -1177,7 +1277,20 @@ private BigQueryReadClient getBigQueryReadClientConnection() throws IOException } TransportChannelProvider activeProvider = this.transportChannelProvider; if (activeProvider == null) { - activeProvider = BigQueryReadSettings.defaultGrpcTransportProviderBuilder().build(); + InstantiatingGrpcChannelProvider.Builder builder = + BigQueryReadSettings.defaultGrpcTransportProviderBuilder(); + if (this.enableGcpTraceExporter + || this.customOpenTelemetry != null + || this.useGlobalOpenTelemetry) { + GrpcOpenTelemetry grpcOpenTelemetry = + GrpcOpenTelemetry.newBuilder().sdk(this.openTelemetry).build(); + builder.setChannelConfigurator( + b -> { + grpcOpenTelemetry.configureChannelBuilder((ManagedChannelBuilder) b); + return b; + }); + } + activeProvider = builder.build(); } if (activeProvider instanceof InstantiatingGrpcChannelProvider) { @@ -1192,6 +1305,13 @@ private BigQueryReadClient getBigQueryReadClientConnection() throws IOException bigQueryReadSettings.setTransportChannelProvider(activeProvider); + if (this.enableGcpTraceExporter + || this.customOpenTelemetry != null + || this.useGlobalOpenTelemetry) { + bigQueryReadSettings.setOpenTelemetryTracerProvider(this.openTelemetry.getTracerProvider()); + bigQueryReadSettings.setEnableOpenTelemetryTracing(true); + } + return BigQueryReadClient.create(bigQueryReadSettings.build()); } @@ -1291,6 +1411,18 @@ public CallableStatement prepareCall( return prepareCall(sql); } + public Tracer getTracer() { + return this.tracer; + } + + public Context getOtelContext() { + return this.otelContext; + } + + public String getPartnerToken() { + return this.partnerToken; + } + public boolean isReadOnlyTokenUsed() { return this.isReadOnlyTokenUsed; } 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 32ed62d91fd6..7d5c1263e03a 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 @@ -43,6 +43,11 @@ import com.google.cloud.bigquery.TableId; import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.jdbc.BigQueryJdbcTypeMappings.ColumnTypeInfo; +import com.google.common.annotations.VisibleForTesting; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; import java.io.BufferedReader; import java.io.IOException; import java.io.InputStream; @@ -857,7 +862,8 @@ public ResultSet getProcedures( procedureNamePattern, procedureNameRegex, LOG); - Future> apiFuture = apiExecutor.submit(apiCallable); + Future> apiFuture = + apiExecutor.submit(Context.current().wrap(apiCallable)); apiFutures.add(apiFuture); } LOG.fine("Finished submitting " + apiFutures.size() + " findMatchingRoutines tasks."); @@ -881,9 +887,13 @@ public ResultSet getProcedures( final Routine finalRoutine = routine; Future processFuture = routineProcessorExecutor.submit( - () -> - processProcedureInfo( - finalRoutine, collectedResults, localResultSchemaFields)); + Context.current() + .wrap( + () -> + processProcedureInfo( + finalRoutine, + collectedResults, + localResultSchemaFields))); processingTaskFutures.add(processFuture); } else { LOG.finer("Skipping non-procedure routine: " + routine.getRoutineId()); @@ -1273,7 +1283,7 @@ List listMatchingProcedureIdsFromDatasets( procedureNamePattern, procedureNameRegex, logger); - listRoutineFutures.add(listRoutinesExecutor.submit(listCallable)); + listRoutineFutures.add(listRoutinesExecutor.submit(Context.current().wrap(listCallable))); } logger.fine( "Submitted " @@ -1350,7 +1360,7 @@ List fetchFullRoutineDetailsForIds( return null; } }; - getRoutineFutures.add(getRoutineDetailsExecutor.submit(getCallable)); + getRoutineFutures.add(getRoutineDetailsExecutor.submit(Context.current().wrap(getCallable))); } logger.fine("Submitted " + getRoutineFutures.size() + " getRoutine detail tasks."); @@ -1400,9 +1410,14 @@ void submitProcedureArgumentProcessingJobs( final Routine finalFullRoutine = fullRoutine; Future processFuture = processArgsExecutor.submit( - () -> - processProcedureArguments( - finalFullRoutine, columnNameRegex, collectedResults, resultSchemaFields)); + Context.current() + .wrap( + () -> + processProcedureArguments( + finalFullRoutine, + columnNameRegex, + collectedResults, + resultSchemaFields))); outArgumentProcessingFutures.add(processFuture); } else { logger.warning( @@ -1708,8 +1723,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(); @@ -1727,161 +1750,195 @@ public ResultSet getTables( "getTables called for catalog: %s, schemaPattern: %s, tableNamePattern: %s, types: %s", effectiveCatalog, effectiveSchemaPattern, tableNamePattern, Arrays.toString(types)); + final Schema resultSchema = defineGetTablesSchema(); + final BlockingQueue queue = + new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); + + Thread fetcherThread = + runGetTablesTaskAsync( + effectiveCatalog, effectiveSchemaPattern, tableNamePattern, types, resultSchema, queue); + + BigQueryJsonResultSet resultSet = + BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread}); + + LOG.info("Started background thread for getTables"); + return resultSet; + } + + @VisibleForTesting + Thread runGetTablesTaskAsync( + String effectiveCatalog, + String effectiveSchemaPattern, + String tableNamePattern, + String[] types, + Schema resultSchema, + BlockingQueue queue) + throws SQLException { + 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(); - - final BlockingQueue queue = - new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); final List collectedResults = Collections.synchronizedList(new ArrayList<>()); final String catalogParam = effectiveCatalog; final String schemaParam = effectiveSchemaPattern; - + 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 = + this.connection + .getTracer() + .spanBuilder("BigQueryDatabaseMetaData.getTables.background") + .setNoParent() + .addLink(parentSpanContext) + .startSpan(); + + try (Scope scope = 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); - } - - if (!Thread.currentThread().isInterrupted()) { - populateQueue(collectedResults, queue, localResultSchemaFields); + } finally { + signalEndOfData(queue, localResultSchemaFields); + shutdownExecutor(apiExecutor); + shutdownExecutor(tableProcessorExecutor); + LOG.info("Table fetcher thread finished."); } - - } 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); - BigQueryJsonResultSet resultSet = - BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread}); - + Runnable wrappedTableFetcher = Context.current().wrap(tableFetcher); + Thread fetcherThread = new Thread(wrappedTableFetcher, "getTables-fetcher-" + effectiveCatalog); fetcherThread.start(); - LOG.info("Started background thread for getTables"); - return resultSet; + return fetcherThread; } Schema defineGetTablesSchema() { @@ -1995,16 +2052,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(); @@ -2073,8 +2133,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(); @@ -2094,120 +2162,160 @@ public ResultSet getColumns( + " columnNamePattern: %s", effectiveCatalog, effectiveSchemaPattern, tableNamePattern, columnNamePattern); + final Schema resultSchema = defineGetColumnsSchema(); + final BlockingQueue queue = + new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); + + Thread fetcherThread = + runGetColumnsTaskAsync( + effectiveCatalog, + effectiveSchemaPattern, + tableNamePattern, + columnNamePattern, + resultSchema, + queue); + + BigQueryJsonResultSet resultSet = + BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread}); + + LOG.info("Started background thread for getColumns"); + return resultSet; + } + + @VisibleForTesting + Thread runGetColumnsTaskAsync( + String effectiveCatalog, + String effectiveSchemaPattern, + String tableNamePattern, + String columnNamePattern, + Schema resultSchema, + BlockingQueue queue) + throws SQLException { + Pattern schemaRegex = compileSqlLikePattern(effectiveSchemaPattern); Pattern tableNameRegex = compileSqlLikePattern(tableNamePattern); Pattern columnNameRegex = compileSqlLikePattern(columnNamePattern); - final Schema resultSchema = defineGetColumnsSchema(); final FieldList resultSchemaFields = resultSchema.getFields(); - final BlockingQueue queue = - new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); final List collectedResults = Collections.synchronizedList(new ArrayList<>()); final String catalogParam = effectiveCatalog; final String schemaParam = effectiveSchemaPattern; + SpanContext parentSpanContext = Span.current().getSpanContext(); Runnable columnFetcher = () -> { - ExecutorService columnExecutor = null; - final List> taskFutures = new ArrayList<>(); - final FieldList localResultSchemaFields = resultSchemaFields; + Span backgroundSpan = + this.connection + .getTracer() + .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; - 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()); - - 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); - BigQueryJsonResultSet resultSet = - BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread}); - + Runnable wrappedColumnFetcher = Context.current().wrap(columnFetcher); + Thread fetcherThread = + new Thread(wrappedColumnFetcher, "getColumns-fetcher-" + effectiveCatalog); fetcherThread.start(); - LOG.info("Started background thread for getColumns"); - return resultSet; + return fetcherThread; } private void processTableColumns( @@ -2274,7 +2382,7 @@ private void processTableColumns( } } - private Schema defineGetColumnsSchema() { + Schema defineGetColumnsSchema() { List fields = new ArrayList<>(24); fields.add( Field.newBuilder("TABLE_CAT", StandardSQLTypeName.STRING) @@ -3616,7 +3724,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."); @@ -3625,99 +3738,126 @@ public ResultSet getSchemas(String catalog, String schemaPattern) { 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(); - final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY); + + Thread fetcherThread = runGetSchemasTaskAsync(catalog, schemaPattern, resultSchema, queue); + + BigQueryJsonResultSet resultSet = + BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread}); + + LOG.info("Started background thread for getSchemas"); + return resultSet; + } + + @VisibleForTesting + Thread runGetSchemasTaskAsync( + String catalog, + String schemaPattern, + Schema resultSchema, + BlockingQueue queue) + throws SQLException { + + final Pattern schemaRegex = compileSqlLikePattern(schemaPattern); + final FieldList resultSchemaFields = resultSchema.getFields(); final List collectedResults = Collections.synchronizedList(new ArrayList<>()); final String catalogParam = catalog; + 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 = + this.connection + .getTracer() + .spanBuilder("BigQueryDatabaseMetaData.getSchemas.background") + .setNoParent() + .addLink(parentSpanContext) + .startSpan(); + + try (Scope scope = 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); - BigQueryJsonResultSet resultSet = - BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread}); - + Runnable wrappedFetcher = Context.current().wrap(schemaFetcher); + Thread fetcherThread = new Thread(wrappedFetcher, "getSchemas-fetcher-" + catalog); fetcherThread.start(); - LOG.info("Started background thread for getSchemas"); - return resultSet; + return fetcherThread; } Schema defineGetSchemasSchema() { @@ -3938,7 +4078,8 @@ public ResultSet getFunctions(String catalog, String schemaPattern, String funct functionNameRegex, LOG); }; - Future> apiFuture = apiExecutor.submit(apiCallable); + Future> apiFuture = + apiExecutor.submit(Context.current().wrap(apiCallable)); apiFutures.add(apiFuture); } LOG.fine( @@ -4373,7 +4514,7 @@ List listMatchingFunctionIdsFromDatasets( functionNamePattern, functionNameRegex, logger); - listRoutineFutures.add(listRoutinesExecutor.submit(listCallable)); + listRoutineFutures.add(listRoutinesExecutor.submit(Context.current().wrap(listCallable))); } logger.fine( "Submitted " @@ -5264,4 +5405,14 @@ private void loadDriverVersionProperties() { throw ex; } } + + private interface TracedMetadataOperation { + T run() throws SQLException; + } + + private T withTracing(String spanName, TracedMetadataOperation operation) + throws SQLException { + return BigQueryJdbcOpenTelemetry.withTracing( + spanName, this.connection, null, () -> operation.run()); + } } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java index 15058ff01bd9..b6747f60a246 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java @@ -20,6 +20,7 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; import io.grpc.LoadBalancerRegistry; import io.grpc.internal.PickFirstLoadBalancerProvider; +import io.opentelemetry.api.OpenTelemetry; import java.io.IOException; import java.sql.Connection; import java.sql.Driver; @@ -127,8 +128,12 @@ public Connection connect(String url, Properties info) throws SQLException { LOG.finest("++enter++"); try { if (acceptsURL(url)) { + Properties connectInfo = info == null ? new Properties() : (Properties) info.clone(); + Object customOpenTelemetryObj = connectInfo.remove("customOpenTelemetry"); + String connectionUri = - BigQueryJdbcUrlUtility.appendPropertiesToURL(url.substring(5), this.toString(), info); + BigQueryJdbcUrlUtility.appendPropertiesToURL( + url.substring(5), this.toString(), connectInfo); Level logLevel; String logPath; try { @@ -148,7 +153,19 @@ public Connection connect(String url, Properties info) throws SQLException { if (logPath == null) { logPath = System.getenv(BigQueryJdbcUrlUtility.LOG_PATH_ENV_VAR); } - if (logPath == null) { + + // Fallback to default path only if not specified and not in Cloud-Only mode + String enableGcpLogExporterStr = + BigQueryJdbcUrlUtility.parseUriPropertyWithoutValidation( + connectionUri, BigQueryJdbcUrlUtility.ENABLE_GCP_LOG_EXPORTER_PROPERTY_NAME); + boolean enableGcpLogExporter = false; + if (enableGcpLogExporterStr != null) { + enableGcpLogExporter = + BigQueryJdbcUrlUtility.convertIntToBoolean( + enableGcpLogExporterStr, + BigQueryJdbcUrlUtility.ENABLE_GCP_LOG_EXPORTER_PROPERTY_NAME); + } + if (logPath == null && !enableGcpLogExporter) { logPath = BigQueryJdbcUrlUtility.DEFAULT_LOG_PATH; } @@ -166,7 +183,9 @@ public Connection connect(String url, Properties info) throws SQLException { LOG.log(Level.SEVERE, "Failed to parse connection URL", e); throw new BigQueryJdbcException("Failed to parse connection URL", e); } - + if (customOpenTelemetryObj instanceof OpenTelemetry) { + ds.setCustomOpenTelemetry((OpenTelemetry) customOpenTelemetryObj); + } BigQueryConnection connection = new BigQueryConnection(connectionUri, ds); LOG.info( "Driver info : { {Database Product Name : %s}, " diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtility.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtility.java index 0b9166d4cda2..632e6ae06769 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtility.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtility.java @@ -307,7 +307,7 @@ private static boolean isFileExists(String filename) { } } - private static boolean isJson(byte[] value) { + static boolean isJson(byte[] value) { try { // This is done this way to ensure strict Json parsing // https://github.com/google/gson/issues/1208#issuecomment-2120764686 diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java new file mode 100644 index 000000000000..40d35c880263 --- /dev/null +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java @@ -0,0 +1,415 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.jdbc; + +import com.google.auth.Credentials; +import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; +import com.google.cloud.logging.Logging; +import com.google.cloud.logging.LoggingOptions; +import com.google.common.hash.Hashing; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.baggage.Baggage; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; +import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporter; +import io.opentelemetry.exporter.otlp.trace.OtlpGrpcSpanExporter; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk; +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.sql.SQLException; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; +import java.util.logging.Handler; +import java.util.logging.Logger; + +public class BigQueryJdbcOpenTelemetry { + + static final String INSTRUMENTATION_SCOPE_NAME = "com.google.cloud.bigquery.jdbc"; + static final String BIGQUERY_NAMESPACE = "com.google.cloud.bigquery"; + public static final String CONNECTION_ID_BAGGAGE_KEY = "jdbc.connection_id"; + public static final String DB_SYSTEM_KEY = "db.system"; + public static final String DB_SYSTEM_VALUE = "bigquery"; + public static final String DB_CONNECTION_ID_KEY = "db.connection_id"; + public static final String DB_APPLICATION_KEY = "db.application"; + public static final String DEFAULT_APPLICATION_NAME = "Google-BigQuery-JDBC-Driver"; + public static final String DB_STATEMENT_KEY = "db.statement"; + public static final String DB_STATEMENT_COUNT_KEY = "db.statement.count"; + public static final String DB_BATCH_STATEMENTS_KEY = "db.batch.statements"; + private static final String OTEL_TRACES_EXPORTER = "otel.traces.exporter"; + private static final String OTEL_EXPORTER_OTLP_ENDPOINT = "otel.exporter.otlp.endpoint"; + private static final String OTEL_LOGS_EXPORTER = "otel.logs.exporter"; + private static final String OTEL_METRICS_EXPORTER = "otel.metrics.exporter"; + private static final String GOOGLE_CLOUD_PROJECT = "google.cloud.project"; + private static final String OTLP_ENDPOINT_VALUE = "https://telemetry.googleapis.com:443"; + private static final URI OTLP_ENDPOINT_URI = URI.create(OTLP_ENDPOINT_VALUE); + private static final String EXPORTER_NONE = "none"; + private static final String EXPORTER_OTLP = "otlp"; + private static final String OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT = + "otel.span.attribute.value.length.limit"; + private static final String OTEL_ATTRIBUTE_VALUE_LENGTH_LIMIT = + "otel.attribute.value.length.limit"; + private static final String DEFAULT_ATTRIBUTE_LENGTH_LIMIT = "32768"; + private static final BigQueryJdbcCustomLogger LOG = + new BigQueryJdbcCustomLogger("BigQueryJdbcOpenTelemetry"); + + private static final class SdkCacheKey { + private final String projectId; + private final String credentialsHashOrPath; + private final boolean enableTrace; + + SdkCacheKey(String projectId, String credentialsHashOrPath, boolean enableTrace) { + this.projectId = projectId; + this.credentialsHashOrPath = credentialsHashOrPath; + this.enableTrace = enableTrace; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + SdkCacheKey that = (SdkCacheKey) o; + return enableTrace == that.enableTrace + && Objects.equals(projectId, that.projectId) + && Objects.equals(credentialsHashOrPath, that.credentialsHashOrPath); + } + + @Override + public int hashCode() { + return Objects.hash(projectId, credentialsHashOrPath, enableTrace); + } + } + + private static final ConcurrentHashMap sdkCache = + new ConcurrentHashMap<>(); + + static class TelemetryConfig { + final OpenTelemetry openTelemetry; + final Logging loggingClient; + final boolean useDirectGcpLogging; + + TelemetryConfig( + OpenTelemetry openTelemetry, Logging loggingClient, Boolean useDirectGcpLogging) { + this.openTelemetry = openTelemetry; + this.loggingClient = loggingClient; + this.useDirectGcpLogging = useDirectGcpLogging != null ? useDirectGcpLogging : false; + } + } + + private static final ConcurrentHashMap connectionConfigs = + new ConcurrentHashMap<>(); + + private BigQueryJdbcOpenTelemetry() {} + + static { + ensureGlobalHandlerAttached(); + Runtime.getRuntime() + .addShutdownHook( + new Thread( + () -> { + for (OpenTelemetrySdk sdk : sdkCache.values()) { + try { + sdk.close(); + } catch (Exception e) { + // Ignore failures during shutdown to ensure all SDKs are attempted to be + // closed. Logging is avoided here because the logging system might have + // already been shut down by the JVM. + } + } + })); + } + + public static synchronized void ensureGlobalHandlerAttached() { + Logger logger = Logger.getLogger(BIGQUERY_NAMESPACE); + boolean present = false; + for (Handler h : logger.getHandlers()) { + if (h instanceof OpenTelemetryJulHandler) { + present = true; + break; + } + } + if (!present) { + logger.addHandler(new OpenTelemetryJulHandler()); + } + } + + public static void registerConnection( + String connectionId, + OpenTelemetry openTelemetry, + Logging loggingClient, + Boolean useDirectGcpLogging) { + connectionConfigs.put( + connectionId, new TelemetryConfig(openTelemetry, loggingClient, useDirectGcpLogging)); + } + + public static void unregisterConnection(String connectionId) { + TelemetryConfig config = connectionConfigs.remove(connectionId); + if (config != null && config.loggingClient != null) { + try { + config.loggingClient.close(); + } catch (Exception e) { + LOG.warning("Failed to close Logging client during unregister: %s", e.getMessage()); + } + } + } + + public static Logging createLoggingClient( + boolean enableGcpLogExporter, + OpenTelemetry customOpenTelemetry, + String effectiveCredentials, + String effectiveProjectId, + Credentials fallbackCredentials) { + + if (!enableGcpLogExporter || customOpenTelemetry != null) { + return null; + } + + try { + Credentials credentials; + if (effectiveCredentials != null) { + credentials = resolveCredentialsFromString(effectiveCredentials); + } else { + credentials = fallbackCredentials; + } + + LoggingOptions.Builder loggingOptionsBuilder = + LoggingOptions.newBuilder().setProjectId(effectiveProjectId); + if (credentials != null) { + loggingOptionsBuilder.setCredentials(credentials); + } + return loggingOptionsBuilder.build().getService(); + } catch (Exception e) { + throw new BigQueryJdbcRuntimeException("Failed to initialize Logging client", e); + } + } + + private static Credentials resolveCredentialsFromString(String credsString) { + Map authProperties = new java.util.HashMap<>(); + authProperties.put(BigQueryJdbcUrlUtility.OAUTH_TYPE_PROPERTY_NAME, "0"); // Service Account + + byte[] credsBytes = credsString.getBytes(StandardCharsets.UTF_8); + if (BigQueryJdbcOAuthUtility.isJson(credsBytes)) { + authProperties.put(BigQueryJdbcUrlUtility.OAUTH_PVT_KEY_PROPERTY_NAME, credsString); + } else { + authProperties.put(BigQueryJdbcUrlUtility.OAUTH_PVT_KEY_PATH_PROPERTY_NAME, credsString); + } + + return BigQueryJdbcOAuthUtility.getCredentials( + authProperties, + new java.util.HashMap<>(), + false, + BigQueryJdbcOpenTelemetry.class.getName()); + } + + public static TelemetryConfig getConnectionConfig(String connectionId) { + return connectionConfigs.get(connectionId); + } + + public static Collection getRegisteredConfigs() { + return connectionConfigs.values(); + } + + private static Map getAuthHeaders(Credentials credentials) { + try { + Map> metadata = credentials.getRequestMetadata(OTLP_ENDPOINT_URI); + Map headers = new HashMap<>(); + metadata.forEach( + (headerKey, headerValues) -> { + if (!headerValues.isEmpty()) { + headers.put(headerKey, headerValues.get(0)); + } + }); + return headers; + } catch (Exception e) { + // We log the warning and return an empty map, allowing the exporter to fail gracefully + // with a standard OTLP response code (e.g., 401 Unauthorized) handled by OTel. + LOG.warning(e, "Failed to get auth headers"); + return new HashMap<>(); + } + } + + private static String getCredentialsIdentifier(String credentials) { + if (credentials == null) { + return ""; + } + byte[] credsBytes = credentials.getBytes(StandardCharsets.UTF_8); + if (BigQueryJdbcOAuthUtility.isJson(credsBytes)) { + return Hashing.sha256().hashString(credentials, StandardCharsets.UTF_8).toString(); + } + return credentials; + } + + /** + * Initializes or returns the OpenTelemetry instance based on hybrid logic. Prefer + * customOpenTelemetry if provided; fallback to an auto-configured GCP exporter if requested. + */ + public static OpenTelemetry getOpenTelemetry( + boolean useGlobalOpenTelemetry, + boolean enableGcpTraceExporter, + boolean enableGcpLogExporter, + OpenTelemetry customOpenTelemetry, + String gcpTelemetryCredentials, + String gcpTelemetryProjectId) { + + if (customOpenTelemetry != null) { + return customOpenTelemetry; + } + + if (useGlobalOpenTelemetry) { + return GlobalOpenTelemetry.get(); + } + + if (!enableGcpTraceExporter && !enableGcpLogExporter) { + return OpenTelemetry.noop(); + } + + SdkCacheKey key = + new SdkCacheKey( + gcpTelemetryProjectId, + getCredentialsIdentifier(gcpTelemetryCredentials), + enableGcpTraceExporter); + return sdkCache.computeIfAbsent( + key, + k -> { + Map props = new HashMap<>(); + + if (enableGcpTraceExporter) { + props.put(OTEL_TRACES_EXPORTER, EXPORTER_OTLP); + props.put(OTEL_EXPORTER_OTLP_ENDPOINT, OTLP_ENDPOINT_VALUE); + } else { + props.put(OTEL_TRACES_EXPORTER, EXPORTER_NONE); + } + + // Logs are handled directly via GCP logging + props.put(OTEL_LOGS_EXPORTER, EXPORTER_NONE); + // Metrics are deferred to a future phase + props.put(OTEL_METRICS_EXPORTER, EXPORTER_NONE); + + if (gcpTelemetryProjectId != null) { + props.put(GOOGLE_CLOUD_PROJECT, gcpTelemetryProjectId); + } + + // Set safe, generous default limits on attribute value lengths (32KB) to protect + // customers from GCP Cloud Trace 64KB span ingestion failures when logging massive + // exception stack traces or database schema metadata. + // Respect any existing user configuration overrides. + if (!props.containsKey(OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT)) { + props.put(OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT, DEFAULT_ATTRIBUTE_LENGTH_LIMIT); + } + if (!props.containsKey(OTEL_ATTRIBUTE_VALUE_LENGTH_LIMIT)) { + props.put(OTEL_ATTRIBUTE_VALUE_LENGTH_LIMIT, DEFAULT_ATTRIBUTE_LENGTH_LIMIT); + } + + AutoConfiguredOpenTelemetrySdk autoConfigured = + AutoConfiguredOpenTelemetrySdk.builder() + .addPropertiesSupplier(() -> props) + .addSpanExporterCustomizer( + (spanExporter, configProperties) -> { + if (gcpTelemetryCredentials == null) { + return spanExporter; + } + try { + Credentials credentials = + resolveCredentialsFromString(gcpTelemetryCredentials); + if (spanExporter instanceof OtlpHttpSpanExporter) { + return ((OtlpHttpSpanExporter) spanExporter) + .toBuilder().setHeaders(() -> getAuthHeaders(credentials)).build(); + } + if (spanExporter instanceof OtlpGrpcSpanExporter) { + return ((OtlpGrpcSpanExporter) spanExporter) + .toBuilder().setHeaders(() -> getAuthHeaders(credentials)).build(); + } + } catch (Exception e) { + LOG.warning( + e, + "Failed to resolve telemetry credentials. Telemetry will be exported using default OpenTelemetry configuration (custom authentication headers will not be injected)."); + } + return spanExporter; + }) + .build(); + + OpenTelemetrySdk sdk = autoConfigured.getOpenTelemetrySdk(); + + return sdk; + }); + } + + /** Gets a Tracer for the JDBC driver instrumentation scope. */ + public static Tracer getTracer(OpenTelemetry openTelemetry) { + return openTelemetry.getTracer(INSTRUMENTATION_SCOPE_NAME); + } + + public static T withTracing( + String spanName, BigQueryConnection connection, String sql, Callable operation) + throws SQLException { + + Tracer tracer = connection.getTracer(); + Span span = tracer.spanBuilder(spanName).setSpanKind(SpanKind.CLIENT).startSpan(); + + span.setAttribute(DB_SYSTEM_KEY, DB_SYSTEM_VALUE); + span.setAttribute(DB_CONNECTION_ID_KEY, connection.getConnectionId()); + + String appName = connection.getPartnerToken(); + if (appName == null || appName.isEmpty()) { + appName = DEFAULT_APPLICATION_NAME; + } + span.setAttribute(DB_APPLICATION_KEY, appName); + + if (sql != null) { + span.setAttribute(DB_STATEMENT_KEY, sql); + } + + Baggage updatedBaggage = + Baggage.fromContext(Context.current()).toBuilder() + .put(CONNECTION_ID_BAGGAGE_KEY, connection.getConnectionId()) + .build(); + + // Create full context with new span and updated baggage + Context fullContext = Context.current().with(span).with(updatedBaggage); + + try (Scope scope = fullContext.makeCurrent()) { + return operation.call(); + } catch (Exception ex) { + span.recordException(ex); + span.setStatus(StatusCode.ERROR, ex.getMessage()); + + if (ex instanceof SQLException) { + throw (SQLException) ex; + } + if (ex instanceof RuntimeException) { + throw (RuntimeException) ex; + } + if (ex instanceof InterruptedException) { + Thread.currentThread().interrupt(); + throw new BigQueryJdbcRuntimeException("Operation interrupted", ex); + } + throw new BigQueryJdbcRuntimeException(ex); + } finally { + span.end(); + } + } +} diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcRootLogger.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcRootLogger.java index aeaf5f041693..5897166a3a85 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcRootLogger.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcRootLogger.java @@ -151,8 +151,9 @@ public static Logger getRootLogger() { public static void setLevel(Level level, String logPath) throws IOException { if (level != Level.OFF) { - setPath(logPath, level); - logger.setLevel(level); + if (logPath != null) { + setPath(logPath, level); + } } else { for (Handler h : logger.getHandlers()) { h.close(); @@ -160,6 +161,7 @@ public static void setLevel(Level level, String logPath) throws IOException { } fileHandler = null; } + logger.setLevel(level); } static void setPath(String logPath, Level level) { diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java index 0a19bed7a2c8..390c8b5eb19c 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java @@ -98,6 +98,8 @@ protected boolean removeEldestEntry(Map.Entry> eldes static final String BIGQUERY_ENDPOINT_OVERRIDE_PROPERTY_NAME = "BIGQUERY"; static final String STS_ENDPOINT_OVERRIDE_PROPERTY_NAME = "STS"; static final String OAUTH_ACCESS_TOKEN_PROPERTY_NAME = "OAuthAccessToken"; + static final String GCP_TELEMETRY_PROJECT_ID_PROPERTY_NAME = "gcpTelemetryProjectId"; + static final String GCP_TELEMETRY_CREDENTIALS_PROPERTY_NAME = "gcpTelemetryCredentials"; static final String OAUTH_ACCESS_TOKEN_READONLY_PROPERTY_NAME = "OAuthAccessTokenReadonly"; static final String OAUTH_REFRESH_TOKEN_PROPERTY_NAME = "OAuthRefreshToken"; static final String OAUTH_CLIENT_ID_PROPERTY_NAME = "OAuthClientId"; @@ -162,6 +164,12 @@ protected boolean removeEldestEntry(Map.Entry> eldes static final int DEFAULT_SWA_APPEND_ROW_COUNT_VALUE = 1000; static final String SWA_ACTIVATION_ROW_COUNT_PROPERTY_NAME = "SWA_ActivationRowCount"; static final int DEFAULT_SWA_ACTIVATION_ROW_COUNT_VALUE = 3; + static final String ENABLE_GCP_TRACE_EXPORTER_PROPERTY_NAME = "enableGcpTraceExporter"; + static final boolean DEFAULT_ENABLE_GCP_TRACE_EXPORTER_VALUE = false; + static final String ENABLE_GCP_LOG_EXPORTER_PROPERTY_NAME = "enableGcpLogExporter"; + static final boolean DEFAULT_ENABLE_GCP_LOG_EXPORTER_VALUE = false; + static final String USE_GLOBAL_OTEL_PROPERTY_NAME = "useGlobalOpenTelemetry"; + static final boolean DEFAULT_USE_GLOBAL_OTEL_VALUE = false; private static final BigQueryJdbcCustomLogger LOG = new BigQueryJdbcCustomLogger(BigQueryJdbcUrlUtility.class.getName()); static final String FILTER_TABLES_ON_DEFAULT_DATASET_PROPERTY_NAME = @@ -612,6 +620,32 @@ protected boolean removeEldestEntry(Map.Entry> eldes .setDescription( "Reason for the request, which is passed as the x-goog-request-reason" + " header.") + .build(), + BigQueryConnectionProperty.newBuilder() + .setName(ENABLE_GCP_TRACE_EXPORTER_PROPERTY_NAME) + .setDescription( + "Enables or disables GCP OpenTelemetry Trace exporter. Disabled by default.") + .setDefaultValue(String.valueOf(DEFAULT_ENABLE_GCP_TRACE_EXPORTER_VALUE)) + .build(), + BigQueryConnectionProperty.newBuilder() + .setName(ENABLE_GCP_LOG_EXPORTER_PROPERTY_NAME) + .setDescription( + "Enables or disables GCP OpenTelemetry Log exporter. Disabled by default.") + .setDefaultValue(String.valueOf(DEFAULT_ENABLE_GCP_LOG_EXPORTER_VALUE)) + .build(), + BigQueryConnectionProperty.newBuilder() + .setName(GCP_TELEMETRY_CREDENTIALS_PROPERTY_NAME) + .setDescription("Path or raw JSON of credentials for OTel exporter.") + .build(), + BigQueryConnectionProperty.newBuilder() + .setName(GCP_TELEMETRY_PROJECT_ID_PROPERTY_NAME) + .setDescription("GCP Project ID for OTel exporter.") + .build(), + BigQueryConnectionProperty.newBuilder() + .setName(USE_GLOBAL_OTEL_PROPERTY_NAME) + .setDescription( + "Enables usage of the Global OpenTelemetry instance when true. Default is false.") + .setDefaultValue(String.valueOf(DEFAULT_USE_GLOBAL_OTEL_VALUE)) .build()))); private static final List NETWORK_PROPERTIES = diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatement.java index c153b1935bb6..5dcd6dae4f2a 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatement.java @@ -90,6 +90,16 @@ private int getParameterCount(String query) { @Override public ResultSet executeQuery() throws SQLException { + LOG.finest("++enter++"); + checkClosed(); + return BigQueryJdbcOpenTelemetry.withTracing( + "BigQueryPreparedStatement.executeQuery", + this.connection, + this.currentQuery, + () -> executeQueryImpl()); + } + + private ResultSet executeQueryImpl() throws SQLException { logQueryExecutionStart(this.currentQuery); try { QueryJobConfiguration.Builder jobConfiguration = getJobConfig(this.currentQuery); @@ -104,6 +114,16 @@ public ResultSet executeQuery() throws SQLException { @Override public long executeLargeUpdate() throws SQLException { + LOG.finest("++enter++"); + checkClosed(); + return BigQueryJdbcOpenTelemetry.withTracing( + "BigQueryPreparedStatement.executeLargeUpdate", + this.connection, + this.currentQuery, + () -> executeLargeUpdateImpl()); + } + + private long executeLargeUpdateImpl() throws SQLException { logQueryExecutionStart(this.currentQuery); try { QueryJobConfiguration.Builder jobConfiguration = getJobConfig(this.currentQuery); @@ -123,6 +143,16 @@ public int executeUpdate() throws SQLException { @Override public boolean execute() throws SQLException { + LOG.finest("++enter++"); + checkClosed(); + return BigQueryJdbcOpenTelemetry.withTracing( + "BigQueryPreparedStatement.execute", + this.connection, + this.currentQuery, + () -> executeImpl()); + } + + private boolean executeImpl() throws SQLException { logQueryExecutionStart(this.currentQuery); try { QueryJobConfiguration.Builder jobConfiguration = getJobConfig(this.currentQuery); diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 0d4d94175b8d..0466f7302689 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -60,6 +60,13 @@ import com.google.common.util.concurrent.Uninterruptibles; import io.grpc.Status; import io.grpc.StatusRuntimeException; +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.Tracer; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; import java.lang.ref.ReferenceQueue; import java.sql.Connection; import java.sql.ResultSet; @@ -239,7 +246,8 @@ private BigQuerySettings generateBigQuerySettings() { @Override public ResultSet executeQuery(String sql) throws SQLException { checkClosed(); - return executeQueryImpl(sql); + return BigQueryJdbcOpenTelemetry.withTracing( + "BigQueryStatement.executeQuery", this.connection, sql, () -> executeQueryImpl(sql)); } private ResultSet executeQueryImpl(String sql) throws SQLException { @@ -262,7 +270,11 @@ private ResultSet executeQueryImpl(String sql) throws SQLException { @Override public long executeLargeUpdate(String sql) throws SQLException { checkClosed(); - return executeLargeUpdateImpl(sql); + return BigQueryJdbcOpenTelemetry.withTracing( + "BigQueryStatement.executeLargeUpdate", + this.connection, + sql, + () -> executeLargeUpdateImpl(sql)); } private long executeLargeUpdateImpl(String sql) throws SQLException { @@ -298,7 +310,8 @@ int checkUpdateCount(long updateCount) { @Override public boolean execute(String sql) throws SQLException { checkClosed(); - return executeImpl(sql); + return BigQueryJdbcOpenTelemetry.withTracing( + "BigQueryStatement.execute", this.connection, sql, () -> executeImpl(sql)); } private boolean executeImpl(String sql) throws SQLException { @@ -871,90 +884,85 @@ Thread populateArrowBufferedQueue( LOG.finer("++enter++"); Runnable arrowStreamProcessor = - () -> { - 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 (ApiException e) { - if (e.getStatusCode().getCode() == 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); - } + Context.current() + .wrap( + () -> + 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; } - } catch (InterruptedException e) { + ArrowRecordBatch currentBatch = response.getArrowRecordBatch(); + Uninterruptibles.putUninterruptibly( + arrowBatchWrapperBlockingQueue, BigQueryArrowBatchWrapper.of(currentBatch)); + rowsRead += response.getRowCount(); + } + break; + } catch (ApiException e) { + if (e.getStatusCode().getCode() == 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.WARNING, - "\n" + Thread.currentThread().getName() + " Interrupted @ arrowStreamProcessor", + Level.SEVERE, + "\n" + + Thread.currentThread().getName() + + " Interrupted @ arrowStreamProcessor, max retries exceeded", e); enqueueError(arrowBatchWrapperBlockingQueue, e); - Thread.currentThread().interrupt(); - } catch (Exception e) { - if (e.getCause() instanceof InterruptedException - || Thread.currentThread().isInterrupted()) { - LOG.log( - Level.WARNING, - "\n" + Thread.currentThread().getName() + " Interrupted @ arrowStreamProcessor", - e); - enqueueError(arrowBatchWrapperBlockingQueue, e); - Thread.currentThread().interrupt(); - } else { - LOG.log( - Level.WARNING, - "\n" + Thread.currentThread().getName() + " Error @ arrowStreamProcessor", - e); - enqueueError(arrowBatchWrapperBlockingQueue, e); - } - } finally { // logic needed for graceful shutdown - enqueueEndOfStream(arrowBatchWrapperBlockingQueue); + break; } - }; + retryCount++; + LOG.warning( + "Connection interrupted during arrow stream read, retrying. attempt: %d", retryCount); + Thread.sleep(RETRY_DELAY_MS); + } + } - Thread populateBufferWorker = JDBC_THREAD_FACTORY.newThread(arrowStreamProcessor); - populateBufferWorker.start(); - return populateBufferWorker; + } 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 */ @@ -1124,58 +1132,21 @@ Thread runNextPageTaskAsync( JobId jobId, BlockingQueue> rpcResponseQueue, BlockingQueue bigQueryFieldValueListWrapperBlockingQueue) { - LOG.finer("++enter++"); + LOG.finest("++enter++"); // parse and put the first page in the pageCache before the other pages are parsed from the RPC // calls populateFirstPage(result, rpcResponseQueue); - // This thread makes the RPC calls and paginates Runnable nextPageTask = - () -> { - String currentPageToken = firstPageToken; - TableResult currentResults = result; - TableId destinationTable = null; - if (firstPageToken != null) { - destinationTable = getDestinationTable(jobId); - } - - try { - while (currentPageToken != null) { - // do not process further pages and shutdown - if (Thread.currentThread().isInterrupted() || queryTaskExecutor.isShutdown()) { - LOG.warning( - "%s Interrupted @ runNextPageTaskAsync", Thread.currentThread().getName()); - break; - } - - long startTime = System.nanoTime(); - currentResults = - this.bigQuery.listTableData( - destinationTable, - TableDataListOption.pageSize(querySettings.getMaxResultPerPage()), - TableDataListOption.pageToken(currentPageToken)); - - currentPageToken = currentResults.getNextPageToken(); - // this will be parsed asynchronously without blocking the current - // thread - Uninterruptibles.putUninterruptibly(rpcResponseQueue, Tuple.of(currentResults, true)); - LOG.fine( - "Fetched %d results from the server in %d ms.", - querySettings.getMaxResultPerPage(), - (int) ((System.nanoTime() - startTime) / 1000000)); - } - } catch (Exception ex) { - Uninterruptibles.putUninterruptibly( - bigQueryFieldValueListWrapperBlockingQueue, - BigQueryFieldValueListWrapper.ofError(new BigQueryJdbcRuntimeException(ex))); - } finally { - // this will stop the parseDataTask as well when the pagination - // completes - Uninterruptibles.putUninterruptibly(rpcResponseQueue, Tuple.of(null, false)); - } - // We cannot do queryTaskExecutor.shutdownNow() here as populate buffer method may not - // have finished processing the records and even that will be interrupted - }; + Context.current() + .wrap( + () -> + fetchNextPages( + firstPageToken, + jobId, + rpcResponseQueue, + bigQueryFieldValueListWrapperBlockingQueue, + result)); Thread nextPageWorker = JDBC_THREAD_FACTORY.newThread(nextPageTask); nextPageWorker.start(); @@ -1194,73 +1165,11 @@ Thread parseAndPopulateRpcDataAsync( LOG.finer("++enter++"); Runnable populateBufferRunnable = - () -> { // producer thread populating the buffer - try { - Iterable fieldValueLists; - // as we have to process the first page - boolean hasRows = true; - while (hasRows) { - try { - Tuple nextPageTuple = rpcResponseQueue.take(); - if (nextPageTuple.x() != null) { - fieldValueLists = nextPageTuple.x().getValues(); - } else { - fieldValueLists = null; - } - hasRows = nextPageTuple.y(); - - } catch (InterruptedException e) { - LOG.log(Level.WARNING, "\n" + Thread.currentThread().getName() + " Interrupted", e); - // Thread might get interrupted while calling the Cancel method, which is - // expected, so logging this instead of throwing the exception back - break; - } - - if (Thread.currentThread().isInterrupted() - || queryTaskExecutor.isShutdown() - || fieldValueLists == null) { - // do not process further pages and shutdown (outerloop) - break; - } - - long startTime = System.nanoTime(); - long results = 0; - for (FieldValueList fieldValueList : fieldValueLists) { - - if (Thread.currentThread().isInterrupted() || queryTaskExecutor.isShutdown()) { - // do not process further pages and shutdown (inner loop) - break; - } - Uninterruptibles.putUninterruptibly( - bigQueryFieldValueListWrapperBlockingQueue, - BigQueryFieldValueListWrapper.of(schema.getFields(), fieldValueList)); - results += 1; - } - LOG.fine( - "Processed %d results in %d ms.", - results, (int) ((System.nanoTime() - startTime) / 1000000)); - } - - } catch (Exception ex) { - if (ex.getCause() instanceof InterruptedException - || Thread.currentThread().isInterrupted()) { - LOG.log( - Level.WARNING, - "\n" + Thread.currentThread().getName() + " Interrupted @ populateBufferAsync", - ex); - enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); - Thread.currentThread().interrupt(); - } else { - LOG.log( - Level.WARNING, - "\n" + Thread.currentThread().getName() + " Error @ populateBufferAsync", - ex); - enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); - } - } finally { - enqueueBufferEndOfStream(bigQueryFieldValueListWrapperBlockingQueue); - } - }; + Context.current() + .wrap( + () -> + parseAndPopulateRpcData( + schema, bigQueryFieldValueListWrapperBlockingQueue, rpcResponseQueue)); Thread populateBufferWorker = JDBC_THREAD_FACTORY.newThread(populateBufferRunnable); populateBufferWorker.start(); @@ -1514,13 +1423,32 @@ public void clearBatch() { @Override public int[] executeBatch() throws SQLException { + LOG.finest("++enter++"); + return BigQueryJdbcOpenTelemetry.withTracing( + "BigQueryStatement.executeBatch", + this.connection, + null, + () -> { + Span span = Span.current(); + span.setAttribute( + BigQueryJdbcOpenTelemetry.DB_STATEMENT_COUNT_KEY, this.batchQueries.size()); + span.setAttribute( + AttributeKey.stringArrayKey(BigQueryJdbcOpenTelemetry.DB_BATCH_STATEMENTS_KEY), + new ArrayList<>(this.batchQueries)); + + String combinedQueries = String.join("", this.batchQueries); + + return executeBatchImpl(combinedQueries); + }); + } + + private int[] executeBatchImpl(String combinedQueries) throws SQLException { int[] result = new int[this.batchQueries.size()]; if (this.batchQueries.isEmpty()) { return result; } try { - String combinedQueries = String.join("", this.batchQueries); QueryJobConfiguration.Builder jobConfiguration = getJobConfig(combinedQueries); jobConfiguration.setPriority(QueryJobConfiguration.Priority.BATCH); runQuery(combinedQueries, jobConfiguration.build()); @@ -1690,4 +1618,124 @@ private void enqueueBufferError(BlockingQueue que private void enqueueBufferEndOfStream(BlockingQueue queue) { Uninterruptibles.putUninterruptibly(queue, BigQueryFieldValueListWrapper.of(null, null, true)); } + + private void fetchNextPages( + String firstPageToken, + JobId jobId, + BlockingQueue> rpcResponseQueue, + BlockingQueue bigQueryFieldValueListWrapperBlockingQueue, + TableResult result) { + SpanContext parentSpanContext = Span.current().getSpanContext(); + String currentPageToken = firstPageToken; + TableResult currentResults = result; + TableId destinationTable = null; + if (firstPageToken != null) { + destinationTable = getDestinationTable(jobId); + } + + try { + Tracer tracer = this.connection.getTracer(); + while (currentPageToken != null) { + if (Thread.currentThread().isInterrupted() || queryTaskExecutor.isShutdown()) { + LOG.warning("%s Interrupted @ runNextPageTaskAsync", Thread.currentThread().getName()); + break; + } + + 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); + + long startTime = System.nanoTime(); + currentResults = + this.bigQuery.listTableData( + destinationTable, + TableDataListOption.pageSize(querySettings.getMaxResultPerPage()), + TableDataListOption.pageToken(currentPageToken)); + + long duration = (System.nanoTime() - startTime) / 1000000; + paginationSpan.setAttribute("db.pagination.duration_ms", duration); + paginationSpan.setAttribute( + "db.pagination.rows_fetched", querySettings.getMaxResultPerPage()); + + currentPageToken = currentResults.getNextPageToken(); + Uninterruptibles.putUninterruptibly(rpcResponseQueue, Tuple.of(currentResults, true)); + LOG.fine( + "Fetched %d results from the server in %d ms.", + querySettings.getMaxResultPerPage(), (int) duration); + } catch (Exception e) { + paginationSpan.recordException(e); + paginationSpan.setStatus(io.opentelemetry.api.trace.StatusCode.ERROR, e.getMessage()); + throw e; + } finally { + paginationSpan.end(); + } + } + } catch (Exception ex) { + Uninterruptibles.putUninterruptibly( + bigQueryFieldValueListWrapperBlockingQueue, + BigQueryFieldValueListWrapper.ofError(new BigQueryJdbcRuntimeException(ex))); + } finally { + Uninterruptibles.putUninterruptibly(rpcResponseQueue, Tuple.of(null, false)); + } + } + + private void parseAndPopulateRpcData( + Schema schema, + BlockingQueue bigQueryFieldValueListWrapperBlockingQueue, + BlockingQueue> rpcResponseQueue) { + try { + Iterable fieldValueLists; + boolean hasRows = true; + while (hasRows) { + try { + Tuple nextPageTuple = rpcResponseQueue.take(); + if (nextPageTuple.x() != null) { + fieldValueLists = nextPageTuple.x().getValues(); + } else { + fieldValueLists = null; + } + hasRows = nextPageTuple.y(); + + } catch (InterruptedException e) { + LOG.log(Level.WARNING, "\n" + Thread.currentThread().getName() + " Interrupted", e); + break; + } + + if (Thread.currentThread().isInterrupted() + || queryTaskExecutor.isShutdown() + || fieldValueLists == null) { + break; + } + + long startTime = System.nanoTime(); + long results = 0; + for (FieldValueList fieldValueList : fieldValueLists) { + + if (Thread.currentThread().isInterrupted() || queryTaskExecutor.isShutdown()) { + break; + } + Uninterruptibles.putUninterruptibly( + bigQueryFieldValueListWrapperBlockingQueue, + BigQueryFieldValueListWrapper.of(schema.getFields(), fieldValueList)); + results += 1; + } + LOG.fine( + "Processed %d results in %d ms.", + results, (int) ((System.nanoTime() - startTime) / 1000000)); + } + + } catch (Exception ex) { + LOG.log( + Level.WARNING, + "\n" + Thread.currentThread().getName() + " Error @ populateBufferAsync", + ex); + enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); + } finally { + enqueueBufferEndOfStream(bigQueryFieldValueListWrapperBlockingQueue); + } + } } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/DataSource.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/DataSource.java index d7de950a935f..6a7ccaf487fc 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/DataSource.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/DataSource.java @@ -21,6 +21,7 @@ import com.google.common.base.Joiner; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import io.opentelemetry.api.OpenTelemetry; import java.io.PrintWriter; import java.sql.Connection; import java.sql.DriverManager; @@ -58,6 +59,8 @@ public class DataSource implements javax.sql.DataSource { private String logLevel; private Boolean enableSession; private String logPath; + private String gcpTelemetryProjectId; + private String gcpTelemetryCredentials; private Integer oAuthType; private String oAuthServiceAcctEmail; private String oAuthPvtKeyPath; @@ -115,6 +118,12 @@ public class DataSource implements javax.sql.DataSource { private String privateServiceConnect; private Long connectionPoolSize; private Long listenerPoolSize; + private boolean enableGcpTraceExporter = + BigQueryJdbcUrlUtility.DEFAULT_ENABLE_GCP_TRACE_EXPORTER_VALUE; + private boolean enableGcpLogExporter = + BigQueryJdbcUrlUtility.DEFAULT_ENABLE_GCP_LOG_EXPORTER_VALUE; + private OpenTelemetry customOpenTelemetry; + private boolean useGlobalOpenTelemetry = BigQueryJdbcUrlUtility.DEFAULT_USE_GLOBAL_OTEL_VALUE; // Make sure the JDBC driver class is loaded. static { @@ -131,6 +140,12 @@ public class DataSource implements javax.sql.DataSource { .put(BigQueryJdbcUrlUtility.PROJECT_ID_PROPERTY_NAME, DataSource::setProjectId) .put(BigQueryJdbcUrlUtility.DEFAULT_DATASET_PROPERTY_NAME, DataSource::setDefaultDataset) .put(BigQueryJdbcUrlUtility.LOCATION_PROPERTY_NAME, DataSource::setLocation) + .put( + BigQueryJdbcUrlUtility.GCP_TELEMETRY_PROJECT_ID_PROPERTY_NAME, + DataSource::setGcpTelemetryProjectId) + .put( + BigQueryJdbcUrlUtility.GCP_TELEMETRY_CREDENTIALS_PROPERTY_NAME, + DataSource::setGcpTelemetryCredentials) .put( BigQueryJdbcUrlUtility.ENABLE_HTAPI_PROPERTY_NAME, (ds, val) -> @@ -332,6 +347,24 @@ public class DataSource implements javax.sql.DataSource { .put( BigQueryJdbcUrlUtility.LISTENER_POOL_SIZE_PROPERTY_NAME, (ds, val) -> ds.setListenerPoolSize(Long.parseLong(val))) + .put( + BigQueryJdbcUrlUtility.ENABLE_GCP_TRACE_EXPORTER_PROPERTY_NAME, + (ds, val) -> + ds.setEnableGcpTraceExporter( + BigQueryJdbcUrlUtility.convertIntToBoolean( + val, BigQueryJdbcUrlUtility.ENABLE_GCP_TRACE_EXPORTER_PROPERTY_NAME))) + .put( + BigQueryJdbcUrlUtility.ENABLE_GCP_LOG_EXPORTER_PROPERTY_NAME, + (ds, val) -> + ds.setEnableGcpLogExporter( + BigQueryJdbcUrlUtility.convertIntToBoolean( + val, BigQueryJdbcUrlUtility.ENABLE_GCP_LOG_EXPORTER_PROPERTY_NAME))) + .put( + BigQueryJdbcUrlUtility.USE_GLOBAL_OTEL_PROPERTY_NAME, + (ds, val) -> + ds.setUseGlobalOpenTelemetry( + BigQueryJdbcUrlUtility.convertIntToBoolean( + val, BigQueryJdbcUrlUtility.USE_GLOBAL_OTEL_PROPERTY_NAME))) .build(); public static DataSource fromUrl(String url) { @@ -389,7 +422,11 @@ public Connection getConnection() throws SQLException { throw new BigQueryJdbcException( "The URL " + getURL() + " is invalid. Please specify a valid Connection URL. "); } - return DriverManager.getConnection(getURL(), createProperties()); + Properties props = createProperties(); + if (this.customOpenTelemetry != null) { + props.put("customOpenTelemetry", this.customOpenTelemetry); + } + return DriverManager.getConnection(getURL(), props); } Properties createProperties() { @@ -635,6 +672,21 @@ Properties createProperties() { BigQueryJdbcUrlUtility.LISTENER_POOL_SIZE_PROPERTY_NAME, String.valueOf(this.listenerPoolSize)); } + if (this.enableGcpTraceExporter) { + connectionProperties.setProperty( + BigQueryJdbcUrlUtility.ENABLE_GCP_TRACE_EXPORTER_PROPERTY_NAME, + String.valueOf(this.enableGcpTraceExporter)); + } + if (this.enableGcpLogExporter) { + connectionProperties.setProperty( + BigQueryJdbcUrlUtility.ENABLE_GCP_LOG_EXPORTER_PROPERTY_NAME, + String.valueOf(this.enableGcpLogExporter)); + } + if (this.useGlobalOpenTelemetry) { + connectionProperties.setProperty( + BigQueryJdbcUrlUtility.USE_GLOBAL_OTEL_PROPERTY_NAME, + String.valueOf(this.useGlobalOpenTelemetry)); + } return connectionProperties; } @@ -768,6 +820,38 @@ public void setListenerPoolSize(Long listenerPoolSize) { this.listenerPoolSize = listenerPoolSize; } + public boolean getEnableGcpTraceExporter() { + return enableGcpTraceExporter; + } + + public void setEnableGcpTraceExporter(boolean enableGcpTraceExporter) { + this.enableGcpTraceExporter = enableGcpTraceExporter; + } + + public boolean getEnableGcpLogExporter() { + return enableGcpLogExporter; + } + + public void setEnableGcpLogExporter(boolean enableGcpLogExporter) { + this.enableGcpLogExporter = enableGcpLogExporter; + } + + public OpenTelemetry getCustomOpenTelemetry() { + return customOpenTelemetry; + } + + public void setCustomOpenTelemetry(OpenTelemetry customOpenTelemetry) { + this.customOpenTelemetry = customOpenTelemetry; + } + + public boolean getUseGlobalOpenTelemetry() { + return useGlobalOpenTelemetry; + } + + public void setUseGlobalOpenTelemetry(boolean useGlobalOpenTelemetry) { + this.useGlobalOpenTelemetry = useGlobalOpenTelemetry; + } + public void setHighThroughputMinTableSize(Integer highThroughputMinTableSize) { if (highThroughputMinTableSize != null) { validateNonNegative( @@ -837,6 +921,22 @@ public void setLogPath(String logPath) { this.logPath = logPath; } + public String getGcpTelemetryProjectId() { + return gcpTelemetryProjectId; + } + + public void setGcpTelemetryProjectId(String gcpTelemetryProjectId) { + this.gcpTelemetryProjectId = gcpTelemetryProjectId; + } + + public String getGcpTelemetryCredentials() { + return gcpTelemetryCredentials; + } + + public void setGcpTelemetryCredentials(String gcpTelemetryCredentials) { + this.gcpTelemetryCredentials = gcpTelemetryCredentials; + } + public String getUniverseDomain() { return universeDomain != null ? universeDomain diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/OpenTelemetryJulHandler.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/OpenTelemetryJulHandler.java new file mode 100644 index 000000000000..f652d91b3f2e --- /dev/null +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/OpenTelemetryJulHandler.java @@ -0,0 +1,191 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.jdbc; + +import com.google.cloud.logging.LogEntry; +import com.google.cloud.logging.Logging; +import com.google.cloud.logging.Payload; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.baggage.Baggage; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.logs.LogRecordBuilder; +import io.opentelemetry.api.logs.Logger; +import io.opentelemetry.api.logs.Severity; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.context.Context; +import java.time.Instant; +import java.util.Collections; +import java.util.logging.Handler; +import java.util.logging.Level; +import java.util.logging.LogRecord; +import java.util.regex.Pattern; + +/** + * Custom logging handler that bridges java.util.logging records to OpenTelemetry or Google Cloud + * Logging. Extracts TraceId, SpanId, and Connection UUID from context. + */ +public class OpenTelemetryJulHandler extends Handler { + private static final Pattern UNSAFE_LOG_CHARACTERS = Pattern.compile("[^a-zA-Z0-9./_-]"); + + public OpenTelemetryJulHandler() { + setLevel(Level.ALL); + } + + @Override + public void publish(LogRecord record) { + if (!isLoggable(record)) { + return; + } + + try { + // Extract connection ID from baggage + String connectionId = + Baggage.fromContext(Context.current()) + .getEntryValue(BigQueryJdbcOpenTelemetry.CONNECTION_ID_BAGGAGE_KEY); + + if (connectionId == null) { + return; + } + + BigQueryJdbcOpenTelemetry.TelemetryConfig config = + BigQueryJdbcOpenTelemetry.getConnectionConfig(connectionId); + if (config == null) { + return; + } + + if (config.useDirectGcpLogging && config.loggingClient != null) { + publishToGcp(record, connectionId, config.loggingClient); + } else if (config.openTelemetry != null) { + publishToOTel(record, connectionId, config.openTelemetry); + } + } catch (Throwable t) { + // Ignore exceptions to prevent breaking application logging or other handlers + } + } + + private void publishToGcp(LogRecord record, String connectionId, Logging loggingClient) { + Context context = Context.current(); + SpanContext spanContext = Span.fromContext(context).getSpanContext(); + String traceId = spanContext.isValid() ? spanContext.getTraceId() : null; + String spanId = spanContext.isValid() ? spanContext.getSpanId() : null; + + String logId = record.getLoggerName(); + if (logId == null) { + logId = BigQueryJdbcOpenTelemetry.INSTRUMENTATION_SCOPE_NAME; + } else { + logId = UNSAFE_LOG_CHARACTERS.matcher(logId).replaceAll("_"); + } + + LogEntry.Builder builder = + LogEntry.newBuilder(Payload.StringPayload.of(formatMessage(record))) + .setSeverity(mapGcpSeverity(record.getLevel())) + .setTimestamp(record.getMillis()) + .setLogName(logId); + + if (traceId != null) { + builder.setTrace(traceId); + } + if (spanId != null) { + builder.setSpanId(spanId); + } + if (connectionId != null) { + builder.addLabel(BigQueryJdbcOpenTelemetry.CONNECTION_ID_BAGGAGE_KEY, connectionId); + } + + loggingClient.write(Collections.singleton(builder.build())); + } + + private com.google.cloud.logging.Severity mapGcpSeverity(Level level) { + if (level == Level.SEVERE) return com.google.cloud.logging.Severity.ERROR; + if (level == Level.WARNING) return com.google.cloud.logging.Severity.WARNING; + if (level == Level.INFO) return com.google.cloud.logging.Severity.INFO; + if (level == Level.CONFIG) return com.google.cloud.logging.Severity.INFO; + if (level == Level.FINE) return com.google.cloud.logging.Severity.DEBUG; + return com.google.cloud.logging.Severity.DEBUG; + } + + private void publishToOTel(LogRecord record, String connectionId, OpenTelemetry openTelemetry) { + String loggerName = record.getLoggerName(); + Logger logger = + openTelemetry + .getLogsBridge() + .get( + loggerName != null + ? loggerName + : BigQueryJdbcOpenTelemetry.INSTRUMENTATION_SCOPE_NAME); + + LogRecordBuilder builder = + logger + .logRecordBuilder() + .setBody(formatMessage(record)) + .setSeverity(mapSeverity(record.getLevel())) + .setTimestamp(Instant.ofEpochMilli(record.getMillis())) + .setContext(Context.current()); + + if (connectionId != null) { + builder.setAttribute( + AttributeKey.stringKey(BigQueryJdbcOpenTelemetry.CONNECTION_ID_BAGGAGE_KEY), + connectionId); + } + + builder.emit(); + } + + private Severity mapSeverity(Level level) { + if (level == Level.SEVERE) return Severity.ERROR; + if (level == Level.WARNING) return Severity.WARN; + if (level == Level.INFO) return Severity.INFO; + if (level == Level.CONFIG) return Severity.INFO; + if (level == Level.FINE) return Severity.DEBUG; + if (level == Level.FINER) return Severity.TRACE; + if (level == Level.FINEST) return Severity.TRACE; + return Severity.TRACE; + } + + private String formatMessage(LogRecord record) { + String message = record.getMessage(); + Object[] params = record.getParameters(); + if (params != null && params.length > 0) { + try { + return java.text.MessageFormat.format(message, params); + } catch (IllegalArgumentException e) { + return message; + } + } + return message; + } + + @Override + public void flush() { + for (BigQueryJdbcOpenTelemetry.TelemetryConfig config : + BigQueryJdbcOpenTelemetry.getRegisteredConfigs()) { + if (config.useDirectGcpLogging && config.loggingClient != null) { + try { + config.loggingClient.flush(); + } catch (Exception e) { + // Ignore failures during flush to protect other connections + } + } + } + } + + @Override + public void close() throws SecurityException { + flush(); + } +} diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowStructTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowStructTest.java index c3e15a3457d9..e1afd98f9f99 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowStructTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowStructTest.java @@ -60,7 +60,6 @@ import java.util.List; import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; -import java.util.stream.Stream; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.JsonStringHashMap; import org.apache.arrow.vector.util.Text; @@ -141,8 +140,7 @@ BIGNUMERIC, new BigDecimal("11.2657"), new BigDecimal("33.4657")), LocalDateTime.parse("2023-03-30T11:15:19.820227")), arrowArraySchemaAndValue( GEOGRAPHY, new Text("POINT(-122 47)"), new Text("POINT(-122 48)")), - arrowArraySchemaAndValue( - BYTES, Stream.of("one", "two").map(String::getBytes).toArray(byte[][]::new))); + arrowArraySchemaAndValue(BYTES, "one".getBytes(), "two".getBytes())); List orderedSchemas = schemaAndValues.stream().map(Tuple::x).collect(Collectors.toList()); diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java index 94cde20fa400..8dc23053818e 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java @@ -16,25 +16,36 @@ package com.google.cloud.bigquery.jdbc; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.when; import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider; import com.google.api.gax.rpc.HeaderProvider; import com.google.api.gax.rpc.TransportChannelProvider; +import com.google.auth.oauth2.GoogleCredentials; import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode; import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.storage.v1.BigQueryReadClient; import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient; +import com.google.cloud.logging.Logging; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.sdk.testing.junit5.OpenTelemetryExtension; +import io.opentelemetry.sdk.trace.data.SpanData; import java.io.IOException; import java.io.InputStream; import java.sql.SQLException; +import java.util.List; import java.util.Optional; import java.util.Properties; import java.util.logging.Level; @@ -42,11 +53,16 @@ import java.util.logging.Logger; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; +import org.mockito.MockedStatic; public class BigQueryConnectionTest extends BigQueryJdbcLoggingBaseTest { + @RegisterExtension + static final OpenTelemetryExtension otelTesting = OpenTelemetryExtension.create(); + private static final String DEFAULT_VERSION = "0.0.0"; private static final String DEFAULT_JDBC_TOKEN_VALUE = "Google-BigQuery-JDBC-Driver"; private static final String BASE_URL = @@ -467,6 +483,27 @@ public void testIsReadOnlyTokenProvided(String readonlyProp, boolean expectedIsR } } + @Test + public void testConnect_withCustomOpenTelemetry_usesCustomInstance() throws Exception { + DataSource ds = DataSource.fromUrl(BASE_URL); + ds.setCustomOpenTelemetry(otelTesting.getOpenTelemetry()); + + try (BigQueryConnection connection = new BigQueryConnection(BASE_URL, ds)) { + assertNotNull(connection); + assertFalse(connection.isClosed()); + + Tracer tracer = connection.getTracer(); + assertNotNull(tracer); + + Span span = tracer.spanBuilder("custom-otel-span").startSpan(); + span.end(); + + List spans = otelTesting.getSpans(); + assertEquals(1, spans.size()); + assertEquals("custom-otel-span", spans.get(0).getName()); + } + } + @Test public void testConnectionPropertiesLoggingAndMasking() throws IOException, SQLException { Logger rootLogger = BigQueryJdbcRootLogger.getRootLogger(); @@ -499,6 +536,135 @@ public void testConnectionPropertiesLoggingAndMasking() throws IOException, SQLE } } + @ParameterizedTest( + name = + "Case {index}: custom={0}, global={1}, trace={2}, log={3} -> expectTrace={4}, expectLog={5}") + @CsvSource({ + // hasCustom, useGlobal, enableTrace, enableLog, expectTrace, expectLog + "true, true, true, true, CUSTOM, CUSTOM", + "true, false, true, true, CUSTOM, CUSTOM", + "false, true, true, true, GLOBAL, GLOBAL", + "false, true, false, false, GLOBAL, GLOBAL", + "false, false, true, false, DRIVER_MANAGED, NONE", + "false, false, false, true, NONE, DRIVER_MANAGED", + "false, false, true, true, DRIVER_MANAGED, DRIVER_MANAGED", + "false, false, false, false, NONE, NONE" + }) + public void testOpenTelemetryPrecedenceHierarchy( + boolean hasCustom, + boolean useGlobal, + boolean enableTrace, + boolean enableLog, + String expectTrace, + String expectLog) + throws Exception { + + DataSource ds = DataSource.fromUrl(BASE_URL); + ds.setUseGlobalOpenTelemetry(useGlobal); + ds.setEnableGcpTraceExporter(enableTrace); + ds.setEnableGcpLogExporter(enableLog); + + OpenTelemetry mockCustomOtel = mock(OpenTelemetry.class); + OpenTelemetry mockGlobalOtel = mock(OpenTelemetry.class); + OpenTelemetry mockDriverManagedOtel = mock(OpenTelemetry.class); + Logging mockLogging = mock(Logging.class); + when(mockCustomOtel.getTracer(anyString())).thenReturn(mock(Tracer.class)); + when(mockGlobalOtel.getTracer(anyString())).thenReturn(mock(Tracer.class)); + when(mockDriverManagedOtel.getTracer(anyString())).thenReturn(mock(Tracer.class)); + + if (hasCustom) { + ds.setCustomOpenTelemetry(mockCustomOtel); + } + + try (MockedStatic mockedOtel = + mockStatic(BigQueryJdbcOpenTelemetry.class); + MockedStatic mockedAuth = + mockStatic(BigQueryJdbcOAuthUtility.class); + MockedStatic mockedCreds = mockStatic(GoogleCredentials.class)) { + + mockedCreds + .when(GoogleCredentials::getApplicationDefault) + .thenReturn(mock(GoogleCredentials.class)); + + // Mock parseOAuthProperties to always return ADC type to bypass validation + mockedAuth + .when(() -> BigQueryJdbcOAuthUtility.parseOAuthProperties(any(), anyString())) + .thenAnswer( + invocation -> { + java.util.Map props = new java.util.HashMap<>(); + props.put( + BigQueryJdbcUrlUtility.OAUTH_TYPE_PROPERTY_NAME, + "APPLICATION_DEFAULT_CREDENTIALS"); + return props; + }); + + mockedAuth + .when(() -> BigQueryJdbcOAuthUtility.getCredentials(any(), any(), any(), any())) + .thenReturn(mock(GoogleCredentials.class)); + + mockedOtel + .when( + () -> + BigQueryJdbcOpenTelemetry.createLoggingClient( + anyBoolean(), any(), any(), any(), any())) + .thenReturn(mockLogging); + + // Stub getOpenTelemetry to return the expected mock based on inputs + mockedOtel + .when( + () -> + BigQueryJdbcOpenTelemetry.getOpenTelemetry( + eq(useGlobal), + eq(enableTrace), + eq(enableLog), + hasCustom ? eq(mockCustomOtel) : isNull(), + any(), + any())) + .thenAnswer( + invocation -> { + if (hasCustom) return mockCustomOtel; + if (useGlobal) return mockGlobalOtel; + if (enableTrace || enableLog) return mockDriverManagedOtel; + return OpenTelemetry.noop(); + }); + + try (BigQueryConnection connection = new BigQueryConnection(BASE_URL, ds)) { + + boolean shouldBeRegistered = enableLog || hasCustom || useGlobal; + + if (!shouldBeRegistered) { + mockedOtel.verify( + () -> + BigQueryJdbcOpenTelemetry.registerConnection( + anyString(), any(), any(), anyBoolean()), + never()); + } else { + final OpenTelemetry expectedOtelInstance; + if ("CUSTOM".equals(expectTrace) || "CUSTOM".equals(expectLog)) { + expectedOtelInstance = mockCustomOtel; + } else if ("GLOBAL".equals(expectTrace) || "GLOBAL".equals(expectLog)) { + expectedOtelInstance = mockGlobalOtel; + } else if ("DRIVER_MANAGED".equals(expectTrace) || "DRIVER_MANAGED".equals(expectLog)) { + expectedOtelInstance = mockDriverManagedOtel; + } else { + expectedOtelInstance = OpenTelemetry.noop(); + } + + boolean expectUseDirectGcp = "DRIVER_MANAGED".equals(expectLog); + Logging expectedLogClient = expectUseDirectGcp ? mockLogging : null; + + mockedOtel.verify( + () -> + BigQueryJdbcOpenTelemetry.registerConnection( + anyString(), + eq(expectedOtelInstance), + eq(expectedLogClient), + eq(expectUseDirectGcp))); + } + } + } + } + @Test public void testWrapperMethods() throws Exception { try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) { 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 58a5a7212066..c279ef78d34f 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 @@ -34,6 +34,14 @@ import com.google.cloud.bigquery.BigQuery.RoutineListOption; import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.jdbc.BigQueryJdbcTypeMappings.ColumnTypeInfo; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; +import io.opentelemetry.sdk.testing.junit5.OpenTelemetryExtension; +import io.opentelemetry.sdk.trace.data.SpanData; import java.io.IOException; import java.io.InputStream; import java.sql.DatabaseMetaData; @@ -43,18 +51,30 @@ import java.sql.Statement; import java.sql.Types; import java.util.*; +import java.util.concurrent.BlockingQueue; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; +import java.util.concurrent.LinkedBlockingQueue; import java.util.regex.Pattern; +import java.util.stream.Stream; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.MethodSource; +@TestInstance(TestInstance.Lifecycle.PER_CLASS) public class BigQueryDatabaseMetaDataTest { + @RegisterExtension + public static final OpenTelemetryExtension otelTesting = OpenTelemetryExtension.create(); + private BigQueryConnection bigQueryConnection; private BigQueryDatabaseMetaData dbMetadata; private BigQuery bigqueryClient; @@ -68,6 +88,22 @@ public void setUp() throws SQLException { when(bigQueryConnection.getConnectionUrl()).thenReturn("jdbc:bigquery://test-project"); when(bigQueryConnection.getBigQuery()).thenReturn(bigqueryClient); when(bigQueryConnection.createStatement()).thenReturn(mockStatement); + when(bigQueryConnection.getConnectionId()).thenReturn("test-connection-id"); + when(bigQueryConnection.getTracer()) + .thenReturn( + otelTesting + .getOpenTelemetry() + .getTracer(BigQueryJdbcOpenTelemetry.INSTRUMENTATION_SCOPE_NAME)); + when(bigQueryConnection.getOtelContext()).thenReturn(Context.current()); + + Page datasetPageMock = mock(Page.class); + when(bigqueryClient.listDatasets(anyString(), any())).thenReturn(datasetPageMock); + + Page
tablePageMock = mock(Page.class); + when(bigqueryClient.listTables(any(DatasetId.class), any())).thenReturn(tablePageMock); + + Table mockTable = mock(Table.class); + when(bigqueryClient.getTable(any(TableId.class))).thenReturn(mockTable); dbMetadata = new BigQueryDatabaseMetaData(bigQueryConnection); } @@ -2951,7 +2987,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); @@ -3241,6 +3277,107 @@ public void testGetSQLStateType() throws SQLException { assertEquals(DatabaseMetaData.sqlStateSQL, dbMetadata.getSQLStateType()); } + @ParameterizedTest + @MethodSource("metadataOperationProvider") + public void testMetadataOperation_generatesSpan( + MetadataOperation operation, String expectedSpanName) throws Exception { + operation.run(); + + SpanData span = + OpenTelemetryTestUtility.findSpanByName(otelTesting.getSpans(), expectedSpanName); + OpenTelemetryTestUtility.assertSpanStatus(span, StatusCode.UNSET); + + OpenTelemetryTestUtility.assertSpanHasAttribute( + span, + AttributeKey.stringKey(BigQueryJdbcOpenTelemetry.DB_SYSTEM_KEY), + BigQueryJdbcOpenTelemetry.DB_SYSTEM_VALUE); + OpenTelemetryTestUtility.assertSpanHasAttribute( + span, + AttributeKey.stringKey(BigQueryJdbcOpenTelemetry.DB_CONNECTION_ID_KEY), + "test-connection-id"); + } + + @FunctionalInterface + interface MetadataOperation { + void run() throws SQLException; + } + + Stream metadataOperationProvider() { + return Stream.of( + Arguments.of( + (MetadataOperation) () -> dbMetadata.getCatalogs(), + "BigQueryDatabaseMetaData.getCatalogs"), + Arguments.of( + (MetadataOperation) () -> dbMetadata.getSchemas("catalog", "schema"), + "BigQueryDatabaseMetaData.getSchemas"), + Arguments.of( + (MetadataOperation) + () -> dbMetadata.getTables("catalog", "schema", "table", new String[] {"TABLE"}), + "BigQueryDatabaseMetaData.getTables"), + Arguments.of( + (MetadataOperation) () -> dbMetadata.getColumns("catalog", "schema", "table", "column"), + "BigQueryDatabaseMetaData.getColumns")); + } + + @ParameterizedTest + @MethodSource("asyncMetadataOperationProvider") + public void testAsyncMetadataOperation_createsDetachedLinkedSpan( + AsyncMetadataOperation operation, String expectedSpanName) throws Exception { + BlockingQueue queue = new LinkedBlockingQueue<>(); + + Tracer testTracer = otelTesting.getOpenTelemetry().getTracer("test"); + Span parentSpan = testTracer.spanBuilder("parent-span").startSpan(); + + try (Scope scope = parentSpan.makeCurrent()) { + Thread workerThread = operation.run(queue); + + Assertions.assertNotNull(workerThread, "Worker thread should not be null"); + workerThread.join(); + + OpenTelemetryTestUtility.assertSpanLinkedToParent( + otelTesting.getSpans(), expectedSpanName, parentSpan); + } finally { + parentSpan.end(); + } + } + + @FunctionalInterface + interface AsyncMetadataOperation { + Thread run(BlockingQueue queue) throws Exception; + } + + Stream asyncMetadataOperationProvider() { + return Stream.of( + Arguments.of( + (AsyncMetadataOperation) + (q) -> + dbMetadata.runGetTablesTaskAsync( + "catalog", + "schema", + "table", + new String[] {"TABLE"}, + dbMetadata.defineGetTablesSchema(), + q), + "BigQueryDatabaseMetaData.getTables.background"), + Arguments.of( + (AsyncMetadataOperation) + (q) -> + dbMetadata.runGetColumnsTaskAsync( + "catalog", + "schema", + "table", + "column", + dbMetadata.defineGetColumnsSchema(), + q), + "BigQueryDatabaseMetaData.getColumns.background"), + Arguments.of( + (AsyncMetadataOperation) + (q) -> + dbMetadata.runGetSchemasTaskAsync( + "catalog", "schema", dbMetadata.defineGetSchemasSchema(), q), + "BigQueryDatabaseMetaData.getSchemas.background")); + } + @Test public void testWrapperMethods() throws SQLException { assertTrue(dbMetadata.isWrapperFor(DatabaseMetaData.class)); diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDriverTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDriverTest.java index 2a5ad5c4767c..ec25a1c61185 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDriverTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDriverTest.java @@ -16,7 +16,9 @@ package com.google.cloud.bigquery.jdbc; import static com.google.common.truth.Truth.assertThat; +import static org.mockito.Mockito.mock; +import io.opentelemetry.api.OpenTelemetry; import java.sql.Connection; import java.sql.DriverPropertyInfo; import java.sql.SQLException; @@ -107,6 +109,23 @@ public void testConnectWithInvalidUrlChainsNoException() throws SQLException { assertThat(connection.isClosed()).isFalse(); } + @Test + public void testConnect_withCustomOpenTelemetry_injectsIntoDataSource() throws SQLException { + OpenTelemetry mockOtel = mock(OpenTelemetry.class); + Properties props = new Properties(); + props.put("customOpenTelemetry", mockOtel); + + // Connect using standard URL setup but pass the SDK via Properties + Connection connection = + bigQueryDriver.connect( + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" + + "OAuthType=2;ProjectId=MyBigQueryProject;OAuthAccessToken=redacted;", + props); + + assertThat(connection).isNotNull(); + assertThat(connection.isClosed()).isFalse(); + } + @Test public void testUnknownPropertyWarningIsLogged() throws SQLException { Connection connection = diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcLoggingBaseTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcLoggingBaseTest.java index 31c8cd7f3a5d..6ed3231eed15 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcLoggingBaseTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcLoggingBaseTest.java @@ -30,10 +30,13 @@ public abstract class BigQueryJdbcLoggingBaseTest extends BigQueryJdbcBaseTest { private Handler handler; private Logger logger; private long threadId; + private java.util.logging.Level originalLevel; @BeforeEach public void setUpLogValidator() { logger = BigQueryJdbcRootLogger.getRootLogger(); + originalLevel = logger.getLevel(); + logger.setLevel(java.util.logging.Level.ALL); capturedLogs.clear(); threadId = Thread.currentThread().getId(); handler = @@ -58,6 +61,7 @@ public void close() throws SecurityException {} public void tearDownLogValidator() { if (logger != null && handler != null) { logger.removeHandler(handler); + logger.setLevel(originalLevel); } try { BigQueryJdbcRootLogger.setLevel(java.util.logging.Level.OFF, null); diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetryTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetryTest.java new file mode 100644 index 000000000000..6f07b0e220d5 --- /dev/null +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetryTest.java @@ -0,0 +1,132 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.jdbc; + +import static com.google.common.truth.Truth.assertThat; +import static org.mockito.Mockito.mock; + +import com.google.auth.oauth2.GoogleCredentials; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.trace.Tracer; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +public class BigQueryJdbcOpenTelemetryTest { + + private MockedStatic mockedCredentials; + + @BeforeEach + public void setUp() { + mockedCredentials = Mockito.mockStatic(GoogleCredentials.class); + mockedCredentials + .when(GoogleCredentials::getApplicationDefault) + .thenReturn(mock(GoogleCredentials.class)); + } + + @AfterEach + public void tearDown() { + mockedCredentials.close(); + } + + @Test + public void testGetOpenTelemetry_withCustomSdk_returnsCustom() { + OpenTelemetry mockCustomOtel = mock(OpenTelemetry.class); + + OpenTelemetry result = + BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, false, false, mockCustomOtel, null, null); + + assertThat(result).isSameInstanceAs(mockCustomOtel); + } + + @Test + public void testGetOpenTelemetry_withCustomSdkAndFlags_returnsCustom() { + OpenTelemetry mockCustomOtel = mock(OpenTelemetry.class); + + // Custom SDK always takes precedence over individual flags + OpenTelemetry result = + BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, true, true, mockCustomOtel, null, null); + + assertThat(result).isSameInstanceAs(mockCustomOtel); + } + + @Test + public void testGetOpenTelemetry_noFlags_returnsNoop() { + OpenTelemetry result = + BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, false, false, null, null, null); + + assertThat(result).isSameInstanceAs(OpenTelemetry.noop()); + } + + @Test + public void testGetTracer_respectsScopeName() { + Tracer result = BigQueryJdbcOpenTelemetry.getTracer(OpenTelemetry.noop()); + + assertThat(result).isNotNull(); + } + + @Test + public void testGetOpenTelemetry_cachesSdkInstances() { + OpenTelemetry result1 = + BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, true, false, null, null, "project1"); + OpenTelemetry result2 = + BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, true, false, null, null, "project1"); + + assertThat(result1).isSameInstanceAs(result2); + } + + @Test + public void testGetOpenTelemetry_createsNewInstanceForDifferentKey() { + OpenTelemetry result1 = + BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, true, false, null, null, "project1"); + OpenTelemetry result2 = + BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, true, false, null, null, "project2"); + + assertThat(result1).isNotSameInstanceAs(result2); + } + + @Test + public void testGetOpenTelemetry_createsNewInstanceForDifferentTraceFlag() { + OpenTelemetry result1 = + BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, true, true, null, null, "project1"); + OpenTelemetry result2 = + BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, false, true, null, null, "project1"); + + assertThat(result1).isNotSameInstanceAs(result2); + } + + @Test + public void testGetOpenTelemetry_ignoresEnableLogFlagInCacheKey() { + OpenTelemetry result1 = + BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, true, true, null, null, "project1"); + OpenTelemetry result2 = + BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, true, false, null, null, "project1"); + + assertThat(result1).isSameInstanceAs(result2); + } + + @Test + public void testGetOpenTelemetry_withUseGlobalOTel_returnsGlobal() { + OpenTelemetry result = + BigQueryJdbcOpenTelemetry.getOpenTelemetry(true, false, false, null, null, null); + + assertThat(result).isSameInstanceAs(GlobalOpenTelemetry.get()); + } +} diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java index 674eb0df64e0..752a32300546 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java @@ -31,8 +31,11 @@ import com.google.api.gax.rpc.ApiException; import com.google.api.gax.rpc.StatusCode; import com.google.cloud.ServiceOptions; +import com.google.cloud.Tuple; import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.BigQuery.QueryResultsOption; +import com.google.cloud.bigquery.BigQuery.TableDataListOption; +import com.google.cloud.bigquery.BigQueryException; import com.google.cloud.bigquery.BigQueryOptions; import com.google.cloud.bigquery.Field; import com.google.cloud.bigquery.FieldList; @@ -58,30 +61,50 @@ import com.google.cloud.bigquery.storage.v1.ReadSession; import com.google.common.collect.ImmutableList; import com.google.common.collect.Maps; +import io.opentelemetry.api.baggage.Baggage; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; +import io.opentelemetry.sdk.testing.junit5.OpenTelemetryExtension; +import io.opentelemetry.sdk.trace.data.SpanData; import java.io.IOException; import java.sql.ResultSet; import java.sql.SQLException; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.UUID; import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingDeque; +import java.util.stream.Stream; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; import org.junit.jupiter.params.provider.ValueSource; import org.mockito.ArgumentCaptor; import org.mockito.Mockito; +@TestInstance(TestInstance.Lifecycle.PER_CLASS) public class BigQueryStatementTest { + @RegisterExtension + public static final OpenTelemetryExtension otelTesting = OpenTelemetryExtension.create(); + private BigQueryConnection bigQueryConnection; private static final String PROJECT = "project"; @@ -133,9 +156,39 @@ private Job getJobMock( return job; } + private TableResult setupMockQueryResults(JobId jobId, StatementType type, Long affectedRows) + throws Exception { + doReturn(true).when(bigQueryConnection).getUseStatelessQueryMode(); + TableResult tableResultMock = mock(TableResult.class); + doReturn(jobId).when(tableResultMock).getJobId(); + doReturn(Schema.of()).when(tableResultMock).getSchema(); + doReturn(tableResultMock) + .when(bigquery) + .queryWithTimeout(any(QueryJobConfiguration.class), any(), any()); + + Job jobMock = getJobMock(tableResultMock, null, type); + if (affectedRows != null) { + JobStatistics.QueryStatistics stats = (JobStatistics.QueryStatistics) jobMock.getStatistics(); + doReturn(affectedRows).when(stats).getNumDmlAffectedRows(); + } + doReturn(jobMock).when(bigquery).getJob(any(JobId.class)); + doReturn(jobMock).when(jobMock).waitFor(); + + Job dryRunJobMock = getJobMock(null, null, type); + doReturn(dryRunJobMock).when(bigquery).create(any(JobInfo.class)); + return tableResultMock; + } + @BeforeEach public void setUp() throws IOException, SQLException { bigQueryConnection = mock(BigQueryConnection.class); + doReturn( + otelTesting + .getOpenTelemetry() + .getTracer(BigQueryJdbcOpenTelemetry.INSTRUMENTATION_SCOPE_NAME)) + .when(bigQueryConnection) + .getTracer(); + doReturn(Context.current()).when(bigQueryConnection).getOtelContext(); rpcFactoryMock = mock(BigQueryRpcFactory.class); bigquery = mock(BigQuery.class); bigQueryConnection.bigQuery = bigquery; @@ -143,6 +196,7 @@ public void setUp() throws IOException, SQLException { jobId = JobId.newBuilder().setJob(jobIdVal).build(); doReturn(bigquery).when(bigQueryConnection).getBigQuery(); + doReturn("test-connection-id").when(bigQueryConnection).getConnectionId(); doReturn(10L).when(bigQueryConnection).getJobTimeoutInSeconds(); doReturn(10L).when(bigQueryConnection).getMaxBytesBilled(); doReturn(LABELS).when(bigQueryConnection).getLabels(); @@ -452,21 +506,11 @@ public void testCloseCancelsJob() throws SQLException, InterruptedException { } @Test - public void testCancelWithJoblessQuery() throws SQLException, InterruptedException { - doReturn(true).when(bigQueryConnection).getUseStatelessQueryMode(); + public void testCancelWithJoblessQuery() throws Exception { + TableResult tableResultMock = setupMockQueryResults(null, StatementType.SELECT, null); BigQueryStatement joblessStatement = new BigQueryStatement(bigQueryConnection); BigQueryStatement joblessStatementSpy = Mockito.spy(joblessStatement); - TableResult tableResultMock = mock(TableResult.class); - doReturn(null).when(tableResultMock).getJobId(); - - doReturn(tableResultMock) - .when(bigquery) - .queryWithTimeout(any(QueryJobConfiguration.class), any(), any()); - - Job dryRunJobMock = getJobMock(null, null, StatementType.SELECT); - doReturn(dryRunJobMock).when(bigquery).create(any(JobInfo.class)); - BigQueryJsonResultSet resultSetMock = mock(BigQueryJsonResultSet.class); doReturn(resultSetMock) .when(joblessStatementSpy) @@ -486,6 +530,126 @@ public void testCancelWithJoblessQuery() throws SQLException, InterruptedExcepti verify(bigquery, Mockito.never()).cancel(any(JobId.class)); } + @Test + public void testFetchNextPages_addsLinkToParent() throws Exception { + Tracer testTracer = otelTesting.getOpenTelemetry().getTracer("test"); + Span parentSpan = testTracer.spanBuilder("parent-span").startSpan(); + + try (Scope scope = parentSpan.makeCurrent()) { + + BlockingQueue> rpcResponseQueue = new LinkedBlockingDeque<>(); + BlockingQueue bigQueryFieldValueListWrapperBlockingQueue = + new LinkedBlockingDeque<>(); + TableResult mockResult = mock(TableResult.class); + JobId mockJobId = JobId.of("job"); + + Job mockJob = mock(Job.class); + QueryJobConfiguration realConfig = + QueryJobConfiguration.newBuilder("SELECT 1") + .setDestinationTable(TableId.of("project", "dataset", "table")) + .build(); + doReturn(mockJob).when(bigquery).getJob(any(JobId.class)); + doReturn(realConfig).when(mockJob).getConfiguration(); + + TableResult mockNextResult = mock(TableResult.class); + doReturn(mockNextResult) + .when(bigquery) + .listTableData( + any(TableId.class), any(TableDataListOption.class), any(TableDataListOption.class)); + doReturn(null).when(mockNextResult).getNextPageToken(); + + Thread workerThread = + bigQueryStatement.runNextPageTaskAsync( + mockResult, + "token", + mockJobId, + rpcResponseQueue, + bigQueryFieldValueListWrapperBlockingQueue); + + Assertions.assertNotNull(workerThread, "Worker thread should not be null"); + workerThread.join(); + + OpenTelemetryTestUtility.assertSpanLinkedToParent( + otelTesting.getSpans(), "BigQueryStatement.pagination", parentSpan); + } finally { + parentSpan.end(); + } + } + + @ParameterizedTest + @MethodSource("statementOperationProvider") + public void testExecuteOperation_generatesSpan( + StatementOperation operation, + String expectedSpanName, + StatementType type, + Map, Object> expectedAttributes) + throws Exception { + setupMockQueryResults(JobId.of("job"), type, 1L); + operation.run(); + + SpanData span = + OpenTelemetryTestUtility.findSpanByName(otelTesting.getSpans(), expectedSpanName); + OpenTelemetryTestUtility.assertSpanStatus(span, io.opentelemetry.api.trace.StatusCode.UNSET); + + if (expectedAttributes != null) { + for (Map.Entry, Object> entry : expectedAttributes.entrySet()) { + OpenTelemetryTestUtility.assertSpanHasAttribute( + span, (AttributeKey) entry.getKey(), entry.getValue()); + } + } + + OpenTelemetryTestUtility.assertSpanHasAttribute( + span, + AttributeKey.stringKey(BigQueryJdbcOpenTelemetry.DB_SYSTEM_KEY), + BigQueryJdbcOpenTelemetry.DB_SYSTEM_VALUE); + OpenTelemetryTestUtility.assertSpanHasAttribute( + span, + AttributeKey.stringKey(BigQueryJdbcOpenTelemetry.DB_CONNECTION_ID_KEY), + "test-connection-id"); + } + + Stream statementOperationProvider() { + return Stream.of( + Arguments.of( + (StatementOperation) () -> bigQueryStatement.executeQuery("SELECT 1"), + "BigQueryStatement.executeQuery", + StatementType.SELECT, + Collections.singletonMap(AttributeKey.stringKey("db.statement"), "SELECT 1")), + Arguments.of( + (StatementOperation) () -> bigQueryStatement.execute("SELECT 1"), + "BigQueryStatement.execute", + StatementType.SELECT, + Collections.singletonMap(AttributeKey.stringKey("db.statement"), "SELECT 1")), + Arguments.of( + (StatementOperation) + () -> bigQueryStatement.executeLargeUpdate("UPDATE table SET col = 1"), + "BigQueryStatement.executeLargeUpdate", + StatementType.UPDATE, + Collections.singletonMap( + AttributeKey.stringKey("db.statement"), "UPDATE table SET col = 1")), + Arguments.of( + (StatementOperation) + () -> { + bigQueryStatement.addBatch("UPDATE table SET col = 1"); + bigQueryStatement.executeBatch(); + }, + "BigQueryStatement.executeBatch", + StatementType.UPDATE, + new HashMap, Object>() { + { + put(AttributeKey.longKey("db.statement.count"), 1L); + put( + AttributeKey.stringArrayKey("db.batch.statements"), + Collections.singletonList("UPDATE table SET col = 1; ")); + } + })); + } + + @FunctionalInterface + interface StatementOperation { + void run() throws Exception; + } + @ParameterizedTest @ValueSource(booleans = {true, false}) public void testGetStatementType(boolean isReadOnlyTokenUsed) throws Exception { @@ -677,6 +841,72 @@ public void testUseReadAPI_ZeroPageSizeDivisionByZeroSafeguard() throws SQLExcep assertThat(useReadApi).isTrue(); // ratio = 500 / 1 = 500 > 2 -> true } + @Test + public void testExecute_registersException() throws Exception { + // Mock bigquery to throw a backend exception + BigQueryException expectedException = new BigQueryException(500, "Backend Error"); + Mockito.doThrow(expectedException) + .when(bigquery) + .queryWithTimeout(Mockito.any(QueryJobConfiguration.class), Mockito.any(), Mockito.any()); + + BigQueryStatement spiedStatement = Mockito.spy(bigQueryStatement); + + // Execute and expect the exception to be propagated to JDBC + Assertions.assertThrows(SQLException.class, () -> spiedStatement.executeQuery("SELECT 1")); + + // Retrieve the exported span from OTel extension + List spans = otelTesting.getSpans(); + SpanData span = + OpenTelemetryTestUtility.findSpanByName(spans, "BigQueryStatement.executeQuery"); + + // Assert that the span recorded the error correctly + OpenTelemetryTestUtility.assertSpanStatus(span, io.opentelemetry.api.trace.StatusCode.ERROR); + OpenTelemetryTestUtility.assertSpanHasException(span, BigQueryJdbcException.class); + } + + @Test + public void testExecute_propagatesContextAndBaggage() throws Exception { + // Mock bigquery using thenAnswer to hook into the call and assert Context/Baggage + Mockito.doAnswer( + invocation -> { + // This code runs on the execution thread during the SDK call + String connectionIdBaggage = + Baggage.current() + .getEntryValue(BigQueryJdbcOpenTelemetry.CONNECTION_ID_BAGGAGE_KEY); + assertEquals("test-connection-id", connectionIdBaggage); + + Span currentSpan = Span.current(); + assertTrue(currentSpan.getSpanContext().isValid()); + + // Return a mock TableResult to allow the execution to proceed + TableResult tableResultMock = mock(TableResult.class); + doReturn(jobId).when(tableResultMock).getJobId(); + doReturn(Schema.of()).when(tableResultMock).getSchema(); + return tableResultMock; + }) + .when(bigquery) + .queryWithTimeout(Mockito.any(QueryJobConfiguration.class), Mockito.any(), Mockito.any()); + + BigQueryStatement spiedStatement = Mockito.spy(bigQueryStatement); + + // Setup connection mocks to allow the statement to execute successfully + doReturn(true).when(bigQueryConnection).getUseStatelessQueryMode(); + Job dryRunJobMock = getJobMock(null, null, StatementType.SELECT); + doReturn(dryRunJobMock).when(bigquery).create(Mockito.any(JobInfo.class)); + + BigQueryJsonResultSet resultSetMock = mock(BigQueryJsonResultSet.class); + doReturn(resultSetMock) + .when(spiedStatement) + .processJsonResultSet(Mockito.any(TableResult.class), Mockito.any()); + + // Execute query + spiedStatement.executeQuery("SELECT 1"); + + // Verify the SDK call actually occurred + verify(bigquery) + .queryWithTimeout(Mockito.any(QueryJobConfiguration.class), Mockito.any(), Mockito.any()); + } + @Test public void testWrapperMethods() throws SQLException { assertTrue(bigQueryStatement.isWrapperFor(java.sql.Statement.class)); diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/OpenTelemetryJulHandlerTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/OpenTelemetryJulHandlerTest.java new file mode 100644 index 000000000000..e923dd011778 --- /dev/null +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/OpenTelemetryJulHandlerTest.java @@ -0,0 +1,134 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.jdbc; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +import com.google.cloud.logging.LogEntry; +import com.google.cloud.logging.Logging; +import com.google.cloud.logging.Payload; +import io.opentelemetry.api.baggage.Baggage; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.logs.Severity; +import io.opentelemetry.context.Scope; +import io.opentelemetry.sdk.logs.data.LogRecordData; +import io.opentelemetry.sdk.testing.junit5.OpenTelemetryExtension; +import java.util.List; +import java.util.logging.Logger; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.mockito.ArgumentCaptor; + +public class OpenTelemetryJulHandlerTest { + + @RegisterExtension + static final OpenTelemetryExtension otelTesting = OpenTelemetryExtension.create(); + + private static final Logger logger = Logger.getLogger("com.google.cloud.bigquery"); + private java.util.logging.Level originalLevel; + + @BeforeEach + public void setUp() { + originalLevel = logger.getLevel(); + logger.setLevel(java.util.logging.Level.ALL); + BigQueryJdbcOpenTelemetry.ensureGlobalHandlerAttached(); + } + + @AfterEach + public void tearDown() { + BigQueryJdbcOpenTelemetry.unregisterConnection("test-uuid"); + BigQueryJdbcOpenTelemetry.unregisterConnection("wrong-uuid"); + BigQueryJdbcOpenTelemetry.unregisterConnection("gcp-uuid"); + logger.setLevel(originalLevel); + } + + @Test + public void testPublishToOTel() { + BigQueryJdbcOpenTelemetry.registerConnection( + "test-uuid", otelTesting.getOpenTelemetry(), null, false); + + BigQueryConnection mockConnection = mock(BigQueryConnection.class); + Baggage baggage = + Baggage.builder() + .put(BigQueryJdbcOpenTelemetry.CONNECTION_ID_BAGGAGE_KEY, "test-uuid") + .build(); + try (Scope scope = baggage.makeCurrent(); + BigQueryJdbcMdc.MdcCloseable mdcScope = BigQueryJdbcMdc.registerInstance("test-uuid")) { + logger.info("Test message"); + } + + List logs = otelTesting.getLogRecords(); + assertEquals(1, logs.size()); + LogRecordData log = logs.get(0); + assertEquals("Test message", log.getBody().asString()); + assertEquals(Severity.INFO, log.getSeverity()); + assertEquals( + "test-uuid", + log.getAttributes() + .get(AttributeKey.stringKey(BigQueryJdbcOpenTelemetry.CONNECTION_ID_BAGGAGE_KEY))); + } + + @Test + public void testPublishWithFiltering() { + // Register for "test-uuid" + BigQueryJdbcOpenTelemetry.registerConnection( + "test-uuid", otelTesting.getOpenTelemetry(), null, false); + + // Log with WRONG connection ID + Baggage baggage = + Baggage.builder() + .put(BigQueryJdbcOpenTelemetry.CONNECTION_ID_BAGGAGE_KEY, "wrong-uuid") + .build(); + try (Scope scope = baggage.makeCurrent()) { + logger.info("Test message"); + } + + List logs = otelTesting.getLogRecords(); + assertTrue(logs.isEmpty()); // Should be filtered out because "wrong-uuid" has no config + } + + @Test + public void testPublishToGcp() { + Logging loggingClient = mock(Logging.class); + BigQueryJdbcOpenTelemetry.registerConnection( + "gcp-uuid", otelTesting.getOpenTelemetry(), loggingClient, true); + + BigQueryConnection mockConnection = mock(BigQueryConnection.class); + Baggage baggage = + Baggage.builder() + .put(BigQueryJdbcOpenTelemetry.CONNECTION_ID_BAGGAGE_KEY, "gcp-uuid") + .build(); + try (Scope scope = baggage.makeCurrent(); + BigQueryJdbcMdc.MdcCloseable mdcScope = BigQueryJdbcMdc.registerInstance("gcp-uuid")) { + logger.info("Test message"); + } + + ArgumentCaptor> captor = ArgumentCaptor.forClass(Iterable.class); + verify(loggingClient).write(captor.capture()); + + Iterable entries = captor.getValue(); + LogEntry entry = entries.iterator().next(); + + assertEquals("Test message", ((Payload.StringPayload) entry.getPayload()).getData()); + assertEquals(com.google.cloud.logging.Severity.INFO, entry.getSeverity()); + } +} diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/OpenTelemetryTestUtility.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/OpenTelemetryTestUtility.java new file mode 100644 index 000000000000..df8f025a0f80 --- /dev/null +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/OpenTelemetryTestUtility.java @@ -0,0 +1,141 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.jdbc; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.sdk.trace.data.SpanData; +import java.util.List; +import java.util.Optional; + +public class OpenTelemetryTestUtility { + + /** + * Asserts that a span with the given name exists in the provided list and returns it. + * + * @param spans The list of exported spans. + * @param spanName The name of the span to find. + * @return The found SpanData. + * @throws AssertionError if the span is not found. + */ + public static SpanData findSpanByName(List spans, String spanName) { + Optional span = spans.stream().filter(s -> s.getName().equals(spanName)).findFirst(); + assertTrue(span.isPresent(), "Span with name '" + spanName + "' not found"); + return span.get(); + } + + /** + * Asserts that a span exists in the list. + * + * @param spans The list of exported spans. + * @param spanName The name of the span to find. + */ + public static void assertSpanExists(List spans, String spanName) { + findSpanByName(spans, spanName); + } + + /** + * Asserts that a span has a specific attribute key and value. + * + * @param span The span to check. + * @param key The attribute key. + * @param expectedValue The expected value of the attribute. + * @param The type of the attribute value. + */ + public static void assertSpanHasAttribute( + SpanData span, AttributeKey key, T expectedValue) { + T actualValue = span.getAttributes().get(key); + assertNotNull( + actualValue, "Attribute '" + key.getKey() + "' not found on span '" + span.getName() + "'"); + assertEquals( + expectedValue, + actualValue, + "Attribute '" + key.getKey() + "' value mismatch on span '" + span.getName() + "'"); + } + + /** + * Asserts the status of a span. + * + * @param span The span to check. + * @param expectedStatus The expected StatusCode. + */ + public static void assertSpanStatus(SpanData span, StatusCode expectedStatus) { + assertEquals( + expectedStatus, + span.getStatus().getStatusCode(), + "Status code mismatch on span '" + span.getName() + "'"); + } + + /** + * Asserts that an exception of a specific type was recorded on the span. + * + * @param span The span to check. + * @param exceptionClass The class of the expected exception. + */ + public static void assertSpanHasException( + SpanData span, Class exceptionClass) { + boolean found = + span.getEvents().stream() + .anyMatch( + event -> + event.getName().equals("exception") + && exceptionClass + .getName() + .equals( + event + .getAttributes() + .get(AttributeKey.stringKey("exception.type")))); + assertTrue( + found, + "Exception of type " + + exceptionClass.getName() + + " not found in events of span '" + + span.getName() + + "'"); + } + + /** + * Asserts that a span is linked to a parent span. + * + * @param spans The list of exported spans. + * @param spanName The name of the span that should have the link. + * @param parentSpan The parent span it should be linked to. + */ + public static void assertSpanLinkedToParent( + List spans, String spanName, Span parentSpan) { + boolean found = + spans.stream() + .anyMatch( + span -> + span.getName().equals(spanName) + && span.getLinks().stream() + .anyMatch( + link -> + link.getSpanContext() + .getTraceId() + .equals(parentSpan.getSpanContext().getTraceId()) + && link.getSpanContext() + .getSpanId() + .equals(parentSpan.getSpanContext().getSpanId()))); + assertTrue(found, "Span " + spanName + " not found or not linked to parent"); + } +} diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandlerTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandlerTest.java index fda5112703fd..a1149733b113 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandlerTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandlerTest.java @@ -45,12 +45,15 @@ public class PerConnectionFileHandlerTest { private PerConnectionFileHandler handler; private BigQueryConnection mockConnection; + private Level originalLevel; @BeforeEach public void setUp() { handler = new PerConnectionFileHandler(tempDir.toString(), Level.INFO); mockConnection = Mockito.mock(BigQueryConnection.class); BigQueryJdbcMdc.clear(); + originalLevel = BigQueryJdbcRootLogger.getRootLogger().getLevel(); + BigQueryJdbcRootLogger.getRootLogger().setLevel(Level.ALL); } @AfterEach @@ -59,6 +62,7 @@ public void tearDown() { handler.close(); } BigQueryJdbcMdc.clear(); + BigQueryJdbcRootLogger.getRootLogger().setLevel(originalLevel); } private Optional findLogFile(String suffix) throws IOException { diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITAuthTests.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITAuthTests.java index 0877553e42c0..2a3a6bfac799 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITAuthTests.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITAuthTests.java @@ -24,15 +24,12 @@ import com.google.auth.oauth2.GoogleCredentials; import com.google.cloud.ServiceOptions; import com.google.gson.JsonObject; -import com.google.gson.JsonParser; import java.io.ByteArrayInputStream; import java.io.File; import java.io.IOException; import java.io.InputStream; -import java.io.InputStreamReader; import java.nio.charset.StandardCharsets; import java.nio.file.Files; -import java.nio.file.Paths; import java.sql.Connection; import java.sql.DriverManager; import java.sql.ResultSet; @@ -48,25 +45,6 @@ public class ITAuthTests extends ITBase { static final String PROJECT_ID = ServiceOptions.getDefaultProjectId(); - private JsonObject getAuthJson() throws IOException { - final String secret = requireEnvVar("SA_SECRET"); - JsonObject authJson; - // Supporting both formats of SA_SECRET: - // - Local runs can point to a json file - // - Cloud Build has JSON value - try { - InputStream stream = Files.newInputStream(Paths.get(secret)); - InputStreamReader reader = new InputStreamReader(stream); - authJson = JsonParser.parseReader(reader).getAsJsonObject(); - } catch (IOException e) { - authJson = JsonParser.parseString(secret).getAsJsonObject(); - } - assertTrue(authJson.has("client_email")); - assertTrue(authJson.has("private_key")); - assertTrue(authJson.has("project_id")); - return authJson; - } - private void validateConnection(String connection_uri) throws SQLException { Connection connection = DriverManager.getConnection(connection_uri); assertNotNull(connection); diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBase.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBase.java index 5b4d36fac4fe..098dcdf75c03 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBase.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBase.java @@ -17,12 +17,23 @@ package com.google.cloud.bigquery.jdbc.it; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import com.google.auth.oauth2.GoogleCredentials; import com.google.cloud.ServiceOptions; import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.BigQueryOptions; import com.google.cloud.bigquery.QueryJobConfiguration; import com.google.cloud.bigquery.jdbc.BigQueryJdbcBaseTest; +import com.google.gson.JsonObject; +import com.google.gson.JsonParser; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Paths; import java.sql.Connection; import java.sql.ResultSet; import java.sql.SQLException; @@ -291,6 +302,31 @@ protected static String requireEnvVar(String varName) { return value; } + protected static JsonObject getAuthJson() throws IOException { + final String secret = requireEnvVar("SA_SECRET"); + JsonObject authJson; + // Supporting both formats of SA_SECRET: + // - Local runs can point to a json file + // - Cloud Build has JSON value + try { + InputStream stream = Files.newInputStream(Paths.get(secret)); + InputStreamReader reader = new InputStreamReader(stream); + authJson = JsonParser.parseReader(reader).getAsJsonObject(); + } catch (IOException e) { + authJson = JsonParser.parseString(secret).getAsJsonObject(); + } + assertTrue(authJson.has("client_email")); + assertTrue(authJson.has("private_key")); + assertTrue(authJson.has("project_id")); + return authJson; + } + + protected static GoogleCredentials getCredentials() throws IOException { + JsonObject authJson = getAuthJson(); + return GoogleCredentials.fromStream( + new ByteArrayInputStream(authJson.toString().getBytes(StandardCharsets.UTF_8))); + } + protected int resultSetRowCount(ResultSet resultSet) throws SQLException { int rowCount = 0; while (resultSet.next()) { diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITOpenTelemetryTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITOpenTelemetryTest.java new file mode 100644 index 000000000000..792ae5071841 --- /dev/null +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITOpenTelemetryTest.java @@ -0,0 +1,364 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.jdbc.it; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.google.api.gax.core.FixedCredentialsProvider; +import com.google.api.gax.paging.Page; +import com.google.auth.oauth2.GoogleCredentials; +import com.google.cloud.ServiceOptions; +import com.google.cloud.bigquery.jdbc.BigQueryConnection; +import com.google.cloud.bigquery.jdbc.DataSource; +import com.google.cloud.logging.LogEntry; +import com.google.cloud.logging.Logging; +import com.google.cloud.logging.LoggingOptions; +import com.google.cloud.trace.v1.TraceServiceClient; +import com.google.cloud.trace.v1.TraceServiceSettings; +import com.google.devtools.cloudtrace.v1.Trace; +import com.google.devtools.cloudtrace.v1.TraceSpan; +import com.google.gson.JsonObject; +import java.io.File; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.List; +import org.junit.jupiter.api.Test; + +public class ITOpenTelemetryTest extends ITBase { + + private static final String PROJECT_ID = ServiceOptions.getDefaultProjectId(); + private static final String CONNECTION_URL = connectionUrl; + + @Test + public void testExecute_withOpenTelemetryGcpExporter() throws Exception { + + // Step 1: Connect with GCP Exporters enabled via DataSource + DataSource ds = DataSource.fromUrl(CONNECTION_URL); + ds.setEnableGcpTraceExporter(true); + ds.setEnableGcpLogExporter(true); + ds.setLogLevel("5"); // Triggers FINE log generation + ds.setGcpTelemetryProjectId(PROJECT_ID); + ds.setEnableHighThroughputAPI(false); + ds.setMaxResults(50L); // Forces small page size (50) to trigger pagination + + String connectionUuid = null; + + try (Connection connection = ds.getConnection(); + Statement statement = connection.createStatement()) { + + // Retrieve the Connection UUID programmatically + BigQueryConnection bqConnection = connection.unwrap(BigQueryConnection.class); + connectionUuid = bqConnection.getConnectionId(); + assertNotNull(connectionUuid, "Connection UUID should be generated"); + + // Execute an in-memory array query (scans 0 bytes, extremely fast) and force pagination (3 + // pages) + String paginationQuery = "SELECT * FROM UNNEST(GENERATE_ARRAY(1, 150)) AS id;"; + try (ResultSet paginatedRs = statement.executeQuery(paginationQuery)) { + while (paginatedRs.next()) { + // Drain the result set to trigger pagination fetches + } + } + } + + // Step 2: Retrieve and assert logs, harvesting the TraceId + String traceId = verifyAndFetchLogs(connectionUuid); + + // Step 3: Query Cloud Trace and assert parent-child hierarchy + Trace trace = verifyAndFetchTrace(traceId); + + boolean foundParentExecuteQuery = false; + boolean foundChildSdkSpans = false; + boolean foundPaginationSpans = false; + long parentSpanId = 0; + + for (TraceSpan span : trace.getSpansList()) { + String spanName = span.getName(); + if (spanName.equals("BigQueryStatement.executeQuery")) { + foundParentExecuteQuery = true; + parentSpanId = span.getSpanId(); + } + } + + assertTrue( + foundParentExecuteQuery, + "Traces must contain JDBC parent span 'BigQueryStatement.executeQuery'"); + + // Verify that we captured child spans or linked pagination spans + for (TraceSpan span : trace.getSpansList()) { + if (span.getParentSpanId() == parentSpanId && parentSpanId != 0) { + foundChildSdkSpans = true; + } + if (span.getName().equals("BigQueryStatement.pagination")) { + foundPaginationSpans = true; + } + } + + assertTrue(foundPaginationSpans, "OTel pagination must generate pagination spans"); + assertTrue( + foundChildSdkSpans, + "OTel context must propagate parent to downstream pagination child spans"); + } + + @Test + public void testExecute_withErrorCorrelation() throws Exception { + + // Step 1: Connect with GCP Exporters enabled via DataSource + DataSource ds = DataSource.fromUrl(CONNECTION_URL); + ds.setEnableGcpTraceExporter(true); + ds.setEnableGcpLogExporter(true); + ds.setLogLevel("5"); // Triggers FINE log generation + ds.setGcpTelemetryProjectId(PROJECT_ID); + + String connectionUuid = null; + + try (Connection connection = ds.getConnection(); + Statement statement = connection.createStatement()) { + + // Retrieve the Connection UUID programmatically + BigQueryConnection bqConnection = connection.unwrap(BigQueryConnection.class); + connectionUuid = bqConnection.getConnectionId(); + assertNotNull(connectionUuid, "Connection UUID should be generated"); + + // Execute a query designed to fail instantly due to syntax error (compiler-level failure) + assertThrows(SQLException.class, () -> statement.executeQuery("SELECT * FROM;")); + } + + // Step 2: Retrieve and assert logs, harvesting the TraceId + String traceId = verifyAndFetchLogs(connectionUuid); + + // Step 3: Query Cloud Trace and assert span status is ERROR + Trace trace = verifyAndFetchTrace(traceId); + + boolean foundParentExecuteQuery = false; + + for (TraceSpan span : trace.getSpansList()) { + String spanName = span.getName(); + if (spanName.equals("BigQueryStatement.executeQuery")) { + foundParentExecuteQuery = true; + } + } + + assertTrue( + foundParentExecuteQuery, + "Traces must contain JDBC parent span 'BigQueryStatement.executeQuery'"); + } + + @Test + public void testExecute_withCustomCredentialsJson() throws Exception { + JsonObject authJson = getAuthJson(); + DataSource ds = DataSource.fromUrl(CONNECTION_URL); + ds.setEnableGcpTraceExporter(true); + ds.setGcpTelemetryProjectId(PROJECT_ID); + ds.setGcpTelemetryCredentials(authJson.toString()); + + verifyTraceDelivery(ds); + } + + @Test + public void testExecute_withCustomCredentialsFilePath() throws Exception { + JsonObject authJson = getAuthJson(); + File tempFile = File.createTempFile("auth", ".json"); + tempFile.deleteOnExit(); + Files.write(tempFile.toPath(), authJson.toString().getBytes(StandardCharsets.UTF_8)); + + DataSource ds = DataSource.fromUrl(CONNECTION_URL); + ds.setEnableGcpTraceExporter(true); + ds.setGcpTelemetryProjectId(PROJECT_ID); + ds.setGcpTelemetryCredentials(tempFile.getAbsolutePath()); + + verifyTraceDelivery(ds); + } + + @Test + public void testExecute_withHttpProtocol() throws Exception { + JsonObject authJson = getAuthJson(); + System.setProperty("otel.exporter.otlp.protocol", "http/protobuf"); + + try { + DataSource ds = DataSource.fromUrl(CONNECTION_URL); + ds.setEnableGcpTraceExporter(true); + ds.setGcpTelemetryProjectId(PROJECT_ID); + ds.setGcpTelemetryCredentials(authJson.toString()); + + verifyTraceDelivery(ds); + } finally { + System.clearProperty("otel.exporter.otlp.protocol"); + } + } + + @Test + public void testExecute_withGrpcProtocol() throws Exception { + JsonObject authJson = getAuthJson(); + System.setProperty("otel.exporter.otlp.protocol", "grpc"); + + try { + DataSource ds = DataSource.fromUrl(CONNECTION_URL); + ds.setEnableGcpTraceExporter(true); + ds.setGcpTelemetryProjectId(PROJECT_ID); + ds.setGcpTelemetryCredentials(authJson.toString()); + + verifyTraceDelivery(ds); + } finally { + System.clearProperty("otel.exporter.otlp.protocol"); + } + } + + private void verifyTraceDelivery(DataSource ds) throws Exception { + ds.setEnableGcpLogExporter(true); + ds.setLogLevel("5"); + + String connectionUuid = null; + try (Connection connection = ds.getConnection(); + Statement statement = connection.createStatement()) { + + BigQueryConnection bqConnection = connection.unwrap(BigQueryConnection.class); + connectionUuid = bqConnection.getConnectionId(); + + String query = "SELECT 1;"; + try (ResultSet rs = statement.executeQuery(query)) { + assertTrue(rs.next()); + } + } + + String traceId = verifyAndFetchLogs(connectionUuid); + Trace trace = verifyAndFetchTrace(traceId); + assertNotNull(trace, "Trace must be found"); + } + + private String verifyAndFetchLogs(String connectionUuid) throws Exception { + GoogleCredentials credentials = getCredentials(); + + try (Logging logging = + LoggingOptions.newBuilder() + .setProjectId(PROJECT_ID) + .setCredentials(credentials) + .build() + .getService()) { + String filter = + "logName:\"projects/" + + PROJECT_ID + + "/logs/com.google.cloud.bigquery\" AND labels.\"jdbc.connection_id\"=\"" + + connectionUuid + + "\""; + + List entries = fetchLogsWithRetry(logging, filter); + assertFalse(entries.isEmpty(), "Telemetry logs should be exported to GCP"); + + LogEntry sampleEntry = entries.get(0); + String traceId = sampleEntry.getTrace(); + String hexSpanId = sampleEntry.getSpanId(); + + assertNotNull(traceId, "Log entry must contain TraceId"); + assertNotNull(hexSpanId, "Log entry must contain SpanId"); + + // Verify Connection UUID label correlation on all entries + for (LogEntry entry : entries) { + assertEquals(connectionUuid, entry.getLabels().get("jdbc.connection_id")); + } + + return traceId; + } + } + + private Trace verifyAndFetchTrace(String traceId) throws Exception { + String hexTraceId = traceId; + if (traceId.contains("/traces/")) { + hexTraceId = traceId.substring(traceId.lastIndexOf("/traces/") + 8); + } + + GoogleCredentials credentials = getCredentials(); + + TraceServiceSettings settings = + TraceServiceSettings.newBuilder() + .setCredentialsProvider(FixedCredentialsProvider.create(credentials)) + .build(); + + try (TraceServiceClient traceClient = TraceServiceClient.create(settings)) { + Trace trace = fetchTraceWithRetry(traceClient, PROJECT_ID, hexTraceId); + assertNotNull(trace, "Trace must be found in Cloud Trace API: " + hexTraceId); + return trace; + } + } + + private T pollWithRetry(java.util.concurrent.Callable task) throws InterruptedException { + int attempts = 0; + int maxAttempts = 30; // 30 attempts * 500ms = 15 seconds max delay + long delayMs = 500; // 500ms linear polling + + while (attempts < maxAttempts) { + attempts++; + try { + T result = task.call(); + if (result != null) { + return result; + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Test execution interrupted", e); + } catch (Exception e) { + // Ignore exceptions during remote lookup and retry + } + if (attempts < maxAttempts) { + Thread.sleep(delayMs); + } + } + return null; + } + + private List fetchLogsWithRetry(Logging logging, String filter) + throws InterruptedException { + List result = + pollWithRetry( + () -> { + Page entriesPage = + logging.listLogEntries( + Logging.EntryListOption.filter(filter), Logging.EntryListOption.pageSize(50)); + List entries = new ArrayList<>(); + entriesPage.iterateAll().forEach(entries::add); + return entries.isEmpty() ? null : entries; + }); + return result != null ? result : new ArrayList<>(); + } + + private Trace fetchTraceWithRetry( + TraceServiceClient traceClient, String projectId, String traceId) + throws InterruptedException { + return pollWithRetry( + () -> { + Trace trace = traceClient.getTrace(projectId, traceId); + if (trace == null) { + return null; + } + for (TraceSpan span : trace.getSpansList()) { + if (span.getName().equals("BigQueryStatement.executeQuery")) { + return trace; + } + } + return null; + }); + } +} diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/suites/ITNightlyTests.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/suites/ITNightlyTests.java index 2700c18adbbd..02fc96c17bad 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/suites/ITNightlyTests.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/suites/ITNightlyTests.java @@ -19,9 +19,15 @@ import com.google.cloud.bigquery.jdbc.it.ITAuthTests; import com.google.cloud.bigquery.jdbc.it.ITBigQueryJDBCTest; import com.google.cloud.bigquery.jdbc.it.ITNightlyBigQueryTest; +import com.google.cloud.bigquery.jdbc.it.ITOpenTelemetryTest; import org.junit.platform.suite.api.SelectClasses; import org.junit.platform.suite.api.Suite; @Suite -@SelectClasses({ITAuthTests.class, ITBigQueryJDBCTest.class, ITNightlyBigQueryTest.class}) +@SelectClasses({ + ITAuthTests.class, + ITBigQueryJDBCTest.class, + ITNightlyBigQueryTest.class, + ITOpenTelemetryTest.class +}) public class ITNightlyTests {}