Skip to content

Commit

Permalink
Fix query ui history redirect
Browse files Browse the repository at this point in the history
  • Loading branch information
@regadas
regadas committed Feb 5, 2023
1 parent 69ac14b commit 49dced5
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 21 deletions.
49 changes: 32 additions & 17 deletions gateway-ha/src/main/java/com/lyft/data/gateway/ha/handler/QueryIdCachingProxyHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,9 @@ 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 +116,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);
} 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 +161,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 +176,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,11 +186,22 @@ 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;
}
Expand All @@ -206,14 +220,15 @@ protected static String extractQueryIdIfPresent(String path, String queryParams)
queryId = tokens[3];
}
}
} else if (path.startsWith(PRESTO_UI_PATH)) {
Matcher matcher = QUERY_ID_PATTERN.matcher(path);
} else if (path.startsWith(PRESTO_UI_PATH) && queryParams != null) {
Matcher matcher = QUERY_ID_PATTERN.matcher(queryParams);
if (matcher.matches()) {
queryId = matcher.group(1);
queryId = matcher.group(0);
}
}

log.debug("query id in url [{}]", queryId);
return queryId;
return Optional.ofNullable(queryId);
}

protected void postConnectionHook(
Expand Down
9 changes: 5 additions & 4 deletions ...way-ha/src/test/java/com/lyft/data/gateway/ha/handler/TestQueryIdCachingProxyHandler.java
View file
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,15 @@ 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"));
}
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

0 comments on commit 49dced5

Please sign in to comment.