Skip to content

Commit fb2dc46

Browse files
author
Jaeho Yoo
committed
Add support for request and response compression
1 parent 3b49d6d commit fb2dc46

File tree

3 files changed

+109
-9
lines changed

3 files changed

+109
-9
lines changed

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -167,8 +167,7 @@ private void performRequest(
167167
for (String name : list(servletRequest.getHeaderNames())) {
168168
for (String value : list(servletRequest.getHeaders(name))) {
169169
// TODO: decide what should and shouldn't be forwarded
170-
if (!name.equalsIgnoreCase("Accept-Encoding")
171-
&& !name.equalsIgnoreCase("Host")
170+
if (!name.equalsIgnoreCase("Host")
172171
&& (addXForwardedHeaders || !name.startsWith("X-Forwarded"))) {
173172
requestBuilder.addHeader(name, value);
174173
}
@@ -231,7 +230,7 @@ private static String getRemoteTarget(URI remoteUri)
231230

232231
private Response buildResponse(ProxyResponse response, ImmutableList<NewCookie> cookie)
233232
{
234-
Response.ResponseBuilder builder = Response.status(response.statusCode()).entity(response.body());
233+
Response.ResponseBuilder builder = Response.status(response.statusCode()).entity(response.getRawBody());
235234
response.headers().forEach((headerName, value) -> builder.header(headerName.toString(), value));
236235
cookie.forEach(builder::cookie);
237236
return builder.build();
@@ -270,26 +269,27 @@ private static WebApplicationException badRequest(String message)
270269
private ProxyResponse recordBackendForQueryId(Request request, ProxyResponse response, Optional<String> username,
271270
RoutingDestination routingDestination)
272271
{
273-
log.debug("For Request [%s] got Response [%s]", request.getUri(), response.body());
272+
String body = response.getDecompressedBody();
273+
log.debug("For Request [%s] got Response [%s]", request.getUri(), body);
274274

275275
QueryHistoryManager.QueryDetail queryDetail = getQueryDetailsFromRequest(request, username);
276276

277277
log.debug("Extracting proxy destination : [%s] for request : [%s]", queryDetail.getBackendUrl(), request.getUri());
278278

279279
if (response.statusCode() == OK.getStatusCode()) {
280280
try {
281-
HashMap<String, String> results = OBJECT_MAPPER.readValue(response.body(), HashMap.class);
281+
HashMap<String, String> results = OBJECT_MAPPER.readValue(body, HashMap.class);
282282
queryDetail.setQueryId(results.get("id"));
283283
routingManager.setBackendForQueryId(queryDetail.getQueryId(), queryDetail.getBackendUrl());
284284
routingManager.setRoutingGroupForQueryId(queryDetail.getQueryId(), routingDestination.routingGroup());
285285
log.debug("QueryId [%s] mapped with proxy [%s]", queryDetail.getQueryId(), queryDetail.getBackendUrl());
286286
}
287287
catch (IOException e) {
288-
log.error("Failed to get QueryId from response [%s] , Status code [%s]", response.body(), response.statusCode());
288+
log.error("Failed to get QueryId from response [%s] , Status code [%s]", body, response.statusCode());
289289
}
290290
}
291291
else {
292-
log.error("Non OK HTTP Status code with response [%s] , Status code [%s], user: [%s]", response.body(), response.statusCode(), username.orElse(null));
292+
log.error("Non OK HTTP Status code with response [%s] , Status code [%s], user: [%s]", body, response.statusCode(), username.orElse(null));
293293
}
294294
queryDetail.setRoutingGroup(routingDestination.routingGroup());
295295
queryDetail.setExternalUrl(routingDestination.externalUrl());

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

Lines changed: 49 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,11 @@
2222
import io.trino.gateway.ha.config.ProxyResponseConfiguration;
2323
import io.trino.gateway.proxyserver.ProxyResponseHandler.ProxyResponse;
2424

25+
import java.io.ByteArrayInputStream;
2526
import java.io.IOException;
27+
import java.io.InputStream;
2628
import java.nio.charset.StandardCharsets;
29+
import java.util.zip.GZIPInputStream;
2730

2831
import static java.util.Objects.requireNonNull;
2932

@@ -47,7 +50,9 @@ public ProxyResponse handleException(Request request, Exception exception)
4750
public ProxyResponse handle(Request request, Response response)
4851
{
4952
try {
50-
return new ProxyResponse(response.getStatusCode(), response.getHeaders(), new String(response.getInputStream().readNBytes((int) responseSize.toBytes()), StandardCharsets.UTF_8));
53+
// Store raw bytes to preserve compression
54+
byte[] responseBodyBytes = response.getInputStream().readNBytes((int) responseSize.toBytes());
55+
return new ProxyResponse(response.getStatusCode(), response.getHeaders(), responseBodyBytes);
5156
}
5257
catch (IOException e) {
5358
throw new ProxyException("Failed reading response from remote Trino server", e);
@@ -57,11 +62,53 @@ public ProxyResponse handle(Request request, Response response)
5762
public record ProxyResponse(
5863
int statusCode,
5964
ListMultimap<HeaderName, String> headers,
60-
String body)
65+
byte[] body)
6166
{
6267
public ProxyResponse
6368
{
6469
requireNonNull(headers, "headers is null");
70+
requireNonNull(body, "body is null");
71+
}
72+
73+
/**
74+
* Get the response body as raw bytes for sending to clients (preserves
75+
* compression)
76+
*/
77+
public byte[] getRawBody()
78+
{
79+
return body;
80+
}
81+
82+
/**
83+
* Get the response body as a decompressed string for JSON parsing and logging.
84+
* Only call this when you need to parse the content, not when passing through
85+
* to clients.
86+
*/
87+
public String getDecompressedBody()
88+
{
89+
// Check if the response is gzip-compressed
90+
String contentEncoding = null;
91+
for (HeaderName headerName : headers.keySet()) {
92+
if (headerName.toString().equalsIgnoreCase("Content-Encoding")) {
93+
contentEncoding = headers.get(headerName).iterator().next();
94+
break;
95+
}
96+
}
97+
98+
if ("gzip".equalsIgnoreCase(contentEncoding)) {
99+
try {
100+
try (InputStream inputStream = new GZIPInputStream(new ByteArrayInputStream(body))) {
101+
return new String(inputStream.readAllBytes(), StandardCharsets.UTF_8);
102+
}
103+
}
104+
catch (IOException e) {
105+
// If decompression fails, return the body as UTF-8 string
106+
return new String(body, StandardCharsets.UTF_8);
107+
}
108+
}
109+
110+
// Not compressed, convert bytes to string
111+
return new String(body, StandardCharsets.UTF_8);
65112
}
66113
}
67114
}

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

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,14 @@ public MockResponse dispatch(RecordedRequest request)
7878
.setBody("{\"starting\": false}");
7979
}
8080

81+
if (request.getPath().equals(healthCheckEndpoint + "?test-compression")) {
82+
// Return the Accept-Encoding header value for compression testing
83+
String acceptEncoding = request.getHeader("Accept-Encoding");
84+
return new MockResponse().setResponseCode(200)
85+
.setHeader(CONTENT_TYPE, JSON_UTF_8)
86+
.setBody(acceptEncoding != null ? acceptEncoding : "null");
87+
}
88+
8189
if (request.getMethod().equals("PUT") && request.getPath().equals(customPutEndpoint)) {
8290
return new MockResponse().setResponseCode(200)
8391
.setHeader(CONTENT_TYPE, JSON_UTF_8)
@@ -157,4 +165,49 @@ void testGetQueryDetailsFromRequest()
157165
assertThat(queryDetail.getSource()).isEqualTo("trino-cli");
158166
assertThat(queryDetail.getBackendUrl()).isEqualTo("http://localhost:" + routerPort);
159167
}
168+
169+
@Test
170+
void testAcceptEncodingHeaderForwarding()
171+
throws Exception
172+
{
173+
// Test that Accept-Encoding header is properly forwarded to backends
174+
String url = "http://localhost:" + routerPort + healthCheckEndpoint + "?test-compression";
175+
String expectedAcceptEncoding = "gzip, deflate, br";
176+
177+
Request request = new Request.Builder()
178+
.url(url)
179+
.get()
180+
.addHeader("Accept-Encoding", expectedAcceptEncoding)
181+
.build();
182+
183+
try (Response response = httpClient.newCall(request).execute()) {
184+
assertThat(response.code()).isEqualTo(200);
185+
assertThat(response.body()).isNotNull();
186+
187+
// The mock backend returns the Accept-Encoding header value in the response body
188+
assertThat(response.body().string()).isEqualTo(expectedAcceptEncoding);
189+
}
190+
}
191+
192+
@Test
193+
void testDefaultAcceptEncodingHeaderForwarding()
194+
throws Exception
195+
{
196+
// Test that requests without explicit Accept-Encoding header work correctly
197+
// Note: OkHttp automatically adds "Accept-Encoding: gzip" when none is specified
198+
String url = "http://localhost:" + routerPort + healthCheckEndpoint + "?test-compression";
199+
200+
Request request = new Request.Builder()
201+
.url(url)
202+
.get()
203+
.build(); // No explicit Accept-Encoding header
204+
205+
try (Response response = httpClient.newCall(request).execute()) {
206+
assertThat(response.code()).isEqualTo(200);
207+
assertThat(response.body()).isNotNull();
208+
209+
// OkHttp automatically adds "Accept-Encoding: gzip" when none is specified
210+
assertThat(response.body().string()).isEqualTo("gzip");
211+
}
212+
}
160213
}

0 commit comments

Comments
 (0)