Skip to content

Commit 6de1277

Browse files
keshavdandevakeshavdandeva
committed
feat: add globalOtel support
1 parent aef912c commit 6de1277

6 files changed

Lines changed: 206 additions & 15 deletions

File tree

‎java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
216216
boolean enableGcpTraceExporter;
217217
boolean enableGcpLogExporter;
218218
OpenTelemetry customOpenTelemetry;
219+
boolean useGlobalOpenTelemetry;
219220
private OpenTelemetry openTelemetry;
220221
private Context otelContext;
221222
Tracer tracer =
@@ -367,6 +368,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
367368
this.enableGcpTraceExporter = ds.getEnableGcpTraceExporter();
368369
this.enableGcpLogExporter = ds.getEnableGcpLogExporter();
369370
this.customOpenTelemetry = ds.getCustomOpenTelemetry();
371+
this.useGlobalOpenTelemetry = ds.getUseGlobalOpenTelemetry();
370372
this.openTelemetry = getOpenTelemetryInstance();
371373
this.bigQuery = getBigQueryConnection();
372374
}
@@ -1048,22 +1050,27 @@ private OpenTelemetry getOpenTelemetryInstance() {
10481050

10491051
OpenTelemetry openTelemetry =
10501052
BigQueryJdbcOpenTelemetry.getOpenTelemetry(
1053+
this.useGlobalOpenTelemetry,
10511054
this.enableGcpTraceExporter,
10521055
this.enableGcpLogExporter,
10531056
this.customOpenTelemetry,
10541057
effectiveCredentials,
10551058
effectiveProjectId);
10561059

1060+
boolean hasExternalOtel = this.customOpenTelemetry != null || this.useGlobalOpenTelemetry;
10571061
Logging localLoggingClient = null;
1058-
if (this.enableGcpLogExporter && !hasCustomOtel) {
1062+
if (this.enableGcpLogExporter && !hasExternalOtel) {
10591063
localLoggingClient =
10601064
BigQueryJdbcOpenTelemetry.createLoggingClient(
10611065
true, null, effectiveCredentials, effectiveProjectId, this.credentials);
10621066
}
10631067

1064-
if (this.enableGcpLogExporter || hasCustomOtel) {
1068+
if (this.enableGcpLogExporter || hasExternalOtel) {
10651069
BigQueryJdbcOpenTelemetry.registerConnection(
1066-
this.connectionId, openTelemetry, localLoggingClient, this.enableGcpLogExporter);
1070+
this.connectionId,
1071+
openTelemetry,
1072+
localLoggingClient,
1073+
this.enableGcpLogExporter && !hasExternalOtel);
10671074
}
10681075

10691076
return openTelemetry;
@@ -1128,7 +1135,9 @@ private BigQuery getBigQueryConnection() {
11281135
if (this.httpTransportOptions != null) {
11291136
bigQueryOptions.setTransportOptions(this.httpTransportOptions);
11301137
}
1131-
if (this.enableGcpTraceExporter || this.customOpenTelemetry != null) {
1138+
if (this.enableGcpTraceExporter
1139+
|| this.customOpenTelemetry != null
1140+
|| this.useGlobalOpenTelemetry) {
11321141
Tracer sdkTracer = this.openTelemetry.getTracer(BigQueryJdbcOpenTelemetry.BIGQUERY_NAMESPACE);
11331142
bigQueryOptions.setOpenTelemetryTracer(sdkTracer);
11341143
this.tracer =

‎java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import com.google.cloud.logging.Logging;
2222
import com.google.cloud.logging.LoggingOptions;
2323
import com.google.common.hash.Hashing;
24+
import io.opentelemetry.api.GlobalOpenTelemetry;
2425
import io.opentelemetry.api.OpenTelemetry;
2526
import io.opentelemetry.api.baggage.Baggage;
2627
import io.opentelemetry.api.trace.Span;
@@ -245,6 +246,7 @@ private static String getCredentialsIdentifier(String credentials) {
245246
* customOpenTelemetry if provided; fallback to an auto-configured GCP exporter if requested.
246247
*/
247248
public static OpenTelemetry getOpenTelemetry(
249+
boolean useGlobalOpenTelemetry,
248250
boolean enableGcpTraceExporter,
249251
boolean enableGcpLogExporter,
250252
OpenTelemetry customOpenTelemetry,
@@ -255,6 +257,10 @@ public static OpenTelemetry getOpenTelemetry(
255257
return customOpenTelemetry;
256258
}
257259

260+
if (useGlobalOpenTelemetry) {
261+
return GlobalOpenTelemetry.get();
262+
}
263+
258264
// NOTE: Currently, tracing only fully supports Application Default Credentials (ADC).
259265
// Once b/503721589 is completed, Service Account (SA) will work as well.
260266
if (!enableGcpTraceExporter && !enableGcpLogExporter) {

‎java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,8 @@ protected boolean removeEldestEntry(Map.Entry<String, Map<String, String>> eldes
168168
static final boolean DEFAULT_ENABLE_GCP_TRACE_EXPORTER_VALUE = false;
169169
static final String ENABLE_GCP_LOG_EXPORTER_PROPERTY_NAME = "enableGcpLogExporter";
170170
static final boolean DEFAULT_ENABLE_GCP_LOG_EXPORTER_VALUE = false;
171+
static final String USE_GLOBAL_OTEL_PROPERTY_NAME = "useGlobalOpenTelemetry";
172+
static final boolean DEFAULT_USE_GLOBAL_OTEL_VALUE = false;
171173
private static final BigQueryJdbcCustomLogger LOG =
172174
new BigQueryJdbcCustomLogger(BigQueryJdbcUrlUtility.class.getName());
173175
static final String FILTER_TABLES_ON_DEFAULT_DATASET_PROPERTY_NAME =
@@ -638,6 +640,12 @@ protected boolean removeEldestEntry(Map.Entry<String, Map<String, String>> eldes
638640
BigQueryConnectionProperty.newBuilder()
639641
.setName(GCP_TELEMETRY_PROJECT_ID_PROPERTY_NAME)
640642
.setDescription("GCP Project ID for OTel exporter.")
643+
.build(),
644+
BigQueryConnectionProperty.newBuilder()
645+
.setName(USE_GLOBAL_OTEL_PROPERTY_NAME)
646+
.setDescription(
647+
"Enables usage of the Global OpenTelemetry instance when true. Default is false.")
648+
.setDefaultValue(String.valueOf(DEFAULT_USE_GLOBAL_OTEL_VALUE))
641649
.build())));
642650

643651
private static final List<String> NETWORK_PROPERTIES =

‎java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/DataSource.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,7 @@ public class DataSource implements javax.sql.DataSource {
123123
private boolean enableGcpLogExporter =
124124
BigQueryJdbcUrlUtility.DEFAULT_ENABLE_GCP_LOG_EXPORTER_VALUE;
125125
private OpenTelemetry customOpenTelemetry;
126+
private boolean useGlobalOpenTelemetry = BigQueryJdbcUrlUtility.DEFAULT_USE_GLOBAL_OTEL_VALUE;
126127

127128
// Make sure the JDBC driver class is loaded.
128129
static {
@@ -358,6 +359,12 @@ public class DataSource implements javax.sql.DataSource {
358359
ds.setEnableGcpLogExporter(
359360
BigQueryJdbcUrlUtility.convertIntToBoolean(
360361
val, BigQueryJdbcUrlUtility.ENABLE_GCP_LOG_EXPORTER_PROPERTY_NAME)))
362+
.put(
363+
BigQueryJdbcUrlUtility.USE_GLOBAL_OTEL_PROPERTY_NAME,
364+
(ds, val) ->
365+
ds.setUseGlobalOpenTelemetry(
366+
BigQueryJdbcUrlUtility.convertIntToBoolean(
367+
val, BigQueryJdbcUrlUtility.USE_GLOBAL_OTEL_PROPERTY_NAME)))
361368
.build();
362369

363370
public static DataSource fromUrl(String url) {
@@ -675,6 +682,11 @@ Properties createProperties() {
675682
BigQueryJdbcUrlUtility.ENABLE_GCP_LOG_EXPORTER_PROPERTY_NAME,
676683
String.valueOf(this.enableGcpLogExporter));
677684
}
685+
if (this.useGlobalOpenTelemetry) {
686+
connectionProperties.setProperty(
687+
BigQueryJdbcUrlUtility.USE_GLOBAL_OTEL_PROPERTY_NAME,
688+
String.valueOf(this.useGlobalOpenTelemetry));
689+
}
678690
return connectionProperties;
679691
}
680692

@@ -832,6 +844,14 @@ public void setCustomOpenTelemetry(OpenTelemetry customOpenTelemetry) {
832844
this.customOpenTelemetry = customOpenTelemetry;
833845
}
834846

847+
public boolean getUseGlobalOpenTelemetry() {
848+
return useGlobalOpenTelemetry;
849+
}
850+
851+
public void setUseGlobalOpenTelemetry(boolean useGlobalOpenTelemetry) {
852+
this.useGlobalOpenTelemetry = useGlobalOpenTelemetry;
853+
}
854+
835855
public void setHighThroughputMinTableSize(Integer highThroughputMinTableSize) {
836856
if (highThroughputMinTableSize != null) {
837857
validateNonNegative(

‎java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,25 @@
1717
package com.google.cloud.bigquery.jdbc;
1818

1919
import static org.junit.jupiter.api.Assertions.*;
20+
import static org.mockito.ArgumentMatchers.any;
21+
import static org.mockito.ArgumentMatchers.anyBoolean;
22+
import static org.mockito.ArgumentMatchers.anyString;
23+
import static org.mockito.ArgumentMatchers.eq;
24+
import static org.mockito.ArgumentMatchers.isNull;
25+
import static org.mockito.Mockito.mock;
26+
import static org.mockito.Mockito.never;
2027

2128
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
2229
import com.google.api.gax.rpc.HeaderProvider;
2330
import com.google.api.gax.rpc.TransportChannelProvider;
31+
import com.google.auth.oauth2.GoogleCredentials;
2432
import com.google.cloud.bigquery.BigQuery;
2533
import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
2634
import com.google.cloud.bigquery.exception.BigQueryJdbcException;
2735
import com.google.cloud.bigquery.storage.v1.BigQueryReadClient;
2836
import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient;
37+
import com.google.cloud.logging.Logging;
38+
import io.opentelemetry.api.OpenTelemetry;
2939
import io.opentelemetry.api.trace.Span;
3040
import io.opentelemetry.api.trace.Tracer;
3141
import io.opentelemetry.sdk.testing.junit5.OpenTelemetryExtension;
@@ -44,6 +54,8 @@
4454
import org.junit.jupiter.api.extension.RegisterExtension;
4555
import org.junit.jupiter.params.ParameterizedTest;
4656
import org.junit.jupiter.params.provider.CsvSource;
57+
import org.mockito.MockedStatic;
58+
import org.mockito.Mockito;
4759

4860
public class BigQueryConnectionTest extends BigQueryJdbcLoggingBaseTest {
4961

@@ -522,4 +534,131 @@ public void testConnectionPropertiesLoggingAndMasking() throws IOException, SQLE
522534
rootLogger.setLevel(originalLevel);
523535
}
524536
}
537+
538+
@ParameterizedTest(
539+
name =
540+
"Case {index}: custom={0}, global={1}, trace={2}, log={3} -> expectTrace={4}, expectLog={5}")
541+
@CsvSource({
542+
// hasCustom, useGlobal, enableTrace, enableLog, expectTrace, expectLog
543+
"true, true, true, true, CUSTOM, CUSTOM",
544+
"true, false, true, true, CUSTOM, CUSTOM",
545+
"false, true, true, true, GLOBAL, GLOBAL",
546+
"false, true, false, false, GLOBAL, GLOBAL",
547+
"false, false, true, false, DRIVER_MANAGED, NONE",
548+
"false, false, false, true, NONE, DRIVER_MANAGED",
549+
"false, false, true, true, DRIVER_MANAGED, DRIVER_MANAGED",
550+
"false, false, false, false, NONE, NONE"
551+
})
552+
public void testOpenTelemetryPrecedenceHierarchy(
553+
boolean hasCustom,
554+
boolean useGlobal,
555+
boolean enableTrace,
556+
boolean enableLog,
557+
String expectTrace,
558+
String expectLog)
559+
throws Exception {
560+
561+
DataSource ds = DataSource.fromUrl(BASE_URL);
562+
ds.setUseGlobalOpenTelemetry(useGlobal);
563+
ds.setEnableGcpTraceExporter(enableTrace);
564+
ds.setEnableGcpLogExporter(enableLog);
565+
566+
OpenTelemetry mockCustomOtel = mock(OpenTelemetry.class);
567+
OpenTelemetry mockGlobalOtel = mock(OpenTelemetry.class);
568+
OpenTelemetry mockDriverManagedOtel = mock(OpenTelemetry.class);
569+
Logging mockLogging = mock(Logging.class);
570+
571+
if (hasCustom) {
572+
ds.setCustomOpenTelemetry(mockCustomOtel);
573+
}
574+
575+
try (MockedStatic<BigQueryJdbcOpenTelemetry> mockedOtel =
576+
Mockito.mockStatic(BigQueryJdbcOpenTelemetry.class);
577+
MockedStatic<BigQueryJdbcOAuthUtility> mockedAuth =
578+
Mockito.mockStatic(BigQueryJdbcOAuthUtility.class);
579+
MockedStatic<GoogleCredentials> mockedCreds = Mockito.mockStatic(GoogleCredentials.class)) {
580+
581+
mockedCreds
582+
.when(GoogleCredentials::getApplicationDefault)
583+
.thenReturn(mock(GoogleCredentials.class));
584+
585+
// Mock parseOAuthProperties to always return ADC type to bypass validation
586+
mockedAuth
587+
.when(() -> BigQueryJdbcOAuthUtility.parseOAuthProperties(any(), anyString()))
588+
.thenAnswer(
589+
invocation -> {
590+
java.util.Map<String, String> props = new java.util.HashMap<>();
591+
props.put(
592+
BigQueryJdbcUrlUtility.OAUTH_TYPE_PROPERTY_NAME,
593+
"APPLICATION_DEFAULT_CREDENTIALS");
594+
return props;
595+
});
596+
597+
mockedAuth
598+
.when(() -> BigQueryJdbcOAuthUtility.getCredentials(any(), any(), any(), any()))
599+
.thenReturn(mock(GoogleCredentials.class));
600+
601+
mockedOtel
602+
.when(
603+
() ->
604+
BigQueryJdbcOpenTelemetry.createLoggingClient(
605+
anyBoolean(), any(), any(), any(), any()))
606+
.thenReturn(mockLogging);
607+
608+
// Stub getOpenTelemetry to return the expected mock based on inputs
609+
mockedOtel
610+
.when(
611+
() ->
612+
BigQueryJdbcOpenTelemetry.getOpenTelemetry(
613+
eq(useGlobal),
614+
eq(enableTrace),
615+
eq(enableLog),
616+
hasCustom ? eq(mockCustomOtel) : isNull(),
617+
any(),
618+
any()))
619+
.thenAnswer(
620+
invocation -> {
621+
if (hasCustom) return mockCustomOtel;
622+
if (useGlobal) return mockGlobalOtel;
623+
if (enableTrace || enableLog) return mockDriverManagedOtel;
624+
return OpenTelemetry.noop();
625+
});
626+
627+
BigQueryConnection connection = new BigQueryConnection(BASE_URL, ds);
628+
629+
boolean shouldBeRegistered = enableLog || hasCustom || useGlobal;
630+
631+
if (!shouldBeRegistered) {
632+
mockedOtel.verify(
633+
() ->
634+
BigQueryJdbcOpenTelemetry.registerConnection(
635+
anyString(), any(), any(), anyBoolean()),
636+
never());
637+
} else {
638+
final OpenTelemetry expectedOtelInstance;
639+
if ("CUSTOM".equals(expectTrace) || "CUSTOM".equals(expectLog)) {
640+
expectedOtelInstance = mockCustomOtel;
641+
} else if ("GLOBAL".equals(expectTrace) || "GLOBAL".equals(expectLog)) {
642+
expectedOtelInstance = mockGlobalOtel;
643+
} else if ("DRIVER_MANAGED".equals(expectTrace) || "DRIVER_MANAGED".equals(expectLog)) {
644+
expectedOtelInstance = mockDriverManagedOtel;
645+
} else {
646+
expectedOtelInstance = OpenTelemetry.noop();
647+
}
648+
649+
boolean expectUseDirectGcp = "DRIVER_MANAGED".equals(expectLog);
650+
Logging expectedLogClient = expectUseDirectGcp ? mockLogging : null;
651+
652+
mockedOtel.verify(
653+
() ->
654+
BigQueryJdbcOpenTelemetry.registerConnection(
655+
anyString(),
656+
eq(expectedOtelInstance),
657+
eq(expectedLogClient),
658+
eq(expectUseDirectGcp)));
659+
}
660+
661+
connection.close();
662+
}
663+
}
525664
}

0 commit comments

Comments
 (0)