If isOpenTelemetryTracingEnabled == true handles span creation, scope management, and + * cleanup, otherwise only executes the operation. + */ + private T executeWithSpan( + String spanName, + String service, + String method, + Map options, + SpanOperation operation) + throws IOException { + + if (!this.options.isOpenTelemetryTracingEnabled() + || this.options.getOpenTelemetryTracer() == null) { + return operation.execute(null); + } + + SpanBuilder builder = + this.options + .getOpenTelemetryTracer() + .spanBuilder(spanName) + .setSpanKind(SpanKind.CLIENT) + .setAttribute("bq.rpc.service", service) + .setAttribute("bq.rpc.method", method) + .setAttribute("bq.rpc.system", "http"); + + if (options != null) { + builder.setAllAttributes(otelAttributesFromOptions(options)); } - TestIamPermissionsResponse bqTestResponse = bqTestRequest.execute(); - if (testIamPermissions != null) { - testIamPermissions.end(); + Span span = builder.startSpan(); + + try (Scope scope = span.makeCurrent()) { + return operation.execute(span); + } finally { + span.end(); } - return bqTestResponse; + } + + /** Overloaded method for executeWithSpan with null options parameter. */ + private T executeWithSpan( + String spanName, String service, String method, SpanOperation operation) + throws IOException { + return executeWithSpan(spanName, service, method, null, operation); + } + + /** Functional interface for span operations that can throw IOException. */ + @FunctionalInterface + private interface SpanOperation { + T execute(Span span) throws IOException; } private static Attributes otelAttributesFromOptions(Map options) { diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpcTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpcTest.java index 3968cd05e22f..489671461c72 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpcTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpcTest.java @@ -16,11 +16,30 @@ package com.google.cloud.bigquery.spi.v2; import static com.google.common.truth.Truth.assertThat; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +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.json.Json; +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.api.services.bigquery.model.Dataset; import com.google.api.services.bigquery.model.DatasetList; import com.google.api.services.bigquery.model.DatasetReference; +import com.google.cloud.NoCredentials; +import com.google.cloud.bigquery.BigQueryOptions; +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.util.Collections; +import java.util.HashMap; +import java.util.List; import org.junit.jupiter.api.Test; public class HttpBigQueryRpcTest { @@ -46,4 +65,96 @@ public void testListToDataset() { assertThat(dataset.getLabels()).containsExactly("foo", "bar"); assertThat(dataset.getLocation()).isEqualTo("test-region-1"); } + + @Test + public void testOpenTelemetrySpanCreation() throws Exception { + InMemorySpanExporter 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"); + + HttpTransport mockTransport = + new MockHttpTransport() { + @Override + public LowLevelHttpRequest buildRequest(String method, String url) { + return new MockLowLevelHttpRequest() { + @Override + public LowLevelHttpResponse execute() { + MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); + response.setStatusCode(200); + response.setContentType(Json.MEDIA_TYPE); + // Return a minimal valid dataset JSON + response.setContent( + "{" + + "\"kind\":\"bigquery#dataset\"," + + "\"id\":\"test-project:test-dataset\"," + + "\"datasetReference\":{" + + " \"projectId\":\"test-project\"," + + " \"datasetId\":\"test-dataset\"" + + "}" + + "}"); + return response; + } + }; + } + }; + + // Create BigQueryOptions with OpenTelemetry enabled and custom transport + BigQueryOptions options = + BigQueryOptions.newBuilder() + .setProjectId("test-project") + .setCredentials(NoCredentials.getInstance()) + .setEnableOpenTelemetryTracing(true) + .setOpenTelemetryTracer(tracer) + .setTransportOptions( + BigQueryOptions.getDefaultHttpTransportOptions().toBuilder() + .setHttpTransportFactory(() -> mockTransport) + .build()) + .build(); + HttpBigQueryRpc rpc = new HttpBigQueryRpc(options); + + Dataset result = rpc.getDataset("test-project", "test-dataset", new HashMap<>()); + + // Verify the RPC call succeeded + assertNotNull(result); + assertEquals("test-project", result.getDatasetReference().getProjectId()); + assertEquals("test-dataset", result.getDatasetReference().getDatasetId()); + + // Verify that spans were created + List spans = spanExporter.getFinishedSpanItems(); + assertThat(spans).isNotEmpty(); + + // Find the RPC span (should be the one we created in executeWithSpan) + SpanData rpcSpan = + spans.stream() + .filter(span -> span.getName().contains("BigQueryRpc.getDataset")) + .findFirst() + .orElse(null); + + assertNotNull(rpcSpan, "Expected to find a span for BigQueryRpc.getDataset"); + + // Verify span attributes are set correctly + assertEquals( + "DatasetService", + rpcSpan + .getAttributes() + .asMap() + .get(io.opentelemetry.api.common.AttributeKey.stringKey("bq.rpc.service"))); + assertEquals( + "GetDataset", + rpcSpan + .getAttributes() + .asMap() + .get(io.opentelemetry.api.common.AttributeKey.stringKey("bq.rpc.method"))); + assertEquals( + "http", + rpcSpan + .getAttributes() + .asMap() + .get(io.opentelemetry.api.common.AttributeKey.stringKey("bq.rpc.system"))); + } }