diff --git a/docs/routing-rules.md b/docs/routing-rules.md index 8133dbaec..a5055a1d8 100644 --- a/docs/routing-rules.md +++ b/docs/routing-rules.md @@ -101,7 +101,7 @@ return a result with the following criteria: * Response status code of OK (200) * Message in JSON format -* Only one group can be returned +* Only one group or one cluster can be returned * If `errors` is not null, the query is routed to the configured default group #### Request headers modification @@ -126,7 +126,16 @@ or setting client tags before the request reaches the Trino cluster. } } ``` - +```json +{ + "routingCluster": "test-cluster", + "errors": [ + "Error1", + "Error2", + "Error3" + ] +} +``` ### Configure routing rules with a file Rules consist of a name, description, condition, and list @@ -158,10 +167,10 @@ In addition to the default objects, rules may optionally utilize [trinoRequestUser](#trinorequestuser) and [trinoQueryProperties](#trinoqueryproperties) , which provide information about the user and query respectively. -You must include an action of the form `result.put(\"routingGroup\", \"foo\")` -to trigger routing of a request that satisfies the condition to the specific -routing group. Without this action, the configured default group is used and the -whole routing rule is redundant. +You must include an action of the form `result.put(\"routingGroup\", \"foo\")` or +`result.put(\"routingCluster\", \"bar\")` to trigger routing of a request that satisfies +the condition to the specific routing group. Without this action, the configured default +group is used and the whole routing rule is redundant. The condition and actions are written in [MVEL](http://mvel.documentnode.com/), an expression language with Java-like syntax. Classes from `java.util`, data-type @@ -373,6 +382,13 @@ priority: 1 condition: 'request.getHeader("X-Trino-Source") == "airflow" && request.getHeader("X-Trino-Client-Tags") contains "label=special"' actions: - 'result.put("routingGroup", "etl-special")' +--- +name: "airflow cluster" +description: "query can also be pinned to a specific cluster" +priority: 10 +condition: 'request.getHeader("X-Trino-Source") == "airflow" && request.getHeader("X-Trino-Client-Tags") contains "label=airflow-cluster"' +actions: + - 'result.put("routingCluster", "airflow-cluster")' ``` Note that both rules still fire. The difference is that you are guaranteed @@ -380,8 +396,14 @@ that the first rule (priority 0) is fired before the second rule (priority 1). Thus `routingGroup` is set to `etl` and then to `etl-special`, so the `routingGroup` is always `etl-special` in the end. +When mixing cluster and group actions, the same rule priority semantics apply. +If a higher-priority rule (evaluated later) sets `routingCluster`, it overwrites +any previously set group, and vice versa. In practice, the +last assignment wins and `RoutingTargetHandler` inspects the resulting +`RoutingDecision`, using the cluster when both values are present. + More specific rules must be set to a higher priority so they are evaluated last -to set a `routingGroup`. +to set a `routingGroup` or a `routingCluster`. ##### Passing State diff --git a/gateway-ha/config.yaml b/gateway-ha/config.yaml index c9e0c4a29..edc8dc3ee 100644 --- a/gateway-ha/config.yaml +++ b/gateway-ha/config.yaml @@ -4,7 +4,7 @@ serverConfig: routingRules: rulesEngineEnabled: False - # rulesConfigPath: "src/main/resources/rules/routing_rules.yml" + # rulesConfigPath: "gateway-ha/src/main/resources/rules/routing_rules.yml" dataStore: jdbcUrl: jdbc:postgresql://localhost:5432/trino_gateway_db diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/handler/RoutingTargetHandler.java b/gateway-ha/src/main/java/io/trino/gateway/ha/handler/RoutingTargetHandler.java index 16521c44e..fdbe746d9 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/handler/RoutingTargetHandler.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/handler/RoutingTargetHandler.java @@ -21,8 +21,8 @@ import io.trino.gateway.ha.handler.schema.RoutingDestination; import io.trino.gateway.ha.handler.schema.RoutingTargetResponse; import io.trino.gateway.ha.router.GatewayCookie; -import io.trino.gateway.ha.router.RoutingGroupSelector; import io.trino.gateway.ha.router.RoutingManager; +import io.trino.gateway.ha.router.RoutingSelector; import io.trino.gateway.ha.router.schema.RoutingSelectorResponse; import jakarta.servlet.http.HttpServletRequest; import jakarta.servlet.http.HttpServletRequestWrapper; @@ -45,7 +45,7 @@ public class RoutingTargetHandler { private static final Logger log = Logger.get(RoutingTargetHandler.class); private final RoutingManager routingManager; - private final RoutingGroupSelector routingGroupSelector; + private final RoutingSelector routingSelector; private final String defaultRoutingGroup; private final List statementPaths; private final boolean requestAnalyserClientsUseV2Format; @@ -55,11 +55,11 @@ public class RoutingTargetHandler @Inject public RoutingTargetHandler( RoutingManager routingManager, - RoutingGroupSelector routingGroupSelector, + RoutingSelector routingSelector, HaGatewayConfiguration haGatewayConfiguration) { this.routingManager = requireNonNull(routingManager); - this.routingGroupSelector = requireNonNull(routingGroupSelector); + this.routingSelector = requireNonNull(routingSelector); this.defaultRoutingGroup = haGatewayConfiguration.getRouting().getDefaultRoutingGroup(); statementPaths = requireNonNull(haGatewayConfiguration.getStatementPaths()); requestAnalyserClientsUseV2Format = haGatewayConfiguration.getRequestAnalyzerConfig().isClientsUseV2Format(); @@ -73,12 +73,12 @@ public RoutingTargetResponse resolveRouting(HttpServletRequest request) Optional previousCluster = getPreviousCluster(queryId, request); RoutingTargetResponse routingTargetResponse = previousCluster.map(cluster -> { - String routingGroup = queryId.map(routingManager::findRoutingGroupForQueryId) + String routingDecision = queryId.map(routingManager::findRoutingDecisionForQueryId) .orElse(defaultRoutingGroup); String externalUrl = queryId.map(routingManager::findExternalUrlForQueryId) .orElse(cluster); return new RoutingTargetResponse( - new RoutingDestination(routingGroup, cluster, buildUriWithNewCluster(cluster, request), externalUrl), + new RoutingDestination(routingDecision, cluster, buildUriWithNewCluster(cluster, request), externalUrl), request); }).orElse(getRoutingTargetResponse(request)); @@ -88,14 +88,15 @@ public RoutingTargetResponse resolveRouting(HttpServletRequest request) private RoutingTargetResponse getRoutingTargetResponse(HttpServletRequest request) { - RoutingSelectorResponse routingDestination = routingGroupSelector.findRoutingDestination(request); + RoutingSelectorResponse routingDestination = routingSelector.findRoutingDestination(request); String user = request.getHeader(USER_HEADER); // This falls back on default routing group backend if there is no cluster found for the routing group. + String routingCluster = routingDestination.routingCluster(); String routingGroup = !isNullOrEmpty(routingDestination.routingGroup()) ? routingDestination.routingGroup() : defaultRoutingGroup; - ProxyBackendConfiguration backendConfiguration = routingManager.provideBackendConfiguration(routingGroup, user); + ProxyBackendConfiguration backendConfiguration = routingManager.provideBackendConfiguration(routingGroup, routingCluster, user); String clusterHost = backendConfiguration.getProxyTo(); String externalUrl = backendConfiguration.getExternalUrl(); // Apply headers from RoutingDestination if there are any @@ -103,8 +104,10 @@ private RoutingTargetResponse getRoutingTargetResponse(HttpServletRequest reques if (!routingDestination.externalHeaders().isEmpty()) { modifiedRequest = new HeaderModifyingRequestWrapper(request, routingDestination.externalHeaders()); } + // routingCluster and routingGroup are mutually exclusive. If neither is set, fall back to the default routing group. + String routingDecision = !isNullOrEmpty(routingCluster) ? routingCluster : routingGroup; return new RoutingTargetResponse( - new RoutingDestination(routingGroup, clusterHost, buildUriWithNewCluster(clusterHost, request), externalUrl), + new RoutingDestination(routingDecision, clusterHost, buildUriWithNewCluster(clusterHost, request), externalUrl), modifiedRequest); } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/handler/schema/RoutingDestination.java b/gateway-ha/src/main/java/io/trino/gateway/ha/handler/schema/RoutingDestination.java index 7060110d4..3c9e7f679 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/handler/schema/RoutingDestination.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/handler/schema/RoutingDestination.java @@ -15,4 +15,4 @@ import java.net.URI; -public record RoutingDestination(String routingGroup, String clusterHost, URI clusterUri, String externalUrl) {} +public record RoutingDestination(String routingDecision, String clusterHost, URI clusterUri, String externalUrl) {} diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java b/gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java index 2973655b8..df703a901 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java @@ -50,8 +50,8 @@ import io.trino.gateway.ha.router.PathFilter; import io.trino.gateway.ha.router.QueryHistoryManager; import io.trino.gateway.ha.router.ResourceGroupsManager; -import io.trino.gateway.ha.router.RoutingGroupSelector; import io.trino.gateway.ha.router.RoutingManager; +import io.trino.gateway.ha.router.RoutingSelector; import io.trino.gateway.ha.security.ApiAuthenticator; import io.trino.gateway.ha.security.AuthorizationManager; import io.trino.gateway.ha.security.BasicAuthFilter; @@ -212,27 +212,27 @@ public AuthorizationManager getAuthorizationManager() @Provides @Singleton - public RoutingGroupSelector getRoutingGroupSelector(@ForRouter HttpClient httpClient) + public RoutingSelector getRoutingSelector(@ForRouter HttpClient httpClient) { RoutingRulesConfiguration routingRulesConfig = configuration.getRoutingRules(); if (routingRulesConfig.isRulesEngineEnabled()) { try { return switch (routingRulesConfig.getRulesType()) { - case FILE -> RoutingGroupSelector.byRoutingRulesEngine( + case FILE -> RoutingSelector.byRoutingRulesEngine( routingRulesConfig.getRulesConfigPath(), routingRulesConfig.getRulesRefreshPeriod(), configuration.getRequestAnalyzerConfig()); case EXTERNAL -> { RulesExternalConfiguration rulesExternalConfiguration = routingRulesConfig.getRulesExternalConfiguration(); - yield RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, configuration.getRequestAnalyzerConfig()); + yield RoutingSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, configuration.getRequestAnalyzerConfig()); } }; } catch (Exception e) { - return RoutingGroupSelector.byRoutingGroupHeader(); + return RoutingSelector.byRoutingGroupHeader(); } } - return RoutingGroupSelector.byRoutingGroupHeader(); + return RoutingSelector.byRoutingGroupHeader(); } @Provides diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/QueryHistory.java b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/QueryHistory.java index de7e68dcc..1fb3818d1 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/QueryHistory.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/QueryHistory.java @@ -25,7 +25,7 @@ public record QueryHistory( @ColumnName("user_name") @Nullable String userName, @ColumnName("source") @Nullable String source, @ColumnName("created") long created, - @ColumnName("routing_group") String routingGroup, + @ColumnName("routing_decision") String routingDecision, @ColumnName("external_url") String externalUrl) { public QueryHistory diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/QueryHistoryDao.java b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/QueryHistoryDao.java index f0ff07bad..8762d8ce6 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/QueryHistoryDao.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/QueryHistoryDao.java @@ -78,10 +78,10 @@ default List findRecentQueriesByUserName(String userName, boolean String findBackendUrlByQueryId(String queryId); @SqlQuery(""" - SELECT routing_group FROM query_history + SELECT routing_decision FROM query_history WHERE query_id = :queryId """) - String findRoutingGroupByQueryId(String queryId); + String findRoutingDecisionByQueryId(String queryId); @SqlQuery(""" SELECT external_url FROM query_history @@ -116,10 +116,10 @@ GROUP BY FLOOR(created / 1000 / 60), backend_url List> findDistribution(long created); @SqlUpdate(""" - INSERT INTO query_history (query_id, query_text, backend_url, user_name, source, created, routing_group, external_url) - VALUES (:queryId, :queryText, :backendUrl, :userName, :source, :created, :routingGroup, :externalUrl) + INSERT INTO query_history (query_id, query_text, backend_url, user_name, source, created, routing_decision, external_url) + VALUES (:queryId, :queryText, :backendUrl, :userName, :source, :created, :routingDecision, :externalUrl) """) - void insertHistory(String queryId, String queryText, String backendUrl, String userName, String source, long created, String routingGroup, String externalUrl); + void insertHistory(String queryId, String queryText, String backendUrl, String userName, String source, long created, String routingDecision, String externalUrl); @SqlUpdate(""" DELETE FROM query_history diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/BaseRoutingManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/BaseRoutingManager.java index 996ee158c..9a5067971 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/BaseRoutingManager.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/BaseRoutingManager.java @@ -42,6 +42,8 @@ import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; +import static com.google.common.base.Strings.isNullOrEmpty; + /** * This class performs health check, stats counts for each backend and provides a backend given * request object. Default implementation comes here. @@ -56,7 +58,7 @@ public abstract class BaseRoutingManager private final String defaultRoutingGroup; private final QueryHistoryManager queryHistoryManager; private final LoadingCache queryIdBackendCache; - private final LoadingCache queryIdRoutingGroupCache; + private final LoadingCache queryIdRoutingDecisionCache; private final LoadingCache queryIdExternalUrlCache; public BaseRoutingManager(GatewayBackendManager gatewayBackendManager, QueryHistoryManager queryHistoryManager, RoutingConfiguration routingConfiguration) @@ -65,7 +67,7 @@ public BaseRoutingManager(GatewayBackendManager gatewayBackendManager, QueryHist this.defaultRoutingGroup = routingConfiguration.getDefaultRoutingGroup(); this.queryHistoryManager = queryHistoryManager; this.queryIdBackendCache = buildCache(this::findBackendForUnknownQueryId); - this.queryIdRoutingGroupCache = buildCache(this::findRoutingGroupForUnknownQueryId); + this.queryIdRoutingDecisionCache = buildCache(this::findRoutingDecisionForUnknownQueryId); this.queryIdExternalUrlCache = buildCache(this::findExternalUrlForUnknownQueryId); this.backendToStatus = new ConcurrentHashMap<>(); } @@ -82,9 +84,9 @@ public void setBackendForQueryId(String queryId, String backend) } @Override - public void setRoutingGroupForQueryId(String queryId, String routingGroup) + public void setRoutingDecisionForQueryId(String queryId, String routingDecision) { - queryIdRoutingGroupCache.put(queryId, routingGroup); + queryIdRoutingDecisionCache.put(queryId, routingDecision); } /** @@ -99,12 +101,24 @@ public ProxyBackendConfiguration provideDefaultBackendConfiguration(String user) } /** - * Performs routing to a given cluster group. This falls back to a default backend, if no scheduled - * backend is found. + * Selects a backend configuration for the request. + * At most one of `routingCluster` or `routingGroup` may be provided; they are mutually exclusive + * - If `routingCluster` is provided, returns that backend when it is active and healthy; otherwise + * falls back to the default backend. + * - If `routingCluster` is not provided, considers all active backends in `routingGroup`, filters to + * healthy ones, and delegates to `selectBackend(...)` to choose; if none are eligible, falls back + * to the default backend. + * - If neither `routingCluster` nor `routingGroup` is provided, falls back to the default backend. */ @Override - public ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user) + public ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String routingCluster, String user) { + if (!isNullOrEmpty(routingCluster)) { + return gatewayBackendManager.getBackendByName(routingCluster) + .filter(ProxyBackendConfiguration::isActive) + .filter(backEnd -> isBackendHealthy(backEnd.getName())) + .orElseGet(() -> provideDefaultBackendConfiguration(user)); + } List backends = gatewayBackendManager.getActiveBackends(routingGroup).stream() .filter(backEnd -> isBackendHealthy(backEnd.getName())) .toList(); @@ -144,21 +158,21 @@ public String findExternalUrlForQueryId(String queryId) } /** - * Looks up the routing group associated with the queryId in the cache. + * Looks up the routing decision associated with the queryId in the cache. * If it's not in the cache, look up in query history */ @Nullable @Override - public String findRoutingGroupForQueryId(String queryId) + public String findRoutingDecisionForQueryId(String queryId) { - String routingGroup = null; + String routingDecision = null; try { - routingGroup = queryIdRoutingGroupCache.get(queryId); + routingDecision = queryIdRoutingDecisionCache.get(queryId); } catch (ExecutionException e) { - log.warn("Exception while loading queryId from routing group cache %s", e.getLocalizedMessage()); + log.warn("Exception while loading queryId from routing decision cache %s", e.getLocalizedMessage()); } - return routingGroup; + return routingDecision; } @Override @@ -241,13 +255,13 @@ private String searchAllBackendForQuery(String queryId) } /** - * Attempts to look up the routing group associated with the query id from query history table + * Attempts to look up the routing decision associated with the query id from query history table */ - private String findRoutingGroupForUnknownQueryId(String queryId) + private String findRoutingDecisionForUnknownQueryId(String queryId) { - String routingGroup = queryHistoryManager.getRoutingGroupForQueryId(queryId); - setRoutingGroupForQueryId(queryId, routingGroup); - return routingGroup; + String routingDecision = queryHistoryManager.getRoutingDecisionForQueryId(queryId); + setRoutingDecisionForQueryId(queryId, routingDecision); + return routingDecision; } /** diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/ExternalRoutingGroupSelector.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/ExternalRoutingSelector.java similarity index 96% rename from gateway-ha/src/main/java/io/trino/gateway/ha/router/ExternalRoutingGroupSelector.java rename to gateway-ha/src/main/java/io/trino/gateway/ha/router/ExternalRoutingSelector.java index a7821e5b1..eca5285ab 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/ExternalRoutingGroupSelector.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/ExternalRoutingSelector.java @@ -51,10 +51,10 @@ import static java.util.Collections.list; import static java.util.Objects.requireNonNull; -public class ExternalRoutingGroupSelector - implements RoutingGroupSelector +public class ExternalRoutingSelector + implements RoutingSelector { - private static final Logger log = Logger.get(ExternalRoutingGroupSelector.class); + private static final Logger log = Logger.get(ExternalRoutingSelector.class); private final Set excludeHeaders; private final URI uri; private final boolean propagateErrors; @@ -65,7 +65,7 @@ public class ExternalRoutingGroupSelector createJsonResponseHandler(jsonCodec(ExternalRouterResponse.class)); @VisibleForTesting - ExternalRoutingGroupSelector(HttpClient httpClient, RulesExternalConfiguration rulesExternalConfiguration, RequestAnalyzerConfig requestAnalyzerConfig) + ExternalRoutingSelector(HttpClient httpClient, RulesExternalConfiguration rulesExternalConfiguration, RequestAnalyzerConfig requestAnalyzerConfig) { this.httpClient = requireNonNull(httpClient, "httpClient is null"); this.excludeHeaders = ImmutableSet.builder() @@ -128,14 +128,14 @@ else if (response.errors() != null && !response.errors().isEmpty()) { log.info("External routing service modified headers to: %s", filteredHeaders); } } - return new RoutingSelectorResponse(response.routingGroup(), filteredHeaders); + return new RoutingSelectorResponse(response.routingGroup(), response.routingCluster(), filteredHeaders); } catch (Exception e) { throwIfInstanceOf(e, WebApplicationException.class); log.error(e, "Error occurred while retrieving routing group " + "from external routing rules processing at " + uri); } - return new RoutingSelectorResponse(servletRequest.getHeader(ROUTING_GROUP_HEADER)); + return new RoutingSelectorResponse(servletRequest.getHeader(ROUTING_GROUP_HEADER), null); } private RoutingGroupExternalBody createRequestBody(HttpServletRequest request) diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/FileBasedRoutingGroupSelector.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/FileBasedRoutingSelector.java similarity index 84% rename from gateway-ha/src/main/java/io/trino/gateway/ha/router/FileBasedRoutingGroupSelector.java rename to gateway-ha/src/main/java/io/trino/gateway/ha/router/FileBasedRoutingSelector.java index 96106b889..70cd778ef 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/FileBasedRoutingGroupSelector.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/FileBasedRoutingSelector.java @@ -29,6 +29,7 @@ import java.nio.file.Path; import java.util.ArrayList; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -38,18 +39,19 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.Collections.sort; -public class FileBasedRoutingGroupSelector - implements RoutingGroupSelector +public class FileBasedRoutingSelector + implements RoutingSelector { - private static final Logger log = Logger.get(FileBasedRoutingGroupSelector.class); + private static final Logger log = Logger.get(FileBasedRoutingSelector.class); public static final String RESULTS_ROUTING_GROUP_KEY = "routingGroup"; + public static final String RESULTS_ROUTING_CLUSTER_KEY = "routingCluster"; private static final ObjectMapper yamlReader = new ObjectMapper(new YAMLFactory()); private final Supplier> rules; private final boolean analyzeRequest; - public FileBasedRoutingGroupSelector(String rulesPath, Duration rulesRefreshPeriod, RequestAnalyzerConfig requestAnalyzerConfig) + public FileBasedRoutingSelector(String rulesPath, Duration rulesRefreshPeriod, RequestAnalyzerConfig requestAnalyzerConfig) { analyzeRequest = requestAnalyzerConfig.isAnalyzeRequest(); @@ -59,7 +61,13 @@ public FileBasedRoutingGroupSelector(String rulesPath, Duration rulesRefreshPeri @Override public RoutingSelectorResponse findRoutingDestination(HttpServletRequest request) { - Map result = new HashMap<>(); + // Keep only the highest-priority rule result by limiting the map to a single entry. + LinkedHashMap result = new LinkedHashMap<>(1) { @Override + protected boolean removeEldestEntry(Map.Entry eldest) + { + return size() > 1; + } + }; Map state = new HashMap<>(); Map data; @@ -78,7 +86,7 @@ public RoutingSelectorResponse findRoutingDestination(HttpServletRequest request rule.evaluateAction(result, data, state); } }); - return new RoutingSelectorResponse(result.get(RESULTS_ROUTING_GROUP_KEY)); + return new RoutingSelectorResponse(result.get(RESULTS_ROUTING_GROUP_KEY), result.get(RESULTS_ROUTING_CLUSTER_KEY)); } public List readRulesFromPath(Path rulesPath) diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaQueryHistoryManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaQueryHistoryManager.java index b7ea39001..036ed3f98 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaQueryHistoryManager.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaQueryHistoryManager.java @@ -61,7 +61,7 @@ public void submitQueryDetail(QueryDetail queryDetail) queryDetail.getUser(), queryDetail.getSource(), queryDetail.getCaptureTime(), - queryDetail.getRoutingGroup(), + queryDetail.getRoutingDecision(), queryDetail.getExternalUrl()); } @@ -89,7 +89,7 @@ private static List upcast(List q queryDetail.setBackendUrl(dao.backendUrl()); queryDetail.setUser(dao.userName()); queryDetail.setSource(dao.source()); - queryDetail.setRoutingGroup(dao.routingGroup()); + queryDetail.setRoutingDecision(dao.routingDecision()); queryDetail.setExternalUrl(dao.externalUrl()); queryDetails.add(queryDetail); } @@ -103,9 +103,9 @@ public String getBackendForQueryId(String queryId) } @Override - public String getRoutingGroupForQueryId(String queryId) + public String getRoutingDecisionForQueryId(String queryId) { - return dao.findRoutingGroupByQueryId(queryId); + return dao.findRoutingDecisionByQueryId(queryId); } @Override diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/MVELRoutingRule.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/MVELRoutingRule.java index 8b2700460..d51d7b076 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/MVELRoutingRule.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/MVELRoutingRule.java @@ -88,7 +88,7 @@ private void initializeParserContext(ParserContext parserContext) parserContext.addImport(String.class); parserContext.addImport(StringBuffer.class); parserContext.addImport(StringBuilder.class); - parserContext.addImport(FileBasedRoutingGroupSelector.class); + parserContext.addImport(FileBasedRoutingSelector.class); } @Override diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/QueryHistoryManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/QueryHistoryManager.java index d62e750cc..1c0c05729 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/QueryHistoryManager.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/QueryHistoryManager.java @@ -32,7 +32,7 @@ public interface QueryHistoryManager String getBackendForQueryId(String queryId); - String getRoutingGroupForQueryId(String queryId); + String getRoutingDecisionForQueryId(String queryId); String getExternalUrlForQueryId(String queryId); @@ -49,7 +49,7 @@ class QueryDetail private String source; private String backendUrl; private long captureTime; - private String routingGroup; + private String routingDecision; private String externalUrl; public QueryDetail() {} @@ -132,14 +132,14 @@ public void setCaptureTime(long captureTime) } @JsonProperty - public String getRoutingGroup() + public String getRoutingDecision() { - return routingGroup; + return routingDecision; } - public void setRoutingGroup(String routingGroup) + public void setRoutingDecision(String routingDecision) { - this.routingGroup = routingGroup; + this.routingDecision = routingDecision; } @JsonProperty @@ -169,14 +169,14 @@ public boolean equals(Object o) Objects.equals(user, that.user) && Objects.equals(source, that.source) && Objects.equals(backendUrl, that.backendUrl) && - Objects.equals(routingGroup, that.routingGroup) && + Objects.equals(routingDecision, that.routingDecision) && Objects.equals(externalUrl, that.externalUrl); } @Override public int hashCode() { - return Objects.hash(queryId, queryText, user, source, backendUrl, captureTime, routingGroup, externalUrl); + return Objects.hash(queryId, queryText, user, source, backendUrl, captureTime, routingDecision, externalUrl); } @Override @@ -189,7 +189,7 @@ public String toString() .add("source", source) .add("backendUrl", backendUrl) .add("captureTime", captureTime) - .add("routingGroup", routingGroup) + .add("routingDecision", routingDecision) .add("externalUrl", externalUrl) .toString(); } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingManager.java index 2f112f6dd..a575c53e4 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingManager.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingManager.java @@ -45,12 +45,12 @@ public interface RoutingManager void setBackendForQueryId(String queryId, String backend); /** - * Associates a routing group with a specific query ID. + * Associates a routing decision with a specific query ID. * * @param queryId the unique identifier of the query - * @param routingGroup the routing group to associate with the query + * @param routingDecision the routing decision to associate with the query */ - void setRoutingGroupForQueryId(String queryId, String routingGroup); + void setRoutingDecisionForQueryId(String queryId, String routingDecision); /** * Finds the backend cluster associated with a given query ID. @@ -69,19 +69,20 @@ public interface RoutingManager String findExternalUrlForQueryId(String queryId); /** - * Finds the routing group associated with a given query ID. + * Finds the routing decision associated with a given query ID. * * @param queryId the unique identifier of the query - * @return the routing group, or null if not found + * @return the routing decision, or null if not found */ - String findRoutingGroupForQueryId(String queryId); + String findRoutingDecisionForQueryId(String queryId); /** - * Provides the backend configuration for a given routing group and user. + * Provides the backend configuration for a given routing group or cluster and user. * * @param routingGroup the routing group to use for backend selection + * @param routingCluster the routing cluster to use for backend selection * @param user the user requesting the backend * @return the backend configuration for the selected cluster */ - ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user); + ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String routingCluster, String user); } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingGroupSelector.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingSelector.java similarity index 60% rename from gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingGroupSelector.java rename to gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingSelector.java index 9f0877880..b3f9456a3 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingGroupSelector.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingSelector.java @@ -21,44 +21,44 @@ import jakarta.servlet.http.HttpServletRequest; /** - * RoutingGroupSelector provides a way to match an HTTP request to a Gateway routing group. + * RoutingSelector provides a way to match an HTTP request to a Gateway routing group. */ -public interface RoutingGroupSelector +public interface RoutingSelector { String ROUTING_GROUP_HEADER = "X-Trino-Routing-Group"; /** - * Routing group selector that relies on the X-Trino-Routing-Group + * Routing selector that relies on the X-Trino-Routing-Group * header to determine the right routing group. */ - static RoutingGroupSelector byRoutingGroupHeader() + static RoutingSelector byRoutingGroupHeader() { - return request -> new RoutingSelectorResponse(request.getHeader(ROUTING_GROUP_HEADER)); + return request -> new RoutingSelectorResponse(request.getHeader(ROUTING_GROUP_HEADER), null); } /** - * Routing group selector that uses routing engine rules - * to determine the right routing group. + * Routing selector that uses routing engine rules + * to determine the right routing group or cluster. */ - static RoutingGroupSelector byRoutingRulesEngine(String rulesConfigPath, Duration rulesRefreshPeriod, RequestAnalyzerConfig requestAnalyzerConfig) + static RoutingSelector byRoutingRulesEngine(String rulesConfigPath, Duration rulesRefreshPeriod, RequestAnalyzerConfig requestAnalyzerConfig) { - return new FileBasedRoutingGroupSelector(rulesConfigPath, rulesRefreshPeriod, requestAnalyzerConfig); + return new FileBasedRoutingSelector(rulesConfigPath, rulesRefreshPeriod, requestAnalyzerConfig); } /** - * Routing group selector that uses RESTful API + * Routing selector that uses RESTful API * to determine the right routing group. */ - static RoutingGroupSelector byRoutingExternal( + static RoutingSelector byRoutingExternal( HttpClient httpClient, RulesExternalConfiguration rulesExternalConfiguration, RequestAnalyzerConfig requestAnalyzerConfig) { - return new ExternalRoutingGroupSelector(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); + return new ExternalRoutingSelector(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); } /** - * Given an HTTP request find a routing group to direct the request to. If a routing group cannot + * Given an HTTP request find a routing destination to direct the request to. If a routing group or cluster cannot * be determined return null. */ RoutingSelectorResponse findRoutingDestination(HttpServletRequest request); diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/ExternalRouterResponse.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/ExternalRouterResponse.java index 88b6db164..8a65ed65c 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/ExternalRouterResponse.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/ExternalRouterResponse.java @@ -21,15 +21,16 @@ /** * Response from the external routing service that includes: - * - routingGroup: The target routing group for the request (optional) + * - routingDecision: The target routing group for the request (optional) * - errors: Any errors that occurred during routing * - externalHeaders: Headers that can be set in the request */ public record ExternalRouterResponse( @Nullable String routingGroup, + @Nullable String routingCluster, List errors, @Nullable Map externalHeaders) - implements RoutingGroupResponse + implements RoutingResponse { public ExternalRouterResponse { externalHeaders = externalHeaders == null ? ImmutableMap.of() : ImmutableMap.copyOf(externalHeaders); diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/RoutingGroupResponse.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/RoutingResponse.java similarity index 83% rename from gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/RoutingGroupResponse.java rename to gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/RoutingResponse.java index 73b667d3d..b4246ebb7 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/RoutingGroupResponse.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/RoutingResponse.java @@ -18,17 +18,19 @@ import java.util.Map; /** - Interface representing the response from a routing group selector. + Interface representing the response from a routing selector. This interface defines the contract for responses that determine how requests should be routed within the Trino Gateway system. Implementations of this interface are used to: - * Specify the target routing group for a request + * Specify the target routing group or cluster for a request * Provide additional headers that should be added to the request */ -public interface RoutingGroupResponse +public interface RoutingResponse { @Nullable String routingGroup(); + @Nullable String routingCluster(); + Map externalHeaders(); } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/RoutingSelectorResponse.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/RoutingSelectorResponse.java index 60b44900a..bfebe16f4 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/RoutingSelectorResponse.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/RoutingSelectorResponse.java @@ -20,18 +20,19 @@ /** * Response from the routing service that includes: - * - routingGroup: The target routing group for the request (Optional) + * - routingDecision: The target routing group for the request (Optional) + * - routingCluster: The target routing cluster for the request (Optional) * - externalHeaders: Headers that can be set in the request (Currently can only be set in ExternalRoutingGroupSelector) */ -public record RoutingSelectorResponse(@Nullable String routingGroup, Map externalHeaders) - implements RoutingGroupResponse +public record RoutingSelectorResponse(@Nullable String routingGroup, @Nullable String routingCluster, Map externalHeaders) + implements RoutingResponse { public RoutingSelectorResponse { externalHeaders = ImmutableMap.copyOf(externalHeaders); } - public RoutingSelectorResponse(String routingGroup) + public RoutingSelectorResponse(String routingGroup, String routingCluster) { - this(routingGroup, ImmutableMap.of()); + this(routingGroup, routingCluster, ImmutableMap.of()); } } diff --git a/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java b/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java index 8da9713ca..4ce9b84f1 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java +++ b/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java @@ -280,7 +280,7 @@ private ProxyResponse recordBackendForQueryId(Request request, ProxyResponse res HashMap results = OBJECT_MAPPER.readValue(response.body(), HashMap.class); queryDetail.setQueryId(results.get("id")); routingManager.setBackendForQueryId(queryDetail.getQueryId(), queryDetail.getBackendUrl()); - routingManager.setRoutingGroupForQueryId(queryDetail.getQueryId(), routingDestination.routingGroup()); + routingManager.setRoutingDecisionForQueryId(queryDetail.getQueryId(), routingDestination.routingDecision()); log.debug("QueryId [%s] mapped with proxy [%s]", queryDetail.getQueryId(), queryDetail.getBackendUrl()); } catch (IOException e) { @@ -290,7 +290,7 @@ private ProxyResponse recordBackendForQueryId(Request request, ProxyResponse res else { log.error("Non OK HTTP Status code with response [%s] , Status code [%s], user: [%s]", response.body(), response.statusCode(), username.orElse(null)); } - queryDetail.setRoutingGroup(routingDestination.routingGroup()); + queryDetail.setRoutingDecision(routingDestination.routingDecision()); queryDetail.setExternalUrl(routingDestination.externalUrl()); queryHistoryManager.submitQueryDetail(queryDetail); return response; diff --git a/gateway-ha/src/main/resources/mysql/V2__add_routingDecision_to_query_history.sql b/gateway-ha/src/main/resources/mysql/V2__add_routingDecision_to_query_history.sql new file mode 100644 index 000000000..9e40ab261 --- /dev/null +++ b/gateway-ha/src/main/resources/mysql/V2__add_routingDecision_to_query_history.sql @@ -0,0 +1,2 @@ +ALTER TABLE query_history + ADD routing_decision VARCHAR(255); diff --git a/gateway-ha/src/main/resources/mysql/V2__add_routingGroup_to_query_history.sql b/gateway-ha/src/main/resources/mysql/V2__add_routingGroup_to_query_history.sql deleted file mode 100644 index 34f4032ce..000000000 --- a/gateway-ha/src/main/resources/mysql/V2__add_routingGroup_to_query_history.sql +++ /dev/null @@ -1,2 +0,0 @@ -ALTER TABLE query_history - ADD routing_group VARCHAR(255); diff --git a/gateway-ha/src/main/resources/oracle/V2__add_routingDecision_to_query_history.sql b/gateway-ha/src/main/resources/oracle/V2__add_routingDecision_to_query_history.sql new file mode 100644 index 000000000..9e40ab261 --- /dev/null +++ b/gateway-ha/src/main/resources/oracle/V2__add_routingDecision_to_query_history.sql @@ -0,0 +1,2 @@ +ALTER TABLE query_history + ADD routing_decision VARCHAR(255); diff --git a/gateway-ha/src/main/resources/oracle/V2__add_routingGroup_to_query_history.sql b/gateway-ha/src/main/resources/oracle/V2__add_routingGroup_to_query_history.sql deleted file mode 100644 index 34f4032ce..000000000 --- a/gateway-ha/src/main/resources/oracle/V2__add_routingGroup_to_query_history.sql +++ /dev/null @@ -1,2 +0,0 @@ -ALTER TABLE query_history - ADD routing_group VARCHAR(255); diff --git a/gateway-ha/src/main/resources/postgresql/V2__add_routingDecision_to_query_history.sql b/gateway-ha/src/main/resources/postgresql/V2__add_routingDecision_to_query_history.sql new file mode 100644 index 000000000..9e40ab261 --- /dev/null +++ b/gateway-ha/src/main/resources/postgresql/V2__add_routingDecision_to_query_history.sql @@ -0,0 +1,2 @@ +ALTER TABLE query_history + ADD routing_decision VARCHAR(255); diff --git a/gateway-ha/src/main/resources/postgresql/V2__add_routingGroup_to_query_history.sql b/gateway-ha/src/main/resources/postgresql/V2__add_routingGroup_to_query_history.sql deleted file mode 100644 index 34f4032ce..000000000 --- a/gateway-ha/src/main/resources/postgresql/V2__add_routingGroup_to_query_history.sql +++ /dev/null @@ -1,2 +0,0 @@ -ALTER TABLE query_history - ADD routing_group VARCHAR(255); diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestRoutingTargetHandler.java b/gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestRoutingTargetHandler.java index 277e0c2aa..8072bb480 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestRoutingTargetHandler.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestRoutingTargetHandler.java @@ -22,8 +22,8 @@ import io.trino.gateway.ha.config.RequestAnalyzerConfig; import io.trino.gateway.ha.config.RulesExternalConfiguration; import io.trino.gateway.ha.handler.schema.RoutingTargetResponse; -import io.trino.gateway.ha.router.RoutingGroupSelector; import io.trino.gateway.ha.router.RoutingManager; +import io.trino.gateway.ha.router.RoutingSelector; import io.trino.gateway.ha.router.schema.ExternalRouterResponse; import jakarta.servlet.http.HttpServletRequest; import jakarta.ws.rs.HttpMethod; @@ -108,13 +108,13 @@ void setUp() config = provideGatewayConfiguration(); httpClient = Mockito.mock(HttpClient.class); routingManager = Mockito.mock(RoutingManager.class); - when(routingManager.provideBackendConfiguration(any(), any())).thenReturn(new ProxyBackendConfiguration()); + when(routingManager.provideBackendConfiguration(any(), any(), any())).thenReturn(new ProxyBackendConfiguration()); request = prepareMockRequest(); // Initialize the handler with the configuration handler = new RoutingTargetHandler( routingManager, - RoutingGroupSelector.byRoutingExternal(httpClient, config.getRoutingRules().getRulesExternalConfiguration(), config.getRequestAnalyzerConfig()), config); + RoutingSelector.byRoutingExternal(httpClient, config.getRoutingRules().getRulesExternalConfiguration(), config.getRequestAnalyzerConfig()), config); } @Test @@ -127,6 +127,7 @@ void testBasicHeaderModification() "X-New-Header", "new-value"); ExternalRouterResponse mockResponse = new ExternalRouterResponse( "test-group", + null, Collections.emptyList(), modifiedHeaders); when(httpClient.execute(any(), any())).thenReturn(mockResponse); @@ -151,6 +152,7 @@ void testExcludedHeaders() "Cookie", "new-session"); ExternalRouterResponse mockResponse = new ExternalRouterResponse( "test-group", + null, Collections.emptyList(), modifiedHeaders); when(httpClient.execute(any(), any())).thenReturn(mockResponse); @@ -172,6 +174,7 @@ void testNoHeaderModification() // Setup routing group selector response with no header modifications ExternalRouterResponse mockResponse = new ExternalRouterResponse( "test-group", + null, Collections.emptyList(), ImmutableMap.of()); when(httpClient.execute(any(), any())).thenReturn(mockResponse); @@ -194,6 +197,7 @@ void testEmptyHeader() "X-New-Header", "new-value"); ExternalRouterResponse mockResponse = new ExternalRouterResponse( "test-group", + null, Collections.emptyList(), modifiedHeaders); when(httpClient.execute(any(), any())).thenReturn(mockResponse); @@ -217,6 +221,7 @@ void testEmptyRoutingGroup() "X-Empty-Group-Header", "should-be-set"); ExternalRouterResponse mockResponse = new ExternalRouterResponse( "", + null, Collections.emptyList(), modifiedHeaders); when(httpClient.execute(any(), any())).thenReturn(mockResponse); @@ -225,7 +230,7 @@ void testEmptyRoutingGroup() RoutingTargetResponse response = handler.resolveRouting(request); // Verify that when no routing group header is set, we default to "adhoc" - assertThat(response.routingDestination().routingGroup()).isEqualTo("default-group"); + assertThat(response.routingDestination().routingDecision()).isEqualTo("default-group"); assertThat(response.modifiedRequest().getHeader("X-Empty-Group-Header")) .isEqualTo("should-be-set"); } @@ -233,46 +238,46 @@ void testEmptyRoutingGroup() @Test void testResponsePropertiesNull() { - ExternalRouterResponse mockResponse = new ExternalRouterResponse(null, null, ImmutableMap.of()); + ExternalRouterResponse mockResponse = new ExternalRouterResponse(null, null, null, ImmutableMap.of()); when(httpClient.execute(any(), any())).thenReturn(mockResponse); RoutingTargetResponse result = handler.resolveRouting(request); - assertThat(result.routingDestination().routingGroup()).isEqualTo("default-group"); + assertThat(result.routingDestination().routingDecision()).isEqualTo("default-group"); } @Test void testResponseGroupSetResponseErrorsNull() { ExternalRouterResponse mockResponse = new ExternalRouterResponse( - "test-group", null, ImmutableMap.of()); + "test-group", null, null, ImmutableMap.of()); when(httpClient.execute(any(), any())).thenReturn(mockResponse); RoutingTargetResponse result = handler.resolveRouting(request); - assertThat(result.routingDestination().routingGroup()).isEqualTo("test-group"); + assertThat(result.routingDestination().routingDecision()).isEqualTo("test-group"); } @Test void testPropagateErrorsFalseResponseGroupNullResponseErrorsSet() { - ExternalRouterResponse mockResponse = new ExternalRouterResponse(null, List.of("some-error"), ImmutableMap.of()); + ExternalRouterResponse mockResponse = new ExternalRouterResponse(null, null, List.of("some-error"), ImmutableMap.of()); when(httpClient.execute(any(), any())).thenReturn(mockResponse); RoutingTargetResponse result = handler.resolveRouting(request); - assertThat(result.routingDestination().routingGroup()).isEqualTo("default-group"); + assertThat(result.routingDestination().routingDecision()).isEqualTo("default-group"); } @Test void testPropagateErrorsFalseResponseGroupAndErrorsSet() { - ExternalRouterResponse mockResponse = new ExternalRouterResponse("test-group", List.of("some-error"), ImmutableMap.of()); + ExternalRouterResponse mockResponse = new ExternalRouterResponse("test-group", null, List.of("some-error"), ImmutableMap.of()); when(httpClient.execute(any(), any())).thenReturn(mockResponse); RoutingTargetResponse result = handler.resolveRouting(request); - assertThat(result.routingDestination().routingGroup()).isEqualTo("test-group"); + assertThat(result.routingDestination().routingDecision()).isEqualTo("test-group"); } @Test @@ -281,7 +286,7 @@ void testPropagateErrorsTrueResponseGroupNullResponseErrorsSet() RoutingTargetHandler handler = createHandlerWithPropagateErrorsTrue(); config.getRoutingRules().getRulesExternalConfiguration().setPropagateErrors(true); - ExternalRouterResponse mockResponse = new ExternalRouterResponse(null, List.of("some-error"), ImmutableMap.of()); + ExternalRouterResponse mockResponse = new ExternalRouterResponse(null, null, List.of("some-error"), ImmutableMap.of()); when(httpClient.execute(any(), any())).thenReturn(mockResponse); assertThatThrownBy(() -> handler.resolveRouting(request)) @@ -293,7 +298,7 @@ void testPropagateErrorsTrueResponseGroupAndErrorsSet() { RoutingTargetHandler handler = createHandlerWithPropagateErrorsTrue(); - ExternalRouterResponse response = new ExternalRouterResponse("test-group", List.of("some-error"), ImmutableMap.of()); + ExternalRouterResponse response = new ExternalRouterResponse("test-group", null, List.of("some-error"), ImmutableMap.of()); when(httpClient.execute(any(), any())).thenReturn(response); assertThatThrownBy(() -> handler.resolveRouting(request)) @@ -305,6 +310,6 @@ private RoutingTargetHandler createHandlerWithPropagateErrorsTrue() config.getRoutingRules().getRulesExternalConfiguration().setPropagateErrors(true); return new RoutingTargetHandler( routingManager, - RoutingGroupSelector.byRoutingExternal(httpClient, config.getRoutingRules().getRulesExternalConfiguration(), config.getRequestAnalyzerConfig()), config); + RoutingSelector.byRoutingExternal(httpClient, config.getRoutingRules().getRulesExternalConfiguration(), config.getRequestAnalyzerConfig()), config); } } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/BaseExternalUrlQueryHistoryTest.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/BaseExternalUrlQueryHistoryTest.java index 1d34a40e3..82411a527 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/BaseExternalUrlQueryHistoryTest.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/BaseExternalUrlQueryHistoryTest.java @@ -72,7 +72,7 @@ void testSubmitQueryWithExternalUrl() queryDetail.setBackendUrl("http://localhost:8080"); queryDetail.setUser("test-user"); queryDetail.setSource("sqlWorkbench"); - queryDetail.setRoutingGroup("adhoc"); + queryDetail.setRoutingDecision("adhoc"); queryDetail.setExternalUrl("https://external-gateway.example.com"); queryDetail.setCaptureTime(System.currentTimeMillis()); @@ -97,7 +97,7 @@ void testGetExternalUrlByQueryId() queryDetail.setBackendUrl("http://backend:8080"); queryDetail.setUser("admin"); queryDetail.setSource("trino-cli"); - queryDetail.setRoutingGroup("analytics"); + queryDetail.setRoutingDecision("analytics"); queryDetail.setExternalUrl("https://analytics-gateway.company.com"); queryDetail.setCaptureTime(System.currentTimeMillis()); @@ -126,7 +126,7 @@ void testSubmitQueryWithNullExternalUrl() queryDetail.setBackendUrl("http://localhost:8080"); queryDetail.setUser("test-user"); queryDetail.setSource("sqlWorkbench"); - queryDetail.setRoutingGroup("adhoc"); + queryDetail.setRoutingDecision("adhoc"); queryDetail.setExternalUrl(null); queryDetail.setCaptureTime(System.currentTimeMillis()); @@ -154,7 +154,7 @@ void testMultipleQueriesWithDifferentExternalUrls() queryDetail.setBackendUrl("http://backend-" + i + ":8080"); queryDetail.setUser("user-" + i); queryDetail.setSource("source-" + i); - queryDetail.setRoutingGroup("group-" + i); + queryDetail.setRoutingDecision("group-" + i); queryDetail.setExternalUrl("https://external-" + i + ".example.com"); queryDetail.setCaptureTime(System.currentTimeMillis()); @@ -179,7 +179,7 @@ void testQueryDetailEqualsAndHashCodeWithExternalUrl() queryDetail1.setBackendUrl("http://localhost:8080"); queryDetail1.setUser("test-user"); queryDetail1.setSource("sqlWorkbench"); - queryDetail1.setRoutingGroup("adhoc"); + queryDetail1.setRoutingDecision("adhoc"); queryDetail1.setExternalUrl("https://external.example.com"); queryDetail1.setCaptureTime(captureTime); @@ -189,7 +189,7 @@ void testQueryDetailEqualsAndHashCodeWithExternalUrl() queryDetail2.setBackendUrl("http://localhost:8080"); queryDetail2.setUser("test-user"); queryDetail2.setSource("sqlWorkbench"); - queryDetail2.setRoutingGroup("adhoc"); + queryDetail2.setRoutingDecision("adhoc"); queryDetail2.setExternalUrl("https://external.example.com"); queryDetail2.setCaptureTime(captureTime); diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestExternalRoutingGroupSelector.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestExternalRoutingSelector.java similarity index 88% rename from gateway-ha/src/test/java/io/trino/gateway/ha/router/TestExternalRoutingGroupSelector.java rename to gateway-ha/src/test/java/io/trino/gateway/ha/router/TestExternalRoutingSelector.java index d5092f09a..caeb1b78d 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestExternalRoutingGroupSelector.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestExternalRoutingSelector.java @@ -51,7 +51,7 @@ import static io.airlift.json.JsonCodec.jsonCodec; import static io.trino.gateway.ha.handler.HttpUtils.TRINO_QUERY_PROPERTIES; import static io.trino.gateway.ha.handler.HttpUtils.TRINO_REQUEST_USER; -import static io.trino.gateway.ha.router.RoutingGroupSelector.ROUTING_GROUP_HEADER; +import static io.trino.gateway.ha.router.RoutingSelector.ROUTING_GROUP_HEADER; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; @@ -63,7 +63,7 @@ @ExtendWith(MockitoExtension.class) @TestInstance(TestInstance.Lifecycle.PER_CLASS) -final class TestExternalRoutingGroupSelector +final class TestExternalRoutingSelector { RequestAnalyzerConfig requestAnalyzerConfig = new RequestAnalyzerConfig(); private HttpClient httpClient; @@ -94,7 +94,7 @@ void testByRoutingRulesExternalEngine() HttpServletRequest mockRequest = prepareMockRequest(); // Create a mock response - ExternalRouterResponse mockResponse = new ExternalRouterResponse("test-group", null, ImmutableMap.of()); + ExternalRouterResponse mockResponse = new ExternalRouterResponse("test-group", null, null, ImmutableMap.of()); // Create ArgumentCaptor ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(Request.class); @@ -138,8 +138,8 @@ void testFallbackToHeaderOnApiFailure() HttpClient httpClient = mock(HttpClient.class); RulesExternalConfiguration rulesExternalConfiguration = provideRoutingRuleExternalConfig(); - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); + RoutingSelector routingSelector = + RoutingSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); HttpServletRequest mockRequest = prepareMockRequest(); setMockHeaders(mockRequest); when(mockRequest.getHeader(ROUTING_GROUP_HEADER)).thenReturn("default-group-api-failure"); @@ -148,7 +148,7 @@ void testFallbackToHeaderOnApiFailure() when(httpClient.execute(any(), any())).thenThrow(new RuntimeException("Simulated failure")); // Mock the behavior of httpClient.execute - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); // Fallback expected assertThat(routingGroup).isEqualTo("default-group-api-failure"); @@ -161,7 +161,7 @@ void testNullUri() rulesExternalConfiguration.setUrlPath(null); // Assert that a RuntimeException is thrown with message - assertThatThrownBy(() -> RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig)) + assertThatThrownBy(() -> RoutingSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig)) .isInstanceOf(RuntimeException.class) .hasMessage("Invalid URL provided, using routing group header as default."); } @@ -173,8 +173,8 @@ void testExcludeHeader() RulesExternalConfiguration rulesExternalConfiguration = provideRoutingRuleExternalConfig(); rulesExternalConfiguration.setExcludeHeaders(List.of("test-exclude-header")); - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); + RoutingSelector routingSelector = + RoutingSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); // Mock headers to be read by mockRequest HttpServletRequest mockRequest = mock(HttpServletRequest.class); @@ -187,11 +187,11 @@ void testExcludeHeader() when(mockRequest.getHeaders("not-excluded-header")).thenReturn(Collections.enumeration(customValidHeaderValues)); // Use reflection to get valid headers after removing excludeHeaders headers - Method getValidHeaders = ExternalRoutingGroupSelector.class.getDeclaredMethod("getValidHeaders", HttpServletRequest.class); + Method getValidHeaders = ExternalRoutingSelector.class.getDeclaredMethod("getValidHeaders", HttpServletRequest.class); getValidHeaders.setAccessible(true); @SuppressWarnings("unchecked") - Multimap validHeaders = (Multimap) getValidHeaders.invoke(routingGroupSelector, mockRequest); + Multimap validHeaders = (Multimap) getValidHeaders.invoke(routingSelector, mockRequest); assertThat(validHeaders.size()).isEqualTo(1); } @@ -201,7 +201,7 @@ void testFindRoutingDestinationWithHeaderValues() { // Setup RulesExternalConfiguration rulesExternalConfiguration = provideRoutingRuleExternalConfig(); - RoutingGroupSelector selector = RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); + RoutingSelector selector = RoutingSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); HttpServletRequest mockRequest = prepareMockRequest(); setMockHeaders(mockRequest); String headerKey = "X-Header"; @@ -209,6 +209,7 @@ void testFindRoutingDestinationWithHeaderValues() ExternalRouterResponse mockResponse = new ExternalRouterResponse( "test-group", + null, ImmutableList.of(), ImmutableMap.of(headerKey, headerValue)); @@ -231,7 +232,7 @@ void testExcludedHeaders() // Setup RulesExternalConfiguration rulesExternalConfiguration = provideRoutingRuleExternalConfig(); rulesExternalConfiguration.setExcludeHeaders(ImmutableList.of("X-Custom-Header")); - RoutingGroupSelector selector = RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); + RoutingSelector selector = RoutingSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); HttpServletRequest mockRequest = prepareMockRequest(); setMockHeaders(mockRequest); String allowedHeaderKey = "X-Header"; @@ -241,6 +242,7 @@ void testExcludedHeaders() ExternalRouterResponse mockResponse = new ExternalRouterResponse( "test-group", + null, ImmutableList.of(), ImmutableMap.of( allowedHeaderKey, allowedHeaderValue, @@ -263,7 +265,7 @@ void testHeaderModificationWithErrors() { // Setup RulesExternalConfiguration rulesExternalConfiguration = provideRoutingRuleExternalConfig(); - RoutingGroupSelector selector = RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); + RoutingSelector selector = RoutingSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); HttpServletRequest mockRequest = prepareMockRequest(); setMockHeaders(mockRequest); String headerKey = "X-Header"; @@ -271,6 +273,7 @@ void testHeaderModificationWithErrors() ExternalRouterResponse mockResponse = new ExternalRouterResponse( "test-group", + null, ImmutableList.of("Error occurred"), ImmutableMap.of(headerKey, headerValue)); @@ -288,11 +291,11 @@ void testHeaderModificationWithErrors() @Test void testHeaderModificationWithNoExternalHeaders() { - RoutingGroupSelector selector = RoutingGroupSelector.byRoutingExternal(httpClient, provideRoutingRuleExternalConfig(), requestAnalyzerConfig); + RoutingSelector selector = RoutingSelector.byRoutingExternal(httpClient, provideRoutingRuleExternalConfig(), requestAnalyzerConfig); HttpServletRequest mockRequest = prepareMockRequest(); setMockHeaders(mockRequest); - ExternalRouterResponse mockResponse = new ExternalRouterResponse("test-group", ImmutableList.of(), null); + ExternalRouterResponse mockResponse = new ExternalRouterResponse("test-group", null, ImmutableList.of(), null); when(httpClient.execute(any(), any())).thenReturn(mockResponse); @@ -308,7 +311,7 @@ void testHeaderModificationWithEmptyRoutingGroup() { // Setup RulesExternalConfiguration rulesExternalConfiguration = provideRoutingRuleExternalConfig(); - RoutingGroupSelector selector = RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); + RoutingSelector selector = RoutingSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); HttpServletRequest mockRequest = prepareMockRequest(); setMockHeaders(mockRequest); String headerKey = "X-Empty-Group-Header"; @@ -316,6 +319,7 @@ void testHeaderModificationWithEmptyRoutingGroup() ExternalRouterResponse mockResponse = new ExternalRouterResponse( "", + null, ImmutableList.of(), ImmutableMap.of(headerKey, headerValue)); @@ -335,12 +339,12 @@ void testPropagateErrorsFalseResponseWithErrors() // Setup RulesExternalConfiguration rulesExternalConfiguration = provideRoutingRuleExternalConfig(); rulesExternalConfiguration.setPropagateErrors(false); - RoutingGroupSelector selector = RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); + RoutingSelector selector = RoutingSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); HttpServletRequest mockRequest = prepareMockRequest(); setMockHeaders(mockRequest); ExternalRouterResponse mockResponse = new ExternalRouterResponse( - "test-group", List.of("some-error"), ImmutableMap.of()); + "test-group", null, List.of("some-error"), ImmutableMap.of()); when(httpClient.execute(any(), any())).thenReturn(mockResponse); @@ -357,13 +361,13 @@ void testPropagateErrorsTrueResponseWithErrors() // Setup RulesExternalConfiguration rulesExternalConfiguration = provideRoutingRuleExternalConfig(); rulesExternalConfiguration.setPropagateErrors(true); - RoutingGroupSelector selector = RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); + RoutingSelector selector = RoutingSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig); HttpServletRequest mockRequest = prepareMockRequest(); setMockHeaders(mockRequest); ExternalRouterResponse mockResponse = new ExternalRouterResponse( - "test-group", List.of("some-error"), ImmutableMap.of()); + "test-group", null, List.of("some-error"), ImmutableMap.of()); when(httpClient.execute(any(), any())).thenReturn(mockResponse); diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestQueryCountBasedRouter.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestQueryCountBasedRouter.java index 04c3f8b6e..088d4b20a 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestQueryCountBasedRouter.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestQueryCountBasedRouter.java @@ -196,7 +196,7 @@ void testUserWithSameNoOfQueuedQueries() { // The user u1 has same number of queries queued on each cluster // The query needs to be routed to cluster with least number of queries running - ProxyBackendConfiguration proxyConfig = queryCountBasedRouter.provideBackendConfiguration("etl", "u1"); + ProxyBackendConfiguration proxyConfig = queryCountBasedRouter.provideBackendConfiguration("etl", null, "u1"); String proxyTo = proxyConfig.getProxyTo(); assertThat(proxyTo).isEqualTo(BACKEND_URL_3); @@ -212,7 +212,7 @@ void testUserWithSameNoOfQueuedQueries() assertThat(c3Stats.userQueuedCount().getOrDefault("u1", 0)) .isEqualTo(6); - proxyConfig = queryCountBasedRouter.provideBackendConfiguration("etl", "u1"); + proxyConfig = queryCountBasedRouter.provideBackendConfiguration("etl", null, "u1"); proxyTo = proxyConfig.getProxyTo(); assertThat(proxyTo).isEqualTo(BACKEND_URL_1); @@ -224,7 +224,7 @@ void testUserWithDifferentQueueLengthUser1() { // The user u2 has different number of queries queued on each cluster // The query needs to be routed to cluster with least number of queued for that user - ProxyBackendConfiguration proxyConfig = queryCountBasedRouter.provideBackendConfiguration(routingConfiguration.getDefaultRoutingGroup(), "u2"); + ProxyBackendConfiguration proxyConfig = queryCountBasedRouter.provideBackendConfiguration(routingConfiguration.getDefaultRoutingGroup(), null, "u2"); String proxyTo = proxyConfig.getProxyTo(); assertThat(BACKEND_URL_2).isEqualTo(proxyTo); @@ -234,7 +234,7 @@ void testUserWithDifferentQueueLengthUser1() @Test void testUserWithDifferentQueueLengthUser2() { - ProxyBackendConfiguration proxyConfig = queryCountBasedRouter.provideBackendConfiguration(routingConfiguration.getDefaultRoutingGroup(), "u3"); + ProxyBackendConfiguration proxyConfig = queryCountBasedRouter.provideBackendConfiguration(routingConfiguration.getDefaultRoutingGroup(), null, "u3"); String proxyTo = proxyConfig.getProxyTo(); assertThat(BACKEND_URL_1).isEqualTo(proxyTo); @@ -244,7 +244,7 @@ void testUserWithDifferentQueueLengthUser2() @Test void testUserWithNoQueuedQueries() { - ProxyBackendConfiguration proxyConfig = queryCountBasedRouter.provideBackendConfiguration(routingConfiguration.getDefaultRoutingGroup(), "u101"); + ProxyBackendConfiguration proxyConfig = queryCountBasedRouter.provideBackendConfiguration(routingConfiguration.getDefaultRoutingGroup(), null, "u101"); String proxyTo = proxyConfig.getProxyTo(); assertThat(BACKEND_URL_3).isEqualTo(proxyTo); @@ -254,7 +254,7 @@ void testUserWithNoQueuedQueries() void testAdhocRoutingGroupFailOver() { // The ETL routing group doesn't exist - ProxyBackendConfiguration proxyConfig = queryCountBasedRouter.provideBackendConfiguration("NonExisting", "u1"); + ProxyBackendConfiguration proxyConfig = queryCountBasedRouter.provideBackendConfiguration("NonExisting", null, "u1"); String proxyTo = proxyConfig.getProxyTo(); assertThat(BACKEND_URL_3).isEqualTo(proxyTo); @@ -271,7 +271,7 @@ void testClusterWithLeastQueueCount() .build(); queryCountBasedRouter.updateClusterStats(clusters); - ProxyBackendConfiguration proxyConfig = queryCountBasedRouter.provideBackendConfiguration("NonExisting", "u1"); + ProxyBackendConfiguration proxyConfig = queryCountBasedRouter.provideBackendConfiguration("NonExisting", null, "u1"); String proxyTo = proxyConfig.getProxyTo(); assertThat(BACKEND_URL_4).isEqualTo(proxyTo); @@ -290,7 +290,7 @@ void testClusterWithLeastRunningCount() queryCountBasedRouter.updateClusterStats(clusters); - ProxyBackendConfiguration proxyConfig = queryCountBasedRouter.provideBackendConfiguration("NonExisting", "u1"); + ProxyBackendConfiguration proxyConfig = queryCountBasedRouter.provideBackendConfiguration("NonExisting", null, "u1"); String proxyTo = proxyConfig.getProxyTo(); assertThat(BACKEND_URL_5).isEqualTo(proxyTo); diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingManagerNotFound.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingManagerNotFound.java index 8d2b05135..8fb35e71c 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingManagerNotFound.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingManagerNotFound.java @@ -40,7 +40,7 @@ public TestRoutingManagerNotFound() void testNonExistentRoutingGroupThrowsNotFoundException() { // When requesting a non-existent routing group, an IllegalStateException should be thrown - assertThatThrownBy(() -> routingManager.provideBackendConfiguration("non_existent_group", "user")) + assertThatThrownBy(() -> routingManager.provideBackendConfiguration("non_existent_group", null, "user")) .isInstanceOf(IllegalStateException.class) .hasMessageContaining("Number of active backends found zero"); } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingRulesManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingRulesManager.java index 668374c17..c6f4da606 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingRulesManager.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingRulesManager.java @@ -49,14 +49,14 @@ void testGetRoutingRules() "airflow", "if query from airflow, route to etl group", null, - List.of("result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\")"), + List.of("result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\" && (request.getHeader(\"X-Trino-Client-Tags\") == null || request.getHeader(\"X-Trino-Client-Tags\").isEmpty())")); assertThat(result.get(1)).isEqualTo( new RoutingRule( "airflow special", "if query from airflow with special label, route to etl-special group", null, - List.of("result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\")"), + List.of("result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\" && request.getHeader(\"X-Trino-Client-Tags\") contains \"label=special\"")); } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingGroupSelector.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingSelector.java similarity index 79% rename from gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingGroupSelector.java rename to gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingSelector.java index 0ca8f41f7..4a59ce319 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingGroupSelector.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingSelector.java @@ -16,6 +16,7 @@ import com.google.common.collect.ImmutableSet; import io.airlift.units.Duration; import io.trino.gateway.ha.config.RequestAnalyzerConfig; +import io.trino.gateway.ha.router.schema.RoutingSelectorResponse; import io.trino.gateway.ha.util.QueryRequestMock; import io.trino.sql.tree.QualifiedName; import jakarta.servlet.http.HttpServletRequest; @@ -43,7 +44,7 @@ import java.util.stream.Stream; import static io.trino.gateway.ha.handler.HttpUtils.TRINO_QUERY_PROPERTIES; -import static io.trino.gateway.ha.router.RoutingGroupSelector.ROUTING_GROUP_HEADER; +import static io.trino.gateway.ha.router.RoutingSelector.ROUTING_GROUP_HEADER; import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.concurrent.TimeUnit.HOURS; import static java.util.concurrent.TimeUnit.MILLISECONDS; @@ -52,7 +53,7 @@ import static org.mockito.Mockito.when; @TestInstance(Lifecycle.PER_CLASS) -final class TestRoutingGroupSelector +final class TestRoutingSelector { public static final String TRINO_SOURCE_HEADER = "X-Trino-Source"; public static final String TRINO_CLIENT_TAGS_HEADER = "X-Trino-Client-Tags"; @@ -86,15 +87,15 @@ void testByRoutingGroupHeader() // If the header is present the routing group is the value of that header. when(mockRequest.getHeader(ROUTING_GROUP_HEADER)).thenReturn("batch_backend"); - RoutingGroupSelector routingGroupSelector = RoutingGroupSelector.byRoutingGroupHeader(); - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + RoutingSelector routingSelector = RoutingSelector.byRoutingGroupHeader(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isEqualTo("batch_backend"); // If the header is not present just return null. when(mockRequest.getHeader(ROUTING_GROUP_HEADER)).thenReturn(null); - routingGroupSelector = RoutingGroupSelector.byRoutingGroupHeader(); - routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + routingSelector = RoutingSelector.byRoutingGroupHeader(); + routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isNull(); } @@ -103,22 +104,22 @@ void testByRoutingGroupHeader() @MethodSource("provideRoutingRuleConfigFiles") void testByRoutingRulesEngine(String rulesConfigPath) { - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine(rulesConfigPath, oneHourRefreshPeriod, requestAnalyzerConfig); + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine(rulesConfigPath, oneHourRefreshPeriod, requestAnalyzerConfig); HttpServletRequest mockRequest = new QueryRequestMock() .httpHeader(TRINO_SOURCE_HEADER, "airflow") .requestAnalyzerConfig(requestAnalyzerConfig) .getHttpServletRequest(); - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isEqualTo("etl"); } @Test void testGetUserFromBasicAuth() { - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine( "src/test/resources/rules/routing_rules_trino_query_properties.yml", oneHourRefreshPeriod, requestAnalyzerConfig); @@ -129,7 +130,7 @@ void testGetUserFromBasicAuth() .requestAnalyzerConfig(requestAnalyzerConfig) .getHttpServletRequest(); - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isEqualTo("will-group"); } @@ -138,8 +139,8 @@ void testGetUserFromBasicAuth() void testTrinoQueryPropertiesQueryDetails() throws IOException { - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine( "src/test/resources/rules/routing_rules_trino_query_properties.yml", oneHourRefreshPeriod, requestAnalyzerConfig); @@ -150,7 +151,7 @@ void testTrinoQueryPropertiesQueryDetails() .httpHeader(TrinoQueryProperties.TRINO_SCHEMA_HEADER_NAME, "schem_\\\"default") .requestAnalyzerConfig(requestAnalyzerConfig) .getHttpServletRequest(); - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isEqualTo("tbl-group"); } @@ -159,8 +160,8 @@ void testTrinoQueryPropertiesQueryDetails() void testTrinoQueryPropertiesCatalogSchemas() throws IOException { - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine( "src/test/resources/rules/routing_rules_trino_query_properties.yml", oneHourRefreshPeriod, requestAnalyzerConfig); @@ -172,15 +173,15 @@ void testTrinoQueryPropertiesCatalogSchemas() .requestAnalyzerConfig(requestAnalyzerConfig) .getHttpServletRequest(); - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isEqualTo("catalog-schema-group"); } @Test void testTrinoQueryPropertiesSessionDefaults() { - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine( "src/test/resources/rules/routing_rules_trino_query_properties.yml", oneHourRefreshPeriod, requestAnalyzerConfig); @@ -191,7 +192,7 @@ void testTrinoQueryPropertiesSessionDefaults() .requestAnalyzerConfig(requestAnalyzerConfig) .getHttpServletRequest(); - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isEqualTo("defaults-group"); } @@ -199,8 +200,8 @@ void testTrinoQueryPropertiesSessionDefaults() void testTrinoQueryPropertiesQueryType() throws IOException { - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine( "src/test/resources/rules/routing_rules_trino_query_properties.yml", oneHourRefreshPeriod, requestAnalyzerConfig); @@ -210,7 +211,7 @@ void testTrinoQueryPropertiesQueryType() .requestAnalyzerConfig(requestAnalyzerConfig) .getHttpServletRequest(); - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isEqualTo("type-group"); } @@ -218,8 +219,8 @@ void testTrinoQueryPropertiesQueryType() void testTrinoQueryPropertiesResourceGroupQueryType() throws IOException { - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine( "src/test/resources/rules/routing_rules_trino_query_properties.yml", oneHourRefreshPeriod, requestAnalyzerConfig); @@ -229,7 +230,7 @@ void testTrinoQueryPropertiesResourceGroupQueryType() .requestAnalyzerConfig(requestAnalyzerConfig) .getHttpServletRequest(); - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isEqualTo("resource-group-type-group"); } @@ -238,8 +239,8 @@ void testTrinoQueryPropertiesAlternateStatementFormat() throws IOException { requestAnalyzerConfig.setClientsUseV2Format(true); - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine( "src/test/resources/rules/routing_rules_trino_query_properties.yml", oneHourRefreshPeriod, requestAnalyzerConfig); @@ -248,7 +249,7 @@ void testTrinoQueryPropertiesAlternateStatementFormat() .requestAnalyzerConfig(requestAnalyzerConfig) .getHttpServletRequest(); - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isEqualTo("type-group"); } @@ -259,8 +260,8 @@ void testTrinoQueryPropertiesPreparedStatementInHeader() String encodedStatements = "statement1=SELECT+%27s1%27+c%0A%0A,statement2=SELECT+%27s2%27+c%0A%0A,statement3=SELECT%0A++%27%2C%27+comma%0A%2C+%27%3D%27+eq%0A%0A,statement4=SELECT%0A++c1%0A%2C+c2%0AFROM%0A++foo%0A"; String body = "EXECUTE statement4"; - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine( "src/test/resources/rules/routing_rules_trino_query_properties.yml", oneHourRefreshPeriod, requestAnalyzerConfig); @@ -274,7 +275,7 @@ void testTrinoQueryPropertiesPreparedStatementInHeader() .requestAnalyzerConfig(requestAnalyzerConfig) .getHttpServletRequest(); - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isEqualTo("statement-header-group"); } @@ -282,8 +283,8 @@ void testTrinoQueryPropertiesPreparedStatementInHeader() void testTrinoQueryPropertiesParsingError() throws IOException { - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine( "src/test/resources/rules/routing_rules_trino_query_properties.yml", oneHourRefreshPeriod, requestAnalyzerConfig); @@ -296,7 +297,7 @@ void testTrinoQueryPropertiesParsingError() .getHttpServletRequest(); // When parsing fails, the query should route to the default "no-match" group - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isEqualTo("no-match"); // Verify that the TrinoQueryProperties indicates a parsing failure @@ -310,8 +311,8 @@ void testTrinoQueryPropertiesParsingError() @MethodSource("provideRoutingRuleConfigFiles") void testByRoutingRulesEngineSpecialLabel(String rulesConfigPath) { - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine(rulesConfigPath, oneHourRefreshPeriod, requestAnalyzerConfig); + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine(rulesConfigPath, oneHourRefreshPeriod, requestAnalyzerConfig); HttpServletRequest mockRequest = new QueryRequestMock() .httpHeader(TRINO_SOURCE_HEADER, "airflow") @@ -319,7 +320,7 @@ void testByRoutingRulesEngineSpecialLabel(String rulesConfigPath) .requestAnalyzerConfig(requestAnalyzerConfig) .getHttpServletRequest(); - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isEqualTo("etl-special"); } @@ -327,8 +328,8 @@ void testByRoutingRulesEngineSpecialLabel(String rulesConfigPath) @MethodSource("provideRoutingRuleConfigFiles") void testByRoutingRulesEngineNoMatch(String rulesConfigPath) { - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine(rulesConfigPath, oneHourRefreshPeriod, requestAnalyzerConfig); + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine(rulesConfigPath, oneHourRefreshPeriod, requestAnalyzerConfig); // even though special label is present, query is not from airflow. // should return no match @@ -337,7 +338,7 @@ void testByRoutingRulesEngineNoMatch(String rulesConfigPath) .requestAnalyzerConfig(requestAnalyzerConfig) .getHttpServletRequest(); - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isNull(); } @@ -359,15 +360,15 @@ void testByRoutingRulesEngineFileChange() } Duration refreshPeriod = new Duration(1, MILLISECONDS); - RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine(file.getPath(), refreshPeriod, requestAnalyzerConfig); + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine(file.getPath(), refreshPeriod, requestAnalyzerConfig); HttpServletRequest mockRequest = new QueryRequestMock() .httpHeader(TRINO_SOURCE_HEADER, "airflow") .requestAnalyzerConfig(requestAnalyzerConfig) .getHttpServletRequest(); - String routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + String routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isEqualTo("etl"); try (BufferedWriter writer = Files.newBufferedWriter(file.toPath(), UTF_8)) { @@ -382,7 +383,7 @@ void testByRoutingRulesEngineFileChange() Thread.sleep(2 * refreshPeriod.toMillis()); when(mockRequest.getHeader(TRINO_SOURCE_HEADER)).thenReturn("airflow"); - routingGroup = routingGroupSelector.findRoutingDestination(mockRequest).routingGroup(); + routingGroup = routingSelector.findRoutingDestination(mockRequest).routingGroup(); assertThat(routingGroup).isEqualTo("etl2"); @@ -556,4 +557,60 @@ void testLongQuery() TrinoQueryProperties trinoQueryProperties = (TrinoQueryProperties) mockRequest.getAttribute(TRINO_QUERY_PROPERTIES); assertThat(trinoQueryProperties.tablesContains("kat.schem.widetable")).isTrue(); } + + @Test + void testPinByRoutingCluster() { + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine("src/test/resources/rules/routing_rules_group_and_cluster.yml", + oneHourRefreshPeriod, + requestAnalyzerConfig); + + HttpServletRequest mockRequest = new QueryRequestMock() + .httpHeader(TrinoQueryProperties.TRINO_CATALOG_HEADER_NAME, DEFAULT_CATALOG) + .httpHeader(TrinoQueryProperties.TRINO_SCHEMA_HEADER_NAME, DEFAULT_SCHEMA) + .requestAnalyzerConfig(requestAnalyzerConfig) + .getHttpServletRequest(); + + when(mockRequest.getHeader("X-Trino-User")).thenReturn("user1"); + + RoutingSelectorResponse routingSelectorResponse = routingSelector.findRoutingDestination(mockRequest); + + assertThat(routingSelectorResponse.routingGroup()).isNull(); + assertThat(routingSelectorResponse.routingCluster()).isEqualTo("cluster01"); + } + + @Test + void testHigherPriorityRoutingRuleWins() { + RoutingSelector routingSelector = + RoutingSelector.byRoutingRulesEngine("src/test/resources/rules/routing_rules_group_and_cluster.yml", + oneHourRefreshPeriod, + requestAnalyzerConfig); + + HttpServletRequest mockRequestForGroup = new QueryRequestMock() + .httpHeader(TrinoQueryProperties.TRINO_CATALOG_HEADER_NAME, DEFAULT_CATALOG) + .httpHeader(TrinoQueryProperties.TRINO_SCHEMA_HEADER_NAME, DEFAULT_SCHEMA) + .requestAnalyzerConfig(requestAnalyzerConfig) + .getHttpServletRequest(); + when(mockRequestForGroup.getHeader("X-Trino-User")).thenReturn("user2"); + + RoutingSelectorResponse responseForGroup = routingSelector.findRoutingDestination(mockRequestForGroup); + + // For user2: routingCluster has higher priority than routingGroup (see routing_rules_group_and_cluster.yml). + // The capped LinkedHashMap keeps only the last (highest-priority) entry, so the routingGroup is evicted. + assertThat(responseForGroup.routingGroup()).isNull(); + assertThat(responseForGroup.routingCluster()).isEqualTo("adhoc01"); + + HttpServletRequest mockRequestForCluster = new QueryRequestMock() + .httpHeader(TrinoQueryProperties.TRINO_CATALOG_HEADER_NAME, DEFAULT_CATALOG) + .httpHeader(TrinoQueryProperties.TRINO_SCHEMA_HEADER_NAME, DEFAULT_SCHEMA) + .requestAnalyzerConfig(requestAnalyzerConfig) + .getHttpServletRequest(); + when(mockRequestForCluster.getHeader("X-Trino-User")).thenReturn("user3"); + + RoutingSelectorResponse responseForCluster = routingSelector.findRoutingDestination(mockRequestForCluster); + + // For user3: routingGroup has higher priority than routingCluster (see routing_rules_group_and_cluster.yml). + assertThat(responseForCluster.routingGroup()).isEqualTo("adhoc"); + assertThat(responseForCluster.routingCluster()).isNull(); + } } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestStochasticRoutingManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestStochasticRoutingManager.java index 230689c4a..a44b60917 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestStochasticRoutingManager.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestStochasticRoutingManager.java @@ -69,7 +69,7 @@ void testAddMockBackends() haRoutingManager.updateBackEndHealth(backend, TrinoStatus.UNHEALTHY); } - assertThat(haRoutingManager.provideBackendConfiguration(groupName, "").getProxyTo()) + assertThat(haRoutingManager.provideBackendConfiguration(groupName, null, "").getProxyTo()) .isEqualTo("test_group0.trino.example.com"); } } diff --git a/gateway-ha/src/test/resources/add_backends_postgres.sql b/gateway-ha/src/test/resources/add_backends_postgres.sql index 64d485df7..036010f2d 100644 --- a/gateway-ha/src/test/resources/add_backends_postgres.sql +++ b/gateway-ha/src/test/resources/add_backends_postgres.sql @@ -2,4 +2,4 @@ INSERT INTO gateway_backend (name, routing_group, backend_url, external_url, active) VALUES ('trino-1', 'adhoc', 'http://localhost:8081', 'http://localhost:8081', true), -('trino-2', 'adhoc', 'http://localhost:8082', 'http://localhost:8082', true); \ No newline at end of file +('trino-2', 'adhoc', 'http://localhost:8082', 'http://localhost:8082', true); diff --git a/gateway-ha/src/test/resources/rules/routing_rules_atomic.yml b/gateway-ha/src/test/resources/rules/routing_rules_atomic.yml index 89a5a09da..08af4b187 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_atomic.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_atomic.yml @@ -3,10 +3,10 @@ name: "airflow" description: "if query from airflow, route to etl group" condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\" && (request.getHeader(\"X-Trino-Client-Tags\") == null || request.getHeader(\"X-Trino-Client-Tags\").isEmpty())" actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\")" + - "result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\")" --- name: "airflow special" description: "if query from airflow with special label, route to etl-special group" condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\" && request.getHeader(\"X-Trino-Client-Tags\") contains \"label=special\"" actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\")" + - "result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\")" diff --git a/gateway-ha/src/test/resources/rules/routing_rules_group_and_cluster.yml b/gateway-ha/src/test/resources/rules/routing_rules_group_and_cluster.yml new file mode 100644 index 000000000..d8ebb77bf --- /dev/null +++ b/gateway-ha/src/test/resources/rules/routing_rules_group_and_cluster.yml @@ -0,0 +1,32 @@ +--- +name: "routing_cluster_pinning_rule" +description: "Pin user to specific cluster" +condition: "request.getHeader(\"X-Trino-User\") == \"user1\"" +actions: + - "result.put(\"routingCluster\", \"cluster01\")" +--- +name: "routing_cluster_higher_priority_rule" +description: "Higher priority sets cluster" +condition: "request.getHeader(\"X-Trino-User\") == \"user2\"" +actions: + - "result.put(\"routingCluster\", \"adhoc01\")" +--- +name: "routing_group_lower_priority_rule" +description: "Lower priority sets routing group" +priority: -1 +condition: "request.getHeader(\"X-Trino-User\") == \"user2\"" +actions: + - "result.put(\"routingGroup\", \"adhoc\")" +--- +name: "routing_group_higher_priority_rule" +description: "Higher priority sets routing group" +condition: "request.getHeader(\"X-Trino-User\") == \"user3\"" +actions: + - "result.put(\"routingGroup\", \"adhoc\")" +--- +name: "routing_cluster_lower_priority_rule" +description: "Lower priority sets routing cluster" +priority: -1 +condition: "request.getHeader(\"X-Trino-User\") == \"user3\"" +actions: + - "result.put(\"routingGroup\", \"adhoc01\")" diff --git a/gateway-ha/src/test/resources/rules/routing_rules_if_statements.yml b/gateway-ha/src/test/resources/rules/routing_rules_if_statements.yml index e0cf9bdbb..291d1e246 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_if_statements.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_if_statements.yml @@ -4,8 +4,8 @@ description: "if query from airflow, route to etl group" condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\"" actions: - "if (request.getHeader(\"X-Trino-Client-Tags\") contains \"label=special\") { - result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\") + result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\") } else { - result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\") + result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\") }" diff --git a/gateway-ha/src/test/resources/rules/routing_rules_priorities.yml b/gateway-ha/src/test/resources/rules/routing_rules_priorities.yml index 0ca00f1c1..f5fdac6ac 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_priorities.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_priorities.yml @@ -4,11 +4,11 @@ description: "if query from airflow, route to etl group" priority: 0 condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\"" actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\")" + - "result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\")" --- name: "airflow special" description: "if query from airflow with special label, route to etl-special group" priority: 1 condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\" && request.getHeader(\"X-Trino-Client-Tags\") contains \"label=special\"" actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\")" + - "result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\")" diff --git a/gateway-ha/src/test/resources/rules/routing_rules_state.yml b/gateway-ha/src/test/resources/rules/routing_rules_state.yml index 53650058b..09870a8d3 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_state.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_state.yml @@ -17,7 +17,7 @@ condition: | request.getHeader("X-Trino-Source") == "airflow" actions: - | - result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, "etl") + result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, "etl") - | state.get("triggeredRules").add("airflow") --- @@ -28,4 +28,4 @@ condition: | state.get("triggeredRules").contains("airflow") && request.getHeader("X-Trino-Client-Tags") contains "label=special" actions: - | - result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, "etl-special") + result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, "etl-special") diff --git a/gateway-ha/src/test/resources/rules/routing_rules_trino_query_properties.yml b/gateway-ha/src/test/resources/rules/routing_rules_trino_query_properties.yml index b76c65769..fc283812e 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_trino_query_properties.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_trino_query_properties.yml @@ -3,7 +3,7 @@ name: "user" description: "if user is will, route to will-group" condition: "trinoRequestUser.userExistsAndEquals(\"will\")" actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"will-group\")" + - "result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"will-group\")" --- name: "query" description: "test extraction of tables and schemas in conjunction with default catalog and schema" @@ -14,7 +14,7 @@ condition: | && trinoQueryProperties.getSchemas().contains("schemy") && trinoQueryProperties.getCatalogs().contains("catx") actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"tbl-group\")" + - "result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"tbl-group\")" --- name: "catalog-schema" description: "test that catalogSchemas with default catalog and schema" @@ -23,14 +23,14 @@ condition: | && trinoQueryProperties.getCatalogSchemas.contains("caty.default") && !trinoQueryProperties.getCatalogSchemas.contains("catx.default") actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"catalog-schema-group\")" + - "result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"catalog-schema-group\")" --- name: "query-type" description: "test table type" condition: | trinoQueryProperties.getQueryType().toLowerCase.equals("insert") actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"type-group\")" + - "result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"type-group\")" --- name: "resource-group-query-type" description: "test table type" @@ -44,7 +44,7 @@ description: "test execute with multiple prepared statements" condition: | trinoQueryProperties.getQueryType().toLowerCase.equals("query") && trinoQueryProperties.tablesContains("cat.schem.foo") actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"statement-header-group\")" + - "result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"statement-header-group\")" --- name: "defaults-group" description: "test default schema and catalog" @@ -53,14 +53,14 @@ condition: | && trinoQueryProperties.getDefaultSchema().equals(java.util.Optional.of("other_schema")) actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"defaults-group\")" + - "result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"defaults-group\")" --- name: "system-group" description: "capture queries to system catalog" condition: | trinoQueryProperties.getCatalogs().contains("system") actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"system\")" + - "result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"system\")" --- name: "nomatch" @@ -68,4 +68,4 @@ priority: -1 description: "default group to catch if no other rules fired" condition: "true" actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"no-match\")" + - "result.put(FileBasedRoutingSelector.RESULTS_ROUTING_GROUP_KEY, \"no-match\")" diff --git a/webapp/src/components/history.tsx b/webapp/src/components/history.tsx index efbd8addf..989e56c27 100644 --- a/webapp/src/components/history.tsx +++ b/webapp/src/components/history.tsx @@ -93,9 +93,9 @@ export function History() { ); - const routingGroupRender = (_: string, record: HistoryDetail) => { + const routingDecisionRender = (_: string, record: HistoryDetail) => { return ( - {record.routingGroup} + {record.routingDecision} ) } @@ -133,24 +133,25 @@ export function History() { onPageChange: list, }}> - { if (!a || !b) return 0; - return a.routingGroup.localeCompare(b.routingGroup); + return a.routingDecision.localeCompare(b.routingDecision); }} filters={ - [...new Set(backendData?.map(b => b.routingGroup))] - .map(routingGroup => { - return { - text: routingGroup, - value: routingGroup - } - })} + [...new Set((historyData?.rows || []).map(r => r.routingDecision).filter(Boolean))] + .map(decision => { + return { + text: decision, + value: decision + } + }) + } onFilter={(value, record) => { if (!record) return false; - return value === record.routingGroup + return value === record.routingDecision }} - render={routingGroupRender} /> + render={routingDecisionRender} /> {backendMapping[text]}} /> diff --git a/webapp/src/types/history.d.ts b/webapp/src/types/history.d.ts index f48de0530..809d0da8d 100644 --- a/webapp/src/types/history.d.ts +++ b/webapp/src/types/history.d.ts @@ -5,7 +5,7 @@ export interface HistoryDetail { source: string; backendUrl: string; captureTime: number; - routingGroup: string; + routingDecision: string; externalUrl: string; }