Skip to content

Commit

Permalink
Convert byte array to string for logging in ProxyResponse
Browse files Browse the repository at this point in the history
  • Loading branch information
andythsu authored and willmostly committed Jan 21, 2025
1 parent c2ba1ea commit 40337f6
Show file tree
Hide file tree
Showing 5 changed files with 72 additions and 3 deletions.
10 changes: 10 additions & 0 deletions docs/installation.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ public class HaGatewayConfiguration
private GatewayCookieConfiguration gatewayCookieConfiguration = new GatewayCookieConfiguration();
private List<String> statementPaths = ImmutableList.of(V1_STATEMENT_PATH);
private boolean includeClusterHostInResponse;
private ProxyResponseConfiguration proxyResponseConfiguration = new ProxyResponseConfiguration();

private RequestAnalyzerConfig requestAnalyzerConfig = new RequestAnalyzerConfig();

Expand Down Expand Up @@ -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<String> statementPaths)
{
if (statementPath.startsWith(V1_STATEMENT_PATH) ||
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -87,6 +88,7 @@ public class ProxyRequestHandler
private final List<String> statementPaths;
private final boolean includeClusterInfoInResponse;
private final TrinoRequestUser.TrinoRequestUserProvider trinoRequestUserProvider;
private final ProxyResponseConfiguration proxyResponseConfiguration;

@Inject
public ProxyRequestHandler(
Expand All @@ -104,6 +106,7 @@ public ProxyRequestHandler(
addXForwardedHeaders = haGatewayConfiguration.getRouting().isAddXForwardedHeaders();
statementPaths = haGatewayConfiguration.getStatementPaths();
this.includeClusterInfoInResponse = haGatewayConfiguration.isIncludeClusterHostInResponse();
proxyResponseConfiguration = haGatewayConfiguration.getProxyResponseConfiguration();
}

@PreDestroy
Expand Down Expand Up @@ -245,7 +248,7 @@ private void setupAsyncResponse(AsyncResponse asyncResponse, ListenableFuture<Re

private FluentFuture<ProxyResponse> 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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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<ProxyResponse, RuntimeException>
{
private final DataSize responseSize;

public ProxyResponseHandler(ProxyResponseConfiguration proxyResponseConfiguration)
{
this.responseSize = requireNonNull(proxyResponseConfiguration.getResponseSize(), "responseSize is null");
}

@Override
public ProxyResponse handleException(Request request, Exception exception)
{
Expand All @@ -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);
Expand All @@ -47,7 +57,7 @@ public ProxyResponse handle(Request request, Response response)
public record ProxyResponse(
int statusCode,
ListMultimap<HeaderName, String> headers,
byte[] body)
String body)
{
public ProxyResponse
{
Expand Down

0 comments on commit 40337f6

Please sign in to comment.