From 18e0d44f986a13aa76f4d71b5e8f122837e177e3 Mon Sep 17 00:00:00 2001 From: siminyou Date: Wed, 6 Nov 2024 22:48:09 -0500 Subject: [PATCH 1/6] add proxy handler for PUT request --- .../gateway/proxyserver/ProxyRequestHandler.java | 11 +++++++++++ .../gateway/proxyserver/RouteToBackendResource.java | 11 +++++++++++ 2 files changed, 22 insertions(+) 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..3ff67a0aa 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 @@ -140,6 +140,17 @@ public void postRequest( performRequest(remoteUri, servletRequest, asyncResponse, request); } + public void putRequest( + String statement, + HttpServletRequest servletRequest, + AsyncResponse asyncResponse, + URI remoteUri) + { + Request.Builder request = Request.Builder.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..75bd4b5c7 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,14 @@ 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) + { + String remoteUri = routingTargetHandler.getRoutingDestination(servletRequest); + proxyRequestHandler.putRequest(body, servletRequest, asyncResponse, URI.create(remoteUri)); + } } From bea076af2929ea7eff311b1c98c12c70b8acd845 Mon Sep 17 00:00:00 2001 From: siminyou Date: Thu, 7 Nov 2024 11:15:49 -0500 Subject: [PATCH 2/6] static import preparePut --- .../java/io/trino/gateway/proxyserver/ProxyRequestHandler.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 3ff67a0aa..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; @@ -146,7 +147,7 @@ public void putRequest( AsyncResponse asyncResponse, URI remoteUri) { - Request.Builder request = Request.Builder.preparePut() + Request.Builder request = preparePut() .setBodyGenerator(createStaticBodyGenerator(statement, UTF_8)); performRequest(remoteUri, servletRequest, asyncResponse, request); } From 773d4a99d35b86a32dc4c404608882ed07995e60 Mon Sep 17 00:00:00 2001 From: siminyou Date: Thu, 7 Nov 2024 11:17:40 -0500 Subject: [PATCH 3/6] use MultiReadHttpServletRequest in putRequest --- .../io/trino/gateway/proxyserver/RouteToBackendResource.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 75bd4b5c7..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 @@ -93,7 +93,8 @@ public void putHandler( @Context HttpServletRequest servletRequest, @Suspended AsyncResponse asyncResponse) { - String remoteUri = routingTargetHandler.getRoutingDestination(servletRequest); - proxyRequestHandler.putRequest(body, servletRequest, asyncResponse, URI.create(remoteUri)); + MultiReadHttpServletRequest multiReadHttpServletRequest = new MultiReadHttpServletRequest(servletRequest, body); + String remoteUri = routingTargetHandler.getRoutingDestination(multiReadHttpServletRequest); + proxyRequestHandler.putRequest(body, multiReadHttpServletRequest, asyncResponse, URI.create(remoteUri)); } } From 7e0fa5d1541f2c9f56aeefd3ca6439da02d0264f Mon Sep 17 00:00:00 2001 From: siminyou Date: Thu, 7 Nov 2024 12:01:51 -0500 Subject: [PATCH 4/6] add test for put --- .../gateway/ha/TestPutRequestHandler.java | 120 ++++++++++++++++++ 1 file changed, 120 insertions(+) create mode 100644 gateway-ha/src/test/java/io/trino/gateway/ha/TestPutRequestHandler.java diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestPutRequestHandler.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestPutRequestHandler.java new file mode 100644 index 000000000..4eea433e8 --- /dev/null +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/TestPutRequestHandler.java @@ -0,0 +1,120 @@ +/* + * 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; + +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 org.junit.jupiter.api.TestInstance.Lifecycle; + +import static com.google.common.net.HttpHeaders.CONTENT_TYPE; +import static com.google.common.net.MediaType.JSON_UTF_8; +import static org.assertj.core.api.Assertions.assertThat; + +@TestInstance(Lifecycle.PER_CLASS) +public class TestPutRequestHandler +{ + private final OkHttpClient httpClient = new OkHttpClient(); + private MockWebServer mockTrinoServer = new MockWebServer();; + + private final int routerPort = 21001 + (int) (Math.random() * 1000); + private final int customBackendPort = 21000 + (int) (Math.random() * 1000); + + private final static String CUSTOM_PUT_ENDPOINT = "/v1/custom"; // this is enabled in test-config-template.yml + private final static String HEALTH_CHECK_ENDPOINT = "/v1/info"; + private final static String PUT = "PUT"; + + @BeforeAll + public void setup() + throws Exception + { + HaGatewayTestUtils.prepareMockBackend(mockTrinoServer, customBackendPort, "default custom response"); + mockTrinoServer.setDispatcher(new Dispatcher() { + @Override + public MockResponse dispatch(RecordedRequest request) + { + if (request.getPath().equals(HEALTH_CHECK_ENDPOINT)) { + return new MockResponse().setResponseCode(200) + .setHeader(CONTENT_TYPE, JSON_UTF_8) + .setBody("{\"starting\": false}"); + } + + if (request.getMethod().equals(PUT) && request.getPath().equals(CUSTOM_PUT_ENDPOINT)) { + return new MockResponse().setResponseCode(200) + .setHeader(CONTENT_TYPE, JSON_UTF_8) + .setBody("ok"); + } + + return new MockResponse().setResponseCode(404); + } + }); + + // seed database + HaGatewayTestUtils.TestConfig testConfig = + HaGatewayTestUtils.buildGatewayConfigAndSeedDb(routerPort, "test-config-template.yml"); + + // Start Gateway + String[] args = {testConfig.configFilePath()}; + HaGatewayLauncher.main(args); + + // Set up backend + HaGatewayTestUtils.setUpBackend( + "custom", "http://localhost:" + customBackendPort, "externalUrl", true, "adhoc", routerPort); + } + + @AfterAll + public void cleanup() + throws Exception + { + if (mockTrinoServer != null) { + mockTrinoServer.shutdown(); + } + } + + @Test + public void testPutRequestHandler() + throws Exception + { + RequestBody requestBody = + RequestBody.create(MediaType.parse("application/json; charset=utf-8"), "SELECT 1"); + Request putRequest = + new Request.Builder() + .url("http://localhost:" + routerPort + CUSTOM_PUT_ENDPOINT) + .put(requestBody) + .build(); + + try (Response response = httpClient.newCall(putRequest).execute()) { + assertThat(response.isSuccessful()).isTrue(); + } + + Request postRequest = + new Request.Builder() + .url("http://localhost:" + routerPort + CUSTOM_PUT_ENDPOINT) + .post(requestBody) + .build(); + try (Response response = httpClient.newCall(postRequest).execute()) { + assertThat(response.isSuccessful()).isFalse(); + } + } +} From 0ff2fac15499c1d979da824d4ac427c930753aa0 Mon Sep 17 00:00:00 2001 From: siminyou Date: Thu, 7 Nov 2024 12:04:39 -0500 Subject: [PATCH 5/6] add test for put request --- .../trino/gateway/ha/TestPutRequestHandler.java | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestPutRequestHandler.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestPutRequestHandler.java index 4eea433e8..5b67b5bdc 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/TestPutRequestHandler.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/TestPutRequestHandler.java @@ -36,14 +36,13 @@ public class TestPutRequestHandler { private final OkHttpClient httpClient = new OkHttpClient(); - private MockWebServer mockTrinoServer = new MockWebServer();; + private MockWebServer mockTrinoServer = new MockWebServer(); private final int routerPort = 21001 + (int) (Math.random() * 1000); private final int customBackendPort = 21000 + (int) (Math.random() * 1000); - private final static String CUSTOM_PUT_ENDPOINT = "/v1/custom"; // this is enabled in test-config-template.yml - private final static String HEALTH_CHECK_ENDPOINT = "/v1/info"; - private final static String PUT = "PUT"; + private final String customPutEndpoint = "/v1/custom"; // this is enabled in test-config-template.yml + private final String healthCheckEndpoint = "/v1/info"; @BeforeAll public void setup() @@ -54,13 +53,13 @@ public void setup() @Override public MockResponse dispatch(RecordedRequest request) { - if (request.getPath().equals(HEALTH_CHECK_ENDPOINT)) { + 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(CUSTOM_PUT_ENDPOINT)) { + if (request.getMethod().equals("PUT") && request.getPath().equals(customPutEndpoint)) { return new MockResponse().setResponseCode(200) .setHeader(CONTENT_TYPE, JSON_UTF_8) .setBody("ok"); @@ -100,7 +99,7 @@ public void testPutRequestHandler() RequestBody.create(MediaType.parse("application/json; charset=utf-8"), "SELECT 1"); Request putRequest = new Request.Builder() - .url("http://localhost:" + routerPort + CUSTOM_PUT_ENDPOINT) + .url("http://localhost:" + routerPort + customPutEndpoint) .put(requestBody) .build(); @@ -110,7 +109,7 @@ public void testPutRequestHandler() Request postRequest = new Request.Builder() - .url("http://localhost:" + routerPort + CUSTOM_PUT_ENDPOINT) + .url("http://localhost:" + routerPort + customPutEndpoint) .post(requestBody) .build(); try (Response response = httpClient.newCall(postRequest).execute()) { From c58ac4863a8d30c5754a3127fb98de4fc88b47d3 Mon Sep 17 00:00:00 2001 From: siminyou Date: Thu, 14 Nov 2024 11:02:00 -0500 Subject: [PATCH 6/6] update test to check body --- .../io/trino/gateway/ha/TestPutRequestHandler.java | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestPutRequestHandler.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestPutRequestHandler.java index 5b67b5bdc..a72a06b25 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/TestPutRequestHandler.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/TestPutRequestHandler.java @@ -41,6 +41,9 @@ public class TestPutRequestHandler 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 final String customPutEndpoint = "/v1/custom"; // this is enabled in test-config-template.yml private final String healthCheckEndpoint = "/v1/info"; @@ -62,10 +65,10 @@ public MockResponse dispatch(RecordedRequest request) if (request.getMethod().equals("PUT") && request.getPath().equals(customPutEndpoint)) { return new MockResponse().setResponseCode(200) .setHeader(CONTENT_TYPE, JSON_UTF_8) - .setBody("ok"); + .setBody(OK); } - return new MockResponse().setResponseCode(404); + return new MockResponse().setResponseCode(NOT_FOUND); } }); @@ -104,7 +107,7 @@ public void testPutRequestHandler() .build(); try (Response response = httpClient.newCall(putRequest).execute()) { - assertThat(response.isSuccessful()).isTrue(); + assertThat(response.body().string()).isEqualTo(OK); } Request postRequest = @@ -113,7 +116,7 @@ public void testPutRequestHandler() .post(requestBody) .build(); try (Response response = httpClient.newCall(postRequest).execute()) { - assertThat(response.isSuccessful()).isFalse(); + assertThat(response.code()).isEqualTo(NOT_FOUND); } } }