Skip to content

Commit 016c21f

Browse files
Chaho12Jaeho Yoo
authored andcommitted
Reformat to match Airlift style
1 parent a37a345 commit 016c21f

File tree

2 files changed

+57
-55
lines changed

2 files changed

+57
-55
lines changed

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

Lines changed: 40 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,40 @@ public ProxyRequestHandler(
109109
proxyResponseConfiguration = haGatewayConfiguration.getProxyResponseConfiguration();
110110
}
111111

112+
private static String getRemoteTarget(URI remoteUri)
113+
{
114+
return format("%s://%s", remoteUri.getScheme(), remoteUri.getAuthority());
115+
}
116+
117+
private static Response handleProxyException(Request request, ProxyException e)
118+
{
119+
log.warn(e, "Proxy request failed: %s %s", request.getMethod(), request.getUri());
120+
throw badRequest(e.getMessage());
121+
}
122+
123+
private static WebApplicationException badRequest(String message)
124+
{
125+
throw new WebApplicationException(
126+
Response.status(Response.Status.BAD_GATEWAY)
127+
.type(TEXT_PLAIN_TYPE)
128+
.entity(message)
129+
.build());
130+
}
131+
132+
public static QueryHistoryManager.QueryDetail getQueryDetailsFromRequest(Request request,
133+
Optional<String> username)
134+
{
135+
QueryHistoryManager.QueryDetail queryDetail = new QueryHistoryManager.QueryDetail();
136+
queryDetail.setBackendUrl(getRemoteTarget(request.getUri()));
137+
queryDetail.setCaptureTime(System.currentTimeMillis());
138+
username.ifPresent(queryDetail::setUser);
139+
queryDetail.setSource(request.getHeader(SOURCE_HEADER));
140+
141+
String queryText = new String(((StaticBodyGenerator) request.getBodyGenerator()).getBody(), UTF_8);
142+
queryDetail.setQueryText(queryText);
143+
return queryDetail;
144+
}
145+
112146
@PreDestroy
113147
public void shutdown()
114148
{
@@ -188,9 +222,11 @@ private void performRequest(
188222

189223
FluentFuture<ProxyResponse> future = executeHttp(request);
190224

191-
if (statementPaths.stream().anyMatch(request.getUri().getPath()::startsWith) && request.getMethod().equals(HttpMethod.POST)) {
225+
if (statementPaths.stream().anyMatch(request.getUri().getPath()::startsWith)
226+
&& request.getMethod().equals(HttpMethod.POST)) {
192227
Optional<String> username = trinoRequestUserProvider.getInstance(servletRequest).getUser();
193-
future = future.transform(response -> recordBackendForQueryId(request, response, username, routingDestination), executor);
228+
future = future.transform(
229+
response -> recordBackendForQueryId(request, response, username, routingDestination), executor);
194230
if (includeClusterInfoInResponse) {
195231
cookieBuilder.add(new NewCookie.Builder("trinoClusterHost").value(remoteUri.getHost()).build());
196232
}
@@ -224,11 +260,6 @@ else if (servletRequest.getCookies() != null) {
224260
return ImmutableList.of();
225261
}
226262

227-
private static String getRemoteTarget(URI remoteUri)
228-
{
229-
return format("%s://%s", remoteUri.getScheme(), remoteUri.getAuthority());
230-
}
231-
232263
private Response buildResponse(ProxyResponse response, ImmutableList<NewCookie> cookie)
233264
{
234265
Response.ResponseBuilder builder = Response.status(response.statusCode()).entity(response.body());
@@ -249,22 +280,8 @@ private void setupAsyncResponse(AsyncResponse asyncResponse, ListenableFuture<Re
249280

250281
private FluentFuture<ProxyResponse> executeHttp(Request request)
251282
{
252-
return FluentFuture.from(httpClient.executeAsync(request, new ProxyResponseHandler(proxyResponseConfiguration)));
253-
}
254-
255-
private static Response handleProxyException(Request request, ProxyException e)
256-
{
257-
log.warn(e, "Proxy request failed: %s %s", request.getMethod(), request.getUri());
258-
throw badRequest(e.getMessage());
259-
}
260-
261-
private static WebApplicationException badRequest(String message)
262-
{
263-
throw new WebApplicationException(
264-
Response.status(Response.Status.BAD_GATEWAY)
265-
.type(TEXT_PLAIN_TYPE)
266-
.entity(message)
267-
.build());
283+
return FluentFuture
284+
.from(httpClient.executeAsync(request, new ProxyResponseHandler(proxyResponseConfiguration)));
268285
}
269286

270287
private ProxyResponse recordBackendForQueryId(Request request, ProxyResponse response, Optional<String> username,
@@ -297,19 +314,6 @@ private ProxyResponse recordBackendForQueryId(Request request, ProxyResponse res
297314
return response;
298315
}
299316

300-
public static QueryHistoryManager.QueryDetail getQueryDetailsFromRequest(Request request, Optional<String> username)
301-
{
302-
QueryHistoryManager.QueryDetail queryDetail = new QueryHistoryManager.QueryDetail();
303-
queryDetail.setBackendUrl(getRemoteTarget(request.getUri()));
304-
queryDetail.setCaptureTime(System.currentTimeMillis());
305-
username.ifPresent(queryDetail::setUser);
306-
queryDetail.setSource(request.getHeader(SOURCE_HEADER));
307-
308-
String queryText = new String(((StaticBodyGenerator) request.getBodyGenerator()).getBody(), UTF_8);
309-
queryDetail.setQueryText(queryText);
310-
return queryDetail;
311-
}
312-
313317
private void addXForwardedHeaders(HttpServletRequest servletRequest, Request.Builder requestBuilder)
314318
{
315319
requestBuilder.addHeader(X_FORWARDED_FOR, servletRequest.getRemoteAddr());

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

Lines changed: 17 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -50,17 +50,14 @@
5050
@TestInstance(PER_CLASS)
5151
final class TestProxyRequestHandler
5252
{
53+
private static final String OK = "OK";
54+
private static final int NOT_FOUND = 404;
55+
private static final MediaType MEDIA_TYPE = MediaType.parse("application/json; charset=utf-8");
5356
private final OkHttpClient httpClient = new OkHttpClient();
5457
private final MockWebServer mockTrinoServer = new MockWebServer();
5558
private final PostgreSQLContainer postgresql = new PostgreSQLContainer("postgres:17");
56-
5759
private final int routerPort = 21001 + (int) (Math.random() * 1000);
5860
private final int customBackendPort = 21000 + (int) (Math.random() * 1000);
59-
60-
private static final String OK = "OK";
61-
private static final int NOT_FOUND = 404;
62-
private static final MediaType MEDIA_TYPE = MediaType.parse("application/json; charset=utf-8");
63-
6461
private final String customPutEndpoint = "/v1/custom"; // this is enabled in test-config-template.yml
6562
private final String healthCheckEndpoint = "/v1/info";
6663

@@ -69,7 +66,8 @@ void setup()
6966
throws Exception
7067
{
7168
prepareMockBackend(mockTrinoServer, customBackendPort, "default custom response");
72-
mockTrinoServer.setDispatcher(new Dispatcher() {
69+
mockTrinoServer.setDispatcher(new Dispatcher()
70+
{
7371
@Override
7472
public MockResponse dispatch(RecordedRequest request)
7573
{
@@ -130,18 +128,18 @@ void testGetQueryDetailsFromRequest()
130128
{
131129
// A sample query longer than 200 characters to test against truncation.
132130
String longQuery = """
133-
SELECT
134-
c.customer_name,
135-
c.customer_region,
136-
COUNT(o.order_id) AS total_orders,
137-
SUM(o.order_value) AS total_revenue
138-
FROM
139-
hive.sales_data.customers AS c
140-
JOIN
141-
hive.sales_data.orders AS o
142-
ON c.customer_id = o.customer_id
143-
WHERE
144-
o.order_date >= date '2023-01-01'""";
131+
SELECT
132+
c.customer_name,
133+
c.customer_region,
134+
COUNT(o.order_id) AS total_orders,
135+
SUM(o.order_value) AS total_revenue
136+
FROM
137+
hive.sales_data.customers AS c
138+
JOIN
139+
hive.sales_data.orders AS o
140+
ON c.customer_id = o.customer_id
141+
WHERE
142+
o.order_date >= date '2023-01-01'""";
145143

146144
io.airlift.http.client.Request request = preparePost()
147145
.setUri(URI.create("http://localhost:" + routerPort + V1_STATEMENT_PATH))

0 commit comments

Comments
 (0)