From a11e5a4967463fe054f2cee498c9cc88fc491903 Mon Sep 17 00:00:00 2001 From: siminyou Date: Wed, 6 Nov 2024 22:48:09 -0500 Subject: [PATCH] Add proxy handler for PUT request --- .../proxyserver/ProxyRequestHandler.java | 12 ++ .../proxyserver/RouteToBackendResource.java | 12 ++ .../proxyserver/TestProxyRequestHandler.java | 114 ++++++++++++++++++ 3 files changed, 138 insertions(+) create mode 100644 gateway-ha/src/test/java/io/trino/gateway/proxyserver/TestProxyRequestHandler.java 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 859c79325..3e1675d7c 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 @@ -58,6 +58,7 @@ import static io.airlift.http.client.Request.Builder.prepareDelete; import static io.airlift.http.client.Request.Builder.prepareGet; import static io.airlift.http.client.Request.Builder.preparePost; +import static io.airlift.http.client.Request.Builder.preparePut; import static io.airlift.http.client.StaticBodyGenerator.createStaticBodyGenerator; import static io.airlift.jaxrs.AsyncResponseHandler.bindAsyncResponse; import static io.trino.gateway.ha.handler.ProxyUtils.QUERY_TEXT_LENGTH_FOR_HISTORY; @@ -140,6 +141,17 @@ public void postRequest( performRequest(remoteUri, servletRequest, asyncResponse, request); } + public void putRequest( + String statement, + HttpServletRequest servletRequest, + AsyncResponse asyncResponse, + URI remoteUri) + { + Request.Builder request = preparePut() + .setBodyGenerator(createStaticBodyGenerator(statement, UTF_8)); + performRequest(remoteUri, servletRequest, asyncResponse, request); + } + private void performRequest( URI remoteUri, HttpServletRequest servletRequest, diff --git a/gateway-ha/src/main/java/io/trino/gateway/proxyserver/RouteToBackendResource.java b/gateway-ha/src/main/java/io/trino/gateway/proxyserver/RouteToBackendResource.java index 3e3dfd2ee..37ca73d3d 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/proxyserver/RouteToBackendResource.java +++ b/gateway-ha/src/main/java/io/trino/gateway/proxyserver/RouteToBackendResource.java @@ -20,6 +20,7 @@ import jakarta.ws.rs.DELETE; import jakarta.ws.rs.GET; import jakarta.ws.rs.POST; +import jakarta.ws.rs.PUT; import jakarta.ws.rs.Path; import jakarta.ws.rs.container.AsyncResponse; import jakarta.ws.rs.container.Suspended; @@ -85,4 +86,15 @@ public void deleteHandler( String remoteUri = routingTargetHandler.getRoutingDestination(servletRequest); proxyRequestHandler.deleteRequest(servletRequest, asyncResponse, URI.create(remoteUri)); } + + @PUT + public void putHandler( + String body, + @Context HttpServletRequest servletRequest, + @Suspended AsyncResponse asyncResponse) + { + MultiReadHttpServletRequest multiReadHttpServletRequest = new MultiReadHttpServletRequest(servletRequest, body); + String remoteUri = routingTargetHandler.getRoutingDestination(multiReadHttpServletRequest); + proxyRequestHandler.putRequest(body, multiReadHttpServletRequest, asyncResponse, URI.create(remoteUri)); + } } diff --git a/gateway-ha/src/test/java/io/trino/gateway/proxyserver/TestProxyRequestHandler.java b/gateway-ha/src/test/java/io/trino/gateway/proxyserver/TestProxyRequestHandler.java new file mode 100644 index 000000000..c1c2b4b01 --- /dev/null +++ b/gateway-ha/src/test/java/io/trino/gateway/proxyserver/TestProxyRequestHandler.java @@ -0,0 +1,114 @@ +/* + * 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.proxyserver; + +import io.trino.gateway.ha.HaGatewayLauncher; +import io.trino.gateway.ha.HaGatewayTestUtils; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.RequestBody; +import okhttp3.Response; +import okhttp3.mockwebserver.Dispatcher; +import okhttp3.mockwebserver.MockResponse; +import okhttp3.mockwebserver.MockWebServer; +import okhttp3.mockwebserver.RecordedRequest; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; + +import static com.google.common.net.HttpHeaders.CONTENT_TYPE; +import static com.google.common.net.MediaType.JSON_UTF_8; +import static io.trino.gateway.ha.HaGatewayTestUtils.buildGatewayConfigAndSeedDb; +import static io.trino.gateway.ha.HaGatewayTestUtils.prepareMockBackend; +import static io.trino.gateway.ha.HaGatewayTestUtils.setUpBackend; +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS; + +@TestInstance(PER_CLASS) +final class TestProxyRequestHandler +{ + private final OkHttpClient httpClient = new OkHttpClient(); + private final MockWebServer mockTrinoServer = new MockWebServer(); + + private final int routerPort = 21001 + (int) (Math.random() * 1000); + private final int customBackendPort = 21000 + (int) (Math.random() * 1000); + + private static final String OK = "OK"; + private static final int NOT_FOUND = 404; + private static final MediaType MEDIA_TYPE = MediaType.parse("application/json; charset=utf-8"); + + private final String customPutEndpoint = "/v1/custom"; // this is enabled in test-config-template.yml + private final String healthCheckEndpoint = "/v1/info"; + + @BeforeAll + void setup() + throws Exception + { + prepareMockBackend(mockTrinoServer, customBackendPort, "default custom response"); + mockTrinoServer.setDispatcher(new Dispatcher() { + @Override + public MockResponse dispatch(RecordedRequest request) + { + if (request.getPath().equals(healthCheckEndpoint)) { + return new MockResponse().setResponseCode(200) + .setHeader(CONTENT_TYPE, JSON_UTF_8) + .setBody("{\"starting\": false}"); + } + + if (request.getMethod().equals("PUT") && request.getPath().equals(customPutEndpoint)) { + return new MockResponse().setResponseCode(200) + .setHeader(CONTENT_TYPE, JSON_UTF_8) + .setBody(OK); + } + + return new MockResponse().setResponseCode(NOT_FOUND); + } + }); + + HaGatewayTestUtils.TestConfig testConfig = buildGatewayConfigAndSeedDb(routerPort, "test-config-template.yml"); + + String[] args = {testConfig.configFilePath()}; + HaGatewayLauncher.main(args); + + setUpBackend("custom", "http://localhost:" + customBackendPort, "externalUrl", true, "adhoc", routerPort); + } + + @AfterAll + void cleanup() + throws Exception + { + mockTrinoServer.shutdown(); + } + + @Test + void testPutRequestHandler() + throws Exception + { + String url = "http://localhost:" + routerPort + customPutEndpoint; + RequestBody requestBody = RequestBody.create("SELECT 1", MEDIA_TYPE); + + Request putRequest = new Request.Builder().url(url).put(requestBody).build(); + try (Response response = httpClient.newCall(putRequest).execute()) { + assertThat(response.body()).isNotNull(); + assertThat(response.body().string()).isEqualTo(OK); + } + + Request postRequest = new Request.Builder().url(url).post(requestBody).build(); + try (Response response = httpClient.newCall(postRequest).execute()) { + assertThat(response.code()).isEqualTo(NOT_FOUND); + } + } +}