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/spi/v2/HttpBigQueryRpc.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java index 16737dc4b7..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 @@ -107,10 +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"); + + if (options.isOpenTelemetryTracingEnabled() && options.getOpenTelemetryTracer() != null) { + 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 new file mode 100644 index 0000000000..6e0a95b63e --- /dev/null +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializer.java @@ -0,0 +1,286 @@ +/* + * 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 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 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; + +/** + * 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"); + + @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", + // Google uses this as a key in resumable uploads. + "upload_id"))); + + private final HttpRequestInitializer delegate; + private final Tracer tracer; + private final String clientRootUrl; + + public HttpTracingRequestInitializer( + HttpRequestInitializer delegate, Tracer tracer, String clientRootUrl) { + this.delegate = delegate; + this.tracer = tracer; + this.clientRootUrl = clientRootUrl; + } + + @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(); + + Span span = createHttpTraceSpan(httpMethod, url, host, port); + + HttpResponseInterceptor originalInterceptor = request.getResponseInterceptor(); + request.setResponseInterceptor( + response -> { + addCommonResponseAttributesToSpan( + request, response, span, httpMethod, response.getStatusCode()); + span.setStatus(StatusCode.OK); + + try { + if (originalInterceptor != null) { + originalInterceptor.interceptResponse(response); + } + } catch (IOException e) { + addExceptionToSpan(e, span); + throw e; + } finally { + span.end(); + } + }); + + HttpUnsuccessfulResponseHandler originalHandler = request.getUnsuccessfulResponseHandler(); + request.setUnsuccessfulResponseHandler( + (request1, response, supportsRetry) -> { + int statusCode = response.getStatusCode(); + addCommonResponseAttributesToSpan(request, response, span, httpMethod, statusCode); + addErrorResponseToSpan(response, span, statusCode); + try { + if (originalHandler != null) { + return originalHandler.handleResponse(request1, response, supportsRetry); + } + return false; + } catch (IOException e) { + addExceptionToSpan(e, span); + throw e; + } finally { + span.end(); + } + }); + } + + /** Initial HTTP trace span creation with basic attributes from request */ + private Span createHttpTraceSpan(String httpMethod, String url, String host, Integer port) { + // 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) + .setAttribute(URL_FULL, sanitizeUrlFull(url)) + .setAttribute(BigQueryTelemetryTracer.SERVER_ADDRESS, host) + .setAttribute(URL_DOMAIN, resolveUrlDomain(host)) + .setAttribute(BigQueryTelemetryTracer.RPC_SYSTEM_NAME, HTTP_RPC_SYSTEM_NAME) + .startSpan(); + 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) { + // 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); + + addResponseBodySizeToSpan(response, span); + span.setAttribute(HTTP_RESPONSE_STATUS_CODE, statusCode); + } + + 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, statusMessage); + span.setStatus(StatusCode.ERROR, statusMessage); + } + + private static void addErrorResponseToSpan(HttpResponse response, Span span, int 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 addRequestBodySizeToSpan(HttpRequest request, Span span) { + try { + long contentLength = request.getContent().getLength(); + if (contentLength > 0) { + span.setAttribute(HTTP_REQUEST_BODY_SIZE, contentLength); + } + } catch (Exception e) { + // Ignore - body size not available + } + } + + private static void addResponseBodySizeToSpan(HttpResponse response, Span span) { + try { + long contentLength = response.getHeaders().getContentLength(); + if (contentLength > 0) { + span.setAttribute(HTTP_RESPONSE_BODY_SIZE, contentLength); + } + } catch (Exception e) { + // Ignore - body size not available + } + } + + 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); + } + } + + @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(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/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..bdd0dfd06e --- /dev/null +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java @@ -0,0 +1,77 @@ +/* + * 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() {} + + 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"); + 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, 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 new file mode 100644 index 0000000000..811a8b41c1 --- /dev/null +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingIntegrationTest.java @@ -0,0 +1,136 @@ +/* + * 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.telemetry.BigQueryTelemetryTracer; +import com.sun.net.httpserver.HttpServer; +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, "https://example-client-endpoint.test/"); + + // 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); + } + } + + @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( + 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 new file mode 100644 index 0000000000..caf8dc44e4 --- /dev/null +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpTracingRequestInitializerTest.java @@ -0,0 +1,304 @@ +/* + * 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.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.ByteArrayContent; +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.telemetry.BigQueryTelemetryTracer; +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 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; + + @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, CLIENT_ROOT_URL); + } + + @Test + public void testSuccessResponseAttributesAreSet() throws IOException { + 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(); + + List spans = spanExporter.getFinishedSpanItems(); + assertEquals(1, spans.size()); + + SpanData span = spans.get(0); + assertEquals( + 200, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_STATUS_CODE)); + assertEquals(1, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_BODY_SIZE)); + assertEquals( + 123, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_BODY_SIZE)); + verifyGeneralSpanData(span, REQUEST_METHOD_POST); + } + + @Test + public void testErrorAttributesAreSetOn404() throws IOException { + 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 + } + + List spans = spanExporter.getFinishedSpanItems(); + assertEquals(1, spans.size()); + + SpanData span = spans.get(0); + assertEquals( + 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(handlerFailureMessage); + }); + HttpTracingRequestInitializer tracingInitializer = + new HttpTracingRequestInitializer(delegateInitializer, tracer, CLIENT_ROOT_URL); + + HttpTransport transport = createTransport(500, serverErrorMessage, null); + HttpRequest request = buildGetRequest(transport, tracingInitializer, BASE_URL); + request.setThrowExceptionOnExecuteError(false); + + IOException thrown = assertThrows(IOException.class, request::execute); + assertEquals(handlerFailureMessage, thrown.getMessage()); + + List spans = spanExporter.getFinishedSpanItems(); + assertEquals(1, spans.size()); + + SpanData span = spans.get(0); + assertEquals( + 500, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_STATUS_CODE)); + assertEquals( + IOException.class.getName(), + span.getAttributes().get(BigQueryTelemetryTracer.EXCEPTION_TYPE)); + assertEquals( + IOException.class.getSimpleName(), + span.getAttributes().get(BigQueryTelemetryTracer.ERROR_TYPE)); + assertEquals( + handlerFailureMessage, span.getAttributes().get(BigQueryTelemetryTracer.STATUS_MESSAGE)); + assertTrue(span.getEvents().stream().anyMatch(event -> "exception".equals(event.getName()))); + verifyGeneralSpanData(span, REQUEST_METHOD_GET); + } + + @Test + public void testDelegateInitializerIsCalledOnSuccessResponse() throws IOException { + HttpRequestInitializer delegateInitializer = mock(HttpRequestInitializer.class); + HttpTracingRequestInitializer tracingInitializer = + new HttpTracingRequestInitializer(delegateInitializer, tracer, CLIENT_ROOT_URL); + + HttpTransport transport = createTransport(200, null, null); + HttpRequest request = buildGetRequest(transport, tracingInitializer, BASE_URL); + + 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, CLIENT_ROOT_URL); + + HttpTransport transport = createTransport(404, "Not Found", null); + HttpRequest request = buildGetRequest(transport, tracingInitializer, BASE_URL); + + try { + HttpResponse response = request.execute(); + response.disconnect(); + } catch (Exception e) { + // Expected - 404 might throw + } + + 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)); + } + + @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() { + @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)); + } +}