@@ -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 ());
0 commit comments