Show RoutingGroup info in query history

Author: andythsu

gateway-ha/src/main/java/io/trino/gateway/ha/handler/ProxyUtils.java

@@ -21,6 +21,7 @@
 
 import java.io.IOException;
 import java.io.InputStreamReader;
+import java.net.URI;
 import java.util.List;
 import java.util.Optional;
 import java.util.regex.Matcher;
@@ -118,8 +119,8 @@ else if (path.startsWith(TRINO_UI_PATH)) {
         return Optional.empty();
     }
 
-    public static String buildUriWithNewBackend(String backendHost, HttpServletRequest request)
+    public static URI buildUriWithNewCluster(String backendHost, HttpServletRequest request)
     {
-        return backendHost + request.getRequestURI() + (request.getQueryString() != null ? "?" + request.getQueryString() : "");
+        return URI.create(backendHost + request.getRequestURI() + (request.getQueryString() != null ? "?" + request.getQueryString() : ""));
     }
 }

gateway-ha/src/main/java/io/trino/gateway/ha/handler/RoutingDestination.java

@@ -0,0 +1,18 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.gateway.ha.handler;
+
+import java.net.URI;
+
+public record RoutingDestination(String routingGroup, String clusterHost, URI clusterUri) {}

gateway-ha/src/main/java/io/trino/gateway/ha/handler/RoutingTargetHandler.java

@@ -36,7 +36,7 @@
 import static io.trino.gateway.ha.handler.HttpUtils.V1_INFO_PATH;
 import static io.trino.gateway.ha.handler.HttpUtils.V1_NODE_PATH;
 import static io.trino.gateway.ha.handler.HttpUtils.V1_QUERY_PATH;
-import static io.trino.gateway.ha.handler.ProxyUtils.buildUriWithNewBackend;
+import static io.trino.gateway.ha.handler.ProxyUtils.buildUriWithNewCluster;
 import static io.trino.gateway.ha.handler.ProxyUtils.extractQueryIdIfPresent;
 import static java.util.Objects.requireNonNull;
 
@@ -66,13 +66,27 @@ public RoutingTargetHandler(
         cookiesEnabled = GatewayCookieConfigurationPropertiesProvider.getInstance().isEnabled();
     }
 
-    public String getRoutingDestination(HttpServletRequest request)
+    public RoutingDestination getRoutingDestination(HttpServletRequest request)
     {
-        Optional<String> previousBackend = getPreviousBackend(request);
-        String clusterHost = previousBackend.orElseGet(() -> getBackendFromRoutingGroup(request));
-        logRewrite(clusterHost, request);
+        Optional<String> queryId = extractQueryIdIfPresent(request, statementPaths, requestAnalyserClientsUseV2Format, requestAnalyserMaxBodySize);
+        Optional<String> previousCluster = getPreviousCluster(queryId, request);
+        RoutingDestination routingDestination = previousCluster.map(cluster -> {
+            String routingGroup = queryId.map(routingManager::findRoutingGroupForQueryId)
+                    .orElse("adhoc");
+            return new RoutingDestination(routingGroup, cluster, buildUriWithNewCluster(cluster, request));
+        }).orElse(getClusterFromRoutingGroup(request));
+        logRewrite(routingDestination.clusterHost(), request);
+        return routingDestination;
+    }
 
-        return buildUriWithNewBackend(clusterHost, request);
+    private RoutingDestination getClusterFromRoutingGroup(HttpServletRequest request)
+    {
+        Optional<String> routingGroup = routingGroupSelector.findRoutingGroup(request);
+        String user = request.getHeader(USER_HEADER);
+        // This falls back on adhoc routing group if there is no cluster found for the routing group.
+        String group = routingGroup.orElse("adhoc");
+        String clusterHost = routingManager.provideClusterForRoutingGroup(group, user);
+        return new RoutingDestination(group, clusterHost, buildUriWithNewCluster(clusterHost, request));
     }
 
     public boolean isPathWhiteListed(String path)
@@ -87,20 +101,8 @@ public boolean isPathWhiteListed(String path)
                 || extraWhitelistPaths.stream().anyMatch(pattern -> pattern.matcher(path).matches());
     }
 
-    private String getBackendFromRoutingGroup(HttpServletRequest request)
-    {
-        String routingGroup = routingGroupSelector.findRoutingGroup(request);
-        String user = request.getHeader(USER_HEADER);
-        if (!isNullOrEmpty(routingGroup)) {
-            // This falls back on adhoc backend if there is no cluster found for the routing group.
-            return routingManager.provideBackendForRoutingGroup(routingGroup, user);
-        }
-        return routingManager.provideAdhocBackend(user);
-    }
-
-    private Optional<String> getPreviousBackend(HttpServletRequest request)
+    private Optional<String> getPreviousCluster(Optional<String> queryId, HttpServletRequest request)
     {
-        Optional<String> queryId = extractQueryIdIfPresent(request, statementPaths, requestAnalyserClientsUseV2Format, requestAnalyserMaxBodySize);
         if (queryId.isPresent()) {
             return queryId.map(routingManager::findBackendForQueryId);
         }
@@ -129,6 +131,6 @@ private void logRewrite(String newBackend, HttpServletRequest request)
                 request.getServerPort(),
                 request.getRequestURI(),
                 (request.getQueryString() != null ? "?" + request.getQueryString() : ""),
-                buildUriWithNewBackend(newBackend, request));
+                buildUriWithNewCluster(newBackend, request));
     }
 }

gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/QueryHistory.java

@@ -24,7 +24,8 @@ public record QueryHistory(
         @ColumnName("backend_url") String backendUrl,
         @ColumnName("user_name") @Nullable String userName,
         @ColumnName("source") @Nullable String source,
-        @ColumnName("created") long created)
+        @ColumnName("created") long created,
+        @ColumnName("routing_group") String routingGroup)
 {
     public QueryHistory
     {

gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/QueryHistoryDao.java

@@ -77,6 +77,12 @@ default List<QueryHistory> findRecentQueriesByUserName(String userName, boolean
             """)
     String findBackendUrlByQueryId(String queryId);
 
+    @SqlQuery("""
+            SELECT routing_group FROM query_history
+            WHERE query_id = :queryId
+            """)
+    String findRoutingGroupByQueryId(String queryId);
+
     @SqlQuery("""
             SELECT * FROM query_history
             WHERE 1 = 1 <condition>
@@ -104,10 +110,10 @@ GROUP BY FLOOR(created / 1000 / 60), backend_url
     List<Map<String, Object>> findDistribution(long created);
 
     @SqlUpdate("""
-            INSERT INTO query_history (query_id, query_text, backend_url, user_name, source, created)
-            VALUES (:queryId, :queryText, :backendUrl, :userName, :source, :created)
+            INSERT INTO query_history (query_id, query_text, backend_url, user_name, source, created, routing_group)
+            VALUES (:queryId, :queryText, :backendUrl, :userName, :source, :created, :routingGroup)
             """)
-    void insertHistory(String queryId, String queryText, String backendUrl, String userName, String source, long created);
+    void insertHistory(String queryId, String queryText, String backendUrl, String userName, String source, long created, String routingGroup);
 
     @SqlUpdate("""
             DELETE FROM query_history

gateway-ha/src/main/java/io/trino/gateway/ha/router/ExternalRoutingGroupSelector.java

@@ -78,7 +78,7 @@ public class ExternalRoutingGroupSelector
     }
 
     @Override
-    public String findRoutingGroup(HttpServletRequest servletRequest)
+    public Optional<String> findRoutingGroup(HttpServletRequest servletRequest)
     {
         try {
             RoutingGroupExternalBody requestBody = createRequestBody(servletRequest);
@@ -100,13 +100,13 @@ public String findRoutingGroup(HttpServletRequest servletRequest)
             else if (response.errors() != null && !response.errors().isEmpty()) {
                 throw new RuntimeException("Response with error: " + String.join(", ", response.errors()));
             }
-            return response.routingGroup();
+            return Optional.ofNullable(response.routingGroup());
         }
         catch (Exception e) {
             log.error(e, "Error occurred while retrieving routing group "
                     + "from external routing rules processing at " + uri);
         }
-        return servletRequest.getHeader(ROUTING_GROUP_HEADER);
+        return Optional.ofNullable(servletRequest.getHeader(ROUTING_GROUP_HEADER));
     }
 
     private RoutingGroupExternalBody createRequestBody(HttpServletRequest request)

gateway-ha/src/main/java/io/trino/gateway/ha/router/FileBasedRoutingGroupSelector.java

@@ -30,6 +30,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 
 import static com.google.common.base.Suppliers.memoizeWithExpiration;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -60,7 +61,7 @@ public FileBasedRoutingGroupSelector(String rulesPath, Duration rulesRefreshPeri
     }
 
     @Override
-    public String findRoutingGroup(HttpServletRequest request)
+    public Optional<String> findRoutingGroup(HttpServletRequest request)
     {
         Map<String, String> result = new HashMap<>();
         Map<String, Object> state = new HashMap<>();
@@ -84,7 +85,7 @@ public String findRoutingGroup(HttpServletRequest request)
                 rule.evaluateAction(result, data, state);
             }
         });
-        return result.get(RESULTS_ROUTING_GROUP_KEY);
+        return Optional.ofNullable(result.get(RESULTS_ROUTING_GROUP_KEY));
     }
 
     public List<RoutingRule> readRulesFromPath(Path rulesPath)

gateway-ha/src/main/java/io/trino/gateway/ha/router/HaQueryHistoryManager.java

@@ -60,7 +60,8 @@ public void submitQueryDetail(QueryDetail queryDetail)
                 queryDetail.getBackendUrl(),
                 queryDetail.getUser(),
                 queryDetail.getSource(),
-                queryDetail.getCaptureTime());
+                queryDetail.getCaptureTime(),
+                queryDetail.getRoutingGroup());
     }
 
     @Override
@@ -87,6 +88,7 @@ private static List<QueryHistoryManager.QueryDetail> upcast(List<QueryHistory> q
             queryDetail.setBackendUrl(dao.backendUrl());
             queryDetail.setUser(dao.userName());
             queryDetail.setSource(dao.source());
+            queryDetail.setRoutingGroup(dao.routingGroup());
             queryDetails.add(queryDetail);
         }
         return queryDetails;
@@ -98,6 +100,12 @@ public String getBackendForQueryId(String queryId)
         return dao.findBackendUrlByQueryId(queryId);
     }
 
+    @Override
+    public String getRoutingGroupForQueryId(String queryId)
+    {
+        return dao.findRoutingGroupByQueryId(queryId);
+    }
+
     @Override
     public TableData<QueryDetail> findQueryHistory(QueryHistoryRequest query)
     {

gateway-ha/src/main/java/io/trino/gateway/ha/router/QueryCountBasedRouter.java

@@ -222,16 +222,16 @@ private synchronized Optional<String> getBackendForRoutingGroup(String routingGr
     }
 
     @Override
-    public String provideAdhocBackend(String user)
+    public String provideAdhocCluster(String user)
     {
         return getBackendForRoutingGroup("adhoc", user).orElseThrow(() -> new RouterException("did not find any cluster for the adhoc routing group"));
     }
 
     @Override
-    public String provideBackendForRoutingGroup(String routingGroup, String user)
+    public String provideClusterForRoutingGroup(String routingGroup, String user)
     {
         return getBackendForRoutingGroup(routingGroup, user)
-                .orElseGet(() -> provideAdhocBackend(user));
+                .orElse(provideAdhocCluster(user));
     }
 
     @Override

gateway-ha/src/main/java/io/trino/gateway/ha/router/QueryHistoryManager.java

@@ -32,6 +32,8 @@ public interface QueryHistoryManager
 
     String getBackendForQueryId(String queryId);
 
+    String getRoutingGroupForQueryId(String queryId);
+
     TableData<QueryDetail> findQueryHistory(QueryHistoryRequest query);
 
     List<DistributionResponse.LineChart> findDistribution(Long ts);
@@ -45,6 +47,7 @@ class QueryDetail
         private String source;
         private String backendUrl;
         private long captureTime;
+        private String routingGroup;
 
         public QueryDetail() {}
 
@@ -125,6 +128,17 @@ public void setCaptureTime(long captureTime)
             this.captureTime = captureTime;
         }
 
+        @JsonProperty
+        public String getRoutingGroup()
+        {
+            return this.routingGroup;
+        }
+
+        public void setRoutingGroup(String routingGroup)
+        {
+            this.routingGroup = routingGroup;
+        }
+
         @Override
         public boolean equals(Object o)
         {
@@ -140,13 +154,14 @@ public boolean equals(Object o)
                     Objects.equals(queryText, that.queryText) &&
                     Objects.equals(user, that.user) &&
                     Objects.equals(source, that.source) &&
-                    Objects.equals(backendUrl, that.backendUrl);
+                    Objects.equals(backendUrl, that.backendUrl) &&
+                    Objects.equals(routingGroup, that.routingGroup);
         }
 
         @Override
         public int hashCode()
         {
-            return Objects.hash(queryId, queryText, user, source, backendUrl, captureTime);
+            return Objects.hash(queryId, queryText, user, source, backendUrl, captureTime, routingGroup);
         }
 
         @Override
@@ -159,6 +174,7 @@ public String toString()
                     .add("source", source)
                     .add("backendUrl", backendUrl)
                     .add("captureTime", captureTime)
+                    .add("routingGroup", routingGroup)
                     .toString();
         }
     }

gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingGroupSelector.java

@@ -19,6 +19,8 @@
 import io.trino.gateway.ha.config.RulesExternalConfiguration;
 import jakarta.servlet.http.HttpServletRequest;
 
+import java.util.Optional;
+
 /**
  * RoutingGroupSelector provides a way to match an HTTP request to a Gateway routing group.
  */
@@ -32,7 +34,7 @@ public interface RoutingGroupSelector
      */
     static RoutingGroupSelector byRoutingGroupHeader()
     {
-        return request -> request.getHeader(ROUTING_GROUP_HEADER);
+        return request -> Optional.ofNullable(request.getHeader(ROUTING_GROUP_HEADER));
     }
 
     /**
@@ -60,5 +62,5 @@ static RoutingGroupSelector byRoutingExternal(
      * Given an HTTP request find a routing group to direct the request to. If a routing group cannot
      * be determined return null.
      */
-    String findRoutingGroup(HttpServletRequest request);
+    Optional<String> findRoutingGroup(HttpServletRequest request);
 }