8000 fix(bigquery): Allow users to supply custom TracerProvider (#2995) · googleapis/java-bigquerystorage@88095e6 · GitHub
[go: up one dir, main page]
Skip to content

Commit 88095e6

Browse files
whuffman36whuffman36
authored
fix(bigquery): Allow users to supply custom TracerProvider (#2995)
* fix(bigquery): Change RegisterGlobal() to GlobalOtel.set() * fix style * Allow users to supply custome TracerProvider
1 parent 7b2ba0c commit 88095e6

File tree

4 files changed

+64
-23
lines changed

4 files changed

+64
-23
lines changed

‎google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/BigQueryReadClient.java

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import com.google.cloud.bigquery.storage.v1.stub.EnhancedBigQueryReadStub;
2323
import io.opentelemetry.api.common.Attributes;
2424
import io.opentelemetry.api.trace.Span;
25+
import io.opentelemetry.api.trace.TracerProvider;
2526
import io.opentelemetry.context.Scope;
2627
import java.io.IOException;
2728
import java.util.concurrent.TimeUnit;
@@ -133,7 +134,8 @@ protected BigQueryReadClient(BigQueryReadSettings settings) throws IOException {
133134
EnhancedBigQueryReadStub.create(
134135
settings.getTypedStubSettings(),
135136
settings.getReadRowsRetryAttemptListener(),
136-
settings.isOpenTelemetryEnabled());
137+
settings.isOpenTelemetryEnabled(),
138+
settings.getOpenTelemetryTracerProvider());
137139
}
138140

139141
@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
527529
}
528530

529531
public void disableOpenTelemetryTracing() {
530-
settings.setEnableOpenTelemetryTracing(false);
532+
settings.setEnableOpenTelemetryTracing(false, null);
531533
}
532534

533-
public void enableOpenTelemetryTracing() {
534-
settings.setEnableOpenTelemetryTracing(true);
535+
public void enableDefaultOpenTelemetryTracing() {
536+
settings.setEnableOpenTelemetryTracing(true, null);
537+
}
538+
539+
public void enableOpenTelemetryTracing(TracerProvider tracerProvider) {
540+
settings.setEnableOpenTelemetryTracing(true, tracerProvider);
535541
}
536542

537543
private static String getFieldAsString(Object field) {

‎google-cloud-bigquerystorage 8000 /src/main/java/com/google/cloud/bigquery/storage/v1/BigQueryReadSettings.java

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
import io.grpc.Metadata;
3131
import io.grpc.Status;
3232
import io.opentelemetry.api.trace.Tracer;
33+
import io.opentelemetry.api.trace.TracerProvider;
3334
import java.io.IOException;
3435
import java.util.List;
3536

@@ -79,6 +80,7 @@ public static interface RetryAttemptListener {
7980
private RetryAttemptListener readRowsRetryAttemptListener = null;
8081
private boolean enableOpenTelemetryTracing = false;
8182
private Tracer openTelemetryTracer = null;
83+
private TracerProvider openTelemetryTracerProvider = null;
8284

8385
/**
8486
* If a non null readRowsRetryAttemptListener is provided, client will call onRetryAttempt
@@ -90,14 +92,23 @@ public void setReadRowsRetryAttemptListener(RetryAttemptListener readRowsRetryAt
9092
this.readRowsRetryAttemptListener = readRowsRetryAttemptListener;
9193
} 8000
9294

93-
public void setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing) {
95+
public void setEnableOpenTelemetryTracing(
96+
boolean enableOpenTelemetryTracing, TracerProvider openTelemetryTracerProvider) {
9497
this.enableOpenTelemetryTracing = enableOpenTelemetryTracing;
9598
if (enableOpenTelemetryTracing) {
96-
this.openTelemetryTracer =
97-
Singletons.getOpenTelemetry()
98-
.getTracerProvider()
99-
.tracerBuilder("com.google.cloud.bigquery.storage.v1.read")
100-
.build();
99+
if (openTelemetryTracerProvider == null) {
100+
this.openTelemetryTracer =
101+
Singletons.getOpenTelemetry()
102+
.getTracerProvider()
103+
.tracerBuilder("com.google.cloud.bigquery.storage.v1.read")
104+
.build();
105+
} else {
106+
this.openTelemetryTracerProvider = openTelemetryTracerProvider;
107+
this.openTelemetryTracer =
108+
openTelemetryTracerProvider
109+
.tracerBuilder("com.google.cloud.bigquery.storage.v1.read")
110+
.build();
111+
}
101112
}
102113
}
103114

@@ -113,6 +124,10 @@ public Tracer getOpenTelemetryTracer() {
113124
return this.openTelemetryTracer;
114125
}
115126

127+
public TracerProvider getOpenTelemetryTracerProvider() {
128+
return this.openTelemetryTracerProvider;
129+
}
130+
116131
/** Returns the object with the settings used for calls to splitReadStream. */
117132
public UnaryCallSettings<SplitReadStreamRequest, SplitReadStreamResponse>
118133
splitReadStreamSettings() {
@@ -222,6 +237,7 @@ public Builder applyToAllUnaryMethods(
222237

223238
private RetryAttemptListener readRowsRetryAttemptListener = null;
224239
private boolean enableOpenTelemetryTracing = false;
240+
private TracerProvider openTelemetryTracerProvider = null;
225241

226242
public Builder setReadRowsRetryAttemptListener(
227243
RetryAttemptListener readRowsRetryAttemptListener) {
@@ -234,6 +250,11 @@ public Builder setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing)
234250
return this;
235251
}
236252

253+
public Builder setOpenTelemetryTracerProvider(TracerProvider tracerProvider) {
254+
this.openTelemetryTracerProvider = tracerProvider;
255+
return this;
256+
}
257+
237258
/** Returns the builder for the settings used for calls to createReadSession. */
238259
public UnaryCallSettings.Builder<CreateReadSessionRequest, ReadSession>
239260
createReadSessionSettings() {
@@ -256,7 +277,8 @@ public Builder setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing)
256277
public BigQueryReadSettings build() throws IOException {
257278
BigQueryReadSettings settings = new BigQueryReadSettings(this);
258279
settings.setReadRowsRetryAttemptListener(readRowsRetryAttemptListener);
259-
settings.setEnableOpenTelemetryTracing(enableOpenTelemetryTracing);
280+
settings.setEnableOpenTelemetryTracing(
281+
enableOpenTelemetryTracing, openTelemetryTracerProvider);
260282
return settings;
261283
}
262284
}

‎google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/stub/EnhancedBigQueryReadStub.java

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
import com.google.common.collect.ImmutableMap;
4545
import io.opentelemetry.api.trace.Span;
4646
import io.opentelemetry.api.trace.Tracer;
47+
import io.opentelemetry.api.trace.TracerProvider;
4748
import io.opentelemetry.context.Scope;
4849
import java.io.IOException;
4950
import java.util.Map;
@@ -73,13 +74,14 @@ public static EnhancedBigQueryReadStub create(
7374
EnhancedBigQueryReadStubSettings settings,
7475
BigQueryReadSettings.RetryAttemptListener readRowsRetryAttemptListener)
7576
throws IOException {
76-
return create(settings, readRowsRetryAttemptListener, false);
77+
return create(settings, readRowsRetryAttemptListener, false, null);
7778
}
7879

7980
public static EnhancedBigQueryReadStub create(
8081
EnhancedBigQueryReadStubSettings settings,
8182
BigQueryReadSettings.RetryAttemptListener readRowsRetryAttemptListener,
82-
boolean enableOpenTelemetryTracing)
83+
boolean enableOpenTelemetryTracing,
84+
TracerProvider openTelemetryTracerProvider)
8385
throws IOException {
8486
// Configure the base settings.
8587
BigQueryReadStubSettings.Builder baseSettingsBuilder =
@@ -118,7 +120,8 @@ public static EnhancedBigQueryReadStub create(
118120
baseSettings,
119121
readRowsRetryAttemptListener,
120122
clientContext,
121-
enableOpenTelemetryTracing);
123+
enableOpenTelemetryTracing,
124+
openTelemetryTracerProvider);
122125
}
123126

124127
@InternalApi("Visible for testing")
@@ -127,18 +130,26 @@ public static EnhancedBigQueryReadStub create(
127130
BigQueryReadStubSettings stubSettings,
128131
BigQueryReadSettings.RetryAttemptListener readRowsRetryAttemptListener,
129132
ClientContext context,
130-
boolean enableOpenTelemetryTracing) {
133+
boolean enableOpenTelemetryTracing,
134+
TracerProvider openTelemetryTracerProvider) {
131135
this.stub = stub;
132136
this.stubSettings = stubSettings;
133137
this.readRowsRetryAttemptListener = readRowsRetryAttemptListener;
134138
this.context = context;
135139
this.enableOpenTelemetryTracing = enableOpenTelemetryTracing;
136140
if (enableOpenTelemetryTracing) {
137-
this.openTelemetryTracer =
138-
Singletons.getOpenTelemetry()
139-
.getTracerProvider()
140-
.tracerBuilder("com.google.cloud.bigquery.storage.v1.read.stub")
141-
.build();
141+
if (openTelemetryTracerProvider == null) {
142+
this.openTelemetryTracer =
143+
Singletons.getOpenTelemetry()
144+
.getTracerProvider()
145+
.tracerBuilder("com.google.cloud.bigquery.storage.v1.read.stub")
146+
.build();
147+
} else {
148+
this.openTelemetryTracer =
149+
openTelemetryTracerProvider
150+
.tracerBuilder("com.google.cloud.bigquery.storage.v1.read.stub")
151+
.build();
152+
}
142153
}
143154
}
144155

‎google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageTest.java

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1598,11 +1598,13 @@ public void testSimpleReadWithOtelTracing() throws IOException {
15981598
.addSpanProcessor(SimpleSpanProcessor.create(new TestSpanExporter()))
15991599
.setSampler(Sampler.alwaysOn())
16001600
.build();
1601-
OpenTelemetry otel =
1602-
OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).buildAndRegisterGlobal();
1601+
OpenTelemetry otel = OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).build();
16031602

16041603
BigQueryReadSettings otelSettings =
1605-
BigQueryReadSettings.newBuilder().setEnableOpenTelemetryTracing(true).build();
1604+
BigQueryReadSettings.newBuilder()
1605+
.setEnableOpenTelemetryTracing(true)
1606+
.setOpenTelemetryTracerProvider(tracerProvider)
1607+
.build();
16061608
BigQueryReadClient otelClient = BigQueryReadClient.create(otelSettings);
16071609

16081610
String table =

0 commit comments

Comments
 (0)
0