Add support for request and response compression

Author: Jaeho Yoo

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

@@ -200,8 +200,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);
                 }
@@ -261,7 +260,7 @@ else if (servletRequest.getCookies() != null) {
 
     private Response buildResponse(ProxyResponse response, ImmutableList<NewCookie> cookie)
     {
-        Response.ResponseBuilder builder = Response.status(response.statusCode()).entity(response.body());
+        Response.ResponseBuilder builder = Response.status(response.statusCode()).entity(response.getRawBody());
         response.headers().forEach((headerName, value) -> builder.header(headerName.toString(), value));
         cookie.forEach(builder::cookie);
         return builder.build();
@@ -286,26 +285,27 @@ private FluentFuture<ProxyResponse> executeHttp(Request request)
     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.getDecompressedBody();
+        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,8 +22,11 @@
 import io.trino.gateway.ha.config.ProxyResponseConfiguration;
 import io.trino.gateway.proxyserver.ProxyResponseHandler.ProxyResponse;
 
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
+import java.util.zip.GZIPInputStream;
 
 import static java.util.Objects.requireNonNull;
 
@@ -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,53 @@ 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 raw bytes for sending to clients (preserves
+         * compression)
+         */
+        public byte[] getRawBody()
+        {
+            return body;
+        }
+
+        /**
+         * 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 getDecompressedBody()
+        {
+            // Check if the response is gzip-compressed
+            String contentEncoding = null;
+            for (HeaderName headerName : headers.keySet()) {
+                if (headerName.toString().equalsIgnoreCase("Content-Encoding")) {
+                    contentEncoding = headers.get(headerName).iterator().next();
+                    break;
+                }
+            }
+
+            if ("gzip".equalsIgnoreCase(contentEncoding)) {
+                try {
+                    try (InputStream inputStream = new GZIPInputStream(new ByteArrayInputStream(body))) {
+                        return new String(inputStream.readAllBytes(), StandardCharsets.UTF_8);
+                    }
+                }
+                catch (IOException e) {
+                    // If decompression fails, return the body as UTF-8 string
+                    return new String(body, StandardCharsets.UTF_8);
+                }
+            }
+
+            // Not compressed, convert bytes to string
+            return new String(body, StandardCharsets.UTF_8);
         }
     }
 }

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

@@ -78,6 +78,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)
@@ -157,4 +165,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");
+        }
+    }
 }