feat(bigquery): add response body tracking for compressed response

Author: ldetmer

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

@@ -0,0 +1,112 @@
+/*
+ * 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 static com.google.cloud.bigquery.telemetry.HttpTracingRequestInitializer.HTTP_RESPONSE_BODY_SIZE;
+
+import com.google.api.client.http.LowLevelHttpResponse;
+import com.google.common.io.CountingInputStream;
+import io.opentelemetry.api.trace.Span;
+import java.io.FilterInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+
+/**
+ * Class that wraps LowLevelHttpResponse to be able to inject an InputStream wrapper that allows us
+ * to track response size in trace spans.
+ */
+class HttpTracingLowLevelHttpResponse extends LowLevelHttpResponse {
+  private final LowLevelHttpResponse delegate;
+  private InputStream countingInputStream;
+  private final Span span;
+
+  public HttpTracingLowLevelHttpResponse(LowLevelHttpResponse delegate, Span span) {
+    this.delegate = delegate;
+    this.span = span;
+  }
+
+  @Override
+  public InputStream getContent() throws IOException {
+    if (countingInputStream == null) {
+      if (span != null && span.getSpanContext().isValid()) {
+        countingInputStream = getWrappedInputStream(span, delegate.getContent());
+      }
+    }
+    return countingInputStream;
+  }
+
+  @Override
+  public String getContentEncoding() throws IOException {
+    return delegate.getContentEncoding();
+  }
+
+  @Override
+  public long getContentLength() throws IOException {
+    return delegate.getContentLength();
+  }
+
+  @Override
+  public String getContentType() throws IOException {
+    return delegate.getContentType();
+  }
+
+  @Override
+  public String getStatusLine() throws IOException {
+    return delegate.getStatusLine();
+  }
+
+  @Override
+  public int getStatusCode() throws IOException {
+    return delegate.getStatusCode();
+  }
+
+  @Override
+  public String getReasonPhrase() throws IOException {
+    return delegate.getReasonPhrase();
+  }
+
+  @Override
+  public int getHeaderCount() throws IOException {
+    return delegate.getHeaderCount();
+  }
+
+  @Override
+  public String getHeaderName(int index) throws IOException {
+    return delegate.getHeaderName(index);
+  }
+
+  @Override
+  public String getHeaderValue(int index) throws IOException {
+    return delegate.getHeaderValue(index);
+  }
+
+  @Override
+  public void disconnect() throws IOException {
+    delegate.disconnect();
+  }
+
+  private InputStream getWrappedInputStream(Span span, InputStream content) {
+    CountingInputStream counter = new CountingInputStream(content);
+    return new FilterInputStream(counter) {
+      @Override
+      public void close() throws IOException {
+        super.close();
+        span.setAttribute(HTTP_RESPONSE_BODY_SIZE, counter.getCount());
+      }
+    };
+  }
+}

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

@@ -20,15 +20,11 @@
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.io.CountingInputStream;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.Tracer;
-import java.io.FilterInputStream;
 import java.io.IOException;
-import java.io.InputStream;
 import java.lang.reflect.Field;
-import org.jspecify.annotations.NonNull;
 
 /**
  * HttpRequestInitializer that wraps a delegate initializer, intercepts all HTTP requests, adds
@@ -151,36 +147,23 @@ static void addResponseBodySizeToSpan(HttpResponse response, Span span) {
   }
 
   /**
-   * Wraps the response's input stream with a CountingInputStream to track bytes read. This handles
-   * compressed transfer encoding where Content-Length is not available.
+   * Wraps the LowLevelHttpResponse with a HttpTracingLowLevelHttpResponse to give us access to get
+   * response body size for compressed responses.
    */
   private static void getResponseBodySizeForCompressedResponse(HttpResponse response, Span span) {
     try {
-      InputStream content = response.getContent();
-      if (content == null) {
-        return;
-      }
-
-      InputStream wrappedStream = getWrappedInputStream(span, content);
-      Field contentField = HttpResponse.class.getDeclaredField("content");
-      contentField.setAccessible(true);
-      contentField.set(response, wrappedStream);
-    } catch (Exception e) {
+      Field lowlevelField = HttpResponse.class.getDeclaredField("response");
+      lowlevelField.setAccessible(true);
+      LowLevelHttpResponse lowLevelHttpResponse =
+          (LowLevelHttpResponse) lowlevelField.get(response);
+      HttpTracingLowLevelHttpResponse wrappedResponse =
+          new HttpTracingLowLevelHttpResponse(lowLevelHttpResponse, span);
+      lowlevelField.set(response, wrappedResponse);
+    } catch (NoSuchFieldException | SecurityException | IllegalAccessException e) {
       // Ignore - stream wrapping failed, we will not track response size
     }
   }
 
-  private static @NonNull InputStream getWrappedInputStream(Span span, InputStream content) {
-    CountingInputStream counter = new CountingInputStream(content);
-    return new FilterInputStream(counter) {
-      @Override
-      public void close() throws IOException {
-        super.close();
-        span.setAttribute(HTTP_RESPONSE_BODY_SIZE, counter.getCount());
-      }
-    };
-  }
-
   /** Removes credentials from URL. */
   private static String getSanitizedUrl(HttpRequest request) {
     GenericUrl clone = request.getUrl().clone();

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

@@ -53,13 +53,11 @@
 import io.opentelemetry.sdk.trace.SdkTracerProvider;
 import io.opentelemetry.sdk.trace.data.SpanData;
 import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
-
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.zip.GZIPOutputStream;
-
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -475,7 +473,7 @@ public LowLevelHttpResponse execute() {
     SpanData span = spans.get(0);
 
     assertEquals(
-        (long) responseBody.length(),
+        compressedData.length,
         span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_BODY_SIZE));
   }
 }