From f5fa8ea9a6de6bfcaa4b1214861c853a2332bd9c Mon Sep 17 00:00:00 2001 From: ldetmer Date: Mon, 2 Mar 2026 15:07:46 -0500 Subject: [PATCH 1/6] feat: add initial opentelemetry tracing to big query HTTP requests --- .../bigquery/spi/v2/HttpBigQueryRpc.java | 6 + .../spi/v2/HttpTracingRequestInitializer.java | 195 ++++++++++ .../telemetry/BigQueryTelemetryTracer.java | 71 ++++ .../spi/v2/HttpTracingIntegrationTest.java | 109 ++++++ .../v2/HttpTracingRequestInitializerTest.java | 363 ++++++++++++++++++ 5 files changed, 744 insertions(+) create mode 100644 google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java create mode 100644 google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java create mode 100644 google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingIntegrationTest.java create mode 100644 google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java index 16737dc4b7..6a6851a807 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java @@ -107,6 +107,12 @@ public HttpBigQueryRpc(BigQueryOptions options) { HttpTransportOptions transportOptions = (HttpTransportOptions) options.getTransportOptions(); HttpTransport transport = transportOptions.getHttpTransportFactory().create(); HttpRequestInitializer initializer = transportOptions.getHttpRequestInitializer(options); + + // Wrap with tracing initializer if OpenTelemetry is enabled + if (options.isOpenTelemetryTracingEnabled() && options.getOpenTelemetryTracer() != null) { + initializer = new HttpTracingRequestInitializer(initializer, options.getOpenTelemetryTracer()); + } + this.options = options; bigquery = new Bigquery.Builder(transport, new GsonFactory(), initializer) diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java new file mode 100644 index 0000000000..50bb112ff4 --- /dev/null +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java @@ -0,0 +1,195 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.spi.v2; + +import com.google.api.client.http.*; +import com.google.api.core.InternalApi; +import com.google.cloud.bigquery.telemetry.BigQueryTelemetryTracer; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.api.trace.Tracer; +import org.checkerframework.checker.nullness.qual.Nullable; + +import java.io.IOException; + +/** + * HttpRequestInitializer that wraps a delegate initializer, intercepts all HTTP requests, + * adds OpenTelemetry tracing and then invokes delegate interceptor. + */ +@InternalApi +public class HttpTracingRequestInitializer implements HttpRequestInitializer { + + // HTTP Specific Telemetry Attributes + public static final AttributeKey HTTP_REQUEST_METHOD = + AttributeKey.stringKey("http.request.method"); + public static final AttributeKey URL_FULL = AttributeKey.stringKey("url.full"); + public static final AttributeKey URL_TEMPLATE = AttributeKey.stringKey("url.template"); + public static final AttributeKey URL_DOMAIN = AttributeKey.stringKey("url.domain"); + public static final AttributeKey HTTP_RESPONSE_STATUS_CODE = + AttributeKey.longKey("http.response.status_code"); + public static final AttributeKey HTTP_REQUEST_RESEND_COUNT = + AttributeKey.longKey("http.request.resend_count"); + public static final AttributeKey HTTP_REQUEST_BODY_SIZE = + AttributeKey.longKey("http.request.body.size"); + public static final AttributeKey HTTP_RESPONSE_BODY_SIZE = + AttributeKey.longKey("http.response.body.size"); + + private final HttpRequestInitializer delegate; + private final Tracer tracer; + private static final String BIGQUERY_DOMAIN = "bigquery.googleapis.com"; + + public HttpTracingRequestInitializer(HttpRequestInitializer delegate, Tracer tracer) { + this.delegate = delegate; + this.tracer = tracer; + } + + @Override + public void initialize(HttpRequest request) throws IOException { + if (delegate != null) { + delegate.initialize(request); + } + + if (tracer == null) { + return; + } + + String httpMethod = request.getRequestMethod(); + String url = request.getUrl().build(); + String host = request.getUrl().getHost(); + Integer port = request.getUrl().getPort(); + + Long requestBodySize = getRequestBodySize(request); + + Span span = getSpan(httpMethod, url, host, port, requestBodySize); + + // Wrap the existing response interceptor + HttpResponseInterceptor originalInterceptor = request.getResponseInterceptor(); + request.setResponseInterceptor( + response -> { + try { + addSuccessResponseToSpan(response, httpMethod, span); + if (originalInterceptor != null) { + originalInterceptor.interceptResponse(response); + } + } finally { + span.end(); + } + }); + +// Wrap the existing unsuccessful response handler +HttpUnsuccessfulResponseHandler originalHandler = request.getUnsuccessfulResponseHandler(); +request.setUnsuccessfulResponseHandler( + (request1, response, supportsRetry) -> { + addErrorResponseToSpan(response, span); + try { + if (originalHandler != null) { + return originalHandler.handleResponse(request1, response, supportsRetry); + } + return false; + } catch (IOException e) { + addExceptionToSpan(e, span); + throw e; + } finally { + span.end(); + } + }); + } + + private static void addExceptionToSpan(IOException e, Span span) { + span.recordException(e); + span.setAttribute(BigQueryTelemetryTracer.EXCEPTION_TYPE, e.getClass().getName()); + span.setAttribute(BigQueryTelemetryTracer.ERROR_TYPE, e.getClass().getSimpleName()); + span.setAttribute(BigQueryTelemetryTracer.STATUS_MESSAGE, e.getMessage() != null ? e.getMessage() : e.getClass().getName()); + span.setStatus(StatusCode.ERROR, e.getMessage()); + } + + private static void addErrorResponseToSpan(HttpResponse response, Span span) { + int statusCode = response.getStatusCode(); + span.setAttribute(HTTP_RESPONSE_STATUS_CODE, statusCode); + String errorMessage = "HTTP " + statusCode; + try { + String statusMessage = response.getStatusMessage(); + if (statusMessage != null && !statusMessage.isEmpty()) { + errorMessage = statusMessage; + } + } catch (Exception ex) { + // Ignore + } + span.setAttribute(BigQueryTelemetryTracer.STATUS_MESSAGE, errorMessage); + span.setAttribute(BigQueryTelemetryTracer.ERROR_TYPE, String.valueOf(statusCode)); + span.setStatus(StatusCode.ERROR, errorMessage); + } + + private static void addSuccessResponseToSpan(HttpResponse response, String httpMethod, Span span) { + String actualMethod = response.getRequest().getRequestMethod(); + if (actualMethod != null && httpMethod == null) { + span.updateName(actualMethod); + span.setAttribute(HTTP_REQUEST_METHOD, actualMethod); + } + int statusCode = response.getStatusCode(); + span.setAttribute(HTTP_RESPONSE_STATUS_CODE, statusCode); + try { + long contentLength = response.getHeaders().getContentLength(); + if (contentLength > 0) { + span.setAttribute(HTTP_RESPONSE_BODY_SIZE, contentLength); + } + } catch (Exception e) { + // Ignore - body size not available + } + if (statusCode >= 400) { + addErrorResponseToSpan(response, span); + } else { + span.setStatus(StatusCode.OK); + } + } + + private Span getSpan(String httpMethod, String url, String host, Integer port, Long requestBodySize) { + //TODO: Determine span name: {method} {url.template} or {method} + Span span = + BigQueryTelemetryTracer.newSpanBuilder(tracer, httpMethod) + // OpenTelemetry semantic convention attributes + .setAttribute(HTTP_REQUEST_METHOD, httpMethod) + .setAttribute(URL_FULL, url) + .setAttribute(BigQueryTelemetryTracer.SERVER_ADDRESS, host) + .setAttribute(URL_DOMAIN, BIGQUERY_DOMAIN) + .setAttribute(BigQueryTelemetryTracer.RPC_SYSTEM_NAME , "http") + .startSpan(); + + // TODO: add url template && resource name + if (port != null && port > 0) { + span.setAttribute(BigQueryTelemetryTracer.SERVER_PORT, port.longValue()); + } + if (requestBodySize != null && requestBodySize > 0) { + span.setAttribute(HTTP_REQUEST_BODY_SIZE, requestBodySize); + } + return span; + } + + private static @Nullable Long getRequestBodySize(HttpRequest request) { + Long requestBodySize = null; + try { + HttpContent content = request.getContent(); + if (content != null) { + requestBodySize = content.getLength(); + } + } catch (Exception e) { + // Ignore - body size not available + } + return requestBodySize; + } +} diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java new file mode 100644 index 0000000000..c73099ef63 --- /dev/null +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java @@ -0,0 +1,71 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.telemetry; + +import com.google.api.core.InternalApi; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.SpanBuilder; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.Tracer; + +/** General BigQuery Telemetry class that stores generic telemetry attributes + * and any associated logic to calculate. + */ +@InternalApi +public final class BigQueryTelemetryTracer { + + private BigQueryTelemetryTracer() {} + + // Common GCP Attributes + public static final AttributeKey GCP_CLIENT_SERVICE = + AttributeKey.stringKey("gcp.client.service"); + public static final AttributeKey GCP_CLIENT_VERSION = + AttributeKey.stringKey("gcp.client.version"); + public static final AttributeKey GCP_CLIENT_REPO = + AttributeKey.stringKey("gcp.client.repo"); + public static final AttributeKey GCP_CLIENT_ARTIFACT = + AttributeKey.stringKey("gcp.client.artifact"); + public static final AttributeKey GCP_CLIENT_LANGUAGE = + AttributeKey.stringKey("gcp.client.language"); + public static final AttributeKey GCP_RESOURCE_NAME = + AttributeKey.stringKey("gcp.resource.name"); + public static final AttributeKey RPC_SYSTEM_NAME = AttributeKey.stringKey("rpc.system.name"); + + + // Common Error Attributes + public static final AttributeKey ERROR_TYPE = AttributeKey.stringKey("error.type"); + public static final AttributeKey EXCEPTION_TYPE = + AttributeKey.stringKey("exception.type"); + public static final AttributeKey STATUS_MESSAGE = + AttributeKey.stringKey("status.message"); + + // Common Server Attributes + public static final AttributeKey SERVER_ADDRESS = + AttributeKey.stringKey("server.address"); + public static final AttributeKey SERVER_PORT = AttributeKey.longKey("server.port"); + + public static SpanBuilder newSpanBuilder(Tracer tracer, String spanName) { + return tracer + .spanBuilder(spanName) + .setSpanKind(SpanKind.CLIENT) + .setAttribute(GCP_CLIENT_SERVICE, "bigquery") + .setAttribute(GCP_CLIENT_REPO, "googleapis/java-bigquery") + .setAttribute(GCP_CLIENT_ARTIFACT, "google-cloud-bigquery") + .setAttribute(GCP_CLIENT_LANGUAGE, "java"); + // TODO: add version + } +} diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingIntegrationTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingIntegrationTest.java new file mode 100644 index 0000000000..4ed9f87675 --- /dev/null +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingIntegrationTest.java @@ -0,0 +1,109 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.spi.v2; + +import static org.junit.jupiter.api.Assertions.assertEquals; + + +import com.google.api.client.http.GenericUrl; +import com.google.api.client.http.HttpRequest; +import com.google.api.client.http.HttpRequestFactory; +import com.google.api.client.http.HttpResponse; +import com.google.api.client.http.javanet.NetHttpTransport; +import com.google.cloud.bigquery.RetryContext; +import com.sun.net.httpserver.HttpServer; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import java.io.IOException; +import java.io.OutputStream; +import java.net.InetSocketAddress; +import java.util.List; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** + * Integration test for HTTP tracing with real HTTP transport and server. + * This test verifies that OpenTelemetry tracing works correctly with actual network calls. + */ +public class HttpTracingIntegrationTest { + + private InMemorySpanExporter spanExporter; + private HttpTracingRequestInitializer initializer; + private HttpServer testServer; + private int serverPort; + + @BeforeEach + public void setUp() throws IOException { + // Set up OpenTelemetry with in-memory exporter + spanExporter = InMemorySpanExporter.create(); + SdkTracerProvider tracerProvider = + SdkTracerProvider.builder() + .addSpanProcessor(SimpleSpanProcessor.create(spanExporter)) + .build(); + OpenTelemetrySdk openTelemetry = + OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).build(); + Tracer tracer = openTelemetry.getTracer("test-tracer"); + initializer = new HttpTracingRequestInitializer(null, tracer); + + // Start a test HTTP server + testServer = HttpServer.create(new InetSocketAddress("localhost", 0), 0); + serverPort = testServer.getAddress().getPort(); + testServer.start(); + } + + @AfterEach + public void tearDown() { + if (testServer != null) { + testServer.stop(0); + } + RetryContext.clearRetryAttempt(); + } + + @Test + public void testHttpTracingWithRealServer() throws IOException { + testServer.createContext("/api/test", exchange -> { + String response = "{\"status\": \"ok\"}"; + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, response.getBytes().length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(response.getBytes()); + } + }); + + NetHttpTransport transport = new NetHttpTransport(); + HttpRequestFactory requestFactory = transport.createRequestFactory(initializer); + HttpRequest request = requestFactory.buildGetRequest( + new GenericUrl("http://localhost:" + serverPort + "/api/test")); + + HttpResponse response = request.execute(); + assertEquals(200, response.getStatusCode()); + response.disconnect(); + + List spans = spanExporter.getFinishedSpanItems(); + assertEquals(1, spans.size()); + SpanData span = spans.get(0); + assertEquals("GET", span.getAttributes().get(AttributeKey.stringKey("http.request.method"))); + assertEquals(200L, span.getAttributes().get(AttributeKey.longKey("http.response.status_code"))); + assertEquals("localhost", span.getAttributes().get(AttributeKey.stringKey("server.address"))); + } +} diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java new file mode 100644 index 0000000000..c680024011 --- /dev/null +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java @@ -0,0 +1,363 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.spi.v2; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +import com.google.api.client.http.GenericUrl; +import com.google.api.client.http.HttpRequest; +import com.google.api.client.http.HttpRequestFactory; +import com.google.api.client.http.HttpRequestInitializer; +import com.google.api.client.http.HttpResponse; +import com.google.api.client.http.HttpTransport; +import com.google.api.client.http.LowLevelHttpRequest; +import com.google.api.client.http.LowLevelHttpResponse; +import com.google.api.client.testing.http.MockHttpTransport; +import com.google.api.client.testing.http.MockLowLevelHttpRequest; +import com.google.api.client.testing.http.MockLowLevelHttpResponse; +import com.google.cloud.bigquery.RetryContext; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import java.io.IOException; +import java.util.List; +import java.util.stream.Collectors; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Unit tests for TracingHttpRequestInitializer */ +public class HttpTracingRequestInitializerTest { + + private InMemorySpanExporter spanExporter; + private Tracer tracer; + private HttpTracingRequestInitializer initializer; + + @BeforeEach + public void setUp() { + spanExporter = InMemorySpanExporter.create(); + SdkTracerProvider tracerProvider = + SdkTracerProvider.builder() + .addSpanProcessor(SimpleSpanProcessor.create(spanExporter)) + .build(); + OpenTelemetrySdk openTelemetry = + OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).build(); + tracer = openTelemetry.getTracer("test-tracer"); + initializer = new HttpTracingRequestInitializer(null, tracer); + } + + @AfterEach + public void tearDown() { + // Clean up thread-local retry context after each test + RetryContext.clearRetryAttempt(); + } + + @Test + public void testSuccessResponseAttributesAreSet() throws IOException { + HttpTransport transport = + new MockHttpTransport() { + @Override + public LowLevelHttpRequest buildRequest(String method, String url) { + return new MockLowLevelHttpRequest() { + @Override + public LowLevelHttpResponse execute() { + MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); + response.setStatusCode(200); + return response; + } + }; + } + }; + + HttpRequestFactory requestFactory = transport.createRequestFactory(initializer); + String urlString = "https://bigquery.googleapis.com:443/bigquery/v2/projects/test/datasets"; + HttpRequest request = + requestFactory.buildGetRequest( + new GenericUrl(urlString)); + + HttpResponse response = request.execute(); + response.disconnect(); + + List spans = spanExporter.getFinishedSpanItems(); + assertEquals(1, spans.size()); + + SpanData span = spans.get(0); + assertEquals( + "GET", + span.getAttributes().get(AttributeKey.stringKey("http.request.method"))); + assertEquals( + "bigquery.googleapis.com", + span.getAttributes().get(AttributeKey.stringKey("server.address"))); + assertEquals( + Long.valueOf(443L), + span.getAttributes().get(AttributeKey.longKey("server.port"))); + assertEquals( + "bigquery.googleapis.com", + span.getAttributes().get(AttributeKey.stringKey("url.domain"))); + assertEquals( + urlString, + span.getAttributes().get(AttributeKey.stringKey("url.full"))); + assertEquals( + Long.valueOf(200L), + span.getAttributes().get(AttributeKey.longKey("http.response.status_code"))); + assertEquals( + "bigquery", + span.getAttributes().get(AttributeKey.stringKey("gcp.client.service"))); + assertEquals( + "googleapis/java-bigquery", + span.getAttributes().get(AttributeKey.stringKey("gcp.client.repo"))); + assertEquals( + "google-cloud-bigquery", + span.getAttributes().get(AttributeKey.stringKey("gcp.client.artifact"))); + assertEquals( + "java", + span.getAttributes().get(AttributeKey.stringKey("gcp.client.language"))); + } + + @Test + public void testErrorAttributesAreSetOn404() throws IOException { + HttpTransport transport = + new MockHttpTransport() { + @Override + public LowLevelHttpRequest buildRequest(String method, String url) { + return new MockLowLevelHttpRequest() { + @Override + public LowLevelHttpResponse execute() { + MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); + response.setStatusCode(404); + response.setReasonPhrase("Not Found"); + return response; + } + }; + } + }; + + HttpRequestFactory requestFactory = transport.createRequestFactory(initializer); + HttpRequest request = + requestFactory.buildGetRequest( + new GenericUrl("https://bigquery.googleapis.com/bigquery/v2/projects/test/datasets/notfound")); + + try { + HttpResponse response = request.execute(); + response.disconnect(); + } catch (Exception e) { + // Expected - 404 might throw + } + + List spans = spanExporter.getFinishedSpanItems(); + assertEquals(1, spans.size()); + + SpanData span = spans.get(0); + assertEquals( + Long.valueOf(404L), + span.getAttributes().get(AttributeKey.longKey("http.response.status_code"))); + assertEquals( + "404", + span.getAttributes().get(AttributeKey.stringKey("error.type"))); + assertNotNull(span.getAttributes().get(AttributeKey.stringKey("status.message"))); + } + + @Test + public void testExceptionAttributesAreSetWhenOriginalUnsuccessfulHandlerThrowsIOException() + throws IOException { + HttpRequestInitializer delegateInitializer = + request -> + request.setUnsuccessfulResponseHandler( + (request1, response, supportsRetry) -> { + throw new IOException("handler failure"); + }); + HttpTracingRequestInitializer tracingInitializer = + new HttpTracingRequestInitializer(delegateInitializer, tracer); + + HttpTransport transport = + new MockHttpTransport() { + @Override + public LowLevelHttpRequest buildRequest(String method, String url) { + return new MockLowLevelHttpRequest() { + @Override + public LowLevelHttpResponse execute() { + MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); + response.setStatusCode(500); + response.setReasonPhrase("Internal Server Error"); + return response; + } + }; + } + }; + + HttpRequestFactory requestFactory = transport.createRequestFactory(tracingInitializer); + HttpRequest request = + requestFactory.buildGetRequest( + new GenericUrl("https://bigquery.googleapis.com/bigquery/v2/projects/test/datasets")); + request.setThrowExceptionOnExecuteError(false); + + IOException thrown = assertThrows(IOException.class, request::execute); + assertEquals("handler failure", thrown.getMessage()); + + List spans = spanExporter.getFinishedSpanItems(); + assertEquals(1, spans.size()); + + SpanData span = spans.get(0); + assertEquals( + Long.valueOf(500L), + span.getAttributes().get(AttributeKey.longKey("http.response.status_code"))); + assertEquals( + IOException.class.getName(), + span.getAttributes().get(AttributeKey.stringKey("exception.type"))); + assertEquals( + IOException.class.getSimpleName(), + span.getAttributes().get(AttributeKey.stringKey("error.type"))); + assertEquals( + "handler failure", + span.getAttributes().get(AttributeKey.stringKey("status.message"))); + } + + @Test + public void testExceptionIsRecordedWhenOriginalUnsuccessfulHandlerThrowsIOException() + throws IOException { + HttpRequestInitializer delegateInitializer = + request -> + request.setUnsuccessfulResponseHandler( + (request1, response, supportsRetry) -> { + throw new IOException("handler failure"); + }); + HttpTracingRequestInitializer tracingInitializer = + new HttpTracingRequestInitializer(delegateInitializer, tracer); + + HttpTransport transport = + new MockHttpTransport() { + @Override + public LowLevelHttpRequest buildRequest(String method, String url) { + return new MockLowLevelHttpRequest() { + @Override + public LowLevelHttpResponse execute() { + MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); + response.setStatusCode(500); + response.setReasonPhrase("Internal Server Error"); + return response; + } + }; + } + }; + + HttpRequestFactory requestFactory = transport.createRequestFactory(tracingInitializer); + HttpRequest request = + requestFactory.buildGetRequest( + new GenericUrl("https://bigquery.googleapis.com/bigquery/v2/projects/test/datasets")); + request.setThrowExceptionOnExecuteError(false); + + assertThrows(IOException.class, request::execute); + + List spans = spanExporter.getFinishedSpanItems(); + assertEquals(1, spans.size()); + SpanData span = spans.get(0); + + assertTrue(span.getEvents().stream().anyMatch(event -> "exception".equals(event.getName()))); + assertTrue( + span.getEvents().stream() + .filter(event -> "exception".equals(event.getName())) + .anyMatch( + event -> + "handler failure" + .equals( + event + .getAttributes() + .get(AttributeKey.stringKey("exception.message"))))); + } + + @Test + public void testDelegateInitializerIsCalledOnSuccessResponse() throws IOException { + HttpRequestInitializer delegateInitializer = mock(HttpRequestInitializer.class); + HttpTracingRequestInitializer tracingInitializer = + new HttpTracingRequestInitializer(delegateInitializer, tracer); + + HttpTransport transport = + new MockHttpTransport() { + @Override + public LowLevelHttpRequest buildRequest(String method, String url) { + return new MockLowLevelHttpRequest() { + @Override + public LowLevelHttpResponse execute() { + MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); + response.setStatusCode(200); + return response; + } + }; + } + }; + + HttpRequestFactory requestFactory = transport.createRequestFactory(tracingInitializer); + HttpRequest request = + requestFactory.buildGetRequest( + new GenericUrl("https://bigquery.googleapis.com/bigquery/v2/projects/test/datasets")); + + HttpResponse response = request.execute(); + response.disconnect(); + + verify(delegateInitializer, times(1)).initialize(any(HttpRequest.class)); + } + + @Test + public void testDelegateInitializerIsCalledOnErrorResponse() throws IOException { + HttpRequestInitializer delegateInitializer = mock(HttpRequestInitializer.class); + HttpTracingRequestInitializer tracingInitializer = + new HttpTracingRequestInitializer(delegateInitializer, tracer); + + HttpTransport transport = + new MockHttpTransport() { + @Override + public LowLevelHttpRequest buildRequest(String method, String url) { + return new MockLowLevelHttpRequest() { + @Override + public LowLevelHttpResponse execute() { + MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); + response.setStatusCode(404); + response.setReasonPhrase("Not Found"); + return response; + } + }; + } + }; + + HttpRequestFactory requestFactory = transport.createRequestFactory(tracingInitializer); + HttpRequest request = + requestFactory.buildGetRequest( + new GenericUrl("https://bigquery.googleapis.com/bigquery/v2/projects/test/datasets/notfound")); + + try { + HttpResponse response = request.execute(); + response.disconnect(); + } catch (Exception e) { + // Expected - 404 might throw + } + + verify(delegateInitializer, times(1)).initialize(any(HttpRequest.class)); + } + +} From 4faaa8252c45f9fbfb19d8074f50d023a41e639e Mon Sep 17 00:00:00 2001 From: ldetmer Date: Mon, 2 Mar 2026 16:08:04 -0500 Subject: [PATCH 2/6] formated code --- google-cloud-bigquery/pom.xml | 5 ++ .../google/cloud/bigquery/RetryContext.java | 76 +++++++++++++++++++ .../spi/v2/HttpTracingRequestInitializer.java | 56 +++++++------- .../telemetry/BigQueryTelemetryTracer.java | 9 ++- .../v2/HttpTracingRequestInitializerTest.java | 42 ++++------ 5 files changed, 130 insertions(+), 58 deletions(-) create mode 100644 google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RetryContext.java diff --git a/google-cloud-bigquery/pom.xml b/google-cloud-bigquery/pom.xml index a74654d1a8..ea455a0975 100644 --- a/google-cloud-bigquery/pom.xml +++ b/google-cloud-bigquery/pom.xml @@ -213,6 +213,11 @@ opentelemetry-sdk-trace test + + io.opentelemetry + opentelemetry-sdk-testing + test + diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RetryContext.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RetryContext.java new file mode 100644 index 0000000000..3b1cfd4558 --- /dev/null +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RetryContext.java @@ -0,0 +1,76 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery; + +/** + * Utility class for storing and retrieving retry attempt count using thread-local storage. + * This allows retry information to be accessed by HTTP/gRPC tracing interceptors during + * retry operations. + * + *

Thread-local storage is used because GAX retry framework executes retries on the same + * thread, making this mechanism both simple and reliable. + */ +public class RetryContext { + + /** + * Thread-local storage for the current retry attempt number. + * Value 0 means the first attempt (no retries yet). + * Value 1 means first retry (second attempt total). + * Value N means Nth retry (N+1 attempts total). + */ + private static final ThreadLocal RETRY_ATTEMPT = ThreadLocal.withInitial(() -> 0); + + /** + * Stores the retry attempt count in thread-local storage. + * + * @param attemptCount The attempt number (0 = first attempt, 1 = first retry, etc.) + */ + public static void setRetryAttempt(int attemptCount) { + RETRY_ATTEMPT.set(attemptCount); + } + + /** + * Retrieves the retry attempt count from thread-local storage. + * + * @return The retry attempt count, or 0 if not set (meaning first attempt) + */ + public static int getRetryAttempt() { + Integer attemptCount = RETRY_ATTEMPT.get(); + return attemptCount != null ? attemptCount : 0; + } + + /** + * Clears the retry attempt count from thread-local storage. + * Should be called when retry operations complete to prevent memory leaks. + */ + public static void clearRetryAttempt() { + RETRY_ATTEMPT.remove(); + } + + /** + * Checks if the current thread is processing a retry (attempt count > 0). + * + * @return true if this is a retry attempt, false if this is the first attempt + */ + public static boolean isRetry() { + return getRetryAttempt() > 0; + } + + private RetryContext() { + // Utility class, prevent instantiation + } +} diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java index 50bb112ff4..31fae4488a 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java @@ -28,26 +28,26 @@ import java.io.IOException; /** - * HttpRequestInitializer that wraps a delegate initializer, intercepts all HTTP requests, - * adds OpenTelemetry tracing and then invokes delegate interceptor. + * HttpRequestInitializer that wraps a delegate initializer, intercepts all HTTP requests, adds + * OpenTelemetry tracing and then invokes delegate interceptor. */ @InternalApi public class HttpTracingRequestInitializer implements HttpRequestInitializer { // HTTP Specific Telemetry Attributes public static final AttributeKey HTTP_REQUEST_METHOD = - AttributeKey.stringKey("http.request.method"); + AttributeKey.stringKey("http.request.method"); public static final AttributeKey URL_FULL = AttributeKey.stringKey("url.full"); public static final AttributeKey URL_TEMPLATE = AttributeKey.stringKey("url.template"); public static final AttributeKey URL_DOMAIN = AttributeKey.stringKey("url.domain"); public static final AttributeKey HTTP_RESPONSE_STATUS_CODE = - AttributeKey.longKey("http.response.status_code"); + AttributeKey.longKey("http.response.status_code"); public static final AttributeKey HTTP_REQUEST_RESEND_COUNT = - AttributeKey.longKey("http.request.resend_count"); + AttributeKey.longKey("http.request.resend_count"); public static final AttributeKey HTTP_REQUEST_BODY_SIZE = - AttributeKey.longKey("http.request.body.size"); + AttributeKey.longKey("http.request.body.size"); public static final AttributeKey HTTP_RESPONSE_BODY_SIZE = - AttributeKey.longKey("http.response.body.size"); + AttributeKey.longKey("http.response.body.size"); private final HttpRequestInitializer delegate; private final Tracer tracer; @@ -80,20 +80,20 @@ public void initialize(HttpRequest request) throws IOException { // Wrap the existing response interceptor HttpResponseInterceptor originalInterceptor = request.getResponseInterceptor(); request.setResponseInterceptor( - response -> { - try { - addSuccessResponseToSpan(response, httpMethod, span); - if (originalInterceptor != null) { - originalInterceptor.interceptResponse(response); - } - } finally { - span.end(); - } - }); - -// Wrap the existing unsuccessful response handler -HttpUnsuccessfulResponseHandler originalHandler = request.getUnsuccessfulResponseHandler(); -request.setUnsuccessfulResponseHandler( + response -> { + try { + addSuccessResponseToSpan(response, httpMethod, span); + if (originalInterceptor != null) { + originalInterceptor.interceptResponse(response); + } + } finally { + span.end(); + } + }); + + // Wrap the existing unsuccessful response handler + HttpUnsuccessfulResponseHandler originalHandler = request.getUnsuccessfulResponseHandler(); + request.setUnsuccessfulResponseHandler( (request1, response, supportsRetry) -> { addErrorResponseToSpan(response, span); try { @@ -114,7 +114,9 @@ private static void addExceptionToSpan(IOException e, Span span) { span.recordException(e); span.setAttribute(BigQueryTelemetryTracer.EXCEPTION_TYPE, e.getClass().getName()); span.setAttribute(BigQueryTelemetryTracer.ERROR_TYPE, e.getClass().getSimpleName()); - span.setAttribute(BigQueryTelemetryTracer.STATUS_MESSAGE, e.getMessage() != null ? e.getMessage() : e.getClass().getName()); + span.setAttribute( + BigQueryTelemetryTracer.STATUS_MESSAGE, + e.getMessage() != null ? e.getMessage() : e.getClass().getName()); span.setStatus(StatusCode.ERROR, e.getMessage()); } @@ -135,7 +137,8 @@ private static void addErrorResponseToSpan(HttpResponse response, Span span) { span.setStatus(StatusCode.ERROR, errorMessage); } - private static void addSuccessResponseToSpan(HttpResponse response, String httpMethod, Span span) { + private static void addSuccessResponseToSpan( + HttpResponse response, String httpMethod, Span span) { String actualMethod = response.getRequest().getRequestMethod(); if (actualMethod != null && httpMethod == null) { span.updateName(actualMethod); @@ -158,8 +161,9 @@ private static void addSuccessResponseToSpan(HttpResponse response, String httpM } } - private Span getSpan(String httpMethod, String url, String host, Integer port, Long requestBodySize) { - //TODO: Determine span name: {method} {url.template} or {method} + private Span getSpan( + String httpMethod, String url, String host, Integer port, Long requestBodySize) { + // TODO: Determine span name: {method} {url.template} or {method} Span span = BigQueryTelemetryTracer.newSpanBuilder(tracer, httpMethod) // OpenTelemetry semantic convention attributes @@ -167,7 +171,7 @@ private Span getSpan(String httpMethod, String url, String host, Integer port, L .setAttribute(URL_FULL, url) .setAttribute(BigQueryTelemetryTracer.SERVER_ADDRESS, host) .setAttribute(URL_DOMAIN, BIGQUERY_DOMAIN) - .setAttribute(BigQueryTelemetryTracer.RPC_SYSTEM_NAME , "http") + .setAttribute(BigQueryTelemetryTracer.RPC_SYSTEM_NAME, "http") .startSpan(); // TODO: add url template && resource name diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java index c73099ef63..fd9c521043 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java @@ -22,8 +22,9 @@ import io.opentelemetry.api.trace.SpanKind; import io.opentelemetry.api.trace.Tracer; -/** General BigQuery Telemetry class that stores generic telemetry attributes - * and any associated logic to calculate. +/** + * General BigQuery Telemetry class that stores generic telemetry attributes and any associated + * logic to calculate. */ @InternalApi public final class BigQueryTelemetryTracer { @@ -43,8 +44,8 @@ private BigQueryTelemetryTracer() {} AttributeKey.stringKey("gcp.client.language"); public static final AttributeKey GCP_RESOURCE_NAME = AttributeKey.stringKey("gcp.resource.name"); - public static final AttributeKey RPC_SYSTEM_NAME = AttributeKey.stringKey("rpc.system.name"); - + public static final AttributeKey RPC_SYSTEM_NAME = + AttributeKey.stringKey("rpc.system.name"); // Common Error Attributes public static final AttributeKey ERROR_TYPE = AttributeKey.stringKey("error.type"); diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java index c680024011..0e19c86e65 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java @@ -96,10 +96,8 @@ public LowLevelHttpResponse execute() { }; HttpRequestFactory requestFactory = transport.createRequestFactory(initializer); - String urlString = "https://bigquery.googleapis.com:443/bigquery/v2/projects/test/datasets"; - HttpRequest request = - requestFactory.buildGetRequest( - new GenericUrl(urlString)); + String urlString = "https://bigquery.googleapis.com:443/bigquery/v2/projects/test/datasets"; + HttpRequest request = requestFactory.buildGetRequest(new GenericUrl(urlString)); HttpResponse response = request.execute(); response.disconnect(); @@ -108,36 +106,26 @@ public LowLevelHttpResponse execute() { assertEquals(1, spans.size()); SpanData span = spans.get(0); - assertEquals( - "GET", - span.getAttributes().get(AttributeKey.stringKey("http.request.method"))); + assertEquals("GET", span.getAttributes().get(AttributeKey.stringKey("http.request.method"))); assertEquals( "bigquery.googleapis.com", span.getAttributes().get(AttributeKey.stringKey("server.address"))); + assertEquals(Long.valueOf(443L), span.getAttributes().get(AttributeKey.longKey("server.port"))); assertEquals( - Long.valueOf(443L), - span.getAttributes().get(AttributeKey.longKey("server.port"))); - assertEquals( - "bigquery.googleapis.com", - span.getAttributes().get(AttributeKey.stringKey("url.domain"))); - assertEquals( - urlString, - span.getAttributes().get(AttributeKey.stringKey("url.full"))); + "bigquery.googleapis.com", span.getAttributes().get(AttributeKey.stringKey("url.domain"))); + assertEquals(urlString, span.getAttributes().get(AttributeKey.stringKey("url.full"))); assertEquals( Long.valueOf(200L), span.getAttributes().get(AttributeKey.longKey("http.response.status_code"))); assertEquals( - "bigquery", - span.getAttributes().get(AttributeKey.stringKey("gcp.client.service"))); + "bigquery", span.getAttributes().get(AttributeKey.stringKey("gcp.client.service"))); assertEquals( "googleapis/java-bigquery", span.getAttributes().get(AttributeKey.stringKey("gcp.client.repo"))); assertEquals( "google-cloud-bigquery", span.getAttributes().get(AttributeKey.stringKey("gcp.client.artifact"))); - assertEquals( - "java", - span.getAttributes().get(AttributeKey.stringKey("gcp.client.language"))); + assertEquals("java", span.getAttributes().get(AttributeKey.stringKey("gcp.client.language"))); } @Test @@ -161,7 +149,8 @@ public LowLevelHttpResponse execute() { HttpRequestFactory requestFactory = transport.createRequestFactory(initializer); HttpRequest request = requestFactory.buildGetRequest( - new GenericUrl("https://bigquery.googleapis.com/bigquery/v2/projects/test/datasets/notfound")); + new GenericUrl( + "https://bigquery.googleapis.com/bigquery/v2/projects/test/datasets/notfound")); try { HttpResponse response = request.execute(); @@ -177,9 +166,7 @@ public LowLevelHttpResponse execute() { assertEquals( Long.valueOf(404L), span.getAttributes().get(AttributeKey.longKey("http.response.status_code"))); - assertEquals( - "404", - span.getAttributes().get(AttributeKey.stringKey("error.type"))); + assertEquals("404", span.getAttributes().get(AttributeKey.stringKey("error.type"))); assertNotNull(span.getAttributes().get(AttributeKey.stringKey("status.message"))); } @@ -234,8 +221,7 @@ public LowLevelHttpResponse execute() { IOException.class.getSimpleName(), span.getAttributes().get(AttributeKey.stringKey("error.type"))); assertEquals( - "handler failure", - span.getAttributes().get(AttributeKey.stringKey("status.message"))); + "handler failure", span.getAttributes().get(AttributeKey.stringKey("status.message"))); } @Test @@ -348,7 +334,8 @@ public LowLevelHttpResponse execute() { HttpRequestFactory requestFactory = transport.createRequestFactory(tracingInitializer); HttpRequest request = requestFactory.buildGetRequest( - new GenericUrl("https://bigquery.googleapis.com/bigquery/v2/projects/test/datasets/notfound")); + new GenericUrl( + "https://bigquery.googleapis.com/bigquery/v2/projects/test/datasets/notfound")); try { HttpResponse response = request.execute(); @@ -359,5 +346,4 @@ public LowLevelHttpResponse execute() { verify(delegateInitializer, times(1)).initialize(any(HttpRequest.class)); } - } From 4202fb330c0a4cbd046cbbb30b7493090f22ce06 Mon Sep 17 00:00:00 2001 From: ldetmer Date: Tue, 3 Mar 2026 14:39:39 -0500 Subject: [PATCH 3/6] cleaned up code + added suggestions from gemini --- .../bigquery/spi/v2/HttpBigQueryRpc.java | 11 +- .../spi/v2/HttpTracingRequestInitializer.java | 169 +++++---- .../telemetry/BigQueryTelemetryTracer.java | 13 +- .../spi/v2/HttpTracingIntegrationTest.java | 69 ++-- .../v2/HttpTracingRequestInitializerTest.java | 326 +++++++----------- 5 files changed, 298 insertions(+), 290 deletions(-) diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java index 6a6851a807..3d1e37dd2a 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java @@ -107,16 +107,19 @@ public HttpBigQueryRpc(BigQueryOptions options) { HttpTransportOptions transportOptions = (HttpTransportOptions) options.getTransportOptions(); HttpTransport transport = transportOptions.getHttpTransportFactory().create(); HttpRequestInitializer initializer = transportOptions.getHttpRequestInitializer(options); - + + String resolvedBigQueryRootUrl = options.getResolvedApiaryHost("bigquery"); // Wrap with tracing initializer if OpenTelemetry is enabled if (options.isOpenTelemetryTracingEnabled() && options.getOpenTelemetryTracer() != null) { - initializer = new HttpTracingRequestInitializer(initializer, options.getOpenTelemetryTracer()); + initializer = + new HttpTracingRequestInitializer( + initializer, options.getOpenTelemetryTracer(), resolvedBigQueryRootUrl); } - + this.options = options; bigquery = new Bigquery.Builder(transport, new GsonFactory(), initializer) - .setRootUrl(options.getResolvedApiaryHost("bigquery")) + .setRootUrl(resolvedBigQueryRootUrl) .setApplicationName(options.getApplicationName()) .build(); } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java index 31fae4488a..6aa9a3ff09 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java @@ -19,13 +19,13 @@ import com.google.api.client.http.*; import com.google.api.core.InternalApi; import com.google.cloud.bigquery.telemetry.BigQueryTelemetryTracer; +import com.google.common.annotations.VisibleForTesting; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.StatusCode; import io.opentelemetry.api.trace.Tracer; -import org.checkerframework.checker.nullness.qual.Nullable; - import java.io.IOException; +import org.jspecify.annotations.Nullable; /** * HttpRequestInitializer that wraps a delegate initializer, intercepts all HTTP requests, adds @@ -49,13 +49,17 @@ public class HttpTracingRequestInitializer implements HttpRequestInitializer { public static final AttributeKey HTTP_RESPONSE_BODY_SIZE = AttributeKey.longKey("http.response.body.size"); + @VisibleForTesting static final String HTTP_RPC_SYSTEM_NAME = "http"; + private final HttpRequestInitializer delegate; private final Tracer tracer; - private static final String BIGQUERY_DOMAIN = "bigquery.googleapis.com"; + private final @Nullable String clientRootUrl; - public HttpTracingRequestInitializer(HttpRequestInitializer delegate, Tracer tracer) { + public HttpTracingRequestInitializer( + HttpRequestInitializer delegate, Tracer tracer, @Nullable String clientRootUrl) { this.delegate = delegate; this.tracer = tracer; + this.clientRootUrl = clientRootUrl; } @Override @@ -73,21 +77,25 @@ public void initialize(HttpRequest request) throws IOException { String host = request.getUrl().getHost(); Integer port = request.getUrl().getPort(); - Long requestBodySize = getRequestBodySize(request); - - Span span = getSpan(httpMethod, url, host, port, requestBodySize); + Span span = createHttpTraceSpan(httpMethod, url, host, port); // Wrap the existing response interceptor HttpResponseInterceptor originalInterceptor = request.getResponseInterceptor(); request.setResponseInterceptor( response -> { - try { - addSuccessResponseToSpan(response, httpMethod, span); - if (originalInterceptor != null) { - originalInterceptor.interceptResponse(response); + if (span.isRecording()) { + try { + int statusCode = response.getStatusCode(); + addCommonResponseAttributesToSpan(request, response, span, httpMethod, statusCode); + addSuccessResponseToSpan(response, span, statusCode); + if (originalInterceptor != null) { + originalInterceptor.interceptResponse(response); + } + } finally { + span.end(); } - } finally { - span.end(); + } else if (originalInterceptor != null) { + originalInterceptor.interceptResponse(response); } }); @@ -95,7 +103,9 @@ public void initialize(HttpRequest request) throws IOException { HttpUnsuccessfulResponseHandler originalHandler = request.getUnsuccessfulResponseHandler(); request.setUnsuccessfulResponseHandler( (request1, response, supportsRetry) -> { - addErrorResponseToSpan(response, span); + int statusCode = response.getStatusCode(); + addCommonResponseAttributesToSpan(request, response, span, httpMethod, statusCode); + addErrorResponseToSpan(response, span, statusCode); try { if (originalHandler != null) { return originalHandler.handleResponse(request1, response, supportsRetry); @@ -105,24 +115,75 @@ public void initialize(HttpRequest request) throws IOException { addExceptionToSpan(e, span); throw e; } finally { - span.end(); + if (span.isRecording()) { + span.end(); + } } }); } + /** Initial HTTP trace span creation with basic attributes from request */ + private Span createHttpTraceSpan(String httpMethod, String url, String host, Integer port) { + // TODO: Determine span name: {method} {url.template} or {method} + Span span = + BigQueryTelemetryTracer.newSpanBuilder(tracer, httpMethod) + .setAttribute(HTTP_REQUEST_METHOD, httpMethod) + .setAttribute(URL_FULL, url) + .setAttribute(BigQueryTelemetryTracer.SERVER_ADDRESS, host) + .setAttribute(URL_DOMAIN, resolveUrlDomain(host)) + .setAttribute(BigQueryTelemetryTracer.RPC_SYSTEM_NAME, HTTP_RPC_SYSTEM_NAME) + .startSpan(); + + // TODO: add url template && resource name + if (port != null && port > 0) { + span.setAttribute(BigQueryTelemetryTracer.SERVER_PORT, port.longValue()); + } + return span; + } + + private String resolveUrlDomain(String requestHost) { + if (clientRootUrl != null) { + try { + String configuredHost = new GenericUrl(clientRootUrl).getHost(); + if (configuredHost != null && !configuredHost.isEmpty()) { + return configuredHost; + } + } catch (IllegalArgumentException ex) { + // Ignore malformed configured root URL and fall back to request host. + } + } + return requestHost; + } + + private static void addCommonResponseAttributesToSpan( + HttpRequest request, HttpResponse response, Span span, String httpMethod, int statusCode) { + // This is called after we get a response as sometimes the request body size isn't available + // before the response is received. + addRequestBodySizeToSpan(request, span); + checkForUpdatedRequestMethod(response, httpMethod, span); + setResponseBodySize(response, span); + span.setAttribute(HTTP_RESPONSE_STATUS_CODE, statusCode); + } + + private static void addSuccessResponseToSpan(HttpResponse response, Span span, int statusCode) { + if (statusCode >= 400) { + addErrorResponseToSpan(response, span, statusCode); + } else { + span.setStatus(StatusCode.OK); + } + } + private static void addExceptionToSpan(IOException e, Span span) { span.recordException(e); + String message = e.getMessage(); + String statusMessage = message != null ? message : e.getClass().getName(); span.setAttribute(BigQueryTelemetryTracer.EXCEPTION_TYPE, e.getClass().getName()); span.setAttribute(BigQueryTelemetryTracer.ERROR_TYPE, e.getClass().getSimpleName()); - span.setAttribute( - BigQueryTelemetryTracer.STATUS_MESSAGE, - e.getMessage() != null ? e.getMessage() : e.getClass().getName()); - span.setStatus(StatusCode.ERROR, e.getMessage()); + span.setAttribute(BigQueryTelemetryTracer.STATUS_MESSAGE, statusMessage); + span.setStatus(StatusCode.ERROR, statusMessage); } - private static void addErrorResponseToSpan(HttpResponse response, Span span) { - int statusCode = response.getStatusCode(); - span.setAttribute(HTTP_RESPONSE_STATUS_CODE, statusCode); + private static void addErrorResponseToSpan(HttpResponse response, Span span, int statusCode) { String errorMessage = "HTTP " + statusCode; try { String statusMessage = response.getStatusMessage(); @@ -137,63 +198,39 @@ private static void addErrorResponseToSpan(HttpResponse response, Span span) { span.setStatus(StatusCode.ERROR, errorMessage); } - private static void addSuccessResponseToSpan( - HttpResponse response, String httpMethod, Span span) { - String actualMethod = response.getRequest().getRequestMethod(); - if (actualMethod != null && httpMethod == null) { - span.updateName(actualMethod); - span.setAttribute(HTTP_REQUEST_METHOD, actualMethod); - } - int statusCode = response.getStatusCode(); - span.setAttribute(HTTP_RESPONSE_STATUS_CODE, statusCode); + private static void addRequestBodySizeToSpan(HttpRequest request, Span span) { + Long requestBodySize = null; try { - long contentLength = response.getHeaders().getContentLength(); - if (contentLength > 0) { - span.setAttribute(HTTP_RESPONSE_BODY_SIZE, contentLength); + HttpContent content = request.getContent(); + + if (content != null) { + requestBodySize = content.getLength(); } } catch (Exception e) { // Ignore - body size not available } - if (statusCode >= 400) { - addErrorResponseToSpan(response, span); - } else { - span.setStatus(StatusCode.OK); - } - } - - private Span getSpan( - String httpMethod, String url, String host, Integer port, Long requestBodySize) { - // TODO: Determine span name: {method} {url.template} or {method} - Span span = - BigQueryTelemetryTracer.newSpanBuilder(tracer, httpMethod) - // OpenTelemetry semantic convention attributes - .setAttribute(HTTP_REQUEST_METHOD, httpMethod) - .setAttribute(URL_FULL, url) - .setAttribute(BigQueryTelemetryTracer.SERVER_ADDRESS, host) - .setAttribute(URL_DOMAIN, BIGQUERY_DOMAIN) - .setAttribute(BigQueryTelemetryTracer.RPC_SYSTEM_NAME, "http") - .startSpan(); - - // TODO: add url template && resource name - if (port != null && port > 0) { - span.setAttribute(BigQueryTelemetryTracer.SERVER_PORT, port.longValue()); - } - if (requestBodySize != null && requestBodySize > 0) { + if (requestBodySize != null) { span.setAttribute(HTTP_REQUEST_BODY_SIZE, requestBodySize); } - return span; } - private static @Nullable Long getRequestBodySize(HttpRequest request) { - Long requestBodySize = null; + private static void setResponseBodySize(HttpResponse response, Span span) { try { - HttpContent content = request.getContent(); - if (content != null) { - requestBodySize = content.getLength(); + long contentLength = response.getHeaders().getContentLength(); + if (contentLength > 0) { + span.setAttribute(HTTP_RESPONSE_BODY_SIZE, contentLength); } } catch (Exception e) { // Ignore - body size not available } - return requestBodySize; + } + + private static void checkForUpdatedRequestMethod( + HttpResponse response, String httpMethod, Span span) { + String actualMethod = response.getRequest().getRequestMethod(); + if (actualMethod != null && httpMethod == null) { + span.updateName(actualMethod); + span.setAttribute(HTTP_REQUEST_METHOD, actualMethod); + } } } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java index fd9c521043..bdd0dfd06e 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java @@ -31,6 +31,11 @@ public final class BigQueryTelemetryTracer { private BigQueryTelemetryTracer() {} + public static final String BQ_GCP_CLIENT_SERVICE = "bigquery"; + public static final String BQ_GCP_CLIENT_REPO = "googleapis/java-bigquery"; + public static final String BQ_GCP_CLIENT_ARTIFACT = "google-cloud-bigquery"; + public static final String BQ_GCP_CLIENT_LANGUAGE = "java"; + // Common GCP Attributes public static final AttributeKey GCP_CLIENT_SERVICE = AttributeKey.stringKey("gcp.client.service"); @@ -63,10 +68,10 @@ public static SpanBuilder newSpanBuilder(Tracer tracer, String spanName) { return tracer .spanBuilder(spanName) .setSpanKind(SpanKind.CLIENT) - .setAttribute(GCP_CLIENT_SERVICE, "bigquery") - .setAttribute(GCP_CLIENT_REPO, "googleapis/java-bigquery") - .setAttribute(GCP_CLIENT_ARTIFACT, "google-cloud-bigquery") - .setAttribute(GCP_CLIENT_LANGUAGE, "java"); + .setAttribute(GCP_CLIENT_SERVICE, BQ_GCP_CLIENT_SERVICE) + .setAttribute(GCP_CLIENT_REPO, BQ_GCP_CLIENT_REPO) + .setAttribute(GCP_CLIENT_ARTIFACT, BQ_GCP_CLIENT_ARTIFACT) + .setAttribute(GCP_CLIENT_LANGUAGE, BQ_GCP_CLIENT_LANGUAGE); // TODO: add version } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingIntegrationTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingIntegrationTest.java index 4ed9f87675..ad2b31a961 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingIntegrationTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingIntegrationTest.java @@ -18,15 +18,14 @@ import static org.junit.jupiter.api.Assertions.assertEquals; - import com.google.api.client.http.GenericUrl; import com.google.api.client.http.HttpRequest; import com.google.api.client.http.HttpRequestFactory; import com.google.api.client.http.HttpResponse; import com.google.api.client.http.javanet.NetHttpTransport; import com.google.cloud.bigquery.RetryContext; +import com.google.cloud.bigquery.telemetry.BigQueryTelemetryTracer; import com.sun.net.httpserver.HttpServer; -import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.sdk.OpenTelemetrySdk; import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; @@ -42,13 +41,13 @@ import org.junit.jupiter.api.Test; /** - * Integration test for HTTP tracing with real HTTP transport and server. - * This test verifies that OpenTelemetry tracing works correctly with actual network calls. + * Integration test for HTTP tracing with real HTTP transport and server. This test verifies that + * OpenTelemetry tracing works correctly with actual network calls. */ public class HttpTracingIntegrationTest { private InMemorySpanExporter spanExporter; - private HttpTracingRequestInitializer initializer; + private HttpTracingRequestInitializer initializer; private HttpServer testServer; private int serverPort; @@ -62,8 +61,9 @@ public void setUp() throws IOException { .build(); OpenTelemetrySdk openTelemetry = OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).build(); - Tracer tracer = openTelemetry.getTracer("test-tracer"); - initializer = new HttpTracingRequestInitializer(null, tracer); + Tracer tracer = openTelemetry.getTracer("test-tracer"); + initializer = + new HttpTracingRequestInitializer(null, tracer, "https://example-client-endpoint.test/"); // Start a test HTTP server testServer = HttpServer.create(new InetSocketAddress("localhost", 0), 0); @@ -81,19 +81,22 @@ public void tearDown() { @Test public void testHttpTracingWithRealServer() throws IOException { - testServer.createContext("/api/test", exchange -> { - String response = "{\"status\": \"ok\"}"; - exchange.getResponseHeaders().add("Content-Type", "application/json"); - exchange.sendResponseHeaders(200, response.getBytes().length); - try (OutputStream os = exchange.getResponseBody()) { - os.write(response.getBytes()); - } - }); + testServer.createContext( + "/api/test", + exchange -> { + String response = "{\"status\": \"ok\"}"; + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, response.getBytes().length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(response.getBytes()); + } + }); NetHttpTransport transport = new NetHttpTransport(); HttpRequestFactory requestFactory = transport.createRequestFactory(initializer); - HttpRequest request = requestFactory.buildGetRequest( - new GenericUrl("http://localhost:" + serverPort + "/api/test")); + HttpRequest request = + requestFactory.buildGetRequest( + new GenericUrl("http://localhost:" + serverPort + "/api/test")); HttpResponse response = request.execute(); assertEquals(200, response.getStatusCode()); @@ -102,8 +105,34 @@ public void testHttpTracingWithRealServer() throws IOException { List spans = spanExporter.getFinishedSpanItems(); assertEquals(1, spans.size()); SpanData span = spans.get(0); - assertEquals("GET", span.getAttributes().get(AttributeKey.stringKey("http.request.method"))); - assertEquals(200L, span.getAttributes().get(AttributeKey.longKey("http.response.status_code"))); - assertEquals("localhost", span.getAttributes().get(AttributeKey.stringKey("server.address"))); + assertEquals( + 200, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_STATUS_CODE)); + assertEquals( + "GET", span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_METHOD)); + assertEquals("localhost", span.getAttributes().get(BigQueryTelemetryTracer.SERVER_ADDRESS)); + assertEquals(serverPort, span.getAttributes().get(BigQueryTelemetryTracer.SERVER_PORT)); + assertEquals( + "example-client-endpoint.test", + span.getAttributes().get(HttpTracingRequestInitializer.URL_DOMAIN)); + assertEquals( + "http://localhost:" + serverPort + "/api/test", + span.getAttributes().get(HttpTracingRequestInitializer.URL_FULL)); + assertEquals( + BigQueryTelemetryTracer.BQ_GCP_CLIENT_SERVICE, + span.getAttributes().get(BigQueryTelemetryTracer.GCP_CLIENT_SERVICE)); + assertEquals( + BigQueryTelemetryTracer.BQ_GCP_CLIENT_REPO, + span.getAttributes().get(BigQueryTelemetryTracer.GCP_CLIENT_REPO)); + assertEquals( + BigQueryTelemetryTracer.BQ_GCP_CLIENT_ARTIFACT, + span.getAttributes().get(BigQueryTelemetryTracer.GCP_CLIENT_ARTIFACT)); + assertEquals( + BigQueryTelemetryTracer.BQ_GCP_CLIENT_LANGUAGE, + span.getAttributes().get(BigQueryTelemetryTracer.GCP_CLIENT_LANGUAGE)); + assertEquals( + HttpTracingRequestInitializer.HTTP_RPC_SYSTEM_NAME, + span.getAttributes().get(BigQueryTelemetryTracer.RPC_SYSTEM_NAME)); + assertEquals( + 16, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_BODY_SIZE)); } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java index 0e19c86e65..c8f97a21d5 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java @@ -18,7 +18,6 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; @@ -26,6 +25,7 @@ import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; +import com.google.api.client.http.ByteArrayContent; import com.google.api.client.http.GenericUrl; import com.google.api.client.http.HttpRequest; import com.google.api.client.http.HttpRequestFactory; @@ -37,8 +37,7 @@ import com.google.api.client.testing.http.MockHttpTransport; import com.google.api.client.testing.http.MockLowLevelHttpRequest; import com.google.api.client.testing.http.MockLowLevelHttpResponse; -import com.google.cloud.bigquery.RetryContext; -import io.opentelemetry.api.common.AttributeKey; +import com.google.cloud.bigquery.telemetry.BigQueryTelemetryTracer; import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.sdk.OpenTelemetrySdk; import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; @@ -47,14 +46,19 @@ import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; import java.io.IOException; import java.util.List; -import java.util.stream.Collectors; -import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Unit tests for TracingHttpRequestInitializer */ public class HttpTracingRequestInitializerTest { + private static final String BASE_URL = + "https://bigquery.googleapis.com:443/bigquery/v2/projects/test/datasets"; + private static final String REQUEST_METHOD_GET = "GET"; + private static final String REQUEST_METHOD_POST = "POST"; + private static final String BIGQUERY_DOMAIN = "bigquery.googleapis.com"; + private static final String CLIENT_ROOT_URL = "https://bigquery.googleapis.com:443"; + private InMemorySpanExporter spanExporter; private Tracer tracer; private HttpTracingRequestInitializer initializer; @@ -69,35 +73,18 @@ public void setUp() { OpenTelemetrySdk openTelemetry = OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).build(); tracer = openTelemetry.getTracer("test-tracer"); - initializer = new HttpTracingRequestInitializer(null, tracer); - } - - @AfterEach - public void tearDown() { - // Clean up thread-local retry context after each test - RetryContext.clearRetryAttempt(); + initializer = new HttpTracingRequestInitializer(null, tracer, CLIENT_ROOT_URL); } @Test public void testSuccessResponseAttributesAreSet() throws IOException { - HttpTransport transport = - new MockHttpTransport() { - @Override - public LowLevelHttpRequest buildRequest(String method, String url) { - return new MockLowLevelHttpRequest() { - @Override - public LowLevelHttpResponse execute() { - MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); - response.setStatusCode(200); - return response; - } - }; - } - }; - - HttpRequestFactory requestFactory = transport.createRequestFactory(initializer); - String urlString = "https://bigquery.googleapis.com:443/bigquery/v2/projects/test/datasets"; - HttpRequest request = requestFactory.buildGetRequest(new GenericUrl(urlString)); + HttpTransport transport = createTransport(200, null, 123L); + HttpRequest request = + buildPostRequest( + transport, + initializer, + BASE_URL, + new ByteArrayContent("application/json", new byte[] {1})); HttpResponse response = request.execute(); response.disconnect(); @@ -106,57 +93,24 @@ public LowLevelHttpResponse execute() { assertEquals(1, spans.size()); SpanData span = spans.get(0); - assertEquals("GET", span.getAttributes().get(AttributeKey.stringKey("http.request.method"))); - assertEquals( - "bigquery.googleapis.com", - span.getAttributes().get(AttributeKey.stringKey("server.address"))); - assertEquals(Long.valueOf(443L), span.getAttributes().get(AttributeKey.longKey("server.port"))); - assertEquals( - "bigquery.googleapis.com", span.getAttributes().get(AttributeKey.stringKey("url.domain"))); - assertEquals(urlString, span.getAttributes().get(AttributeKey.stringKey("url.full"))); - assertEquals( - Long.valueOf(200L), - span.getAttributes().get(AttributeKey.longKey("http.response.status_code"))); - assertEquals( - "bigquery", span.getAttributes().get(AttributeKey.stringKey("gcp.client.service"))); assertEquals( - "googleapis/java-bigquery", - span.getAttributes().get(AttributeKey.stringKey("gcp.client.repo"))); + 200, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_STATUS_CODE)); + assertEquals(1, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_BODY_SIZE)); assertEquals( - "google-cloud-bigquery", - span.getAttributes().get(AttributeKey.stringKey("gcp.client.artifact"))); - assertEquals("java", span.getAttributes().get(AttributeKey.stringKey("gcp.client.language"))); + 123, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_BODY_SIZE)); + verifyGeneralSpanData(span, REQUEST_METHOD_POST); } @Test public void testErrorAttributesAreSetOn404() throws IOException { - HttpTransport transport = - new MockHttpTransport() { - @Override - public LowLevelHttpRequest buildRequest(String method, String url) { - return new MockLowLevelHttpRequest() { - @Override - public LowLevelHttpResponse execute() { - MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); - response.setStatusCode(404); - response.setReasonPhrase("Not Found"); - return response; - } - }; - } - }; - - HttpRequestFactory requestFactory = transport.createRequestFactory(initializer); - HttpRequest request = - requestFactory.buildGetRequest( - new GenericUrl( - "https://bigquery.googleapis.com/bigquery/v2/projects/test/datasets/notfound")); + HttpTransport transport = createTransport(404, "Not Found", null); + HttpRequest request = buildGetRequest(transport, initializer, BASE_URL); try { HttpResponse response = request.execute(); response.disconnect(); } catch (Exception e) { - // Expected - 404 might throw + // Expected } List spans = spanExporter.getFinishedSpanItems(); @@ -164,144 +118,60 @@ public LowLevelHttpResponse execute() { SpanData span = spans.get(0); assertEquals( - Long.valueOf(404L), - span.getAttributes().get(AttributeKey.longKey("http.response.status_code"))); - assertEquals("404", span.getAttributes().get(AttributeKey.stringKey("error.type"))); - assertNotNull(span.getAttributes().get(AttributeKey.stringKey("status.message"))); + 404, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_STATUS_CODE)); + assertEquals("404", span.getAttributes().get(BigQueryTelemetryTracer.ERROR_TYPE)); + assertNotNull(span.getAttributes().get(BigQueryTelemetryTracer.STATUS_MESSAGE)); + verifyGeneralSpanData(span, REQUEST_METHOD_GET); } @Test public void testExceptionAttributesAreSetWhenOriginalUnsuccessfulHandlerThrowsIOException() throws IOException { + String handlerFailureMessage = "handler failure"; + String serverErrorMessage = "Internal Server Error"; + HttpRequestInitializer delegateInitializer = request -> request.setUnsuccessfulResponseHandler( (request1, response, supportsRetry) -> { - throw new IOException("handler failure"); + throw new IOException(handlerFailureMessage); }); HttpTracingRequestInitializer tracingInitializer = - new HttpTracingRequestInitializer(delegateInitializer, tracer); - - HttpTransport transport = - new MockHttpTransport() { - @Override - public LowLevelHttpRequest buildRequest(String method, String url) { - return new MockLowLevelHttpRequest() { - @Override - public LowLevelHttpResponse execute() { - MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); - response.setStatusCode(500); - response.setReasonPhrase("Internal Server Error"); - return response; - } - }; - } - }; + new HttpTracingRequestInitializer(delegateInitializer, tracer, CLIENT_ROOT_URL); - HttpRequestFactory requestFactory = transport.createRequestFactory(tracingInitializer); - HttpRequest request = - requestFactory.buildGetRequest( - new GenericUrl("https://bigquery.googleapis.com/bigquery/v2/projects/test/datasets")); + HttpTransport transport = createTransport(500, serverErrorMessage, null); + HttpRequest request = buildGetRequest(transport, tracingInitializer, BASE_URL); request.setThrowExceptionOnExecuteError(false); IOException thrown = assertThrows(IOException.class, request::execute); - assertEquals("handler failure", thrown.getMessage()); + assertEquals(handlerFailureMessage, thrown.getMessage()); List spans = spanExporter.getFinishedSpanItems(); assertEquals(1, spans.size()); SpanData span = spans.get(0); assertEquals( - Long.valueOf(500L), - span.getAttributes().get(AttributeKey.longKey("http.response.status_code"))); + 500, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_STATUS_CODE)); assertEquals( IOException.class.getName(), - span.getAttributes().get(AttributeKey.stringKey("exception.type"))); + span.getAttributes().get(BigQueryTelemetryTracer.EXCEPTION_TYPE)); assertEquals( IOException.class.getSimpleName(), - span.getAttributes().get(AttributeKey.stringKey("error.type"))); + span.getAttributes().get(BigQueryTelemetryTracer.ERROR_TYPE)); assertEquals( - "handler failure", span.getAttributes().get(AttributeKey.stringKey("status.message"))); - } - - @Test - public void testExceptionIsRecordedWhenOriginalUnsuccessfulHandlerThrowsIOException() - throws IOException { - HttpRequestInitializer delegateInitializer = - request -> - request.setUnsuccessfulResponseHandler( - (request1, response, supportsRetry) -> { - throw new IOException("handler failure"); - }); - HttpTracingRequestInitializer tracingInitializer = - new HttpTracingRequestInitializer(delegateInitializer, tracer); - - HttpTransport transport = - new MockHttpTransport() { - @Override - public LowLevelHttpRequest buildRequest(String method, String url) { - return new MockLowLevelHttpRequest() { - @Override - public LowLevelHttpResponse execute() { - MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); - response.setStatusCode(500); - response.setReasonPhrase("Internal Server Error"); - return response; - } - }; - } - }; - - HttpRequestFactory requestFactory = transport.createRequestFactory(tracingInitializer); - HttpRequest request = - requestFactory.buildGetRequest( - new GenericUrl("https://bigquery.googleapis.com/bigquery/v2/projects/test/datasets")); - request.setThrowExceptionOnExecuteError(false); - - assertThrows(IOException.class, request::execute); - - List spans = spanExporter.getFinishedSpanItems(); - assertEquals(1, spans.size()); - SpanData span = spans.get(0); - + handlerFailureMessage, span.getAttributes().get(BigQueryTelemetryTracer.STATUS_MESSAGE)); assertTrue(span.getEvents().stream().anyMatch(event -> "exception".equals(event.getName()))); - assertTrue( - span.getEvents().stream() - .filter(event -> "exception".equals(event.getName())) - .anyMatch( - event -> - "handler failure" - .equals( - event - .getAttributes() - .get(AttributeKey.stringKey("exception.message"))))); + verifyGeneralSpanData(span, REQUEST_METHOD_GET); } @Test public void testDelegateInitializerIsCalledOnSuccessResponse() throws IOException { HttpRequestInitializer delegateInitializer = mock(HttpRequestInitializer.class); HttpTracingRequestInitializer tracingInitializer = - new HttpTracingRequestInitializer(delegateInitializer, tracer); + new HttpTracingRequestInitializer(delegateInitializer, tracer, CLIENT_ROOT_URL); - HttpTransport transport = - new MockHttpTransport() { - @Override - public LowLevelHttpRequest buildRequest(String method, String url) { - return new MockLowLevelHttpRequest() { - @Override - public LowLevelHttpResponse execute() { - MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); - response.setStatusCode(200); - return response; - } - }; - } - }; - - HttpRequestFactory requestFactory = transport.createRequestFactory(tracingInitializer); - HttpRequest request = - requestFactory.buildGetRequest( - new GenericUrl("https://bigquery.googleapis.com/bigquery/v2/projects/test/datasets")); + HttpTransport transport = createTransport(200, null, null); + HttpRequest request = buildGetRequest(transport, tracingInitializer, BASE_URL); HttpResponse response = request.execute(); response.disconnect(); @@ -313,29 +183,10 @@ public LowLevelHttpResponse execute() { public void testDelegateInitializerIsCalledOnErrorResponse() throws IOException { HttpRequestInitializer delegateInitializer = mock(HttpRequestInitializer.class); HttpTracingRequestInitializer tracingInitializer = - new HttpTracingRequestInitializer(delegateInitializer, tracer); - - HttpTransport transport = - new MockHttpTransport() { - @Override - public LowLevelHttpRequest buildRequest(String method, String url) { - return new MockLowLevelHttpRequest() { - @Override - public LowLevelHttpResponse execute() { - MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); - response.setStatusCode(404); - response.setReasonPhrase("Not Found"); - return response; - } - }; - } - }; + new HttpTracingRequestInitializer(delegateInitializer, tracer, CLIENT_ROOT_URL); - HttpRequestFactory requestFactory = transport.createRequestFactory(tracingInitializer); - HttpRequest request = - requestFactory.buildGetRequest( - new GenericUrl( - "https://bigquery.googleapis.com/bigquery/v2/projects/test/datasets/notfound")); + HttpTransport transport = createTransport(404, "Not Found", null); + HttpRequest request = buildGetRequest(transport, tracingInitializer, BASE_URL); try { HttpResponse response = request.execute(); @@ -346,4 +197,87 @@ public LowLevelHttpResponse execute() { verify(delegateInitializer, times(1)).initialize(any(HttpRequest.class)); } + + @Test + public void testUrlDomainUsesClientRootUrlHost() throws IOException { + HttpTracingRequestInitializer tracingInitializer = + new HttpTracingRequestInitializer(null, tracer, "https://example-client-endpoint.test/"); + + HttpTransport transport = createTransport(200, null, null); + HttpRequest request = buildGetRequest(transport, tracingInitializer, BASE_URL); + + HttpResponse response = request.execute(); + response.disconnect(); + + List spans = spanExporter.getFinishedSpanItems(); + assertEquals(1, spans.size()); + assertEquals( + "example-client-endpoint.test", + spans.get(0).getAttributes().get(HttpTracingRequestInitializer.URL_DOMAIN)); + } + + private static HttpTransport createTransport( + int statusCode, String reasonPhrase, Long contentLength) { + return new MockHttpTransport() { + @Override + public LowLevelHttpRequest buildRequest(String method, String url) { + return new MockLowLevelHttpRequest() { + @Override + public LowLevelHttpResponse execute() { + MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); + response.setStatusCode(statusCode); + if (reasonPhrase != null) { + response.setReasonPhrase(reasonPhrase); + } + if (contentLength != null) { + response.addHeader("Content-Length", String.valueOf(contentLength)); + } + return response; + } + }; + } + }; + } + + private static HttpRequest buildGetRequest( + HttpTransport transport, HttpRequestInitializer requestInitializer, String url) + throws IOException { + HttpRequestFactory requestFactory = transport.createRequestFactory(requestInitializer); + return requestFactory.buildGetRequest(new GenericUrl(url)); + } + + private static HttpRequest buildPostRequest( + HttpTransport transport, + HttpRequestInitializer requestInitializer, + String url, + ByteArrayContent content) + throws IOException { + HttpRequestFactory requestFactory = transport.createRequestFactory(requestInitializer); + return requestFactory.buildPostRequest(new GenericUrl(url), content); + } + + private void verifyGeneralSpanData(SpanData span, String requestMethod) { + assertEquals( + requestMethod, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_METHOD)); + assertEquals(BIGQUERY_DOMAIN, span.getAttributes().get(BigQueryTelemetryTracer.SERVER_ADDRESS)); + assertEquals(443, span.getAttributes().get(BigQueryTelemetryTracer.SERVER_PORT)); + assertEquals( + BIGQUERY_DOMAIN, span.getAttributes().get(HttpTracingRequestInitializer.URL_DOMAIN)); + assertEquals(BASE_URL, span.getAttributes().get(HttpTracingRequestInitializer.URL_FULL)); + assertEquals( + BigQueryTelemetryTracer.BQ_GCP_CLIENT_SERVICE, + span.getAttributes().get(BigQueryTelemetryTracer.GCP_CLIENT_SERVICE)); + assertEquals( + BigQueryTelemetryTracer.BQ_GCP_CLIENT_REPO, + span.getAttributes().get(BigQueryTelemetryTracer.GCP_CLIENT_REPO)); + assertEquals( + BigQueryTelemetryTracer.BQ_GCP_CLIENT_ARTIFACT, + span.getAttributes().get(BigQueryTelemetryTracer.GCP_CLIENT_ARTIFACT)); + assertEquals( + BigQueryTelemetryTracer.BQ_GCP_CLIENT_LANGUAGE, + span.getAttributes().get(BigQueryTelemetryTracer.GCP_CLIENT_LANGUAGE)); + assertEquals( + HttpTracingRequestInitializer.HTTP_RPC_SYSTEM_NAME, + span.getAttributes().get(BigQueryTelemetryTracer.RPC_SYSTEM_NAME)); + } } From 0372db5302ea42c5a085ea4bf2916402cdaad55d Mon Sep 17 00:00:00 2001 From: ldetmer Date: Tue, 3 Mar 2026 15:57:15 -0500 Subject: [PATCH 4/6] added handling of full.url redacted content --- .../google/cloud/bigquery/RetryContext.java | 76 ------------------- .../spi/v2/HttpTracingRequestInitializer.java | 59 +++++++++++++- .../spi/v2/HttpTracingIntegrationTest.java | 2 - .../v2/HttpTracingRequestInitializerTest.java | 21 +++++ 4 files changed, 79 insertions(+), 79 deletions(-) delete mode 100644 google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RetryContext.java diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RetryContext.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RetryContext.java deleted file mode 100644 index 3b1cfd4558..0000000000 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RetryContext.java +++ /dev/null @@ -1,76 +0,0 @@ -/* - * Copyright 2026 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.bigquery; - -/** - * Utility class for storing and retrieving retry attempt count using thread-local storage. - * This allows retry information to be accessed by HTTP/gRPC tracing interceptors during - * retry operations. - * - *

Thread-local storage is used because GAX retry framework executes retries on the same - * thread, making this mechanism both simple and reliable. - */ -public class RetryContext { - - /** - * Thread-local storage for the current retry attempt number. - * Value 0 means the first attempt (no retries yet). - * Value 1 means first retry (second attempt total). - * Value N means Nth retry (N+1 attempts total). - */ - private static final ThreadLocal RETRY_ATTEMPT = ThreadLocal.withInitial(() -> 0); - - /** - * Stores the retry attempt count in thread-local storage. - * - * @param attemptCount The attempt number (0 = first attempt, 1 = first retry, etc.) - */ - public static void setRetryAttempt(int attemptCount) { - RETRY_ATTEMPT.set(attemptCount); - } - - /** - * Retrieves the retry attempt count from thread-local storage. - * - * @return The retry attempt count, or 0 if not set (meaning first attempt) - */ - public static int getRetryAttempt() { - Integer attemptCount = RETRY_ATTEMPT.get(); - return attemptCount != null ? attemptCount : 0; - } - - /** - * Clears the retry attempt count from thread-local storage. - * Should be called when retry operations complete to prevent memory leaks. - */ - public static void clearRetryAttempt() { - RETRY_ATTEMPT.remove(); - } - - /** - * Checks if the current thread is processing a retry (attempt count > 0). - * - * @return true if this is a retry attempt, false if this is the first attempt - */ - public static boolean isRetry() { - return getRetryAttempt() > 0; - } - - private RetryContext() { - // Utility class, prevent instantiation - } -} diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java index 6aa9a3ff09..45e45ccb63 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java @@ -25,6 +25,12 @@ import io.opentelemetry.api.trace.StatusCode; import io.opentelemetry.api.trace.Tracer; import java.io.IOException; +import java.net.URI; +import java.net.URISyntaxException; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.Set; import org.jspecify.annotations.Nullable; /** @@ -50,6 +56,10 @@ public class HttpTracingRequestInitializer implements HttpRequestInitializer { AttributeKey.longKey("http.response.body.size"); @VisibleForTesting static final String HTTP_RPC_SYSTEM_NAME = "http"; + private static final String REDACTED_VALUE = "REDACTED"; + private static final Set SENSITIVE_QUERY_KEYS = + Collections.unmodifiableSet( + new HashSet<>(Arrays.asList("AWSAccessKeyId", "Signature", "sig", "X-Goog-Signature"))); private final HttpRequestInitializer delegate; private final Tracer tracer; @@ -128,7 +138,7 @@ private Span createHttpTraceSpan(String httpMethod, String url, String host, Int Span span = BigQueryTelemetryTracer.newSpanBuilder(tracer, httpMethod) .setAttribute(HTTP_REQUEST_METHOD, httpMethod) - .setAttribute(URL_FULL, url) + .setAttribute(URL_FULL, sanitizeUrlFull(url)) .setAttribute(BigQueryTelemetryTracer.SERVER_ADDRESS, host) .setAttribute(URL_DOMAIN, resolveUrlDomain(host)) .setAttribute(BigQueryTelemetryTracer.RPC_SYSTEM_NAME, HTTP_RPC_SYSTEM_NAME) @@ -233,4 +243,51 @@ private static void checkForUpdatedRequestMethod( span.setAttribute(HTTP_REQUEST_METHOD, actualMethod); } } + + @VisibleForTesting + static String sanitizeUrlFull(String url) { + try { + URI uri = new URI(url); + String sanitizedUserInfo = + uri.getRawUserInfo() != null ? REDACTED_VALUE + ":" + REDACTED_VALUE : null; + String sanitizedQuery = redactSensitiveQueryValues(uri.getRawQuery()); + URI sanitizedUri = + new URI( + uri.getScheme(), + sanitizedUserInfo, + uri.getHost(), + uri.getPort(), + uri.getRawPath(), + sanitizedQuery, + uri.getRawFragment()); + return sanitizedUri.toString(); + } catch (URISyntaxException | IllegalArgumentException ex) { + return url; + } + } + + private static String redactSensitiveQueryValues(@Nullable String rawQuery) { + if (rawQuery == null || rawQuery.isEmpty()) { + return rawQuery; + } + + String[] params = rawQuery.split("&", -1); + for (int i = 0; i < params.length; i++) { + String param = params[i]; + int equalsIndex = param.indexOf('='); + String key = equalsIndex >= 0 ? param.substring(0, equalsIndex) : param; + if (SENSITIVE_QUERY_KEYS.contains(key)) { + params[i] = key + "=" + REDACTED_VALUE; + } + } + + StringBuilder redactedQuery = new StringBuilder(); + for (int i = 0; i < params.length; i++) { + if (i > 0) { + redactedQuery.append('&'); + } + redactedQuery.append(params[i]); + } + return redactedQuery.toString(); + } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingIntegrationTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingIntegrationTest.java index ad2b31a961..811a8b41c1 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingIntegrationTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingIntegrationTest.java @@ -23,7 +23,6 @@ import com.google.api.client.http.HttpRequestFactory; import com.google.api.client.http.HttpResponse; import com.google.api.client.http.javanet.NetHttpTransport; -import com.google.cloud.bigquery.RetryContext; import com.google.cloud.bigquery.telemetry.BigQueryTelemetryTracer; import com.sun.net.httpserver.HttpServer; import io.opentelemetry.api.trace.Tracer; @@ -76,7 +75,6 @@ public void tearDown() { if (testServer != null) { testServer.stop(0); } - RetryContext.clearRetryAttempt(); } @Test diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java index c8f97a21d5..caf8dc44e4 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java @@ -216,6 +216,27 @@ public void testUrlDomainUsesClientRootUrlHost() throws IOException { spans.get(0).getAttributes().get(HttpTracingRequestInitializer.URL_DOMAIN)); } + @Test + public void testUrlFullIsRequestBasedAndRedactsSensitiveContent() throws IOException { + HttpTransport transport = createTransport(200, null, null); + HttpRequest request = + buildGetRequest( + transport, + initializer, + "https://user:password@bigquery.googleapis.com:443/bigquery/v2/projects/test/datasets" + + "?AWSAccessKeyId=abc&Signature=def&sig=ghi&X-Goog-Signature=jkl&safe=value&signature=lower#frag"); + + HttpResponse response = request.execute(); + response.disconnect(); + + List spans = spanExporter.getFinishedSpanItems(); + assertEquals(1, spans.size()); + assertEquals( + "https://REDACTED:REDACTED@bigquery.googleapis.com:443/bigquery/v2/projects/test/datasets" + + "?AWSAccessKeyId=REDACTED&Signature=REDACTED&sig=REDACTED&X-Goog-Signature=REDACTED&safe=value&signature=lower#frag", + spans.get(0).getAttributes().get(HttpTracingRequestInitializer.URL_FULL)); + } + private static HttpTransport createTransport( int statusCode, String reasonPhrase, Long contentLength) { return new MockHttpTransport() { From de0ee8a1f6dc31b144ac45726cff698b5bbe5ba0 Mon Sep 17 00:00:00 2001 From: ldetmer Date: Tue, 3 Mar 2026 16:15:20 -0500 Subject: [PATCH 5/6] remove nullable annotation --- .../bigquery/spi/v2/HttpTracingRequestInitializer.java | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java index 45e45ccb63..ea74d7e83b 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java @@ -31,7 +31,6 @@ import java.util.Collections; import java.util.HashSet; import java.util.Set; -import org.jspecify.annotations.Nullable; /** * HttpRequestInitializer that wraps a delegate initializer, intercepts all HTTP requests, adds @@ -63,10 +62,10 @@ public class HttpTracingRequestInitializer implements HttpRequestInitializer { private final HttpRequestInitializer delegate; private final Tracer tracer; - private final @Nullable String clientRootUrl; + private final String clientRootUrl; public HttpTracingRequestInitializer( - HttpRequestInitializer delegate, Tracer tracer, @Nullable String clientRootUrl) { + HttpRequestInitializer delegate, Tracer tracer, String clientRootUrl) { this.delegate = delegate; this.tracer = tracer; this.clientRootUrl = clientRootUrl; @@ -266,7 +265,7 @@ static String sanitizeUrlFull(String url) { } } - private static String redactSensitiveQueryValues(@Nullable String rawQuery) { + private static String redactSensitiveQueryValues(String rawQuery) { if (rawQuery == null || rawQuery.isEmpty()) { return rawQuery; } From 68704bd9eb2c036dc8bbb7d0480f3166b09a0c05 Mon Sep 17 00:00:00 2001 From: ldetmer Date: Wed, 4 Mar 2026 10:29:27 -0500 Subject: [PATCH 6/6] simplified logic --- .../bigquery/spi/v2/HttpBigQueryRpc.java | 2 +- .../spi/v2/HttpTracingRequestInitializer.java | 76 +++++++++---------- 2 files changed, 36 insertions(+), 42 deletions(-) diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java index 3d1e37dd2a..894db32aad 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java @@ -109,7 +109,7 @@ public HttpBigQueryRpc(BigQueryOptions options) { HttpRequestInitializer initializer = transportOptions.getHttpRequestInitializer(options); String resolvedBigQueryRootUrl = options.getResolvedApiaryHost("bigquery"); - // Wrap with tracing initializer if OpenTelemetry is enabled + if (options.isOpenTelemetryTracingEnabled() && options.getOpenTelemetryTracer() != null) { initializer = new HttpTracingRequestInitializer( diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java index ea74d7e83b..6e0a95b63e 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java @@ -55,10 +55,20 @@ public class HttpTracingRequestInitializer implements HttpRequestInitializer { AttributeKey.longKey("http.response.body.size"); @VisibleForTesting static final String HTTP_RPC_SYSTEM_NAME = "http"; + private static final String REDACTED_VALUE = "REDACTED"; + // Required by OpenTelemetry semantic conventions: + // https://opentelemetry.io/docs/specs/semconv/registry/attributes/url/#url-full private static final Set SENSITIVE_QUERY_KEYS = Collections.unmodifiableSet( - new HashSet<>(Arrays.asList("AWSAccessKeyId", "Signature", "sig", "X-Goog-Signature"))); + new HashSet<>( + Arrays.asList( + "AWSAccessKeyId", + "Signature", + "sig", + "X-Goog-Signature", + // Google uses this as a key in resumable uploads. + "upload_id"))); private final HttpRequestInitializer delegate; private final Tracer tracer; @@ -88,27 +98,25 @@ public void initialize(HttpRequest request) throws IOException { Span span = createHttpTraceSpan(httpMethod, url, host, port); - // Wrap the existing response interceptor HttpResponseInterceptor originalInterceptor = request.getResponseInterceptor(); request.setResponseInterceptor( response -> { - if (span.isRecording()) { - try { - int statusCode = response.getStatusCode(); - addCommonResponseAttributesToSpan(request, response, span, httpMethod, statusCode); - addSuccessResponseToSpan(response, span, statusCode); - if (originalInterceptor != null) { - originalInterceptor.interceptResponse(response); - } - } finally { - span.end(); + addCommonResponseAttributesToSpan( + request, response, span, httpMethod, response.getStatusCode()); + span.setStatus(StatusCode.OK); + + try { + if (originalInterceptor != null) { + originalInterceptor.interceptResponse(response); } - } else if (originalInterceptor != null) { - originalInterceptor.interceptResponse(response); + } catch (IOException e) { + addExceptionToSpan(e, span); + throw e; + } finally { + span.end(); } }); - // Wrap the existing unsuccessful response handler HttpUnsuccessfulResponseHandler originalHandler = request.getUnsuccessfulResponseHandler(); request.setUnsuccessfulResponseHandler( (request1, response, supportsRetry) -> { @@ -124,16 +132,15 @@ public void initialize(HttpRequest request) throws IOException { addExceptionToSpan(e, span); throw e; } finally { - if (span.isRecording()) { - span.end(); - } + span.end(); } }); } /** Initial HTTP trace span creation with basic attributes from request */ private Span createHttpTraceSpan(String httpMethod, String url, String host, Integer port) { - // TODO: Determine span name: {method} {url.template} or {method} + // TODO: add url template && resource name + // TODO: appropriately determine span name using: {method} {url.template} or {method} Span span = BigQueryTelemetryTracer.newSpanBuilder(tracer, httpMethod) .setAttribute(HTTP_REQUEST_METHOD, httpMethod) @@ -142,8 +149,6 @@ private Span createHttpTraceSpan(String httpMethod, String url, String host, Int .setAttribute(URL_DOMAIN, resolveUrlDomain(host)) .setAttribute(BigQueryTelemetryTracer.RPC_SYSTEM_NAME, HTTP_RPC_SYSTEM_NAME) .startSpan(); - - // TODO: add url template && resource name if (port != null && port > 0) { span.setAttribute(BigQueryTelemetryTracer.SERVER_PORT, port.longValue()); } @@ -166,20 +171,14 @@ private String resolveUrlDomain(String requestHost) { private static void addCommonResponseAttributesToSpan( HttpRequest request, HttpResponse response, Span span, String httpMethod, int statusCode) { - // This is called after we get a response as sometimes the request body size isn't available - // before the response is received. + // We add request body size/update request method after we receive response as they sometimes + // the data is + // not available until after the http request execution addRequestBodySizeToSpan(request, span); checkForUpdatedRequestMethod(response, httpMethod, span); - setResponseBodySize(response, span); - span.setAttribute(HTTP_RESPONSE_STATUS_CODE, statusCode); - } - private static void addSuccessResponseToSpan(HttpResponse response, Span span, int statusCode) { - if (statusCode >= 400) { - addErrorResponseToSpan(response, span, statusCode); - } else { - span.setStatus(StatusCode.OK); - } + addResponseBodySizeToSpan(response, span); + span.setAttribute(HTTP_RESPONSE_STATUS_CODE, statusCode); } private static void addExceptionToSpan(IOException e, Span span) { @@ -208,22 +207,17 @@ private static void addErrorResponseToSpan(HttpResponse response, Span span, int } private static void addRequestBodySizeToSpan(HttpRequest request, Span span) { - Long requestBodySize = null; try { - HttpContent content = request.getContent(); - - if (content != null) { - requestBodySize = content.getLength(); + long contentLength = request.getContent().getLength(); + if (contentLength > 0) { + span.setAttribute(HTTP_REQUEST_BODY_SIZE, contentLength); } } catch (Exception e) { // Ignore - body size not available } - if (requestBodySize != null) { - span.setAttribute(HTTP_REQUEST_BODY_SIZE, requestBodySize); - } } - private static void setResponseBodySize(HttpResponse response, Span span) { + private static void addResponseBodySizeToSpan(HttpResponse response, Span span) { try { long contentLength = response.getHeaders().getContentLength(); if (contentLength > 0) {