feat(bigquery-jdbc): OpenTelemetry integration in BQ JDBC

java-bigquery-jdbc/pom.xml

@@ -83,6 +83,10 @@
                 <ignoredUsedUndeclaredDependency>org.apache.httpcomponents.*:*</ignoredUsedUndeclaredDependency>
                 <ignoredUsedUndeclaredDependency>io.grpc:*</ignoredUsedUndeclaredDependency>
             </ignoredUsedUndeclaredDependencies>
+            <ignoredUnusedDeclaredDependencies>
+                <ignoredUnusedDeclaredDependency>io.opentelemetry.contrib:opentelemetry-gcp-auth-extension</ignoredUnusedDeclaredDependency>
+                <ignoredUnusedDeclaredDependency>io.opentelemetry:opentelemetry-exporter-otlp</ignoredUnusedDeclaredDependency>
+            </ignoredUnusedDeclaredDependencies>
         </configuration>
       </plugin>
       <plugin>
@@ -167,6 +171,17 @@
                     <relocation>
                       <pattern>io.opentelemetry</pattern>
                       <shadedPattern>com.google.bqjdbc.shaded.io.opentelemetry</shadedPattern>
+                      <excludes>
+                        <!-- 
+                          OpenTelemetry API and Context must remain unshaded to ensure interoperability.
+                          Shading these would prevent the driver from participating in the application's 
+                          existing tracing context. We are aware that unshaded dependencies can lead to 
+                          version mismatches, but this is a necessary trade-off for the OpenTelemetry 
+                          integration to function correctly across different applications.
+                        -->
+                        <exclude>io.opentelemetry.api.*</exclude>
+                        <exclude>io.opentelemetry.context.*</exclude>
+                      </excludes>
                     </relocation>
                     <relocation>
                       <pattern>io.perfmark</pattern>
@@ -217,6 +232,11 @@
       <artifactId>google-cloud-bigquerystorage</artifactId>
       <version>3.29.0-SNAPSHOT</version><!-- {x-version-update:google-cloud-bigquerystorage:current} -->
     </dependency>
+    <dependency>
+      <groupId>com.google.cloud</groupId>
+      <artifactId>google-cloud-logging</artifactId>
+      <version>3.32.0</version><!-- {x-version-update:google-cloud-logging:current} -->
+    </dependency>
     <dependency>
       <groupId>com.google.http-client</groupId>
       <artifactId>google-http-client-apache-v5</artifactId>
@@ -304,6 +324,16 @@
       <groupId>io.grpc</groupId>
       <artifactId>grpc-netty-shaded</artifactId>
     </dependency>
+    <dependency>
+      <groupId>io.grpc</groupId>
+      <artifactId>grpc-opentelemetry</artifactId>
+      <exclusions>
+        <exclusion>
+          <groupId>io.opentelemetry</groupId>
+          <artifactId>opentelemetry-api</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
 
     <dependency>
       <groupId>org.apache.arrow</groupId>
@@ -322,6 +352,42 @@
       <artifactId>httpcore5</artifactId>
     </dependency>
 
+    <!-- OpenTelemetry APIs (unshaded) -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-api</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-context</artifactId>
+    </dependency>
+    <!-- OpenTelemetry SDK and Exporters (will be shaded via existing 'io' relocation) -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-trace</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-exporter-otlp</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-extension-autoconfigure</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-extension-autoconfigure-spi</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry.contrib</groupId>
+      <artifactId>opentelemetry-gcp-auth-extension</artifactId>
+      <version>1.56.0-alpha</version>
+    </dependency>
+
     <!--  Test Dependencies  -->
 
     <dependency>
@@ -371,6 +437,22 @@
       <artifactId>junit-platform-suite-engine</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-testing</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-logs</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.google.cloud</groupId>
+      <artifactId>google-cloud-trace</artifactId>
+      <version>2.92.0</version><!-- {x-version-update:google-cloud-trace:current} -->
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <profiles>

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java

@@ -27,6 +27,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;
@@ -236,29 +237,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()) {

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java

@@ -28,6 +28,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;
@@ -59,6 +63,7 @@ public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet
   protected boolean isClosed = false;
   protected boolean wasNull = false;
   protected final BigQueryTypeCoercer bigQueryTypeCoercer = BigQueryTypeCoercionUtility.INSTANCE;
+  protected final SpanContext originalSpanContext;
 
   protected BigQueryBaseResultSet(
       BigQuery bigQuery, BigQueryStatement statement, Schema schema, boolean isNested) {
@@ -67,11 +72,16 @@ protected BigQueryBaseResultSet(
     this.schema = schema;
     this.schemaFieldList = schema != null ? schema.getFields() : null;
     this.isNested = isNested;
+    this.originalSpanContext = Span.current().getSpanContext();
     this.LOG =
         BigQueryJdbcResultSetLogger.getLogger(
             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;