Add support for request and response compression

Author: Chaho12

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

@@ -201,8 +201,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);
                 }
@@ -262,7 +261,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();
@@ -287,26 +286,26 @@ 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());
+        log.debug("For Request [%s] got Response [%s]", request.getUri(), response.getDecompressedBody());
 
         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(response.getDecompressedBody(), 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]", response.getDecompressedBody(), 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]", response.getDecompressedBody(), 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);
         }
     }
 }