Extract query ID from all kill_query procedure variations

gateway-ha/src/main/java/io/trino/gateway/ha/config/RequestAnalyzerConfig.java

@@ -17,7 +17,7 @@
 
 public class RequestAnalyzerConfig
 {
-    private Integer maxBodySize = 1_000_000;
+    private int maxBodySize = 1_000_000;
 
     private boolean isClientsUseV2Format;
     private String tokenUserField = "email";
@@ -26,13 +26,13 @@ public class RequestAnalyzerConfig
 
     public RequestAnalyzerConfig() {}
 
-    public Integer getMaxBodySize()
+    public int getMaxBodySize()
     {
         return maxBodySize;
     }
 
     @Max(Integer.MAX_VALUE)
-    public void setMaxBodySize(Integer maxBodySize)
+    public void setMaxBodySize(int maxBodySize)
     {
         this.maxBodySize = maxBodySize;
     }

gateway-ha/src/main/java/io/trino/gateway/ha/handler/ProxyUtils.java

@@ -15,8 +15,11 @@
 
 import com.google.common.io.CharStreams;
 import io.airlift.log.Logger;
+import io.trino.gateway.ha.router.TrinoQueryProperties;
 import jakarta.servlet.http.HttpServletRequest;
+import jakarta.ws.rs.HttpMethod;
 
+import java.io.IOException;
 import java.io.InputStreamReader;
 import java.util.List;
 import java.util.Optional;
@@ -49,51 +52,50 @@ public final class ProxyUtils
      * capitalization.
      */
     private static final Pattern QUERY_ID_PARAM_PATTERN = Pattern.compile(".*(?:%2F|(?i)query_?id(?-i)=|^)(\\d+_\\d+_\\d+_\\w+).*");
-    private static final Pattern EXTRACT_BETWEEN_SINGLE_QUOTES = Pattern.compile("'([^\\s']+)'");
 
     private ProxyUtils() {}
 
-    public static String extractQueryIdIfPresent(HttpServletRequest request, List<String> statementPaths)
+    public static Optional<String> extractQueryIdIfPresent(
+            HttpServletRequest request,
+            List<String> statementPaths,
+            boolean requestAnalyserClientsUseV2Format,
+            int requestAnalyserMaxBodySize)
     {
         String path = request.getRequestURI();
         String queryParams = request.getQueryString();
+        if (!request.getMethod().equals(HttpMethod.POST)) {
+            return extractQueryIdIfPresent(path, queryParams, statementPaths);
+        }
+        String queryText;
         try {
-            String queryText = CharStreams.toString(new InputStreamReader(request.getInputStream(), UTF_8));
-            if (!isNullOrEmpty(queryText)
-                    && queryText.toLowerCase(ENGLISH).contains("system.runtime.kill_query")) {
-                // extract and return the queryId
-                String[] parts = queryText.split(",");
-                for (String part : parts) {
-                    if (part.contains("query_id")) {
-                        Matcher matcher = EXTRACT_BETWEEN_SINGLE_QUOTES.matcher(part);
-                        if (matcher.find()) {
-                            String queryQuoted = matcher.group();
-                            if (!isNullOrEmpty(queryQuoted) && queryQuoted.length() > 0) {
-                                return queryQuoted.substring(1, queryQuoted.length() - 1);
-                            }
-                        }
-                    }
-                }
-            }
+            queryText = CharStreams.toString(new InputStreamReader(request.getInputStream(), UTF_8));
         }
-        catch (Exception e) {
-            log.error(e, "Error extracting query payload from request");
+        catch (IOException e) {
+            throw new RuntimeException("Error reading request body", e);
         }
-
-        return extractQueryIdIfPresent(path, queryParams, statementPaths);
+        if (!isNullOrEmpty(queryText) && queryText.toLowerCase(ENGLISH).contains("kill_query")) {
+            TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(request, requestAnalyserClientsUseV2Format, requestAnalyserMaxBodySize);
+            return trinoQueryProperties.getQueryId();
+        }
+        return Optional.empty();
     }
 
-    public static String extractQueryIdIfPresent(String path, String queryParams, List<String> statementPaths)
+    public static Optional<String> extractQueryIdIfPresent(String path, String queryParams, List<String> statementPaths)
     {
         if (path == null) {
-            return null;
+            return Optional.empty();
         }
-        String queryId = null;
         log.debug("Trying to extract query id from path [%s] or queryString [%s]", path, queryParams);
         // matchingStatementPath should match paths such as /v1/statement/executing/query_id/nonce/sequence_number,
         // and if custom paths are supplied using the statementPaths configuration, paths such as
         // /custom/statement/path/executing/query_id/nonce/sequence_number
         Optional<String> matchingStatementPath = statementPaths.stream().filter(path::startsWith).findAny();
+        if (!isNullOrEmpty(queryParams)) {
+            Matcher matcher = QUERY_ID_PARAM_PATTERN.matcher(queryParams);
+            if (matcher.matches()) {
+                return Optional.of(matcher.group(1));
+            }
+        }
         if (matchingStatementPath.isPresent() || path.startsWith(V1_QUERY_PATH)) {
             path = path.replace(matchingStatementPath.orElse(V1_QUERY_PATH), "");
             String[] tokens = path.split("/");
@@ -102,27 +104,20 @@ public static String extractQueryIdIfPresent(String path, String queryParams, Li
                         || tokens[1].equals("scheduled")
                         || tokens[1].equals("executing")
                         || tokens[1].equals("partialCancel")) {
-                    queryId = tokens[2];
+                    return Optional.of(tokens[2]);
                 }
                 else {
-                    queryId = tokens[1];
+                    return Optional.of(tokens[1]);
                 }
             }
         }
         else if (path.startsWith(TRINO_UI_PATH)) {
             Matcher matcher = QUERY_ID_PATH_PATTERN.matcher(path);
             if (matcher.matches()) {
-                queryId = matcher.group(1);
-            }
-        }
-        if (!isNullOrEmpty(queryParams)) {
-            Matcher matcher = QUERY_ID_PARAM_PATTERN.matcher(queryParams);
-            if (matcher.matches()) {
-                queryId = matcher.group(1);
+                return Optional.of(matcher.group(1));
             }
         }
-        log.debug("Query id in URL [%s]", queryId);
-        return queryId;
+        return Optional.empty();
     }
 
     public static String buildUriWithNewBackend(String backendHost, HttpServletRequest request)

gateway-ha/src/main/java/io/trino/gateway/ha/handler/RoutingTargetHandler.java

@@ -47,6 +47,8 @@ public class RoutingTargetHandler
     private final RoutingGroupSelector routingGroupSelector;
     private final List<String> statementPaths;
     private final List<Pattern> extraWhitelistPaths;
+    private final boolean requestAnalyserClientsUseV2Format;
+    private final int requestAnalyserMaxBodySize;
     private final boolean cookiesEnabled;
 
     @Inject
@@ -57,8 +59,10 @@ public RoutingTargetHandler(
     {
         this.routingManager = requireNonNull(routingManager);
         this.routingGroupSelector = requireNonNull(routingGroupSelector);
-        this.statementPaths = requireNonNull(haGatewayConfiguration.getStatementPaths());
-        this.extraWhitelistPaths = requireNonNull(haGatewayConfiguration.getExtraWhitelistPaths()).stream().map(Pattern::compile).collect(toImmutableList());
+        statementPaths = requireNonNull(haGatewayConfiguration.getStatementPaths());
+        extraWhitelistPaths = requireNonNull(haGatewayConfiguration.getExtraWhitelistPaths()).stream().map(Pattern::compile).collect(toImmutableList());
+        requestAnalyserClientsUseV2Format = haGatewayConfiguration.getRequestAnalyzerConfig().isClientsUseV2Format();
+        requestAnalyserMaxBodySize = haGatewayConfiguration.getRequestAnalyzerConfig().getMaxBodySize();
         cookiesEnabled = GatewayCookieConfigurationPropertiesProvider.getInstance().isEnabled();
     }
 
@@ -96,9 +100,9 @@ private String getBackendFromRoutingGroup(HttpServletRequest request)
 
     private Optional<String> getPreviousBackend(HttpServletRequest request)
     {
-        String queryId = extractQueryIdIfPresent(request, statementPaths);
-        if (!isNullOrEmpty(queryId)) {
-            return Optional.of(routingManager.findBackendForQueryId(queryId));
+        Optional<String> queryId = extractQueryIdIfPresent(request, statementPaths, requestAnalyserClientsUseV2Format, requestAnalyserMaxBodySize);
+        if (queryId.isPresent()) {
+            return queryId.map(routingManager::findBackendForQueryId);
         }
         if (cookiesEnabled && request.getCookies() != null) {
             List<GatewayCookie> cookies = Arrays.stream(request.getCookies())

gateway-ha/src/main/java/io/trino/gateway/ha/router/ExternalRoutingGroupSelector.java

@@ -118,7 +118,7 @@ private RoutingGroupExternalBody createRequestBody(HttpServletRequest request)
         TrinoQueryProperties trinoQueryProperties = null;
         TrinoRequestUser trinoRequestUser = null;
         if (requestAnalyzerConfig.isAnalyzeRequest()) {
-            trinoQueryProperties = new TrinoQueryProperties(request, requestAnalyzerConfig);
+            trinoQueryProperties = new TrinoQueryProperties(request, requestAnalyzerConfig.isClientsUseV2Format(), requestAnalyzerConfig.getMaxBodySize());
             trinoRequestUser = trinoRequestUserProvider.getInstance(request);
         }
 

gateway-ha/src/main/java/io/trino/gateway/ha/router/RuleReloadingRoutingGroupSelector.java

@@ -96,7 +96,10 @@ public String findRoutingGroup(HttpServletRequest request)
 
             facts.put("request", request);
             if (requestAnalyzerConfig.isAnalyzeRequest()) {
-                TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(request, requestAnalyzerConfig);
+                TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(
+                        request,
+                        requestAnalyzerConfig.isClientsUseV2Format(),
+                        requestAnalyzerConfig.getMaxBodySize());
                 TrinoRequestUser trinoRequestUser = trinoRequestUserProvider.getInstance(request);
                 facts.put("trinoQueryProperties", trinoQueryProperties);
                 facts.put("trinoRequestUser", trinoRequestUser);