feat(bigquery-jdbc): support custom OTel credentials and dynamic token refresh

java-bigquery-jdbc/pom.xml

@@ -366,6 +366,10 @@
       <groupId>io.opentelemetry</groupId>
       <artifactId>opentelemetry-sdk</artifactId>
     </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-trace</artifactId>
+    </dependency>
     <dependency>
       <groupId>io.opentelemetry</groupId>
       <artifactId>opentelemetry-exporter-otlp</artifactId>
@@ -439,11 +443,6 @@
       <artifactId>opentelemetry-sdk-logs</artifactId>
       <scope>test</scope>
     </dependency>
-    <dependency>
-      <groupId>io.opentelemetry</groupId>
-      <artifactId>opentelemetry-sdk-trace</artifactId>
-      <scope>test</scope>
-    </dependency>
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>google-cloud-trace</artifactId>

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<TelemetryConfig> getRegisteredConfigs() {
     return connectionConfigs.values();
   }
 
+  private static Map<String, String> getAuthHeaders(Credentials credentials) {
+    try {
+      Map<String, List<String>> metadata =
+          credentials.getRequestMetadata(URI.create(OTLP_ENDPOINT_VALUE));
+      Map<String, String> 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<String, String> 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();
 

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);

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

@@ -17,12 +17,23 @@
 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.auth.oauth2.GoogleCredentials;
 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.ByteArrayInputStream;
+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.ResultSet;
 import java.sql.SQLException;
@@ -291,6 +302,31 @@ 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 static GoogleCredentials getCredentials() throws IOException {
+    JsonObject authJson = getAuthJson();
+    return GoogleCredentials.fromStream(
+        new ByteArrayInputStream(authJson.toString().getBytes(StandardCharsets.UTF_8)));
+  }
+
   protected int resultSetRowCount(ResultSet resultSet) throws SQLException {
     int rowCount = 0;
     while (resultSet.next()) {

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

@@ -22,16 +22,23 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import com.google.api.gax.core.FixedCredentialsProvider;
 import com.google.api.gax.paging.Page;
+import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.bigquery.jdbc.BigQueryConnection;
 import com.google.cloud.bigquery.jdbc.DataSource;
 import com.google.cloud.logging.LogEntry;
 import com.google.cloud.logging.Logging;
 import com.google.cloud.logging.LoggingOptions;
 import com.google.cloud.trace.v1.TraceServiceClient;
+import com.google.cloud.trace.v1.TraceServiceSettings;
 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,13 +47,10 @@
 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 =
-      String.format(
-          "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;ProjectId=%s;OAuthType=3;Timeout=3600;",
-          PROJECT_ID);
+  private static final String CONNECTION_URL = connectionUrl;
 
   @Test
   public void testExecute_withOpenTelemetryGcpExporter() throws Exception {
@@ -163,9 +167,97 @@ 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 {
+    GoogleCredentials credentials = getCredentials();
+
     try (Logging logging =
-        LoggingOptions.newBuilder().setProjectId(PROJECT_ID).build().getService()) {
+        LoggingOptions.newBuilder()
+            .setProjectId(PROJECT_ID)
+            .setCredentials(credentials)
+            .build()
+            .getService()) {
       String filter =
           "logName:\"projects/"
               + PROJECT_ID
@@ -198,7 +290,14 @@ private Trace verifyAndFetchTrace(String traceId) throws Exception {
       hexTraceId = traceId.substring(traceId.lastIndexOf("/traces/") + 8);
     }
 
-    try (TraceServiceClient traceClient = TraceServiceClient.create()) {
+    GoogleCredentials credentials = getCredentials();
+
+    TraceServiceSettings settings =
+        TraceServiceSettings.newBuilder()
+            .setCredentialsProvider(FixedCredentialsProvider.create(credentials))
+            .build();
+
+    try (TraceServiceClient traceClient = TraceServiceClient.create(settings)) {
       Trace trace = fetchTraceWithRetry(traceClient, PROJECT_ID, hexTraceId);
       assertNotNull(trace, "Trace must be found in Cloud Trace API: " + hexTraceId);
       return trace;