feat(bigquery): add HTTP response attribute tracing (#12109)

This adds generic HTTP response attribute capture via open telemetry to
BigQuery. Error attributes are not included in this PR.

[example success
POST](https://screenshot.googleplex.com/9BARUi5tNbLBrNL.png)
[example failed
GET](https://screenshot.googleplex.com/4GUwJFKoVkEVvKD.png)

Author: ldetmer

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/HttpTracingRequestInitializer.java

@@ -77,6 +77,25 @@ public void initialize(HttpRequest request) throws IOException {
     String host = request.getUrl().getHost();
     int port = request.getUrl().getPort();
     addInitialHttpAttributesToSpan(span, host, port);
+
+    HttpResponseInterceptor originalInterceptor = request.getResponseInterceptor();
+    request.setResponseInterceptor(
+        response -> {
+          addCommonResponseAttributesToSpan(request, response, span);
+          if (originalInterceptor != null) {
+            originalInterceptor.interceptResponse(response);
+          }
+        });
+
+    HttpUnsuccessfulResponseHandler originalHandler = request.getUnsuccessfulResponseHandler();
+    request.setUnsuccessfulResponseHandler(
+        (request1, response, supportsRetry) -> {
+          addCommonResponseAttributesToSpan(request1, response, span);
+          if (originalHandler != null) {
+            return originalHandler.handleResponse(request1, response, supportsRetry);
+          }
+          return false;
+        });
   }
 
   /** Add initial HTTP attributes to the existing active span */
@@ -89,4 +108,32 @@ private void addInitialHttpAttributesToSpan(Span span, String host, Integer port
     }
     // TODO add full sanitized url, url domain, request method
   }
+
+  private static void addCommonResponseAttributesToSpan(
+      HttpRequest request, HttpResponse response, Span span) {
+    // We add request body size and update request method after we receive response as
+    // the data is not always available until after the http request execution
+    addRequestBodySizeToSpan(request, span);
+    span.setAttribute(HTTP_REQUEST_METHOD, request.getRequestMethod());
+    addResponseBodySizeToSpan(response, span);
+    span.setAttribute(HTTP_RESPONSE_STATUS_CODE, response.getStatusCode());
+  }
+
+  static void addRequestBodySizeToSpan(HttpRequest request, Span span) {
+    try {
+      if (request.getContent() != null && request.getEncoding() == null) {
+        span.setAttribute(HTTP_REQUEST_BODY_SIZE, request.getContent().getLength());
+      }
+    } catch (IOException e) {
+      // Ignore - body size not available
+    }
+  }
+
+  static void addResponseBodySizeToSpan(HttpResponse response, Span span) {
+    Long contentLength = response.getHeaders().getContentLength();
+    if (contentLength != null && contentLength > 0) {
+      span.setAttribute(HTTP_RESPONSE_BODY_SIZE, contentLength);
+    }
+    // TODO handle chunked responses
+  }
 }

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/telemetry/HttpTracingRequestInitializerTest.java

@@ -19,16 +19,24 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
+import com.google.api.client.http.ByteArrayContent;
 import com.google.api.client.http.GenericUrl;
+import com.google.api.client.http.HttpContent;
+import com.google.api.client.http.HttpEncoding;
 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.HttpResponseException;
+import com.google.api.client.http.HttpResponseInterceptor;
 import com.google.api.client.http.HttpTransport;
+import com.google.api.client.http.HttpUnsuccessfulResponseHandler;
 import com.google.api.client.http.LowLevelHttpRequest;
 import com.google.api.client.http.LowLevelHttpResponse;
 import com.google.api.client.testing.http.MockHttpTransport;
@@ -77,22 +85,23 @@ public void setUp() {
   }
 
   @Test
-  public void testRequestAttributesAreSetIfSpanExists() throws IOException {
-    HttpTransport transport = createTransport();
-    HttpRequest request = buildGetRequest(transport, initializer, BASE_URL);
+  public void testAllAttributesAreSetOnSuccessfulCall() throws IOException {
+    HttpResponseInterceptor originalInterceptor = mock(HttpResponseInterceptor.class);
+    HttpRequestInitializer delegateInitializer =
+        request -> request.setResponseInterceptor(originalInterceptor);
+    HttpTracingRequestInitializer tracingInitializer =
+        new HttpTracingRequestInitializer(delegateInitializer, tracer);
+    String responseContent = "{\"status\": \"ok\"}";
+    HttpTransport transport = createTransport(200, responseContent);
+    HttpContent requestContent =
+        ByteArrayContent.fromString("application/json", "{\"test\": \"data\"}");
+    HttpRequest request = buildPostRequest(transport, tracingInitializer, BASE_URL, requestContent);
 
     HttpResponse response = request.execute();
     response.disconnect();
 
-    // End the span before verifying exported spans, so it appears in the exporter
-    spanScope.close();
-    parentSpan.end();
-
-    List<SpanData> spans = spanExporter.getFinishedSpanItems();
-    assertEquals(1, spans.size());
-
-    SpanData span = spans.get(0);
-    verifyGeneralSpanData(span);
+    verify(originalInterceptor, times(1)).interceptResponse(any(HttpResponse.class));
+    closeAndVerifySpanData(200, "POST", 16, requestContent.getLength());
   }
 
   @Test
@@ -103,14 +112,11 @@ public void testExistingParentAttributesAreNotAffectedByRequestAttributes() thro
 
     HttpResponse response = request.execute();
     response.disconnect();
-
-    // End the span before verifying exported spans, so it appears in the exporter
     spanScope.close();
     parentSpan.end();
 
     List<SpanData> spans = spanExporter.getFinishedSpanItems();
     assertEquals(1, spans.size());
-
     SpanData span = spans.get(0);
     assertEquals("value", span.getAttributes().get(AttributeKey.stringKey("parent_attribute")));
   }
@@ -152,14 +158,160 @@ public void testDelegateInitializerIsCalled() throws IOException {
     verify(delegateInitializer, times(1)).initialize(any(HttpRequest.class));
   }
 
+  @Test
+  public void testUnsuccessfulResponseHandlerSetsAttributesAndCallsOriginal() throws IOException {
+    HttpUnsuccessfulResponseHandler originalHandler = mock(HttpUnsuccessfulResponseHandler.class);
+    when(originalHandler.handleResponse(
+            any(HttpRequest.class), any(HttpResponse.class), anyBoolean()))
+        .thenReturn(false);
+    HttpRequestInitializer delegateInitializer =
+        request -> request.setUnsuccessfulResponseHandler(originalHandler);
+    HttpTracingRequestInitializer tracingInitializer =
+        new HttpTracingRequestInitializer(delegateInitializer, tracer);
+    String responseContent = "{\"error\": \"not found\"}";
+    HttpTransport transport = createTransport(404, responseContent);
+    HttpContent requestContent =
+        ByteArrayContent.fromString("application/json", "{\"test\": \"data\"}");
+    HttpRequest request = buildPostRequest(transport, tracingInitializer, BASE_URL, requestContent);
+
+    try {
+      request.execute();
+    } catch (HttpResponseException e) {
+      // Expected
+    }
+
+    verify(originalHandler, times(1))
+        .handleResponse(any(HttpRequest.class), any(HttpResponse.class), anyBoolean());
+    closeAndVerifySpanData(404, "POST", 16, responseContent.length());
+  }
+
+  @Test
+  public void testUnsuccessfulResponseHandlerSetsErrorIfNoOriginal() throws IOException {
+    HttpTracingRequestInitializer tracingInitializer =
+        new HttpTracingRequestInitializer(null, tracer);
+    HttpTransport transport = createTransport(401);
+    HttpRequest request = buildGetRequest(transport, tracingInitializer, BASE_URL);
+
+    try {
+      request.execute();
+    } catch (HttpResponseException e) {
+      // Expected
+    }
+
+    closeAndVerifySpanData(401, "GET", -1, -1);
+  }
+
+  @Test
+  public void testAddRequestBodySizeToSpan_ExceptionHandled() throws IOException {
+    HttpContent content = mock(HttpContent.class);
+    when(content.getLength()).thenThrow(new IOException("test"));
+    HttpTransport transport = createTransport();
+    HttpRequest request = buildPostRequest(transport, null, BASE_URL, content);
+
+    HttpTracingRequestInitializer.addRequestBodySizeToSpan(request, parentSpan);
+
+    spanScope.close();
+    parentSpan.end();
+
+    List<SpanData> spans = spanExporter.getFinishedSpanItems();
+    assertEquals(1, spans.size());
+    SpanData span = spans.get(0);
+    assertNull(span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_BODY_SIZE));
+  }
+
+  @Test
+  public void testAddRequestBodySizeToSpan_NullBody() throws IOException {
+    HttpTransport transport = createTransport();
+    HttpRequest request = buildPostRequest(transport, null, BASE_URL, null);
+
+    HttpTracingRequestInitializer.addRequestBodySizeToSpan(request, parentSpan);
+
+    spanScope.close();
+    parentSpan.end();
+
+    List<SpanData> spans = spanExporter.getFinishedSpanItems();
+    assertEquals(1, spans.size());
+    SpanData span = spans.get(0);
+    assertNull(span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_BODY_SIZE));
+  }
+
+  @Test
+  public void testAddRequestBodySizeToSpan_WithEncoding() throws IOException {
+    HttpTransport transport = createTransport();
+    HttpContent content = ByteArrayContent.fromString("application/json", "{\"test\": \"data\"}");
+    HttpRequest request = buildPostRequest(transport, null, BASE_URL, content);
+    request.setEncoding(mock(HttpEncoding.class));
+
+    HttpTracingRequestInitializer.addRequestBodySizeToSpan(request, parentSpan);
+
+    spanScope.close();
+    parentSpan.end();
+
+    List<SpanData> spans = spanExporter.getFinishedSpanItems();
+    assertEquals(1, spans.size());
+    SpanData span = spans.get(0);
+    assertNull(span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_BODY_SIZE));
+  }
+
+  @Test
+  public void testAddRequestBodySizeToSpan() throws IOException {
+    HttpTransport transport = createTransport();
+    HttpContent content = ByteArrayContent.fromString("application/json", "{\"test\": \"data\"}");
+    HttpRequest request = buildPostRequest(transport, null, BASE_URL, content);
+
+    HttpTracingRequestInitializer.addRequestBodySizeToSpan(request, parentSpan);
+
+    spanScope.close();
+    parentSpan.end();
+
+    List<SpanData> spans = spanExporter.getFinishedSpanItems();
+    assertEquals(1, spans.size());
+    SpanData span = spans.get(0);
+    assertEquals(
+        content.getLength(),
+        span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_BODY_SIZE));
+  }
+
+  @Test
+  public void testAddResponseBodySizeToSpan_NullLength() throws IOException {
+    HttpTransport transport =
+        createTransport(200, null); // Null response content means null Content-Length header
+    HttpRequest request = buildGetRequest(transport, null, BASE_URL);
+    HttpResponse response = request.execute();
+
+    HttpTracingRequestInitializer.addResponseBodySizeToSpan(response, parentSpan);
+
+    spanScope.close();
+    parentSpan.end();
+
+    List<SpanData> spans = spanExporter.getFinishedSpanItems();
+    assertEquals(1, spans.size());
+    SpanData span = spans.get(0);
+    assertNull(span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_BODY_SIZE));
+  }
+
   private static HttpTransport createTransport() {
+    return createTransport(200, null);
+  }
+
+  private static HttpTransport createTransport(int statusCode) {
+    return createTransport(statusCode, null);
+  }
+
+  private static HttpTransport createTransport(int statusCode, String responseContent) {
     return new MockHttpTransport() {
       @Override
       public LowLevelHttpRequest buildRequest(String method, String url) {
         return new MockLowLevelHttpRequest() {
           @Override
           public LowLevelHttpResponse execute() {
-            return new MockLowLevelHttpResponse();
+            MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
+            response.setStatusCode(statusCode);
+            if (responseContent != null) {
+              response.setContent(responseContent);
+              response.addHeader("Content-Length", String.valueOf(responseContent.length()));
+            }
+            return response;
           }
         };
       }
@@ -173,7 +325,25 @@ private static HttpRequest buildGetRequest(
     return requestFactory.buildGetRequest(new GenericUrl(url));
   }
 
-  private void verifyGeneralSpanData(SpanData span) {
+  private static HttpRequest buildPostRequest(
+      HttpTransport transport,
+      HttpRequestInitializer requestInitializer,
+      String url,
+      HttpContent content)
+      throws IOException {
+    HttpRequestFactory requestFactory = transport.createRequestFactory(requestInitializer);
+    return requestFactory.buildPostRequest(new GenericUrl(url), content);
+  }
+
+  private void closeAndVerifySpanData(
+      long responseCode, String method, long requestBodySize, long responseBodySize) {
+    // close span so it is available for verification
+    spanScope.close();
+    parentSpan.end();
+    List<SpanData> spans = spanExporter.getFinishedSpanItems();
+
+    assertEquals(1, spans.size());
+    SpanData span = spans.get(0);
     assertEquals(SPAN_NAME, span.getName());
     assertEquals(BIGQUERY_DOMAIN, span.getAttributes().get(BigQueryTelemetryTracer.SERVER_ADDRESS));
     assertEquals(443, span.getAttributes().get(BigQueryTelemetryTracer.SERVER_PORT));
@@ -192,5 +362,24 @@ private void verifyGeneralSpanData(SpanData span) {
     assertEquals(
         HttpTracingRequestInitializer.HTTP_RPC_SYSTEM_NAME,
         span.getAttributes().get(BigQueryTelemetryTracer.RPC_SYSTEM_NAME));
+    assertEquals(
+        responseCode,
+        span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_STATUS_CODE));
+    assertEquals(
+        method, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_METHOD));
+    if (requestBodySize >= 0) {
+      assertEquals(
+          requestBodySize,
+          span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_BODY_SIZE));
+    } else {
+      assertNull(span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_BODY_SIZE));
+    }
+    if (responseBodySize >= 0) {
+      assertEquals(
+          responseBodySize,
+          span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_BODY_SIZE));
+    } else {
+      assertNull(span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_BODY_SIZE));
+    }
   }
 }