Fix query id matching in request querystring

trinodb · Mar 6, 2024 · 8bd5e90 · 8bd5e90
1 parent 2717e47
commit 8bd5e90
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 18 deletions.
diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/handler/QueryIdCachingProxyHandler.java b/gateway-ha/src/main/java/io/trino/gateway/ha/handler/QueryIdCachingProxyHandler.java
@@ -60,7 +60,19 @@ public class QueryIdCachingProxyHandler
     public static final String SOURCE_HEADER = "X-Trino-Source";
     public static final String HOST_HEADER = "Host";
     private static final int QUERY_TEXT_LENGTH_FOR_HISTORY = 200;
-    private static final Pattern QUERY_ID_PATTERN = Pattern.compile(".*[/=?](\\d+_\\d+_\\d+_\\w+).*");
+    /**
+     * This regular expression matches query ids as they appear in the path of a URL. The query id must be preceded
+     * by a "/". A query id is defined as three groups of digits separated by underscores, with a final group
+     * consisting of any alphanumeric characters.
+     */
+    private static final Pattern QUERY_ID_PATH_PATTERN = Pattern.compile(".*/(\\d+_\\d+_\\d+_\\w+).*");
+    /**
+     * This regular expression matches query ids as they appear in the query parameters of a URL. The query id is
+     * defined as in QUERY_TEXT_LENGTH_FOR_HISTORY. The query id must either be at the beginning of the query parameter
+     * string, or be preceded by %2F (a URL-encoded "/"), or  "query_id=", with or without the underscore and any
+     * 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']+)'");
 
@@ -113,7 +125,13 @@ protected static String extractQueryIdIfPresent(String path, String queryParams)
             }
         }
         else if (path.startsWith(TRINO_UI_PATH)) {
-            Matcher matcher = QUERY_ID_PATTERN.matcher(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);
             }

diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestQueryIdCachingProxyHandler.java b/gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestQueryIdCachingProxyHandler.java
@@ -23,6 +23,7 @@
 
 import java.io.IOException;
 
+import static io.trino.gateway.ha.handler.QueryIdCachingProxyHandler.extractQueryIdIfPresent;
 import static java.lang.String.format;
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -33,22 +34,31 @@ public class TestQueryIdCachingProxyHandler
     public void testExtractQueryIdFromUrl()
             throws IOException
     {
-        String[] paths = {
-                "/ui/api/query/20200416_160256_03078_6b4yt",
-                "/ui/api/query/20200416_160256_03078_6b4yt?bogus_fictional_param",
-                "/ui/api/query?query_id=20200416_160256_03078_6b4yt",
-                "/ui/api/query.html?20200416_160256_03078_6b4yt"};
-        for (String path : paths) {
-            String queryId = QueryIdCachingProxyHandler.extractQueryIdIfPresent(path, null);
-            assertThat(queryId).isEqualTo("20200416_160256_03078_6b4yt");
-        }
-        String[] nonPaths = {
-                "/ui/api/query/myOtherThing",
-                "/ui/api/query/20200416_blah?bogus_fictional_param"};
-        for (String path : nonPaths) {
-            String queryId = QueryIdCachingProxyHandler.extractQueryIdIfPresent(path, null);
-            assertThat(queryId).isNull();
-        }
+        assertThat(extractQueryIdIfPresent("/v1/statement/executing/20200416_160256_03078_6b4yt/ya7e884929c67cdf86207a80e7a77ab2166fa2e7b/1368", null))
+                .isEqualTo("20200416_160256_03078_6b4yt");
+        assertThat(extractQueryIdIfPresent("/v1/statement/queued/20200416_160256_03078_6b4yt/y0d7620a6941e78d3950798a1085383234258a566/1", null))
+                .isEqualTo("20200416_160256_03078_6b4yt");
+        assertThat(extractQueryIdIfPresent("/ui/api/query/20200416_160256_03078_6b4yt", null))
+                .isEqualTo("20200416_160256_03078_6b4yt");
+        assertThat(extractQueryIdIfPresent("/ui/api/query/20200416_160256_03078_6b4yt/killed", null))
+                .isEqualTo("20200416_160256_03078_6b4yt");
+        assertThat(extractQueryIdIfPresent("/ui/api/query/20200416_160256_03078_6b4yt/preempted", null))
+                .isEqualTo("20200416_160256_03078_6b4yt");
+        assertThat(extractQueryIdIfPresent("/v1/query/20200416_160256_03078_6b4yt", "pretty"))
+                .isEqualTo("20200416_160256_03078_6b4yt");
+        assertThat(extractQueryIdIfPresent("/ui/troubleshooting", "queryId=20200416_160256_03078_6b4yt"))
+                .isEqualTo("20200416_160256_03078_6b4yt");
+        assertThat(extractQueryIdIfPresent("/ui/query.html", "20200416_160256_03078_6b4yt"))
+                .isEqualTo("20200416_160256_03078_6b4yt");
+        assertThat(extractQueryIdIfPresent("/login", "redirect=%2Fui%2Fapi%2Fquery%2F20200416_160256_03078_6b4yt"))
+                .isEqualTo("20200416_160256_03078_6b4yt");
+
+        assertThat(extractQueryIdIfPresent("/ui/api/query/myOtherThing", null))
+                .isNull();
+        assertThat(extractQueryIdIfPresent("/ui/api/query/20200416_blah", "bogus_fictional_param"))
+                .isNull();
+        assertThat(extractQueryIdIfPresent("/ui/", "lang=en&p=1&id=0_1_2_a"))
+                .isNull();
     }
 
     @Test