Add proxy handler for DELETE method

* Needed for query termination from the UI or via API * Includes est
trinodb · Jul 22, 2024 · fdb1aa0 · fdb1aa0
1 parent 523deef
commit fdb1aa0
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 0 deletions.
diff --git a/gateway-ha/pom.xml b/gateway-ha/pom.xml
@@ -327,6 +327,13 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-client</artifactId>
+            <version>444</version>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>org.assertj</groupId>
             <artifactId>assertj-core</artifactId>

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
@@ -53,6 +53,7 @@
 import static com.google.common.net.HttpHeaders.X_FORWARDED_PROTO;
 import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 import static io.airlift.concurrent.Threads.daemonThreadsNamed;
+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.StaticBodyGenerator.createStaticBodyGenerator;
@@ -104,6 +105,15 @@ public void shutdown()
         executor.shutdownNow();
     }
 
+    public void deleteRequest(
+            HttpServletRequest servletRequest,
+            AsyncResponse asyncResponse,
+            URI remoteUri)
+    {
+        Request.Builder request = prepareDelete();
+        performRequest(remoteUri, servletRequest, asyncResponse, request);
+    }
+
     public void getRequest(
             HttpServletRequest servletRequest,
             AsyncResponse asyncResponse,

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
@@ -17,6 +17,7 @@
 import io.trino.gateway.ha.handler.ProxyHandlerStats;
 import io.trino.gateway.ha.handler.RoutingTargetHandler;
 import jakarta.servlet.http.HttpServletRequest;
+import jakarta.ws.rs.DELETE;
 import jakarta.ws.rs.GET;
 import jakarta.ws.rs.POST;
 import jakarta.ws.rs.Path;
@@ -74,4 +75,13 @@ public void getHandler(
         String remoteUri = routingTargetHandler.getRoutingDestination(servletRequest);
         proxyRequestHandler.getRequest(servletRequest, asyncResponse, URI.create(remoteUri));
     }
+
+    @DELETE
+    public void deleteHandler(
+            @Context HttpServletRequest servletRequest,
+            @Suspended AsyncResponse asyncResponse)
+    {
+        String remoteUri = routingTargetHandler.getRoutingDestination(servletRequest);
+        proxyRequestHandler.deleteRequest(servletRequest, asyncResponse, URI.create(remoteUri));
+    }
 }
diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaMultipleBackend.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaMultipleBackend.java
@@ -15,6 +15,8 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableMap;
+import io.airlift.json.JsonCodec;
+import io.trino.client.QueryResults;
 import io.trino.gateway.ha.config.ProxyBackendConfiguration;
 import io.trino.gateway.ha.router.GatewayCookie;
 import io.trino.gateway.ha.router.OAuth2GatewayCookie;
@@ -153,6 +155,31 @@ public void testQueryDeliveryToMultipleRoutingGroups()
         assertThat(response4.body().string()).contains("http://localhost:" + routerPort);
     }
 
+    @Test
+    public void testDeleteQueryId()
+            throws IOException
+    {
+        RequestBody requestBody =
+                RequestBody.create(MediaType.parse("application/json; charset=utf-8"), "SELECT 1");
+        Request request =
+                new Request.Builder()
+                        .url("http://localhost:" + routerPort + "/v1/statement")
+                        .addHeader("X-Trino-User", "test")
+                        .post(requestBody)
+                        .addHeader("X-Trino-Routing-Group", "scheduled")
+                        .build();
+        Response response = httpClient.newCall(request).execute();
+        JsonCodec<QueryResults> responseCodec = JsonCodec.jsonCodec(QueryResults.class);
+        QueryResults queryResults = responseCodec.fromJson(response.body().string());
+
+        Request deleteRequest = new Request.Builder()
+                .url(queryResults.getNextUri().toURL())
+                .delete()
+                .build();
+        Response deleteResponse = httpClient.newCall(deleteRequest).execute();
+        assertThat(deleteResponse.code()).isBetween(200, 204);
+    }
+
     @Test
     public void testBackendConfiguration()
             throws Exception