From 40337f6444362478beba6d555c1425f3490b3cc1 Mon Sep 17 00:00:00 2001 From: Andy Date: Fri, 20 Dec 2024 19:00:08 -0500 Subject: [PATCH] Convert byte array to string for logging in ProxyResponse --- docs/installation.md | 10 ++++++ .../ha/config/HaGatewayConfiguration.java | 11 ++++++ .../ha/config/ProxyResponseConfiguration.java | 35 +++++++++++++++++++ .../proxyserver/ProxyRequestHandler.java | 5 ++- .../proxyserver/ProxyResponseHandler.java | 14 ++++++-- 5 files changed, 72 insertions(+), 3 deletions(-) create mode 100644 gateway-ha/src/main/java/io/trino/gateway/ha/config/ProxyResponseConfiguration.java diff --git a/docs/installation.md b/docs/installation.md index 30a4a4237..d63d7a823 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -212,6 +212,16 @@ serverConfig: http-server.process-forwarded: true ``` +## Configure larger proxy response size + +Trino Gateway reads the response from Trino in bytes (up to 32MB by default). +It can be configured by setting: + +```yaml +proxyResponseConfiguration: + responseSize: 50MB +``` + ## Running Trino Gateway Start Trino Gateway with the following java command in the directory of the diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java b/gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java index 6cc7cd13b..479abcf96 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java @@ -42,6 +42,7 @@ public class HaGatewayConfiguration private GatewayCookieConfiguration gatewayCookieConfiguration = new GatewayCookieConfiguration(); private List statementPaths = ImmutableList.of(V1_STATEMENT_PATH); private boolean includeClusterHostInResponse; + private ProxyResponseConfiguration proxyResponseConfiguration = new ProxyResponseConfiguration(); private RequestAnalyzerConfig requestAnalyzerConfig = new RequestAnalyzerConfig(); @@ -255,6 +256,16 @@ public void setIncludeClusterHostInResponse(boolean includeClusterHostInResponse this.includeClusterHostInResponse = includeClusterHostInResponse; } + public ProxyResponseConfiguration getProxyResponseConfiguration() + { + return this.proxyResponseConfiguration; + } + + public void setProxyResponseConfiguration(ProxyResponseConfiguration proxyResponseConfiguration) + { + this.proxyResponseConfiguration = proxyResponseConfiguration; + } + private void validateStatementPath(String statementPath, List statementPaths) { if (statementPath.startsWith(V1_STATEMENT_PATH) || diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/config/ProxyResponseConfiguration.java b/gateway-ha/src/main/java/io/trino/gateway/ha/config/ProxyResponseConfiguration.java new file mode 100644 index 000000000..1aad37f96 --- /dev/null +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/config/ProxyResponseConfiguration.java @@ -0,0 +1,35 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.gateway.ha.config; + +import io.airlift.units.DataSize; + +import static io.airlift.units.DataSize.Unit.MEGABYTE; + +public class ProxyResponseConfiguration +{ + private DataSize responseSize = DataSize.of(32, MEGABYTE); + + public ProxyResponseConfiguration() {} + + public DataSize getResponseSize() + { + return responseSize; + } + + public void setResponseSize(DataSize responseSize) + { + this.responseSize = responseSize; + } +} diff --git a/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java b/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java index 42d22f7c5..36a01bb13 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java +++ b/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java @@ -25,6 +25,7 @@ import io.airlift.units.Duration; import io.trino.gateway.ha.config.GatewayCookieConfigurationPropertiesProvider; import io.trino.gateway.ha.config.HaGatewayConfiguration; +import io.trino.gateway.ha.config.ProxyResponseConfiguration; import io.trino.gateway.ha.router.GatewayCookie; import io.trino.gateway.ha.router.OAuth2GatewayCookie; import io.trino.gateway.ha.router.QueryHistoryManager; @@ -87,6 +88,7 @@ public class ProxyRequestHandler private final List statementPaths; private final boolean includeClusterInfoInResponse; private final TrinoRequestUser.TrinoRequestUserProvider trinoRequestUserProvider; + private final ProxyResponseConfiguration proxyResponseConfiguration; @Inject public ProxyRequestHandler( @@ -104,6 +106,7 @@ public ProxyRequestHandler( addXForwardedHeaders = haGatewayConfiguration.getRouting().isAddXForwardedHeaders(); statementPaths = haGatewayConfiguration.getStatementPaths(); this.includeClusterInfoInResponse = haGatewayConfiguration.isIncludeClusterHostInResponse(); + proxyResponseConfiguration = haGatewayConfiguration.getProxyResponseConfiguration(); } @PreDestroy @@ -245,7 +248,7 @@ private void setupAsyncResponse(AsyncResponse asyncResponse, ListenableFuture executeHttp(Request request) { - return FluentFuture.from(httpClient.executeAsync(request, new ProxyResponseHandler())); + return FluentFuture.from(httpClient.executeAsync(request, new ProxyResponseHandler(proxyResponseConfiguration))); } private static Response handleProxyException(Request request, ProxyException e) diff --git a/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyResponseHandler.java b/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyResponseHandler.java index f36bb7672..e20c18ada 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyResponseHandler.java +++ b/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyResponseHandler.java @@ -18,15 +18,25 @@ import io.airlift.http.client.Request; import io.airlift.http.client.Response; import io.airlift.http.client.ResponseHandler; +import io.airlift.units.DataSize; +import io.trino.gateway.ha.config.ProxyResponseConfiguration; import io.trino.gateway.proxyserver.ProxyResponseHandler.ProxyResponse; import java.io.IOException; +import java.nio.charset.StandardCharsets; import static java.util.Objects.requireNonNull; public class ProxyResponseHandler implements ResponseHandler { + private final DataSize responseSize; + + public ProxyResponseHandler(ProxyResponseConfiguration proxyResponseConfiguration) + { + this.responseSize = requireNonNull(proxyResponseConfiguration.getResponseSize(), "responseSize is null"); + } + @Override public ProxyResponse handleException(Request request, Exception exception) { @@ -37,7 +47,7 @@ public ProxyResponse handleException(Request request, Exception exception) public ProxyResponse handle(Request request, Response response) { try { - return new ProxyResponse(response.getStatusCode(), response.getHeaders(), response.getInputStream().readAllBytes()); + return new ProxyResponse(response.getStatusCode(), response.getHeaders(), new String(response.getInputStream().readNBytes((int) responseSize.toBytes()), StandardCharsets.UTF_8)); } catch (IOException e) { throw new ProxyException("Failed reading response from remote Trino server", e); @@ -47,7 +57,7 @@ public ProxyResponse handle(Request request, Response response) public record ProxyResponse( int statusCode, ListMultimap headers, - byte[] body) + String body) { public ProxyResponse {