Skip to content

Commit f3b1301

Browse files
keshavdandevakeshavdandeva
committed
feat(bigquery-jdbc): support custom OTel credentials and dynamic token refresh
1 parent 67b84ef commit f3b1301

5 files changed

Lines changed: 162 additions & 41 deletions

File tree

‎java-bigquery-jdbc/pom.xml‎

Lines changed: 4 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>
@@ -443,11 +447,6 @@
443447
<artifactId>opentelemetry-sdk-logs</artifactId>
444448
<scope>test</scope>
445449
</dependency>
446-
<dependency>
447-
<groupId>io.opentelemetry</groupId>
448-
<artifactId>opentelemetry-sdk-trace</artifactId>
449-
<scope>test</scope>
450-
</dependency>
451450
<dependency>
452451
<groupId>com.google.cloud</groupId>
453452
<artifactId>google-cloud-trace</artifactId>

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

Lines changed: 44 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -30,12 +30,17 @@
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.io.IOException;
38+
import java.net.URI;
3539
import java.nio.charset.StandardCharsets;
3640
import java.sql.SQLException;
3741
import java.util.Collection;
3842
import java.util.HashMap;
43+
import java.util.List;
3944
import java.util.Map;
4045
import java.util.Objects;
4146
import java.util.concurrent.Callable;
@@ -61,8 +66,6 @@ public class BigQueryJdbcOpenTelemetry {
6166
private static final String OTEL_LOGS_EXPORTER = "otel.logs.exporter";
6267
private static final String OTEL_METRICS_EXPORTER = "otel.metrics.exporter";
6368
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";
6669
private static final String OTLP_ENDPOINT_VALUE = "https://telemetry.googleapis.com:443";
6770
private static final String EXPORTER_NONE = "none";
6871
private static final String EXPORTER_OTLP = "otlp";
@@ -230,6 +233,26 @@ 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 =
239+
credentials.getRequestMetadata(URI.create(OTLP_ENDPOINT_VALUE));
240+
Map<String, String> headers = new HashMap<>();
241+
metadata.forEach(
242+
(headerKey, headerValues) -> {
243+
if (!headerValues.isEmpty()) {
244+
headers.put(headerKey, headerValues.get(0));
245+
}
246+
});
247+
return headers;
248+
} catch (IOException e) {
249+
// We log the warning and return an empty map, allowing the exporter to fail gracefully
250+
// with a standard OTLP response code (e.g., 401 Unauthorized) handled by OTel.
251+
LOG.warning("Failed to get auth headers: %s", e.getMessage());
252+
return new HashMap<>();
253+
}
254+
}
255+
233256
private static String getCredentialsIdentifier(String credentials) {
234257
if (credentials == null) {
235258
return "";
@@ -261,8 +284,6 @@ public static OpenTelemetry getOpenTelemetry(
261284
return GlobalOpenTelemetry.get();
262285
}
263286

264-
// NOTE: Currently, tracing only fully supports Application Default Credentials (ADC).
265-
// Once b/503721589 is completed, Service Account (SA) will work as well.
266287
if (!enableGcpTraceExporter && !enableGcpLogExporter) {
267288
return OpenTelemetry.noop();
268289
}
@@ -276,14 +297,6 @@ public static OpenTelemetry getOpenTelemetry(
276297
key,
277298
k -> {
278299
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-
}
287300

288301
if (enableGcpTraceExporter) {
289302
props.put(OTEL_TRACES_EXPORTER, EXPORTER_OTLP);
@@ -313,7 +326,25 @@ public static OpenTelemetry getOpenTelemetry(
313326
}
314327

315328
AutoConfiguredOpenTelemetrySdk autoConfigured =
316-
AutoConfiguredOpenTelemetrySdk.builder().addPropertiesSupplier(() -> props).build();
329+
AutoConfiguredOpenTelemetrySdk.builder()
330+
.addPropertiesSupplier(() -> props)
331+
.addSpanExporterCustomizer(
332+
(spanExporter, configProperties) -> {
333+
if (gcpTelemetryCredentials != null) {
334+
Credentials credentials =
335+
resolveCredentialsFromString(gcpTelemetryCredentials);
336+
if (spanExporter instanceof OtlpHttpSpanExporter) {
337+
return ((OtlpHttpSpanExporter) spanExporter)
338+
.toBuilder().setHeaders(() -> getAuthHeaders(credentials)).build();
339+
}
340+
if (spanExporter instanceof OtlpGrpcSpanExporter) {
341+
return ((OtlpGrpcSpanExporter) spanExporter)
342+
.toBuilder().setHeaders(() -> getAuthHeaders(credentials)).build();
343+
}
344+
}
345+
return spanExporter;
346+
})
347+
.build();
317348

318349
OpenTelemetrySdk sdk = autoConfigured.getOpenTelemetrySdk();
319350

‎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: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,20 @@
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

2122
import com.google.cloud.ServiceOptions;
2223
import com.google.cloud.bigquery.BigQuery;
2324
import com.google.cloud.bigquery.BigQueryOptions;
2425
import com.google.cloud.bigquery.QueryJobConfiguration;
2526
import com.google.cloud.bigquery.jdbc.BigQueryJdbcBaseTest;
27+
import com.google.gson.JsonObject;
28+
import com.google.gson.JsonParser;
29+
import java.io.IOException;
30+
import java.io.InputStream;
31+
import java.io.InputStreamReader;
32+
import java.nio.file.Files;
33+
import java.nio.file.Paths;
2634
import java.sql.Connection;
2735
import java.sql.ResultSet;
2836
import java.sql.SQLException;
@@ -291,6 +299,25 @@ protected static String requireEnvVar(String varName) {
291299
return value;
292300
}
293301

302+
protected static JsonObject getAuthJson() throws IOException {
303+
final String secret = requireEnvVar("SA_SECRET");
304+
JsonObject authJson;
305+
// Supporting both formats of SA_SECRET:
306+
// - Local runs can point to a json file
307+
// - Cloud Build has JSON value
308+
try {
309+
InputStream stream = Files.newInputStream(Paths.get(secret));
310+
InputStreamReader reader = new InputStreamReader(stream);
311+
authJson = JsonParser.parseReader(reader).getAsJsonObject();
312+
} catch (IOException e) {
313+
authJson = JsonParser.parseString(secret).getAsJsonObject();
314+
}
315+
assertTrue(authJson.has("client_email"));
316+
assertTrue(authJson.has("private_key"));
317+
assertTrue(authJson.has("project_id"));
318+
return authJson;
319+
}
320+
294321
protected int resultSetRowCount(ResultSet resultSet) throws SQLException {
295322
int rowCount = 0;
296323
while (resultSet.next()) {

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

Lines changed: 87 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,10 @@
3232
import com.google.cloud.trace.v1.TraceServiceClient;
3333
import com.google.devtools.cloudtrace.v1.Trace;
3434
import com.google.devtools.cloudtrace.v1.TraceSpan;
35+
import com.google.gson.JsonObject;
36+
import java.io.File;
37+
import java.nio.charset.StandardCharsets;
38+
import java.nio.file.Files;
3539
import java.sql.Connection;
3640
import java.sql.ResultSet;
3741
import java.sql.SQLException;
@@ -40,7 +44,7 @@
4044
import java.util.List;
4145
import org.junit.jupiter.api.Test;
4246

43-
public class ITOpenTelemetryTest {
47+
public class ITOpenTelemetryTest extends ITBase {
4448

4549
private static final String PROJECT_ID = ServiceOptions.getDefaultProjectId();
4650
private static final String CONNECTION_URL =
@@ -163,6 +167,88 @@ 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 {
167253
try (Logging logging =
168254
LoggingOptions.newBuilder().setProjectId(PROJECT_ID).build().getService()) {

0 commit comments

Comments
 (0)