From e28d4c2a1019ebb666f500529f34f60595942a18 Mon Sep 17 00:00:00 2001 From: Peiyingy Date: Wed, 12 Nov 2025 13:12:58 -0800 Subject: [PATCH] wip Apply local changes from detached HEAD wip wip erroeous ver wip clean build add comments add tests clean up front end code wip update routing setting clean up frontend code --- docs/routing-rules.md | 36 ++++- gateway-ha/config.yaml | 2 +- .../ha/handler/RoutingTargetHandler.java | 21 +-- .../ha/handler/schema/RoutingDestination.java | 2 +- .../ha/module/HaGatewayProviderModule.java | 12 +- .../ha/persistence/dao/QueryHistory.java | 2 +- .../ha/persistence/dao/QueryHistoryDao.java | 10 +- .../gateway/ha/router/BaseRoutingManager.java | 50 +++--- ...ctor.java => ExternalRoutingSelector.java} | 12 +- ...tor.java => FileBasedRoutingSelector.java} | 20 ++- .../ha/router/HaQueryHistoryManager.java | 8 +- .../gateway/ha/router/MVELRoutingRule.java | 2 +- .../ha/router/QueryHistoryManager.java | 18 +-- .../gateway/ha/router/RoutingManager.java | 17 +- ...roupSelector.java => RoutingSelector.java} | 26 +-- .../router/schema/ExternalRouterResponse.java | 5 +- ...roupResponse.java => RoutingResponse.java} | 8 +- .../schema/RoutingSelectorResponse.java | 11 +- .../proxyserver/ProxyRequestHandler.java | 4 +- ...__add_routingDecision_to_query_history.sql | 2 + .../V2__add_routingGroup_to_query_history.sql | 2 - ...__add_routingDecision_to_query_history.sql | 2 + .../V2__add_routingGroup_to_query_history.sql | 2 - ...__add_routingDecision_to_query_history.sql | 2 + .../V2__add_routingGroup_to_query_history.sql | 2 - .../ha/handler/TestRoutingTargetHandler.java | 35 ++-- .../BaseExternalUrlQueryHistoryTest.java | 12 +- ....java => TestExternalRoutingSelector.java} | 46 +++--- .../ha/router/TestQueryCountBasedRouter.java | 16 +- .../ha/router/TestRoutingManagerNotFound.java | 2 +- .../ha/router/TestRoutingRulesManager.java | 4 +- ...Selector.java => TestRoutingSelector.java} | 149 ++++++++++++------ .../router/TestStochasticRoutingManager.java | 2 +- .../test/resources/add_backends_postgres.sql | 2 +- .../resources/rules/routing_rules_atomic.yml | 4 +- .../rules/routing_rules_group_and_cluster.yml | 32 ++++ .../rules/routing_rules_if_statements.yml | 4 +- .../rules/routing_rules_priorities.yml | 4 +- .../resources/rules/routing_rules_state.yml | 4 +- .../routing_rules_trino_query_properties.yml | 16 +- webapp/src/components/history.tsx | 27 ++-- webapp/src/types/history.d.ts | 2 +- 42 files changed, 395 insertions(+), 244 deletions(-) rename gateway-ha/src/main/java/io/trino/gateway/ha/router/{ExternalRoutingGroupSelector.java => ExternalRoutingSelector.java} (96%) rename gateway-ha/src/main/java/io/trino/gateway/ha/router/{FileBasedRoutingGroupSelector.java => FileBasedRoutingSelector.java} (84%) rename gateway-ha/src/main/java/io/trino/gateway/ha/router/{RoutingGroupSelector.java => RoutingSelector.java} (60%) rename gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/{RoutingGroupResponse.java => RoutingResponse.java} (83%) create mode 100644 gateway-ha/src/main/resources/mysql/V2__add_routingDecision_to_query_history.sql delete mode 100644 gateway-ha/src/main/resources/mysql/V2__add_routingGroup_to_query_history.sql create mode 100644 gateway-ha/src/main/resources/oracle/V2__add_routingDecision_to_query_history.sql delete mode 100644 gateway-ha/src/main/resources/oracle/V2__add_routingGroup_to_query_history.sql create mode 100644 gateway-ha/src/main/resources/postgresql/V2__add_routingDecision_to_query_history.sql delete mode 100644 gateway-ha/src/main/resources/postgresql/V2__add_routingGroup_to_query_history.sql rename gateway-ha/src/test/java/io/trino/gateway/ha/router/{TestExternalRoutingGroupSelector.java => TestExternalRoutingSelector.java} (88%) rename gateway-ha/src/test/java/io/trino/gateway/ha/router/{TestRoutingGroupSelector.java => TestRoutingSelector.java} (79%) create mode 100644 gateway-ha/src/test/resources/rules/routing_rules_group_and_cluster.yml 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; }