Update ProxyResponseHandler.java

Author: Chaho12

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);
         }
     }
 }