fix(bigquery): Allow users to supply custom TracerProvider (#2995)

whuffman36 · web-flow · commit 88095e60df7e · 2025-06-04T10:23:14.000-07:00
* fix(bigquery): Change RegisterGlobal() to GlobalOtel.set()

* fix style

* Allow users to supply custome TracerProvider
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/BigQueryReadClient.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/BigQueryReadClient.java
@@ -22,6 +22,7 @@
 import com.google.cloud.bigquery.storage.v1.stub.EnhancedBigQueryReadStub;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.TracerProvider;
 import io.opentelemetry.context.Scope;
 import java.io.IOException;
 import java.util.concurrent.TimeUnit;
@@ -133,7 +134,8 @@ protected BigQueryReadClient(BigQueryReadSettings settings) throws IOException {
         EnhancedBigQueryReadStub.create(
             settings.getTypedStubSettings(),
             settings.getReadRowsRetryAttemptListener(),
-            settings.isOpenTelemetryEnabled());
+            settings.isOpenTelemetryEnabled(),
+            settings.getOpenTelemetryTracerProvider());
   }
 
   @BetaApi("A restructuring of stub classes is planned, so this may break in the future")
@@ -527,11 +529,15 @@ public boolean awaitTermination(long duration, TimeUnit unit) throws Interrupted
   }
 
   public void disableOpenTelemetryTracing() {
-    settings.setEnableOpenTelemetryTracing(false);
+    settings.setEnableOpenTelemetryTracing(false, null);
   }
 
-  public void enableOpenTelemetryTracing() {
-    settings.setEnableOpenTelemetryTracing(true);
+  public void enableDefaultOpenTelemetryTracing() {
+    settings.setEnableOpenTelemetryTracing(true, null);
+  }
+
+  public void enableOpenTelemetryTracing(TracerProvider tracerProvider) {
+    settings.setEnableOpenTelemetryTracing(true, tracerProvider);
   }
 
   private static String getFieldAsString(Object field) {
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/BigQueryReadSettings.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/BigQueryReadSettings.java
@@ -30,6 +30,7 @@
 import io.grpc.Metadata;
 import io.grpc.Status;
 import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.api.trace.TracerProvider;
 import java.io.IOException;
 import java.util.List;
 
@@ -79,6 +80,7 @@ public static interface RetryAttemptListener {
   private RetryAttemptListener readRowsRetryAttemptListener = null;
   private boolean enableOpenTelemetryTracing = false;
   private Tracer openTelemetryTracer = null;
+  private TracerProvider openTelemetryTracerProvider = null;
 
   /**
    * If a non null readRowsRetryAttemptListener is provided, client will call onRetryAttempt
@@ -90,14 +92,23 @@ public void setReadRowsRetryAttemptListener(RetryAttemptListener readRowsRetryAt
     this.readRowsRetryAttemptListener = readRowsRetryAttemptListener;
   }
 
-  public void setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing) {
+  public void setEnableOpenTelemetryTracing(
+      boolean enableOpenTelemetryTracing, TracerProvider openTelemetryTracerProvider) {
     this.enableOpenTelemetryTracing = enableOpenTelemetryTracing;
     if (enableOpenTelemetryTracing) {
-      this.openTelemetryTracer =
-          Singletons.getOpenTelemetry()
-              .getTracerProvider()
-              .tracerBuilder("com.google.cloud.bigquery.storage.v1.read")
-              .build();
+      if (openTelemetryTracerProvider == null) {
+        this.openTelemetryTracer =
+            Singletons.getOpenTelemetry()
+                .getTracerProvider()
+                .tracerBuilder("com.google.cloud.bigquery.storage.v1.read")
+                .build();
+      } else {
+        this.openTelemetryTracerProvider = openTelemetryTracerProvider;
+        this.openTelemetryTracer =
+            openTelemetryTracerProvider
+                .tracerBuilder("com.google.cloud.bigquery.storage.v1.read")
+                .build();
+      }
     }
   }
 
@@ -113,6 +124,10 @@ public Tracer getOpenTelemetryTracer() {
     return this.openTelemetryTracer;
   }
 
+  public TracerProvider getOpenTelemetryTracerProvider() {
+    return this.openTelemetryTracerProvider;
+  }
+
   /** Returns the object with the settings used for calls to splitReadStream. */
   public UnaryCallSettings<SplitReadStreamRequest, SplitReadStreamResponse>
       splitReadStreamSettings() {
@@ -222,6 +237,7 @@ public Builder applyToAllUnaryMethods(
 
     private RetryAttemptListener readRowsRetryAttemptListener = null;
     private boolean enableOpenTelemetryTracing = false;
+    private TracerProvider openTelemetryTracerProvider = null;
 
     public Builder setReadRowsRetryAttemptListener(
         RetryAttemptListener readRowsRetryAttemptListener) {
@@ -234,6 +250,11 @@ public Builder setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing)
       return this;
     }
 
+    public Builder setOpenTelemetryTracerProvider(TracerProvider tracerProvider) {
+      this.openTelemetryTracerProvider = tracerProvider;
+      return this;
+    }
+
     /** Returns the builder for the settings used for calls to createReadSession. */
     public UnaryCallSettings.Builder<CreateReadSessionRequest, ReadSession>
         createReadSessionSettings() {
@@ -256,7 +277,8 @@ public Builder setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing)
     public BigQueryReadSettings build() throws IOException {
       BigQueryReadSettings settings = new BigQueryReadSettings(this);
       settings.setReadRowsRetryAttemptListener(readRowsRetryAttemptListener);
-      settings.setEnableOpenTelemetryTracing(enableOpenTelemetryTracing);
+      settings.setEnableOpenTelemetryTracing(
+          enableOpenTelemetryTracing, openTelemetryTracerProvider);
       return settings;
     }
   }
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/stub/EnhancedBigQueryReadStub.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/stub/EnhancedBigQueryReadStub.java
@@ -44,6 +44,7 @@
 import com.google.common.collect.ImmutableMap;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.api.trace.TracerProvider;
 import io.opentelemetry.context.Scope;
 import java.io.IOException;
 import java.util.Map;
@@ -73,13 +74,14 @@ public static EnhancedBigQueryReadStub create(
       EnhancedBigQueryReadStubSettings settings,
       BigQueryReadSettings.RetryAttemptListener readRowsRetryAttemptListener)
       throws IOException {
-    return create(settings, readRowsRetryAttemptListener, false);
+    return create(settings, readRowsRetryAttemptListener, false, null);
   }
 
   public static EnhancedBigQueryReadStub create(
       EnhancedBigQueryReadStubSettings settings,
       BigQueryReadSettings.RetryAttemptListener readRowsRetryAttemptListener,
-      boolean enableOpenTelemetryTracing)
+      boolean enableOpenTelemetryTracing,
+      TracerProvider openTelemetryTracerProvider)
       throws IOException {
     // Configure the base settings.
     BigQueryReadStubSettings.Builder baseSettingsBuilder =
@@ -118,7 +120,8 @@ public static EnhancedBigQueryReadStub create(
         baseSettings,
         readRowsRetryAttemptListener,
         clientContext,
-        enableOpenTelemetryTracing);
+        enableOpenTelemetryTracing,
+        openTelemetryTracerProvider);
   }
 
   @InternalApi("Visible for testing")
@@ -127,18 +130,26 @@ public static EnhancedBigQueryReadStub create(
       BigQueryReadStubSettings stubSettings,
       BigQueryReadSettings.RetryAttemptListener readRowsRetryAttemptListener,
       ClientContext context,
-      boolean enableOpenTelemetryTracing) {
+      boolean enableOpenTelemetryTracing,
+      TracerProvider openTelemetryTracerProvider) {
     this.stub = stub;
     this.stubSettings = stubSettings;
     this.readRowsRetryAttemptListener = readRowsRetryAttemptListener;
     this.context = context;
     this.enableOpenTelemetryTracing = enableOpenTelemetryTracing;
     if (enableOpenTelemetryTracing) {
-      this.openTelemetryTracer =
-          Singletons.getOpenTelemetry()
-              .getTracerProvider()
-              .tracerBuilder("com.google.cloud.bigquery.storage.v1.read.stub")
-              .build();
+      if (openTelemetryTracerProvider == null) {
+        this.openTelemetryTracer =
+            Singletons.getOpenTelemetry()
+                .getTracerProvider()
+                .tracerBuilder("com.google.cloud.bigquery.storage.v1.read.stub")
+                .build();
+      } else {
+        this.openTelemetryTracer =
+            openTelemetryTracerProvider
+                .tracerBuilder("com.google.cloud.bigquery.storage.v1.read.stub")
+                .build();
+      }
     }
   }
 
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageTest.java
@@ -1598,11 +1598,13 @@ public void testSimpleReadWithOtelTracing() throws IOException {
             .addSpanProcessor(SimpleSpanProcessor.create(new TestSpanExporter()))
             .setSampler(Sampler.alwaysOn())
             .build();
-    OpenTelemetry otel =
-        OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).buildAndRegisterGlobal();
+    OpenTelemetry otel = OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).build();
 
     BigQueryReadSettings otelSettings =
-        BigQueryReadSettings.newBuilder().setEnableOpenTelemetryTracing(true).build();
+        BigQueryReadSettings.newBuilder()
+            .setEnableOpenTelemetryTracing(true)
+            .setOpenTelemetryTracerProvider(tracerProvider)
+            .build();
     BigQueryReadClient otelClient = BigQueryReadClient.create(otelSettings);
 
     String table =
