diff --git a/java-bigquery-jdbc/pom.xml b/java-bigquery-jdbc/pom.xml index ba78d2094996..4c9c367bfa2f 100644 --- a/java-bigquery-jdbc/pom.xml +++ b/java-bigquery-jdbc/pom.xml @@ -84,6 +84,13 @@ org.apache.httpcomponents.*:* io.grpc:* + + io.opentelemetry.contrib:opentelemetry-gcp-auth-extension + io.opentelemetry:opentelemetry-exporter-otlp + + + io.opentelemetry:opentelemetry-sdk-logs + @@ -173,6 +180,17 @@ io.opentelemetry com.google.bqjdbc.shaded.io.opentelemetry + + + io.opentelemetry.api.** + io.opentelemetry.context.** + io.perfmark @@ -223,6 +241,11 @@ google-cloud-bigquerystorage 3.31.0-SNAPSHOT + + com.google.cloud + google-cloud-logging + 3.32.0 + com.google.http-client google-http-client-apache-v5 @@ -310,6 +333,16 @@ io.grpc grpc-netty-shaded + + io.grpc + grpc-opentelemetry + + + io.opentelemetry + opentelemetry-api + + + org.apache.arrow @@ -328,6 +361,46 @@ 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 + opentelemetry-sdk-logs + + + io.opentelemetry.contrib + opentelemetry-gcp-auth-extension + 1.56.0-alpha + + @@ -377,6 +450,17 @@ junit-platform-suite-engine test + + io.opentelemetry + opentelemetry-sdk-testing + 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 e4ba837643f8..3123b6c09b40 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; @@ -264,29 +265,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 ab478fa26c02..9216732b49b2 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 5446df272d74..ee541f146782 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 @@ -44,8 +44,16 @@ 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.ImmutableList; 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; @@ -157,6 +165,8 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { int transactionIsolation; List sqlWarnings; String catalog; + String gcpTelemetryCredentials; + String gcpTelemetryProjectId; int holdability; long retryTimeoutInSeconds; Duration retryTimeoutDuration; @@ -215,6 +225,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(); @@ -228,6 +246,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)) { @@ -250,6 +273,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(); @@ -369,6 +394,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(); // Cached pool executes queries immediately without queueing and reclaims all idle threads // when inactive, minimizing resources. @@ -447,7 +477,8 @@ String getConnectionUrl() { return connectionUrl; } - String getConnectionId() { + @VisibleForTesting + public String getConnectionId() { return this.connectionId; } @@ -1077,6 +1108,9 @@ private void closeImpl() throws SQLException { } finally { BigQueryJdbcMdc.clear(); BigQueryJdbcRootLogger.closeConnectionHandler(this.connectionId); + BigQueryJdbcOpenTelemetry.unregisterConnection(this.connectionId); + BigQueryJdbcOpenTelemetry.releaseSdk(this.openTelemetry); + this.openTelemetry = null; } if (exceptionToThrow != null) { throw exceptionToThrow; @@ -1150,6 +1184,73 @@ 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() { + if (this.gcpTelemetryCredentials != null) { + return this.gcpTelemetryCredentials; + } + + String authTypeStr = this.authProperties.get(BigQueryJdbcUrlUtility.OAUTH_TYPE_PROPERTY_NAME); + if (!BigQueryJdbcOAuthUtility.AuthType.GOOGLE_SERVICE_ACCOUNT.name().equals(authTypeStr)) { + return null; + } + + String pvtKey = this.authProperties.get(BigQueryJdbcUrlUtility.OAUTH_PVT_KEY_PROPERTY_NAME); + if (pvtKey != null) { + return pvtKey; + } + + return this.authProperties.get(BigQueryJdbcUrlUtility.OAUTH_PVT_KEY_PATH_PROPERTY_NAME); + } + + 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 @@ -1186,6 +1287,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( @@ -1220,7 +1330,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) { @@ -1235,6 +1358,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()); } @@ -1334,6 +1464,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 79f6bdb2a2f6..2cbad20c3bc3 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 @@ -52,6 +52,7 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.jdbc.BigQueryJdbcTypeMappings.ColumnTypeInfo; import com.google.cloud.bigquery.jdbc.utils.BigQueryJdbcVersionUtility; +import io.opentelemetry.context.Context; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStream; @@ -837,7 +838,8 @@ public ResultSet getProcedures( procedureNamePattern, procedureNameRegex, false); - Future> apiFuture = apiExecutor.submit(apiCallable); + Future> apiFuture = + apiExecutor.submit(Context.current().wrap(apiCallable)); apiFutures.add(apiFuture); } LOG.fine("Finished submitting " + apiFutures.size() + " findMatchingRoutines tasks."); @@ -1139,7 +1141,7 @@ List listMatchingProcedureIdsFromDatasets( procedureNamePattern, procedureNameRegex, false); - listRoutineFutures.add(listRoutinesExecutor.submit(listCallable)); + listRoutineFutures.add(listRoutinesExecutor.submit(Context.current().wrap(listCallable))); } logger.fine( "Submitted " @@ -1214,7 +1216,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."); @@ -1559,6 +1561,14 @@ Comparator defineGetProcedureColumnsComparator(FieldList resultS @Override public ResultSet getTables( + 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) { if ((catalog != null && catalog.isEmpty()) @@ -1638,7 +1648,8 @@ public ResultSet getTables( tableNamePattern, tableNameRegex, false); - Future> apiFuture = apiExecutor.submit(apiCallable); + Future> apiFuture = + apiExecutor.submit(Context.current().wrap(apiCallable)); apiFutures.add(apiFuture); } LOG.fine("Finished submitting " + apiFutures.size() + " findMatchingTables tasks."); @@ -1682,7 +1693,8 @@ public ResultSet getTables( } }; - Future fetcherFuture = connection.getExecutorService().submit(tableFetcher); + Future fetcherFuture = + connection.getExecutorService().submit(Context.current().wrap(tableFetcher)); BigQueryJsonResultSet resultSet = BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); @@ -1809,6 +1821,10 @@ public ResultSet getSchemas() throws SQLException { @Override public ResultSet getCatalogs() throws SQLException { + return withTracing("BigQueryDatabaseMetaData.getCatalogs", () -> getCatalogsImpl()); + } + + private ResultSet getCatalogsImpl() throws SQLException { LOG.info("getCatalogs() called"); final List accessibleCatalogs = getAccessibleCatalogNames(); @@ -1878,6 +1894,14 @@ static List prepareGetTableTypesRows(Schema schema) { @Override public ResultSet getColumns( + 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) { if ((catalog != null && catalog.isEmpty()) @@ -1992,7 +2016,8 @@ public ResultSet getColumns( } }; - Future fetcherFuture = connection.getExecutorService().submit(columnFetcher); + Future fetcherFuture = + connection.getExecutorService().submit(Context.current().wrap(columnFetcher)); BigQueryJsonResultSet resultSet = BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); @@ -3594,6 +3619,11 @@ public RowIdLifetime getRowIdLifetime() { @Override 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."); @@ -3649,7 +3679,8 @@ public ResultSet getSchemas(String catalog, String schemaPattern) throws SQLExce } }; - Future fetcherFuture = connection.getExecutorService().submit(multiSchemaFetcher); + Future fetcherFuture = + connection.getExecutorService().submit(Context.current().wrap(multiSchemaFetcher)); BigQueryJsonResultSet resultSet = BigQueryJsonResultSet.of(resultSchema, -1, queue, null, fetcherFuture); @@ -3863,7 +3894,8 @@ public ResultSet getFunctions(String catalog, String schemaPattern, String funct functionNameRegex, false); }; - Future> apiFuture = apiExecutor.submit(apiCallable); + Future> apiFuture = + apiExecutor.submit(Context.current().wrap(apiCallable)); apiFutures.add(apiFuture); } LOG.fine( @@ -4213,7 +4245,7 @@ List listMatchingFunctionIdsFromDatasets( functionNamePattern, functionNameRegex, false); - listRoutineFutures.add(listRoutinesExecutor.submit(listCallable)); + listRoutineFutures.add(listRoutinesExecutor.submit(Context.current().wrap(listCallable))); } logger.fine( "Submitted " @@ -5410,4 +5442,14 @@ private void closeStatementIgnoreException(Statement stmt) { // ignore } } + + 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 8c032dc79b14..8c748b1f52bc 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 @@ -21,6 +21,7 @@ import com.google.cloud.bigquery.jdbc.utils.BigQueryJdbcVersionUtility; 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; @@ -125,8 +126,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 { @@ -146,7 +151,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; } @@ -164,7 +181,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 d2236a7ad90b..c7c281d51485 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 @@ -315,7 +315,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..956c2cc02298 --- /dev/null +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java @@ -0,0 +1,463 @@ +/* + * 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.auth.oauth2.GoogleCredentials; +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.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +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 class CachedSdk { + final OpenTelemetrySdk sdk; + final AtomicInteger refCount = new AtomicInteger(1); + + CachedSdk(OpenTelemetrySdk sdk) { + this.sdk = sdk; + } + } + + 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(); + } + + 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); + } + } + + public static void releaseSdk(OpenTelemetry openTelemetry) { + if (openTelemetry == null || openTelemetry == OpenTelemetry.noop()) { + return; + } + + AtomicBoolean shouldClose = new AtomicBoolean(false); + for (Map.Entry entry : sdkCache.entrySet()) { + if (entry.getValue().sdk == openTelemetry) { + sdkCache.computeIfPresent( + entry.getKey(), + (k, cachedSdk) -> { + if (cachedSdk.sdk != openTelemetry) { + return cachedSdk; + } + if (cachedSdk.refCount.decrementAndGet() <= 0) { + shouldClose.set(true); + return null; + } + return cachedSdk; + }); + break; + } + } + + if (shouldClose.get() && openTelemetry instanceof OpenTelemetrySdk) { + try { + ((OpenTelemetrySdk) openTelemetry).close(); + } catch (Exception e) { + // Swallow exceptions so that telemetry shutdown failures (e.g. flush timeouts) + // do not propagate and disrupt the core BigQueryConnection.close() logic. + // Throwing here could prevent the core connection from cleaning up correctly. + LOG.warning("Failed to close OpenTelemetry SDK: %s", e.getMessage()); + } + } + } + + private static Credentials resolveCredentialsFromString(String credsString) { + Map authProperties = new java.util.HashMap<>(); + authProperties.put( + BigQueryJdbcUrlUtility.OAUTH_TYPE_PROPERTY_NAME, + BigQueryJdbcOAuthUtility.AuthType.GOOGLE_SERVICE_ACCOUNT.name()); // 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, + null, + 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.compute( + key, + (k, cachedSdk) -> { + if (cachedSdk != null) { + cachedSdk.refCount.incrementAndGet(); + return cachedSdk; + } + + 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) -> { + try { + Credentials credentials; + if (gcpTelemetryCredentials != null) { + credentials = resolveCredentialsFromString(gcpTelemetryCredentials); + } else { + credentials = GoogleCredentials.getApplicationDefault(); + } + 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 new CachedSdk(sdk); + }) + .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 481df4abf35f..928ed785a795 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 @@ -99,6 +99,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"; @@ -164,6 +166,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 = @@ -636,6 +644,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 bdae55fd2b0c..e88d3a68c8d3 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 @@ -97,12 +97,20 @@ private int getParameterCount(String query) { @Override public ResultSet executeQuery() throws SQLException { - return super.executeQuery(this.currentQuery); + return BigQueryJdbcOpenTelemetry.withTracing( + "BigQueryPreparedStatement.executeQuery", + this.connection, + this.currentQuery, + () -> super.executeQuery(this.currentQuery)); } @Override public long executeLargeUpdate() throws SQLException { - return super.executeLargeUpdate(this.currentQuery); + return BigQueryJdbcOpenTelemetry.withTracing( + "BigQueryPreparedStatement.executeLargeUpdate", + this.connection, + this.currentQuery, + () -> super.executeLargeUpdate(this.currentQuery)); } @Override @@ -112,7 +120,11 @@ public int executeUpdate() throws SQLException { @Override public boolean execute() throws SQLException { - return super.execute(this.currentQuery); + return BigQueryJdbcOpenTelemetry.withTracing( + "BigQueryPreparedStatement.execute", + this.connection, + this.currentQuery, + () -> super.execute(this.currentQuery)); } @Override 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 17373af77b27..30c11d13d263 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,12 @@ 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.context.Context; +import io.opentelemetry.context.Scope; import java.lang.ref.ReferenceQueue; import java.sql.Connection; import java.sql.ResultSet; @@ -237,7 +243,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 { @@ -261,7 +268,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 +309,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 { @@ -898,86 +910,94 @@ Future populateArrowBufferedQueue( ExecutorService executor = connection.getExecutorService(); 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() || executor.isShutdown()) { - break; + Context.current() + .wrap( + () -> { + long rowsRead = 0; + int retryCount = 0; + try { + // Use the first stream to perform reading. + String streamName = readSession.getStreams(0).getName(); + + while (true) { + try { + ReadRowsRequest readRowsRequest = + ReadRowsRequest.newBuilder() + .setReadStream(streamName) + .setOffset(rowsRead) + .build(); + + // Process each block of rows as they arrive and decode using our simple row + // reader. + com.google.api.gax.rpc.ServerStream stream = + bqReadClient.readRowsCallable().call(readRowsRequest); + for (ReadRowsResponse response : stream) { + if (Thread.currentThread().isInterrupted() || executor.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); + } + } + + } catch (InterruptedException e) { + LOG.log( + Level.WARNING, + "\n" + + Thread.currentThread().getName() + + " Interrupted @ arrowStreamProcessor", + 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); } - - 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); - } - } - - } catch (InterruptedException e) { - LOG.log( - Level.WARNING, - "\n" + Thread.currentThread().getName() + " Interrupted @ arrowStreamProcessor", - 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); - } - }; + }); return executor.submit(arrowStreamProcessor); } @@ -1177,60 +1197,79 @@ Future 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); + SpanContext parentSpanContext = Span.current().getSpanContext(); ExecutorService executor = connection.getExecutorService(); // 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); - } + Context.current() + .wrap( + () -> { + 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() || executor.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 - }; + try { + while (currentPageToken != null) { + // do not process further pages and shutdown + if (Thread.currentThread().isInterrupted() || executor.isShutdown()) { + LOG.warning( + "%s Interrupted @ runNextPageTaskAsync", + Thread.currentThread().getName()); + break; + } + + SpanBuilder spanBuilder = + connection.getTracer().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)); + + 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)); + } finally { + paginationSpan.end(); + } + } + } 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 + }); return executor.submit(nextPageTask); } @@ -1249,77 +1288,84 @@ Future parseAndPopulateRpcDataAsync( ExecutorService executor = connection.getExecutorService(); Runnable populateBufferRunnable = - () -> { // producer thread populating the buffer - try { - Iterable fieldValueLists; - boolean[] isComplexColumn = - BigQueryFieldValueListWrapper.createComplexColumnFlags( - schema != null ? schema.getFields() : null); - // 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() - || executor.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() || executor.isShutdown()) { - // do not process further pages and shutdown (inner loop) - break; - } - Uninterruptibles.putUninterruptibly( - bigQueryFieldValueListWrapperBlockingQueue, - BigQueryFieldValueListWrapper.of( - schema.getFields(), fieldValueList, isComplexColumn)); - 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( + () -> { // producer thread populating the buffer + try { + Iterable fieldValueLists; + boolean[] isComplexColumn = + BigQueryFieldValueListWrapper.createComplexColumnFlags( + schema != null ? schema.getFields() : null); + // 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() + || executor.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() || executor.isShutdown()) { + // do not process further pages and shutdown (inner loop) + break; + } + Uninterruptibles.putUninterruptibly( + bigQueryFieldValueListWrapperBlockingQueue, + BigQueryFieldValueListWrapper.of( + schema.getFields(), fieldValueList, isComplexColumn)); + 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); + } + }); return executor.submit(populateBufferRunnable); } @@ -1576,13 +1622,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()); 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 f6489d33dd37..2c07fc483bac 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; @@ -118,6 +121,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 { @@ -134,6 +143,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) -> @@ -347,6 +362,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) { @@ -404,7 +437,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() { @@ -665,6 +702,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; } @@ -798,6 +850,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( @@ -867,6 +951,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 e766ba3bf856..f29bdc553569 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,6 +16,7 @@ package com.google.cloud.bigquery.jdbc; +import static org.junit.jupiter.api.Assertions.*; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -23,7 +24,14 @@ 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.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.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -33,6 +41,7 @@ import com.google.api.gax.paging.Page; 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.BigQueryException; import com.google.cloud.bigquery.Project; @@ -40,6 +49,12 @@ 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; @@ -52,11 +67,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 = @@ -477,6 +497,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(); @@ -509,6 +550,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(), 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 3ef43bf648fb..a41bfe8939ad 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,11 @@ 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.StatusCode; +import io.opentelemetry.context.Context; +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; @@ -50,14 +55,23 @@ import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.regex.Pattern; +import java.util.stream.Stream; import org.junit.jupiter.api.AfterEach; 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; @@ -73,6 +87,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, withSettings().withoutAnnotations()); + when(bigqueryClient.listDatasets(anyString(), any())).thenReturn(datasetPageMock); + + Page tablePageMock = mock(Page.class, withSettings().withoutAnnotations()); + when(bigqueryClient.listTables(any(DatasetId.class), any())).thenReturn(tablePageMock); + + Table mockTable = mock(Table.class); + when(bigqueryClient.getTable(any(TableId.class))).thenReturn(mockTable); when(bigQueryConnection.getMetadataExecutor()).thenReturn(metadataExecutor); when(bigQueryConnection.getExecutorService()).thenReturn(metadataExecutor); @@ -2978,7 +3008,7 @@ public void testPrepareGetCatalogsRows() { } @Test - public void testGetSchemas_NoArgs_DelegatesCorrectly() throws SQLException { + public void testGetSchemas_NoArgs_DelegatesCorrectly() throws Exception { BigQueryDatabaseMetaData spiedDbMetadata = spy(dbMetadata); ResultSet mockResultSet = mock(ResultSet.class); doReturn(mockResultSet).when(spiedDbMetadata).getSchemas(null, null); @@ -3271,6 +3301,48 @@ 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")); + } + @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 4ed3109f9845..8acbc5abb8dc 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,8 +16,10 @@ package com.google.cloud.bigquery.jdbc; import static com.google.common.truth.Truth.assertThat; +import static org.mockito.Mockito.mock; import com.google.cloud.bigquery.jdbc.utils.BigQueryJdbcVersionUtility; +import io.opentelemetry.api.OpenTelemetry; import java.sql.Connection; import java.sql.DriverPropertyInfo; import java.sql.SQLException; @@ -110,6 +112,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 9e7d42969bbc..91ae858a6ab8 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 @@ -32,8 +32,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.DatasetId; import com.google.cloud.bigquery.DatasetInfo; @@ -61,32 +64,54 @@ 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.ExecutorService; +import java.util.concurrent.Executors; import java.util.concurrent.Future; +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.AfterEach; +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"; @@ -138,9 +163,41 @@ 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; + } + + private ExecutorService testExecutorService; + @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; @@ -148,6 +205,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(); @@ -155,8 +213,8 @@ public void setUp() throws IOException, SQLException { .when(bigQueryConnection) .getQueryDialect(); doReturn(1000L).when(bigQueryConnection).getMaxResults(); - ExecutorService executorService = mock(ExecutorService.class); - doReturn(executorService).when(bigQueryConnection).getExecutorService(); + testExecutorService = Executors.newSingleThreadExecutor(); + doReturn(testExecutorService).when(bigQueryConnection).getExecutorService(); bigQueryStatement = new BigQueryStatement(bigQueryConnection); VectorSchemaRoot vectorSchemaRoot = getTestVectorSchemaRoot(); arrowSchema = @@ -167,6 +225,13 @@ public void setUp() throws IOException, SQLException { } + @AfterEach + public void tearDown() { + if (testExecutorService != null) { + testExecutorService.shutdownNow(); + } + } + private VectorSchemaRoot getTestVectorSchemaRoot() { RootAllocator allocator = new RootAllocator(); BitVector boolField = @@ -459,21 +524,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) @@ -493,6 +548,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(); + + Future workerThread = + bigQueryStatement.runNextPageTaskAsync( + mockResult, + "token", + mockJobId, + rpcResponseQueue, + bigQueryFieldValueListWrapperBlockingQueue); + + Assertions.assertNotNull(workerThread, "Worker thread should not be null"); + workerThread.get(); + + 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; + } + @Test public void testCancelDoesNotRollbackTransaction() throws SQLException { doReturn(true).when(bigQueryConnection).isTransactionStarted(); @@ -698,6 +873,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 db77acc58342..cbc4dec030c8 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 @@ -46,12 +46,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 @@ -60,6 +63,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 a95a05229f9b..f1895bd124f9 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,13 +24,9 @@ import com.google.auth.oauth2.GoogleCredentials; import com.google.cloud.ServiceOptions; import com.google.gson.JsonObject; -import com.google.gson.JsonParser; import java.io.File; import java.io.IOException; -import java.io.InputStream; -import java.io.InputStreamReader; import java.nio.file.Files; -import java.nio.file.Paths; import java.sql.Connection; import java.sql.DriverManager; import java.sql.ResultSet; @@ -46,25 +42,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 0224cafab695..366d1f70a353 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.QueryJobConfiguration; import com.google.cloud.bigquery.jdbc.BigQueryJdbcBaseTest; import com.google.cloud.bigquery.jdbc.utils.TestUtilities; +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; @@ -372,6 +383,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..1c0e63458f95 --- /dev/null +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITOpenTelemetryTest.java @@ -0,0 +1,363 @@ +/* + * 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 = "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 = 10; + long delayMs = 10000; + + // 10 second wait for GCP to ingest data + Thread.sleep(10000); + + 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 + e.printStackTrace(); + } + 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 {}