From f3b1301e6d52356326dfcaf3eabd3d35ec9cd098 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Fri, 29 May 2026 17:14:16 +0000 Subject: [PATCH] feat(bigquery-jdbc): support custom OTel credentials and dynamic token refresh --- java-bigquery-jdbc/pom.xml | 9 +- .../jdbc/BigQueryJdbcOpenTelemetry.java | 57 +++++++++--- .../cloud/bigquery/jdbc/it/ITAuthTests.java | 22 ----- .../google/cloud/bigquery/jdbc/it/ITBase.java | 27 ++++++ .../bigquery/jdbc/it/ITOpenTelemetryTest.java | 88 ++++++++++++++++++- 5 files changed, 162 insertions(+), 41 deletions(-) diff --git a/java-bigquery-jdbc/pom.xml b/java-bigquery-jdbc/pom.xml index 19e08d77861b..01539410ebaf 100644 --- a/java-bigquery-jdbc/pom.xml +++ b/java-bigquery-jdbc/pom.xml @@ -366,6 +366,10 @@ io.opentelemetry opentelemetry-sdk + + io.opentelemetry + opentelemetry-sdk-trace + io.opentelemetry opentelemetry-exporter-otlp @@ -443,11 +447,6 @@ opentelemetry-sdk-logs test - - io.opentelemetry - opentelemetry-sdk-trace - test - com.google.cloud google-cloud-trace diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java index 5f35c9cde510..cf832ab9cc26 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java @@ -30,12 +30,17 @@ import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.context.Context; import io.opentelemetry.context.Scope; +import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporter; +import io.opentelemetry.exporter.otlp.trace.OtlpGrpcSpanExporter; import io.opentelemetry.sdk.OpenTelemetrySdk; import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk; +import java.io.IOException; +import java.net.URI; import java.nio.charset.StandardCharsets; import java.sql.SQLException; import java.util.Collection; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.Objects; import java.util.concurrent.Callable; @@ -61,8 +66,6 @@ public class BigQueryJdbcOpenTelemetry { private static final String OTEL_LOGS_EXPORTER = "otel.logs.exporter"; private static final String OTEL_METRICS_EXPORTER = "otel.metrics.exporter"; private static final String GOOGLE_CLOUD_PROJECT = "google.cloud.project"; - private static final String CREDENTIALS_JSON = "google.cloud.credentials.json"; - private static final String CREDENTIALS_PATH = "google.cloud.credentials.path"; private static final String OTLP_ENDPOINT_VALUE = "https://telemetry.googleapis.com:443"; private static final String EXPORTER_NONE = "none"; private static final String EXPORTER_OTLP = "otlp"; @@ -230,6 +233,26 @@ public static Collection getRegisteredConfigs() { return connectionConfigs.values(); } + private static Map getAuthHeaders(Credentials credentials) { + try { + Map> metadata = + credentials.getRequestMetadata(URI.create(OTLP_ENDPOINT_VALUE)); + Map headers = new HashMap<>(); + metadata.forEach( + (headerKey, headerValues) -> { + if (!headerValues.isEmpty()) { + headers.put(headerKey, headerValues.get(0)); + } + }); + return headers; + } catch (IOException e) { + // We log the warning and return an empty map, allowing the exporter to fail gracefully + // with a standard OTLP response code (e.g., 401 Unauthorized) handled by OTel. + LOG.warning("Failed to get auth headers: %s", e.getMessage()); + return new HashMap<>(); + } + } + private static String getCredentialsIdentifier(String credentials) { if (credentials == null) { return ""; @@ -261,8 +284,6 @@ public static OpenTelemetry getOpenTelemetry( return GlobalOpenTelemetry.get(); } - // NOTE: Currently, tracing only fully supports Application Default Credentials (ADC). - // Once b/503721589 is completed, Service Account (SA) will work as well. if (!enableGcpTraceExporter && !enableGcpLogExporter) { return OpenTelemetry.noop(); } @@ -276,14 +297,6 @@ public static OpenTelemetry getOpenTelemetry( key, k -> { Map props = new HashMap<>(); - if (gcpTelemetryCredentials != null) { - byte[] credsBytes = gcpTelemetryCredentials.getBytes(StandardCharsets.UTF_8); - if (BigQueryJdbcOAuthUtility.isJson(credsBytes)) { - props.put(CREDENTIALS_JSON, gcpTelemetryCredentials); - } else { - props.put(CREDENTIALS_PATH, gcpTelemetryCredentials); - } - } if (enableGcpTraceExporter) { props.put(OTEL_TRACES_EXPORTER, EXPORTER_OTLP); @@ -313,7 +326,25 @@ public static OpenTelemetry getOpenTelemetry( } AutoConfiguredOpenTelemetrySdk autoConfigured = - AutoConfiguredOpenTelemetrySdk.builder().addPropertiesSupplier(() -> props).build(); + AutoConfiguredOpenTelemetrySdk.builder() + .addPropertiesSupplier(() -> props) + .addSpanExporterCustomizer( + (spanExporter, configProperties) -> { + if (gcpTelemetryCredentials != null) { + Credentials credentials = + resolveCredentialsFromString(gcpTelemetryCredentials); + if (spanExporter instanceof OtlpHttpSpanExporter) { + return ((OtlpHttpSpanExporter) spanExporter) + .toBuilder().setHeaders(() -> getAuthHeaders(credentials)).build(); + } + if (spanExporter instanceof OtlpGrpcSpanExporter) { + return ((OtlpGrpcSpanExporter) spanExporter) + .toBuilder().setHeaders(() -> getAuthHeaders(credentials)).build(); + } + } + return spanExporter; + }) + .build(); OpenTelemetrySdk sdk = autoConfigured.getOpenTelemetrySdk(); diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITAuthTests.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITAuthTests.java index 0877553e42c0..2a3a6bfac799 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITAuthTests.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITAuthTests.java @@ -24,15 +24,12 @@ import com.google.auth.oauth2.GoogleCredentials; import com.google.cloud.ServiceOptions; import com.google.gson.JsonObject; -import com.google.gson.JsonParser; import java.io.ByteArrayInputStream; import java.io.File; import java.io.IOException; import java.io.InputStream; -import java.io.InputStreamReader; import java.nio.charset.StandardCharsets; import java.nio.file.Files; -import java.nio.file.Paths; import java.sql.Connection; import java.sql.DriverManager; import java.sql.ResultSet; @@ -48,25 +45,6 @@ public class ITAuthTests extends ITBase { static final String PROJECT_ID = ServiceOptions.getDefaultProjectId(); - private JsonObject getAuthJson() throws IOException { - final String secret = requireEnvVar("SA_SECRET"); - JsonObject authJson; - // Supporting both formats of SA_SECRET: - // - Local runs can point to a json file - // - Cloud Build has JSON value - try { - InputStream stream = Files.newInputStream(Paths.get(secret)); - InputStreamReader reader = new InputStreamReader(stream); - authJson = JsonParser.parseReader(reader).getAsJsonObject(); - } catch (IOException e) { - authJson = JsonParser.parseString(secret).getAsJsonObject(); - } - assertTrue(authJson.has("client_email")); - assertTrue(authJson.has("private_key")); - assertTrue(authJson.has("project_id")); - return authJson; - } - private void validateConnection(String connection_uri) throws SQLException { Connection connection = DriverManager.getConnection(connection_uri); assertNotNull(connection); diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBase.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBase.java index 5b4d36fac4fe..21a358ad3f3b 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBase.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBase.java @@ -17,12 +17,20 @@ package com.google.cloud.bigquery.jdbc.it; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; import com.google.cloud.ServiceOptions; import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.BigQueryOptions; import com.google.cloud.bigquery.QueryJobConfiguration; import com.google.cloud.bigquery.jdbc.BigQueryJdbcBaseTest; +import com.google.gson.JsonObject; +import com.google.gson.JsonParser; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.nio.file.Files; +import java.nio.file.Paths; import java.sql.Connection; import java.sql.ResultSet; import java.sql.SQLException; @@ -291,6 +299,25 @@ protected static String requireEnvVar(String varName) { return value; } + protected static JsonObject getAuthJson() throws IOException { + final String secret = requireEnvVar("SA_SECRET"); + JsonObject authJson; + // Supporting both formats of SA_SECRET: + // - Local runs can point to a json file + // - Cloud Build has JSON value + try { + InputStream stream = Files.newInputStream(Paths.get(secret)); + InputStreamReader reader = new InputStreamReader(stream); + authJson = JsonParser.parseReader(reader).getAsJsonObject(); + } catch (IOException e) { + authJson = JsonParser.parseString(secret).getAsJsonObject(); + } + assertTrue(authJson.has("client_email")); + assertTrue(authJson.has("private_key")); + assertTrue(authJson.has("project_id")); + return authJson; + } + protected int resultSetRowCount(ResultSet resultSet) throws SQLException { int rowCount = 0; while (resultSet.next()) { diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITOpenTelemetryTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITOpenTelemetryTest.java index 8e7ffa92c2e9..5cf527f89503 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITOpenTelemetryTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITOpenTelemetryTest.java @@ -32,6 +32,10 @@ import com.google.cloud.trace.v1.TraceServiceClient; import com.google.devtools.cloudtrace.v1.Trace; import com.google.devtools.cloudtrace.v1.TraceSpan; +import com.google.gson.JsonObject; +import java.io.File; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; import java.sql.Connection; import java.sql.ResultSet; import java.sql.SQLException; @@ -40,7 +44,7 @@ import java.util.List; import org.junit.jupiter.api.Test; -public class ITOpenTelemetryTest { +public class ITOpenTelemetryTest extends ITBase { private static final String PROJECT_ID = ServiceOptions.getDefaultProjectId(); private static final String CONNECTION_URL = @@ -163,6 +167,88 @@ public void testExecute_withErrorCorrelation() throws Exception { "Traces must contain JDBC parent span 'BigQueryStatement.executeQuery'"); } + @Test + public void testExecute_withCustomCredentialsJson() throws Exception { + JsonObject authJson = getAuthJson(); + DataSource ds = DataSource.fromUrl(CONNECTION_URL); + ds.setEnableGcpTraceExporter(true); + ds.setGcpTelemetryProjectId(PROJECT_ID); + ds.setGcpTelemetryCredentials(authJson.toString()); + + verifyTraceDelivery(ds); + } + + @Test + public void testExecute_withCustomCredentialsFilePath() throws Exception { + JsonObject authJson = getAuthJson(); + File tempFile = File.createTempFile("auth", ".json"); + tempFile.deleteOnExit(); + Files.write(tempFile.toPath(), authJson.toString().getBytes(StandardCharsets.UTF_8)); + + DataSource ds = DataSource.fromUrl(CONNECTION_URL); + ds.setEnableGcpTraceExporter(true); + ds.setGcpTelemetryProjectId(PROJECT_ID); + ds.setGcpTelemetryCredentials(tempFile.getAbsolutePath()); + + verifyTraceDelivery(ds); + } + + @Test + public void testExecute_withHttpProtocol() throws Exception { + JsonObject authJson = getAuthJson(); + System.setProperty("otel.exporter.otlp.protocol", "http/protobuf"); + + try { + DataSource ds = DataSource.fromUrl(CONNECTION_URL); + ds.setEnableGcpTraceExporter(true); + ds.setGcpTelemetryProjectId(PROJECT_ID); + ds.setGcpTelemetryCredentials(authJson.toString()); + + verifyTraceDelivery(ds); + } finally { + System.clearProperty("otel.exporter.otlp.protocol"); + } + } + + @Test + public void testExecute_withGrpcProtocol() throws Exception { + JsonObject authJson = getAuthJson(); + System.setProperty("otel.exporter.otlp.protocol", "grpc"); + + try { + DataSource ds = DataSource.fromUrl(CONNECTION_URL); + ds.setEnableGcpTraceExporter(true); + ds.setGcpTelemetryProjectId(PROJECT_ID); + ds.setGcpTelemetryCredentials(authJson.toString()); + + verifyTraceDelivery(ds); + } finally { + System.clearProperty("otel.exporter.otlp.protocol"); + } + } + + private void verifyTraceDelivery(DataSource ds) throws Exception { + ds.setEnableGcpLogExporter(true); + ds.setLogLevel("5"); + + String connectionUuid = null; + try (Connection connection = ds.getConnection(); + Statement statement = connection.createStatement()) { + + BigQueryConnection bqConnection = connection.unwrap(BigQueryConnection.class); + connectionUuid = bqConnection.getConnectionId(); + + String query = "SELECT 1;"; + try (ResultSet rs = statement.executeQuery(query)) { + assertTrue(rs.next()); + } + } + + String traceId = verifyAndFetchLogs(connectionUuid); + Trace trace = verifyAndFetchTrace(traceId); + assertNotNull(trace, "Trace must be found"); + } + private String verifyAndFetchLogs(String connectionUuid) throws Exception { try (Logging logging = LoggingOptions.newBuilder().setProjectId(PROJECT_ID).build().getService()) {