cleaned up code + added suggestions from gemini

Author: ldetmer

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

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<Long> 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,29 +77,35 @@ 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);
           }
         });
 
     // Wrap the existing unsuccessful response handler
     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);
+    }
   }
 }

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<String> 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
   }
 }