Skip to content

Commit 115b9b3

Browse files
keshavdandevakeshavdandeva
authored
fix(bigquery-jdbc): pass connection proxy settings to OpenTelemetry exporters (#14011)
b/543818095 ### Changes 1. **Driver Implementation (`BigQueryJdbcOpenTelemetry.java` & `BigQueryConnection.java`)**: * **Trace Exporter Proxy Support**: Passed `proxyProperties` to `getOpenTelemetry(...)`, configuring `ProxyOptions` (via custom `ProxySelector`) on `OtlpHttpSpanExporterBuilder`. * **Cache Key Isolation**: Added `proxyHost` and `proxyPort` to `SdkCacheKey` so different proxy configurations receive dedicated OpenTelemetry SDK instances. * **Telemetry Metadata**: Passed `this.headerProvider` to `createLoggingClient(...)` for telemetry attribution. 2. **Integration Tests (`ITOpenTelemetryTest.java`)**: * Added `testExecute_withHttpProtocol_andDirectTraceVerification()` to verify OpenTelemetry HTTP trace delivery end-to-end through proxies directly via Cloud Trace API (independent of Cloud Logging). * Refactored `verifyAndFetchTrace(...)` to automatically route `TraceServiceClient` requests through proxy tunnels when `ProxyHost`/`ProxyPort` are present on the connection URL. 3. **Unit Tests**: * Updated `BigQueryConnectionTest.java` and `BigQueryJdbcOpenTelemetryTest.java` to match new method signatures.
1 parent 6bf696f commit 115b9b3

5 files changed

Lines changed: 200 additions & 33 deletions

File tree

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

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
200200
List<ConnectionProperty> queryProperties;
201201
Map<String, String> authProperties;
202202
Map<String, String> overrideProperties;
203+
Map<String, String> proxyProperties;
203204
Credentials credentials;
204205
boolean useStatelessQueryMode;
205206
int numBufferedRows;
@@ -299,7 +300,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
299300
String.valueOf(ds.getRequestGoogleDriveScope()),
300301
BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME);
301302

302-
Map<String, String> proxyProperties =
303+
this.proxyProperties =
303304
BigQueryJdbcProxyUtility.parseProxyProperties(ds, this.connectionClassName);
304305

305306
this.sslTrustStorePath = ds.getSSLTrustStorePath();
@@ -1204,14 +1205,20 @@ private OpenTelemetry getOpenTelemetryInstance() {
12041205
this.customOpenTelemetry,
12051206
this.gcpTelemetryCredentials,
12061207
effectiveProjectId,
1207-
this.credentials);
1208+
this.credentials,
1209+
this.proxyProperties);
12081210

12091211
boolean hasExternalOtel = this.customOpenTelemetry != null || this.useGlobalOpenTelemetry;
12101212
Logging localLoggingClient = null;
12111213
if (this.enableGcpLogExporter && !hasExternalOtel) {
12121214
localLoggingClient =
12131215
BigQueryJdbcOpenTelemetry.createLoggingClient(
1214-
true, null, this.gcpTelemetryCredentials, effectiveProjectId, this.credentials);
1216+
true,
1217+
null,
1218+
this.gcpTelemetryCredentials,
1219+
effectiveProjectId,
1220+
this.credentials,
1221+
this.headerProvider);
12151222
}
12161223

12171224
if (this.enableGcpLogExporter || hasExternalOtel) {

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

Lines changed: 65 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package com.google.cloud.bigquery.jdbc;
1818

19+
import com.google.api.gax.rpc.HeaderProvider;
1920
import com.google.auth.Credentials;
2021
import com.google.auth.oauth2.GoogleCredentials;
2122
import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
@@ -32,9 +33,12 @@
3233
import io.opentelemetry.context.Context;
3334
import io.opentelemetry.context.Scope;
3435
import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporter;
36+
import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporterBuilder;
3537
import io.opentelemetry.exporter.otlp.trace.OtlpGrpcSpanExporter;
3638
import io.opentelemetry.sdk.OpenTelemetrySdk;
3739
import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk;
40+
import io.opentelemetry.sdk.common.export.ProxyOptions;
41+
import java.net.InetSocketAddress;
3842
import java.net.URI;
3943
import java.nio.charset.StandardCharsets;
4044
import java.sql.SQLException;
@@ -85,11 +89,25 @@ private static final class SdkCacheKey {
8589
private final String projectId;
8690
private final String credentialsHashOrPath;
8791
private final boolean enableTrace;
88-
89-
SdkCacheKey(String projectId, String credentialsHashOrPath, boolean enableTrace) {
92+
private final String proxyHost;
93+
private final String proxyPort;
94+
95+
SdkCacheKey(
96+
String projectId,
97+
String credentialsHashOrPath,
98+
boolean enableTrace,
99+
Map<String, String> proxyProperties) {
90100
this.projectId = projectId;
91101
this.credentialsHashOrPath = credentialsHashOrPath;
92102
this.enableTrace = enableTrace;
103+
this.proxyHost =
104+
proxyProperties != null
105+
? proxyProperties.get(BigQueryJdbcUrlUtility.PROXY_HOST_PROPERTY_NAME)
106+
: null;
107+
this.proxyPort =
108+
proxyProperties != null
109+
? proxyProperties.get(BigQueryJdbcUrlUtility.PROXY_PORT_PROPERTY_NAME)
110+
: null;
93111
}
94112

95113
@Override
@@ -99,12 +117,14 @@ public boolean equals(Object o) {
99117
SdkCacheKey that = (SdkCacheKey) o;
100118
return enableTrace == that.enableTrace
101119
&& Objects.equals(projectId, that.projectId)
102-
&& Objects.equals(credentialsHashOrPath, that.credentialsHashOrPath);
120+
&& Objects.equals(credentialsHashOrPath, that.credentialsHashOrPath)
121+
&& Objects.equals(proxyHost, that.proxyHost)
122+
&& Objects.equals(proxyPort, that.proxyPort);
103123
}
104124

105125
@Override
106126
public int hashCode() {
107-
return Objects.hash(projectId, credentialsHashOrPath, enableTrace);
127+
return Objects.hash(projectId, credentialsHashOrPath, enableTrace, proxyHost, proxyPort);
108128
}
109129
}
110130

@@ -181,7 +201,8 @@ public static Logging createLoggingClient(
181201
OpenTelemetry customOpenTelemetry,
182202
String effectiveCredentials,
183203
String effectiveProjectId,
184-
Credentials fallbackCredentials) {
204+
Credentials fallbackCredentials,
205+
HeaderProvider headerProvider) {
185206

186207
if (!enableGcpLogExporter || customOpenTelemetry != null) {
187208
return null;
@@ -200,6 +221,9 @@ public static Logging createLoggingClient(
200221
if (credentials != null) {
201222
loggingOptionsBuilder.setCredentials(credentials);
202223
}
224+
if (headerProvider != null) {
225+
loggingOptionsBuilder.setHeaderProvider(headerProvider);
226+
}
203227
return loggingOptionsBuilder.build().getService();
204228
} catch (Exception e) {
205229
throw new BigQueryJdbcRuntimeException("Failed to initialize Logging client", e);
@@ -317,7 +341,8 @@ public static OpenTelemetry getOpenTelemetry(
317341
OpenTelemetry customOpenTelemetry,
318342
String gcpTelemetryCredentials,
319343
String gcpTelemetryProjectId,
320-
Credentials fallbackCredentials) {
344+
Credentials fallbackCredentials,
345+
Map<String, String> proxyProperties) {
321346

322347
if (customOpenTelemetry != null) {
323348
return customOpenTelemetry;
@@ -335,7 +360,8 @@ public static OpenTelemetry getOpenTelemetry(
335360
new SdkCacheKey(
336361
gcpTelemetryProjectId,
337362
getCredentialsIdentifier(gcpTelemetryCredentials),
338-
enableGcpTraceExporter);
363+
enableGcpTraceExporter,
364+
proxyProperties);
339365
CachedSdk fastCheck = sdkCache.get(key);
340366
if (fastCheck != null) {
341367
CachedSdk result =
@@ -415,11 +441,17 @@ public static OpenTelemetry getOpenTelemetry(
415441
final Credentials finalCredentials = credentials;
416442

417443
if (spanExporter instanceof OtlpHttpSpanExporter) {
418-
return ((OtlpHttpSpanExporter) spanExporter)
419-
.toBuilder()
420-
.setHeaders(
421-
() -> getAuthHeaders(finalCredentials, gcpTelemetryProjectId))
422-
.build();
444+
OtlpHttpSpanExporterBuilder builder =
445+
((OtlpHttpSpanExporter) spanExporter).toBuilder();
446+
builder.setHeaders(
447+
() -> getAuthHeaders(finalCredentials, gcpTelemetryProjectId));
448+
449+
ProxyOptions proxyOptions = createProxyOptions(proxyProperties);
450+
if (proxyOptions != null) {
451+
builder.setProxy(proxyOptions);
452+
}
453+
454+
return builder.build();
423455
}
424456
if (spanExporter instanceof OtlpGrpcSpanExporter) {
425457
return ((OtlpGrpcSpanExporter) spanExporter)
@@ -506,4 +538,25 @@ public static <T> T withTracing(
506538
span.end();
507539
}
508540
}
541+
542+
private static ProxyOptions createProxyOptions(Map<String, String> proxyProperties) {
543+
if (proxyProperties == null) {
544+
return null;
545+
}
546+
547+
final String host = proxyProperties.get(BigQueryJdbcUrlUtility.PROXY_HOST_PROPERTY_NAME);
548+
final String portStr = proxyProperties.get(BigQueryJdbcUrlUtility.PROXY_PORT_PROPERTY_NAME);
549+
if (host == null || host.isEmpty() || portStr == null || portStr.isEmpty()) {
550+
return null;
551+
}
552+
553+
int port;
554+
try {
555+
port = Integer.parseInt(portStr);
556+
} catch (NumberFormatException e) {
557+
throw new BigQueryJdbcRuntimeException("Invalid proxy port number: " + portStr, e);
558+
}
559+
560+
return ProxyOptions.create(InetSocketAddress.createUnresolved(host, port));
561+
}
509562
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -620,7 +620,7 @@ public void testOpenTelemetryPrecedenceHierarchy(
620620
.when(
621621
() ->
622622
BigQueryJdbcOpenTelemetry.createLoggingClient(
623-
anyBoolean(), any(), any(), any(), any()))
623+
anyBoolean(), any(), any(), any(), any(), any()))
624624
.thenReturn(mockLogging);
625625

626626
// Stub getOpenTelemetry to return the expected mock based on inputs
@@ -634,6 +634,7 @@ public void testOpenTelemetryPrecedenceHierarchy(
634634
hasCustom ? eq(mockCustomOtel) : isNull(),
635635
any(),
636636
any(),
637+
any(),
637638
any()))
638639
.thenAnswer(
639640
invocation -> {

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

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ public void testGetOpenTelemetry_withCustomSdk_returnsCustom() {
5252

5353
OpenTelemetry result =
5454
BigQueryJdbcOpenTelemetry.getOpenTelemetry(
55-
false, false, false, mockCustomOtel, null, null, null);
55+
false, false, false, mockCustomOtel, null, null, null, null);
5656

5757
assertThat(result).isSameInstanceAs(mockCustomOtel);
5858
}
@@ -64,15 +64,16 @@ public void testGetOpenTelemetry_withCustomSdkAndFlags_returnsCustom() {
6464
// Custom SDK always takes precedence over individual flags
6565
OpenTelemetry result =
6666
BigQueryJdbcOpenTelemetry.getOpenTelemetry(
67-
false, true, true, mockCustomOtel, null, null, null);
67+
false, true, true, mockCustomOtel, null, null, null, null);
6868

6969
assertThat(result).isSameInstanceAs(mockCustomOtel);
7070
}
7171

7272
@Test
7373
public void testGetOpenTelemetry_noFlags_returnsNoop() {
7474
OpenTelemetry result =
75-
BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, false, false, null, null, null, null);
75+
BigQueryJdbcOpenTelemetry.getOpenTelemetry(
76+
false, false, false, null, null, null, null, null);
7677

7778
assertThat(result).isSameInstanceAs(OpenTelemetry.noop());
7879
}
@@ -88,10 +89,10 @@ public void testGetTracer_respectsScopeName() {
8889
public void testGetOpenTelemetry_cachesSdkInstances() {
8990
OpenTelemetry result1 =
9091
BigQueryJdbcOpenTelemetry.getOpenTelemetry(
91-
false, true, false, null, null, "project1", null);
92+
false, true, false, null, null, "project1", null, null);
9293
OpenTelemetry result2 =
9394
BigQueryJdbcOpenTelemetry.getOpenTelemetry(
94-
false, true, false, null, null, "project1", null);
95+
false, true, false, null, null, "project1", null, null);
9596

9697
assertThat(result1).isSameInstanceAs(result2);
9798
}
@@ -100,40 +101,43 @@ public void testGetOpenTelemetry_cachesSdkInstances() {
100101
public void testGetOpenTelemetry_createsNewInstanceForDifferentKey() {
101102
OpenTelemetry result1 =
102103
BigQueryJdbcOpenTelemetry.getOpenTelemetry(
103-
false, true, false, null, null, "project1", null);
104+
false, true, false, null, null, "project1", null, null);
104105
OpenTelemetry result2 =
105106
BigQueryJdbcOpenTelemetry.getOpenTelemetry(
106-
false, true, false, null, null, "project2", null);
107+
false, true, false, null, null, "project2", null, null);
107108

108109
assertThat(result1).isNotSameInstanceAs(result2);
109110
}
110111

111112
@Test
112113
public void testGetOpenTelemetry_createsNewInstanceForDifferentTraceFlag() {
113114
OpenTelemetry result1 =
114-
BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, true, true, null, null, "project1", null);
115+
BigQueryJdbcOpenTelemetry.getOpenTelemetry(
116+
false, true, true, null, null, "project1", null, null);
115117
OpenTelemetry result2 =
116118
BigQueryJdbcOpenTelemetry.getOpenTelemetry(
117-
false, false, true, null, null, "project1", null);
119+
false, false, true, null, null, "project1", null, null);
118120

119121
assertThat(result1).isNotSameInstanceAs(result2);
120122
}
121123

122124
@Test
123125
public void testGetOpenTelemetry_ignoresEnableLogFlagInCacheKey() {
124126
OpenTelemetry result1 =
125-
BigQueryJdbcOpenTelemetry.getOpenTelemetry(false, true, true, null, null, "project1", null);
127+
BigQueryJdbcOpenTelemetry.getOpenTelemetry(
128+
false, true, true, null, null, "project1", null, null);
126129
OpenTelemetry result2 =
127130
BigQueryJdbcOpenTelemetry.getOpenTelemetry(
128-
false, true, false, null, null, "project1", null);
131+
false, true, false, null, null, "project1", null, null);
129132

130133
assertThat(result1).isSameInstanceAs(result2);
131134
}
132135

133136
@Test
134137
public void testGetOpenTelemetry_withUseGlobalOTel_returnsGlobal() {
135138
OpenTelemetry result =
136-
BigQueryJdbcOpenTelemetry.getOpenTelemetry(true, false, false, null, null, null, null);
139+
BigQueryJdbcOpenTelemetry.getOpenTelemetry(
140+
true, false, false, null, null, null, null, null);
137141

138142
assertThat(result).isSameInstanceAs(GlobalOpenTelemetry.get());
139143
}

0 commit comments

Comments
 (0)