Add support for request and response compression

Author: Jaeho Yoo

gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java

@@ -166,8 +166,7 @@ private void performRequest(
         for (String name : list(servletRequest.getHeaderNames())) {
             for (String value : list(servletRequest.getHeaders(name))) {
                 // TODO: decide what should and shouldn't be forwarded
-                if (!name.equalsIgnoreCase("Accept-Encoding")
-                        && !name.equalsIgnoreCase("Host")
+                if (!name.equalsIgnoreCase("Host")
                         && (addXForwardedHeaders || !name.startsWith("X-Forwarded"))) {
                     requestBuilder.addHeader(name, value);
                 }
@@ -269,26 +268,27 @@ private static WebApplicationException badRequest(String message)
     private ProxyResponse recordBackendForQueryId(Request request, ProxyResponse response, Optional<String> username,
             RoutingDestination routingDestination)
     {
-        log.debug("For Request [%s] got Response [%s]", request.getUri(), response.body());
+        String body = response.decompressedBody();
+        log.debug("For Request [%s] got Response [%s]", request.getUri(), body);
 
         QueryHistoryManager.QueryDetail queryDetail = getQueryDetailsFromRequest(request, username);
 
         log.debug("Extracting proxy destination : [%s] for request : [%s]", queryDetail.getBackendUrl(), request.getUri());
 
         if (response.statusCode() == OK.getStatusCode()) {
             try {
-                HashMap<String, String> results = OBJECT_MAPPER.readValue(response.body(), HashMap.class);
+                HashMap<String, String> results = OBJECT_MAPPER.readValue(body, HashMap.class);
                 queryDetail.setQueryId(results.get("id"));
                 routingManager.setBackendForQueryId(queryDetail.getQueryId(), queryDetail.getBackendUrl());
                 routingManager.setRoutingGroupForQueryId(queryDetail.getQueryId(), routingDestination.routingGroup());
                 log.debug("QueryId [%s] mapped with proxy [%s]", queryDetail.getQueryId(), queryDetail.getBackendUrl());
             }
             catch (IOException e) {
-                log.error("Failed to get QueryId from response [%s] , Status code [%s]", response.body(), response.statusCode());
+                log.error("Failed to get QueryId from response [%s] , Status code [%s]", body, response.statusCode());
             }
         }
         else {
-            log.error("Non OK HTTP Status code with response [%s] , Status code [%s], user: [%s]", response.body(), response.statusCode(), username.orElse(null));
+            log.error("Non OK HTTP Status code with response [%s] , Status code [%s], user: [%s]", body, response.statusCode(), username.orElse(null));
         }
         queryDetail.setRoutingGroup(routingDestination.routingGroup());
         queryDetail.setExternalUrl(routingDestination.externalUrl());

gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyResponseHandler.java

@@ -22,9 +22,12 @@
 import io.trino.gateway.ha.config.ProxyResponseConfiguration;
 import io.trino.gateway.proxyserver.ProxyResponseHandler.ProxyResponse;
 
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
-import java.nio.charset.StandardCharsets;
+import java.io.InputStream;
+import java.util.zip.GZIPInputStream;
 
+import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Objects.requireNonNull;
 
 public class ProxyResponseHandler
@@ -47,7 +50,9 @@ public ProxyResponse handleException(Request request, Exception exception)
     public ProxyResponse handle(Request request, Response response)
     {
         try {
-            return new ProxyResponse(response.getStatusCode(), response.getHeaders(), new String(response.getInputStream().readNBytes((int) responseSize.toBytes()), StandardCharsets.UTF_8));
+            // Store raw bytes to preserve compression
+            byte[] responseBodyBytes = response.getInputStream().readNBytes((int) responseSize.toBytes());
+            return new ProxyResponse(response.getStatusCode(), response.getHeaders(), responseBodyBytes);
         }
         catch (IOException e) {
             throw new ProxyException("Failed reading response from remote Trino server", e);
@@ -57,11 +62,36 @@ public ProxyResponse handle(Request request, Response response)
     public record ProxyResponse(
             int statusCode,
             ListMultimap<HeaderName, String> headers,
-            String body)
+            byte[] body)
     {
         public ProxyResponse
         {
             requireNonNull(headers, "headers is null");
+            requireNonNull(body, "body is null");
+        }
+
+        /**
+         * Get the response body as a decompressed string for JSON parsing and logging.
+         * Only call this when you need to parse the content, not when passing through
+         * to clients.
+         */
+        public String decompressedBody()
+        {
+            // Check if the response is gzip-compressed
+            String contentEncoding = headers.get(HeaderName.of("Content-Encoding")).stream().findFirst().orElse(null);
+
+            if ("gzip".equalsIgnoreCase(contentEncoding)) {
+                try (InputStream inputStream = new GZIPInputStream(new ByteArrayInputStream(body))) {
+                    return new String(inputStream.readAllBytes(), UTF_8);
+                }
+                catch (IOException e) {
+                    // If decompression fails, return the body as UTF-8 string
+                    return new String(body, UTF_8);
+                }
+            }
+
+            // Not compressed, convert bytes to string
+            return new String(body, UTF_8);
         }
     }
 }

gateway-ha/src/test/java/io/trino/gateway/proxyserver/TestProxyRequestHandler.java

@@ -81,6 +81,14 @@ public MockResponse dispatch(RecordedRequest request)
                             .setBody("{\"starting\": false}");
                 }
 
+                if (request.getPath().equals(healthCheckEndpoint + "?test-compression")) {
+                    // Return the Accept-Encoding header value for compression testing
+                    String acceptEncoding = request.getHeader("Accept-Encoding");
+                    return new MockResponse().setResponseCode(200)
+                            .setHeader(CONTENT_TYPE, JSON_UTF_8)
+                            .setBody(acceptEncoding != null ? acceptEncoding : "null");
+                }
+
                 if (request.getMethod().equals("PUT") && request.getPath().equals(customPutEndpoint)) {
                     return new MockResponse().setResponseCode(200)
                             .setHeader(CONTENT_TYPE, JSON_UTF_8)
@@ -160,4 +168,49 @@ void testGetQueryDetailsFromRequest()
         assertThat(queryDetail.getSource()).isEqualTo("trino-cli");
         assertThat(queryDetail.getBackendUrl()).isEqualTo("http://localhost:" + routerPort);
     }
+
+    @Test
+    void testAcceptEncodingHeaderForwarding()
+            throws Exception
+    {
+        // Test that Accept-Encoding header is properly forwarded to backends
+        String url = "http://localhost:" + routerPort + healthCheckEndpoint + "?test-compression";
+        String expectedAcceptEncoding = "gzip, deflate, br";
+
+        Request request = new Request.Builder()
+                .url(url)
+                .get()
+                .addHeader("Accept-Encoding", expectedAcceptEncoding)
+                .build();
+
+        try (Response response = httpClient.newCall(request).execute()) {
+            assertThat(response.code()).isEqualTo(200);
+            assertThat(response.body()).isNotNull();
+
+            // The mock backend returns the Accept-Encoding header value in the response body
+            assertThat(response.body().string()).isEqualTo(expectedAcceptEncoding);
+        }
+    }
+
+    @Test
+    void testDefaultAcceptEncodingHeaderForwarding()
+            throws Exception
+    {
+        // Test that requests without explicit Accept-Encoding header work correctly
+        // Note: OkHttp automatically adds "Accept-Encoding: gzip" when none is specified
+        String url = "http://localhost:" + routerPort + healthCheckEndpoint + "?test-compression";
+
+        Request request = new Request.Builder()
+                .url(url)
+                .get()
+                .build(); // No explicit Accept-Encoding header
+
+        try (Response response = httpClient.newCall(request).execute()) {
+            assertThat(response.code()).isEqualTo(200);
+            assertThat(response.body()).isNotNull();
+
+            // OkHttp automatically adds "Accept-Encoding: gzip" when none is specified
+            assertThat(response.body().string()).isEqualTo("gzip");
+        }
+    }
 }