Fix partialCancel statement path redirection

Author: maswin

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

@@ -13,6 +13,7 @@
  */
 package io.trino.gateway.ha.handler;
 
+import com.google.common.collect.ImmutableSet;
 import com.google.common.io.CharStreams;
 import io.airlift.log.Logger;
 import io.trino.gateway.ha.router.TrinoQueryProperties;
@@ -52,6 +53,8 @@ public final class ProxyUtils
      * capitalization.
      */
     private static final Pattern QUERY_ID_PARAM_PATTERN = Pattern.compile(".*(?:%2F|(?i)query_?id(?-i)=|^)(\\d+_\\d+_\\d+_\\w+).*");
+    public static final ImmutableSet<String> QUERY_STATE_PATH = ImmutableSet.of("queued", "scheduled", "executing");
+    public static final String PARTIAL_CANCEL_PATH = "partialCancel";
 
     private ProxyUtils() {}
 
@@ -100,10 +103,10 @@ public static Optional<String> extractQueryIdIfPresent(String path, String query
             path = path.replace(matchingStatementPath.orElse(V1_QUERY_PATH), "");
             String[] tokens = path.split("/");
             if (tokens.length >= 2) {
-                if (tokens[1].equals("queued")
-                        || tokens[1].equals("scheduled")
-                        || tokens[1].equals("executing")
-                        || tokens[1].equals("partialCancel")) {
+                if (tokens.length >= 3 && QUERY_STATE_PATH.contains(tokens[1])) {
+                    if (tokens.length >= 4 && tokens[2].equals(PARTIAL_CANCEL_PATH)) {
+                        return Optional.of(tokens[3]);
+                    }
                     return Optional.of(tokens[2]);
                 }
                 return Optional.of(tokens[1]);