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 6fe251f0e8fa..0fdb650d6e0d 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 @@ -29,12 +29,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; @@ -221,6 +226,23 @@ private static Credentials resolveCredentialsFromString(String credsString) { BigQueryJdbcOpenTelemetry.class.getName()); } + 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) { + throw new RuntimeException("Failed to get auth headers", e); + } + } + public static TelemetryConfig getConnectionConfig(String connectionId) { return connectionConfigs.get(connectionId); } @@ -270,13 +292,9 @@ public static OpenTelemetry getOpenTelemetry( key, k -> { Map props = new HashMap<>(); + Credentials credentials = null; 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); - } + credentials = resolveCredentialsFromString(gcpTelemetryCredentials); } if (enableGcpTraceExporter) { @@ -306,8 +324,30 @@ public static OpenTelemetry getOpenTelemetry( props.put(OTEL_ATTRIBUTE_VALUE_LENGTH_LIMIT, DEFAULT_ATTRIBUTE_LENGTH_LIMIT); } - AutoConfiguredOpenTelemetrySdk autoConfigured = - AutoConfiguredOpenTelemetrySdk.builder().addPropertiesSupplier(() -> props).build(); + final Credentials finalCreds = credentials; + AutoConfiguredOpenTelemetrySdk autoConfigured; + + if (finalCreds != null) { + autoConfigured = + AutoConfiguredOpenTelemetrySdk.builder() + .addPropertiesSupplier(() -> props) + .addSpanExporterCustomizer( + (spanExporter, configProperties) -> { + if (spanExporter instanceof OtlpHttpSpanExporter) { + return ((OtlpHttpSpanExporter) spanExporter) + .toBuilder().setHeaders(() -> getAuthHeaders(finalCreds)).build(); + } + if (spanExporter instanceof OtlpGrpcSpanExporter) { + return ((OtlpGrpcSpanExporter) spanExporter) + .toBuilder().setHeaders(() -> getAuthHeaders(finalCreds)).build(); + } + return spanExporter; + }) + .build(); + } else { + autoConfigured = + AutoConfiguredOpenTelemetrySdk.builder().addPropertiesSupplier(() -> props).build(); + } OpenTelemetrySdk sdk = autoConfigured.getOpenTelemetrySdk(); 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 dff06f976b79..3f31cddc05da 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 @@ -164,6 +164,60 @@ public void testExecute_withErrorCorrelation() throws Exception { "Traces must contain JDBC parent span 'BigQueryStatement.executeQuery'"); } + @Test + public void testExecute_withExplicitCredentialsJson() throws Exception { + // Goal: Verify that passing a raw JSON string in gcpTelemetryCredentials works and invokes our + // customizer. + // How to test: + // If you have a test service account JSON, you can read it as a string and set it in + // props/DataSource: + // ds.setGcpTelemetryCredentials(saJsonString); + // Verify that traces are still successfully delivered to Cloud Trace. + } + + @Test + public void testExecute_withExplicitCredentialsFilePath() throws Exception { + // Goal: Verify that passing a file path works. + // How to test: + // Save the test service account JSON to a temporary file. + // Set gcpTelemetryCredentials to the tempFilePath: + // ds.setGcpTelemetryCredentials(tempFilePath); + // Verify trace delivery. + } + + @Test + public void testExecute_withMultiTenancySdkCaching() throws Exception { + // Goal: Verify that the driver correctly creates and caches separate SDK instances for + // different configurations. + // How to test: + // Create Connection A with gcpTelemetryProjectId = "project-a". + // Create Connection B with gcpTelemetryProjectId = "project-b". + // Even if project-b doesn't exist or fails to export, you can verify that the driver doesn't + // crash and that it attempts to create two separate pipelines. + // To be rigorous, we could add a package-private method in BigQueryJdbcOpenTelemetry to return + // the size of the sdkCache and assert that it is 2 after creating these connections. + } + + @Test + public void testExecute_withExplicitCredentials_HTTP() throws Exception { + // Scenario A: Explicit Credentials + HTTP + // Goal: Verify that our customizer works for OtlpHttpSpanExporter. + // How to test: + // Set gcpTelemetryCredentials (JSON string or path). + // Set EnableHighThroughputAPI = 0 (to force HTTP). + // Verify that traces are delivered. + } + + @Test + public void testExecute_withExplicitCredentials_gRPC() throws Exception { + // Scenario B: Explicit Credentials + gRPC + // Goal: Verify that our customizer works for OtlpGrpcSpanExporter. + // How to test: + // Set gcpTelemetryCredentials (JSON string or path). + // Set EnableHighThroughputAPI = 1 (to force gRPC). + // Verify that traces are delivered. + } + private String verifyAndFetchLogs(String connectionUuid) throws Exception { try (Logging logging = LoggingOptions.newBuilder().setProjectId(PROJECT_ID).build().getService()) {