Skip to content
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

Fix query ui history redirect #193

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Fix query ui history redirect #193

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
81 changes: 53 additions & 28 deletions gateway-ha/src/main/java/com/lyft/data/gateway/ha/handler/QueryIdCachingProxyHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,10 @@ public class QueryIdCachingProxyHandler extends ProxyHandler {
public static final String SOURCE_HEADER = "X-Trino-Source";
public static final String ALTERNATE_SOURCE_HEADER = "X-Presto-Source";
public static final String HOST_HEADER = "Host";
public static final String REFERER_STRING = "Referer";
private static final int QUERY_TEXT_LENGTH_FOR_HISTORY = 200;
private static final Pattern QUERY_ID_PATTERN = Pattern.compile(".*[/=?](\\d+_\\d+_\\d+_\\w+).*");
private static final Pattern QUERY_ID_PATTERN =
Pattern.compile("(.*[/=?])?(\\d+_\\d+_\\d+_\\w+).*");

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

Expand Down Expand Up @@ -115,22 +117,23 @@ public String rewriteTarget(HttpServletRequest request) {

// Only load balance presto query APIs.
if (isPathWhiteListed(request.getRequestURI())) {
String queryId = extractQueryIdIfPresent(request);
Optional<String> queryId = extractQueryIdIfPresent(request);

backendAddress = queryId
// Find query id and get url from cache
if (!Strings.isNullOrEmpty(queryId)) {
backendAddress = routingManager.findBackendForQueryId(queryId);
} else {
.map(routingManager::findBackendForQueryId)
.orElseGet(() -> {
String routingGroup = routingGroupSelector.findRoutingGroup(request);
String user = Optional.ofNullable(request.getHeader(USER_HEADER))
.orElse(request.getHeader(ALTERNATE_USER_HEADER));
if (!Strings.isNullOrEmpty(routingGroup)) {
// This falls back on adhoc backend if there are no cluster found for the routing group.
backendAddress = routingManager.provideBackendForRoutingGroup(routingGroup, user);
return routingManager.provideBackendForRoutingGroup(routingGroup, user);
Copy link
Contributor

@puneetjaiswal puneetjaiswal Feb 8, 2023 •
edited
Loading

Choose a reason for hiding this comment

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

please dont return from here. Below lines are skipped. without this, subsequent calls would require a broadcast to all backends for query id lookup.

       // set target backend so that we could save queryId to backend mapping later.
      ((MultiReadHttpServletRequest) request).addHeader(PROXY_TARGET_HEADER, backendAddress);

would be nice if you could add a test case validating the presence of this header in the response.

Copy link
Author

Choose a reason for hiding this comment

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

Hi @puneetjaiswal thanks for taking a look. The above return is within the orElseGet lambda for the backendAddress assignment. Effectively, L137 doesn't get skipped. I'll add a test.

} else {
backendAddress = routingManager.provideAdhocBackend(user);
return routingManager.provideAdhocBackend(user);
}
}
});

// set target backend so that we could save queryId to backend mapping later.
((MultiReadHttpServletRequest) request).addHeader(PROXY_TARGET_HEADER, backendAddress);
}
Expand Down Expand Up @@ -159,7 +162,7 @@ public String rewriteTarget(HttpServletRequest request) {
return targetLocation;
}

protected String extractQueryIdIfPresent(HttpServletRequest request) {
protected Optional<String> extractQueryIdIfPresent(HttpServletRequest request) {
String path = request.getRequestURI();
String queryParams = request.getQueryString();
try {
Expand All @@ -174,7 +177,8 @@ protected String extractQueryIdIfPresent(HttpServletRequest request) {
if (m.find()) {
String queryQuoted = m.group();
if (!Strings.isNullOrEmpty(queryQuoted) && queryQuoted.length() > 0) {
return queryQuoted.substring(1, queryQuoted.length() - 1);
String res = queryQuoted.substring(1, queryQuoted.length() - 1);
return Optional.of(res).filter(s -> !s.isEmpty());
}
}
}
Expand All @@ -183,37 +187,58 @@ protected String extractQueryIdIfPresent(HttpServletRequest request) {
} catch (Exception e) {
log.error("Error extracting query payload from request", e);
}

return extractQueryIdIfPresent(path, queryParams);

return extractQueryIdIfPresent(path, queryParams).or(() -> {
return Optional.ofNullable(request.getHeader(REFERER_STRING))
.flatMap(referer -> {
try {
URI uri = URI.create(referer);
return extractQueryIdIfPresent(uri.getPath(), uri.getQuery());
} catch (Exception e) {
log.error("Error extracting query id from Referer header", e);
}
return Optional.empty();
});
}).filter(s -> !s.isEmpty());
}

protected static String extractQueryIdIfPresent(String path, String queryParams) {
protected static Optional<String> extractQueryIdIfPresent(String path, String queryParams) {
if (path == null) {
return null;
return Optional.empty();
}
String queryId = null;

log.debug("trying to extract query id from path [{}] or queryString [{}]", path, queryParams);
if (path.startsWith(V1_STATEMENT_PATH) || path.startsWith(V1_QUERY_PATH)) {
String[] tokens = path.split("/");
if (tokens.length >= 4) {
if (path.contains("queued")
|| path.contains("scheduled")
|| path.contains("executing")
|| path.contains("partialCancel")) {
queryId = tokens[4];
} else {
queryId = tokens[3];
}
if (tokens.length < 4) {
return Optional.empty();
}
} else if (path.startsWith(PRESTO_UI_PATH)) {

if (path.contains("queued")
|| path.contains("scheduled")
|| path.contains("executing")
|| path.contains("partialCancel")) {
return Optional.of(tokens[4]);
}

return Optional.of(tokens[3]);
}

if (path.startsWith(PRESTO_UI_PATH)) {
Matcher matcher = QUERY_ID_PATTERN.matcher(path);
if (matcher.matches()) {
queryId = matcher.group(1);
return Optional.of(matcher.group(2));
}
}

if (queryParams != null) {
Matcher matcher = QUERY_ID_PATTERN.matcher(queryParams);
if (matcher.matches()) {
return Optional.of(matcher.group(2));
}
}
log.debug("query id in url [{}]", queryId);
return queryId;

return Optional.empty();
}

protected void postConnectionHook(
Expand Down
15 changes: 11 additions & 4 deletions ...way-ha/src/test/java/com/lyft/data/gateway/ha/handler/TestQueryIdCachingProxyHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import static org.testng.Assert.assertNull;

import java.io.IOException;
import java.util.Optional;

import javax.servlet.http.HttpServletRequest;

Expand All @@ -23,15 +24,21 @@ public void testExtractQueryIdFromUrl() throws IOException {
"/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);
assertEquals(queryId, "20200416_160256_03078_6b4yt");
Optional<String> queryId = QueryIdCachingProxyHandler.extractQueryIdIfPresent(path, null);
assertEquals(queryId, Optional.of("20200416_160256_03078_6b4yt"));
}

Optional<String> queryParamId = QueryIdCachingProxyHandler.extractQueryIdIfPresent(
"/ui/query",
"20200416_160256_03078_6b4yt");
assertEquals(queryParamId, Optional.of("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);
assertNull(queryId);
Optional<String> queryId = QueryIdCachingProxyHandler.extractQueryIdIfPresent(path, null);
assertEquals(queryId, Optional.empty());
}
}

Expand Down