From 8bd5e90f3c92002c9fa4aa0f707cb7c14c3dd3a1 Mon Sep 17 00:00:00 2001 From: Will Morrison Date: Tue, 27 Feb 2024 11:38:05 -0500 Subject: [PATCH] Fix query id matching in request querystring --- .../handler/QueryIdCachingProxyHandler.java | 22 +++++++++- .../TestQueryIdCachingProxyHandler.java | 42 ++++++++++++------- 2 files changed, 46 insertions(+), 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 index 32671bb41..413eb282b 100644 --- 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 index ee4d88751..4e8cfd10f 100644 --- 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