Skip to content

Fix query id extraction from request query #265

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mosabua merged 1 commit into from
Mar 6, 2024
Merged

Fix query id extraction from request query #265

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions gateway-ha/src/main/java/io/trino/gateway/ha/handler/QueryIdCachingProxyHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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+).*");
Copy link
Copy Markdown
Member

@mosabua mosabua Feb 29, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh how I hate regex like this .. should we somehow add some Java doc that explains it a little bit .. or do we just treat the test as the docs?


private static final Pattern EXTRACT_BETWEEN_SINGLE_QUOTES = Pattern.compile("'([^\\s']+)'");

Expand Down Expand Up @@ -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);
Comment thread
mosabua marked this conversation as resolved.
if (matcher.matches()) {
queryId = matcher.group(1);
}
Expand Down
42 changes: 26 additions & 16 deletions gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestQueryIdCachingProxyHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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
Expand Down