Skip to content

Commit a4a8dc0

Browse files
keshavdandevakeshavdandeva
authored
feat(bigquery-jdbc): support custom OTel credentials and dynamic token refresh (#13302)
b/516416076 This PR enables the BigQuery JDBC driver to use custom Service Account credentials (JSON string or file path) for OpenTelemetry tracing, bypassing the ADC-only limitation of the default GCP extension. ### **Key Changes** * **`BigQueryJdbcOpenTelemetry.java`**: Added a customizer to inject dynamic OAuth2 headers into OTLP exporters (supporting both HTTP and gRPC) while preserving auto-configured properties. * **`ITOpenTelemetryTest.java`**: Added 4 integration tests to verify custom credentials and transport protocols. * **`ITBase.java`**: Moved the shared `getAuthJson()` helper here to remove duplication. * **`pom.xml`**: Moved `opentelemetry-sdk-trace` to compile scope to support the implementation.
1 parent 6cb06cf commit a4a8dc0

5 files changed

Lines changed: 200 additions & 47 deletions

File tree

‎java-bigquery-jdbc/pom.xml‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -366,6 +366,10 @@
366366
<groupId>io.opentelemetry</groupId>
367367
<artifactId>opentelemetry-sdk</artifactId>
368368
</dependency>
369+
<dependency>
370+
<groupId>io.opentelemetry</groupId>
371+
<artifactId>opentelemetry-sdk-trace</artifactId>
372+
</dependency>
369373
<dependency>
370374
<groupId>io.opentelemetry</groupId>
371375
<artifactId>opentelemetry-exporter-otlp</artifactId>
@@ -374,6 +378,10 @@
374378
<groupId>io.opentelemetry</groupId>
375379
<artifactId>opentelemetry-sdk-extension-autoconfigure</artifactId>
376380
</dependency>
381+
<dependency>
382+
<groupId>io.opentelemetry</groupId>
383+
<artifactId>opentelemetry-sdk-extension-autoconfigure-spi</artifactId>
384+
</dependency>
377385
<dependency>
378386
<groupId>io.opentelemetry.contrib</groupId>
379387
<artifactId>opentelemetry-gcp-auth-extension</artifactId>
@@ -439,11 +447,6 @@
439447
<artifactId>opentelemetry-sdk-logs</artifactId>
440448
<scope>test</scope>
441449
</dependency>
442-
<dependency>
443-
<groupId>io.opentelemetry</groupId>
444-
<artifactId>opentelemetry-sdk-trace</artifactId>
445-
<scope>test</scope>
446-
</dependency>
447450
<dependency>
448451
<groupId>com.google.cloud</groupId>
449452
<artifactId>google-cloud-trace</artifactId>

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

Lines changed: 50 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -30,12 +30,16 @@
3030
import io.opentelemetry.api.trace.Tracer;
3131
import io.opentelemetry.context.Context;
3232
import io.opentelemetry.context.Scope;
33+
import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporter;
34+
import io.opentelemetry.exporter.otlp.trace.OtlpGrpcSpanExporter;
3335
import io.opentelemetry.sdk.OpenTelemetrySdk;
3436
import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk;
37+
import java.net.URI;
3538
import java.nio.charset.StandardCharsets;
3639
import java.sql.SQLException;
3740
import java.util.Collection;
3841
import java.util.HashMap;
42+
import java.util.List;
3943
import java.util.Map;
4044
import java.util.Objects;
4145
import java.util.concurrent.Callable;
@@ -61,9 +65,8 @@ public class BigQueryJdbcOpenTelemetry {
6165
private static final String OTEL_LOGS_EXPORTER = "otel.logs.exporter";
6266
private static final String OTEL_METRICS_EXPORTER = "otel.metrics.exporter";
6367
private static final String GOOGLE_CLOUD_PROJECT = "google.cloud.project";
64-
private static final String CREDENTIALS_JSON = "google.cloud.credentials.json";
65-
private static final String CREDENTIALS_PATH = "google.cloud.credentials.path";
6668
private static final String OTLP_ENDPOINT_VALUE = "https://telemetry.googleapis.com:443";
69+
private static final URI OTLP_ENDPOINT_URI = URI.create(OTLP_ENDPOINT_VALUE);
6770
private static final String EXPORTER_NONE = "none";
6871
private static final String EXPORTER_OTLP = "otlp";
6972
private static final String OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT =
@@ -230,6 +233,25 @@ public static Collection<TelemetryConfig> getRegisteredConfigs() {
230233
return connectionConfigs.values();
231234
}
232235

236+
private static Map<String, String> getAuthHeaders(Credentials credentials) {
237+
try {
238+
Map<String, List<String>> metadata = credentials.getRequestMetadata(OTLP_ENDPOINT_URI);
239+
Map<String, String> headers = new HashMap<>();
240+
metadata.forEach(
241+
(headerKey, headerValues) -> {
242+
if (!headerValues.isEmpty()) {
243+
headers.put(headerKey, headerValues.get(0));
244+
}
245+
});
246+
return headers;
247+
} catch (Exception e) {
248+
// We log the warning and return an empty map, allowing the exporter to fail gracefully
249+
// with a standard OTLP response code (e.g., 401 Unauthorized) handled by OTel.
250+
LOG.warning(e, "Failed to get auth headers");
251+
return new HashMap<>();
252+
}
253+
}
254+
233255
private static String getCredentialsIdentifier(String credentials) {
234256
if (credentials == null) {
235257
return "";
@@ -261,8 +283,6 @@ public static OpenTelemetry getOpenTelemetry(
261283
return GlobalOpenTelemetry.get();
262284
}
263285

264-
// NOTE: Currently, tracing only fully supports Application Default Credentials (ADC).
265-
// Once b/503721589 is completed, Service Account (SA) will work as well.
266286
if (!enableGcpTraceExporter && !enableGcpLogExporter) {
267287
return OpenTelemetry.noop();
268288
}
@@ -276,14 +296,6 @@ public static OpenTelemetry getOpenTelemetry(
276296
key,
277297
k -> {
278298
Map<String, String> props = new HashMap<>();
279-
if (gcpTelemetryCredentials != null) {
280-
byte[] credsBytes = gcpTelemetryCredentials.getBytes(StandardCharsets.UTF_8);
281-
if (BigQueryJdbcOAuthUtility.isJson(credsBytes)) {
282-
props.put(CREDENTIALS_JSON, gcpTelemetryCredentials);
283-
} else {
284-
props.put(CREDENTIALS_PATH, gcpTelemetryCredentials);
285-
}
286-
}
287299

288300
if (enableGcpTraceExporter) {
289301
props.put(OTEL_TRACES_EXPORTER, EXPORTER_OTLP);
@@ -313,7 +325,32 @@ public static OpenTelemetry getOpenTelemetry(
313325
}
314326

315327
AutoConfiguredOpenTelemetrySdk autoConfigured =
316-
AutoConfiguredOpenTelemetrySdk.builder().addPropertiesSupplier(() -> props).build();
328+
AutoConfiguredOpenTelemetrySdk.builder()
329+
.addPropertiesSupplier(() -> props)
330+
.addSpanExporterCustomizer(
331+
(spanExporter, configProperties) -> {
332+
if (gcpTelemetryCredentials == null) {
333+
return spanExporter;
334+
}
335+
try {
336+
Credentials credentials =
337+
resolveCredentialsFromString(gcpTelemetryCredentials);
338+
if (spanExporter instanceof OtlpHttpSpanExporter) {
339+
return ((OtlpHttpSpanExporter) spanExporter)
340+
.toBuilder().setHeaders(() -> getAuthHeaders(credentials)).build();
341+
}
342+
if (spanExporter instanceof OtlpGrpcSpanExporter) {
343+
return ((OtlpGrpcSpanExporter) spanExporter)
344+
.toBuilder().setHeaders(() -> getAuthHeaders(credentials)).build();
345+
}
346+
} catch (Exception e) {
347+
LOG.warning(
348+
e,
349+
"Failed to resolve telemetry credentials. Telemetry will be exported using default OpenTelemetry configuration (custom authentication headers will not be injected).");
350+
}
351+
return spanExporter;
352+
})
353+
.build();
317354

318355
OpenTelemetrySdk sdk = autoConfigured.getOpenTelemetrySdk();
319356

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

Lines changed: 0 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -24,15 +24,12 @@
2424
import com.google.auth.oauth2.GoogleCredentials;
2525
import com.google.cloud.ServiceOptions;
2626
import com.google.gson.JsonObject;
27-
import com.google.gson.JsonParser;
2827
import java.io.ByteArrayInputStream;
2928
import java.io.File;
3029
import java.io.IOException;
3130
import java.io.InputStream;
32-
import java.io.InputStreamReader;
3331
import java.nio.charset.StandardCharsets;
3432
import java.nio.file.Files;
35-
import java.nio.file.Paths;
3633
import java.sql.Connection;
3734
import java.sql.DriverManager;
3835
import java.sql.ResultSet;
@@ -48,25 +45,6 @@
4845
public class ITAuthTests extends ITBase {
4946
static final String PROJECT_ID = ServiceOptions.getDefaultProjectId();
5047

51-
private JsonObject getAuthJson() throws IOException {
52-
final String secret = requireEnvVar("SA_SECRET");
53-
JsonObject authJson;
54-
// Supporting both formats of SA_SECRET:
55-
// - Local runs can point to a json file
56-
// - Cloud Build has JSON value
57-
try {
58-
InputStream stream = Files.newInputStream(Paths.get(secret));
59-
InputStreamReader reader = new InputStreamReader(stream);
60-
authJson = JsonParser.parseReader(reader).getAsJsonObject();
61-
} catch (IOException e) {
62-
authJson = JsonParser.parseString(secret).getAsJsonObject();
63-
}
64-
assertTrue(authJson.has("client_email"));
65-
assertTrue(authJson.has("private_key"));
66-
assertTrue(authJson.has("project_id"));
67-
return authJson;
68-
}
69-
7048
private void validateConnection(String connection_uri) throws SQLException {
7149
Connection connection = DriverManager.getConnection(connection_uri);
7250
assertNotNull(connection);

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,23 @@
1717
package com.google.cloud.bigquery.jdbc.it;
1818

1919
import static org.junit.jupiter.api.Assertions.assertNotNull;
20+
import static org.junit.jupiter.api.Assertions.assertTrue;
2021

22+
import com.google.auth.oauth2.GoogleCredentials;
2123
import com.google.cloud.ServiceOptions;
2224
import com.google.cloud.bigquery.BigQuery;
2325
import com.google.cloud.bigquery.BigQueryOptions;
2426
import com.google.cloud.bigquery.QueryJobConfiguration;
2527
import com.google.cloud.bigquery.jdbc.BigQueryJdbcBaseTest;
28+
import com.google.gson.JsonObject;
29+
import com.google.gson.JsonParser;
30+
import java.io.ByteArrayInputStream;
31+
import java.io.IOException;
32+
import java.io.InputStream;
33+
import java.io.InputStreamReader;
34+
import java.nio.charset.StandardCharsets;
35+
import java.nio.file.Files;
36+
import java.nio.file.Paths;
2637
import java.sql.Connection;
2738
import java.sql.ResultSet;
2839
import java.sql.SQLException;
@@ -291,6 +302,31 @@ protected static String requireEnvVar(String varName) {
291302
return value;
292303
}
293304

305+
protected static JsonObject getAuthJson() throws IOException {
306+
final String secret = requireEnvVar("SA_SECRET");
307+
JsonObject authJson;
308+
// Supporting both formats of SA_SECRET:
309+
// - Local runs can point to a json file
310+
// - Cloud Build has JSON value
311+
try {
312+
InputStream stream = Files.newInputStream(Paths.get(secret));
313+
InputStreamReader reader = new InputStreamReader(stream);
314+
authJson = JsonParser.parseReader(reader).getAsJsonObject();
315+
} catch (IOException e) {
316+
authJson = JsonParser.parseString(secret).getAsJsonObject();
317+
}
318+
assertTrue(authJson.has("client_email"));
319+
assertTrue(authJson.has("private_key"));
320+
assertTrue(authJson.has("project_id"));
321+
return authJson;
322+
}
323+
324+
protected static GoogleCredentials getCredentials() throws IOException {
325+
JsonObject authJson = getAuthJson();
326+
return GoogleCredentials.fromStream(
327+
new ByteArrayInputStream(authJson.toString().getBytes(StandardCharsets.UTF_8)));
328+
}
329+
294330
protected int resultSetRowCount(ResultSet resultSet) throws SQLException {
295331
int rowCount = 0;
296332
while (resultSet.next()) {

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

Lines changed: 106 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -22,16 +22,23 @@
2222
import static org.junit.jupiter.api.Assertions.assertThrows;
2323
import static org.junit.jupiter.api.Assertions.assertTrue;
2424

25+
import com.google.api.gax.core.FixedCredentialsProvider;
2526
import com.google.api.gax.paging.Page;
27+
import com.google.auth.oauth2.GoogleCredentials;
2628
import com.google.cloud.ServiceOptions;
2729
import com.google.cloud.bigquery.jdbc.BigQueryConnection;
2830
import com.google.cloud.bigquery.jdbc.DataSource;
2931
import com.google.cloud.logging.LogEntry;
3032
import com.google.cloud.logging.Logging;
3133
import com.google.cloud.logging.LoggingOptions;
3234
import com.google.cloud.trace.v1.TraceServiceClient;
35+
import com.google.cloud.trace.v1.TraceServiceSettings;
3336
import com.google.devtools.cloudtrace.v1.Trace;
3437
import com.google.devtools.cloudtrace.v1.TraceSpan;
38+
import com.google.gson.JsonObject;
39+
import java.io.File;
40+
import java.nio.charset.StandardCharsets;
41+
import java.nio.file.Files;
3542
import java.sql.Connection;
3643
import java.sql.ResultSet;
3744
import java.sql.SQLException;
@@ -40,13 +47,10 @@
4047
import java.util.List;
4148
import org.junit.jupiter.api.Test;
4249

43-
public class ITOpenTelemetryTest {
50+
public class ITOpenTelemetryTest extends ITBase {
4451

4552
private static final String PROJECT_ID = ServiceOptions.getDefaultProjectId();
46-
private static final String CONNECTION_URL =
47-
String.format(
48-
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;ProjectId=%s;OAuthType=3;Timeout=3600;",
49-
PROJECT_ID);
53+
private static final String CONNECTION_URL = connectionUrl;
5054

5155
@Test
5256
public void testExecute_withOpenTelemetryGcpExporter() throws Exception {
@@ -163,9 +167,97 @@ public void testExecute_withErrorCorrelation() throws Exception {
163167
"Traces must contain JDBC parent span 'BigQueryStatement.executeQuery'");
164168
}
165169

170+
@Test
171+
public void testExecute_withCustomCredentialsJson() throws Exception {
172+
JsonObject authJson = getAuthJson();
173+
DataSource ds = DataSource.fromUrl(CONNECTION_URL);
174+
ds.setEnableGcpTraceExporter(true);
175+
ds.setGcpTelemetryProjectId(PROJECT_ID);
176+
ds.setGcpTelemetryCredentials(authJson.toString());
177+
178+
verifyTraceDelivery(ds);
179+
}
180+
181+
@Test
182+
public void testExecute_withCustomCredentialsFilePath() throws Exception {
183+
JsonObject authJson = getAuthJson();
184+
File tempFile = File.createTempFile("auth", ".json");
185+
tempFile.deleteOnExit();
186+
Files.write(tempFile.toPath(), authJson.toString().getBytes(StandardCharsets.UTF_8));
187+
188+
DataSource ds = DataSource.fromUrl(CONNECTION_URL);
189+
ds.setEnableGcpTraceExporter(true);
190+
ds.setGcpTelemetryProjectId(PROJECT_ID);
191+
ds.setGcpTelemetryCredentials(tempFile.getAbsolutePath());
192+
193+
verifyTraceDelivery(ds);
194+
}
195+
196+
@Test
197+
public void testExecute_withHttpProtocol() throws Exception {
198+
JsonObject authJson = getAuthJson();
199+
System.setProperty("otel.exporter.otlp.protocol", "http/protobuf");
200+
201+
try {
202+
DataSource ds = DataSource.fromUrl(CONNECTION_URL);
203+
ds.setEnableGcpTraceExporter(true);
204+
ds.setGcpTelemetryProjectId(PROJECT_ID);
205+
ds.setGcpTelemetryCredentials(authJson.toString());
206+
207+
verifyTraceDelivery(ds);
208+
} finally {
209+
System.clearProperty("otel.exporter.otlp.protocol");
210+
}
211+
}
212+
213+
@Test
214+
public void testExecute_withGrpcProtocol() throws Exception {
215+
JsonObject authJson = getAuthJson();
216+
System.setProperty("otel.exporter.otlp.protocol", "grpc");
217+
218+
try {
219+
DataSource ds = DataSource.fromUrl(CONNECTION_URL);
220+
ds.setEnableGcpTraceExporter(true);
221+
ds.setGcpTelemetryProjectId(PROJECT_ID);
222+
ds.setGcpTelemetryCredentials(authJson.toString());
223+
224+
verifyTraceDelivery(ds);
225+
} finally {
226+
System.clearProperty("otel.exporter.otlp.protocol");
227+
}
228+
}
229+
230+
private void verifyTraceDelivery(DataSource ds) throws Exception {
231+
ds.setEnableGcpLogExporter(true);
232+
ds.setLogLevel("5");
233+
234+
String connectionUuid = null;
235+
try (Connection connection = ds.getConnection();
236+
Statement statement = connection.createStatement()) {
237+
238+
BigQueryConnection bqConnection = connection.unwrap(BigQueryConnection.class);
239+
connectionUuid = bqConnection.getConnectionId();
240+
241+
String query = "SELECT 1;";
242+
try (ResultSet rs = statement.executeQuery(query)) {
243+
assertTrue(rs.next());
244+
}
245+
}
246+
247+
String traceId = verifyAndFetchLogs(connectionUuid);
248+
Trace trace = verifyAndFetchTrace(traceId);
249+
assertNotNull(trace, "Trace must be found");
250+
}
251+
166252
private String verifyAndFetchLogs(String connectionUuid) throws Exception {
253+
GoogleCredentials credentials = getCredentials();
254+
167255
try (Logging logging =
168-
LoggingOptions.newBuilder().setProjectId(PROJECT_ID).build().getService()) {
256+
LoggingOptions.newBuilder()
257+
.setProjectId(PROJECT_ID)
258+
.setCredentials(credentials)
259+
.build()
260+
.getService()) {
169261
String filter =
170262
"logName:\"projects/"
171263
+ PROJECT_ID
@@ -198,7 +290,14 @@ private Trace verifyAndFetchTrace(String traceId) throws Exception {
198290
hexTraceId = traceId.substring(traceId.lastIndexOf("/traces/") + 8);
199291
}
200292

201-
try (TraceServiceClient traceClient = TraceServiceClient.create()) {
293+
GoogleCredentials credentials = getCredentials();
294+
295+
TraceServiceSettings settings =
296+
TraceServiceSettings.newBuilder()
297+
.setCredentialsProvider(FixedCredentialsProvider.create(credentials))
298+
.build();
299+
300+
try (TraceServiceClient traceClient = TraceServiceClient.create(settings)) {
202301
Trace trace = fetchTraceWithRetry(traceClient, PROJECT_ID, hexTraceId);
203302
assertNotNull(trace, "Trace must be found in Cloud Trace API: " + hexTraceId);
204303
return trace;

0 commit comments

Comments
 (0)