From d7fe5fc4c6c242bdc9a99d93e0a55829e78d0b38 Mon Sep 17 00:00:00 2001 From: Peiyingy Date: Sun, 23 Nov 2025 23:37:59 -0800 Subject: [PATCH 1/6] flag wip clean up code wip add tests --- .../trino/gateway/ha/domain/RoutingRule.java | 2 + .../ha/handler/RoutingTargetHandler.java | 2 +- .../gateway/ha/router/BaseRoutingManager.java | 16 ++++- .../router/ExternalRoutingGroupSelector.java | 2 +- .../router/FileBasedRoutingGroupSelector.java | 9 ++- .../gateway/ha/router/MVELRoutingRule.java | 10 +++ .../gateway/ha/router/RoutingManager.java | 3 +- .../trino/gateway/ha/router/RoutingRule.java | 2 + .../router/schema/ExternalRouterResponse.java | 4 +- .../router/schema/RoutingGroupResponse.java | 2 + .../schema/RoutingSelectorResponse.java | 6 +- .../ha/handler/TestRoutingTargetHandler.java | 30 +++++---- .../TestExternalRoutingGroupSelector.java | 20 +++--- .../ha/router/TestQueryCountBasedRouter.java | 16 ++--- .../gateway/ha/router/TestRoutingAPI.java | 4 +- .../ha/router/TestRoutingManagerNotFound.java | 2 +- .../ha/router/TestRoutingRulesManager.java | 12 ++-- .../router/TestStochasticRoutingManager.java | 67 +++++++++++++++++-- .../rules/routing_rules_concurrent.yml | 1 + .../resources/rules/routing_rules_update.yml | 2 + 20 files changed, 157 insertions(+), 55 deletions(-) diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java b/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java index db3f905ef..de208ba0c 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java @@ -33,6 +33,7 @@ public record RoutingRule( String name, String description, Integer priority, + Boolean enforceIsolation, List actions, String condition) { @@ -40,6 +41,7 @@ public record RoutingRule( requireNonNull(name, "name is null"); description = requireNonNullElse(description, ""); priority = requireNonNullElse(priority, 0); + enforceIsolation = requireNonNullElse(enforceIsolation, false); actions = ImmutableList.copyOf(actions); requireNonNull(condition, "condition is null"); } 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..c9b5bbc04 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 @@ -95,7 +95,7 @@ private RoutingTargetResponse getRoutingTargetResponse(HttpServletRequest reques String routingGroup = !isNullOrEmpty(routingDestination.routingGroup()) ? routingDestination.routingGroup() : defaultRoutingGroup; - ProxyBackendConfiguration backendConfiguration = routingManager.provideBackendConfiguration(routingGroup, user); + ProxyBackendConfiguration backendConfiguration = routingManager.provideBackendConfiguration(routingGroup, user, routingDestination.enforceIsolation()); String clusterHost = backendConfiguration.getProxyTo(); String externalUrl = backendConfiguration.getExternalUrl(); // Apply headers from RoutingDestination if there are any 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..b93e58712 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 @@ -27,6 +27,8 @@ import jakarta.annotation.Nullable; import jakarta.annotation.PreDestroy; import jakarta.ws.rs.HttpMethod; +import jakarta.ws.rs.WebApplicationException; +import jakarta.ws.rs.core.Response; import java.net.HttpURLConnection; import java.net.URI; @@ -42,6 +44,8 @@ import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; +import static jakarta.ws.rs.core.Response.Status.NOT_FOUND; + /** * This class performs health check, stats counts for each backend and provides a backend given * request object. Default implementation comes here. @@ -99,15 +103,21 @@ 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. + * Performs routing to a given cluster group. This falls back to a default backend if the target group + * has no suitable backend unless {@code enforceIsolation} is true, in which case a 404 is returned. */ @Override - public ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user) + public ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user, Boolean enforceIsolation) { List backends = gatewayBackendManager.getActiveBackends(routingGroup).stream() .filter(backEnd -> isBackendHealthy(backEnd.getName())) .toList(); + if (backends.isEmpty() && enforceIsolation) { + throw new WebApplicationException( + Response.status(NOT_FOUND) + .entity(String.format("No healthy backends available for routing group '%s' under enforced isolation for user '%s'", routingGroup, user)) + .build()); + } return selectBackend(backends, user).orElseGet(() -> provideDefaultBackendConfiguration(user)); } 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/ExternalRoutingGroupSelector.java index a7821e5b1..91ec819ba 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/ExternalRoutingGroupSelector.java @@ -128,7 +128,7 @@ 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(), filteredHeaders, response.enforceIsolation()); } catch (Exception e) { throwIfInstanceOf(e, WebApplicationException.class); 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/FileBasedRoutingGroupSelector.java index 96106b889..4a014579c 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/FileBasedRoutingGroupSelector.java @@ -37,6 +37,7 @@ import static io.trino.gateway.ha.handler.HttpUtils.TRINO_REQUEST_USER; import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.Collections.sort; +import static java.util.Objects.requireNonNull; public class FileBasedRoutingGroupSelector implements RoutingGroupSelector @@ -72,13 +73,15 @@ public RoutingSelectorResponse findRoutingDestination(HttpServletRequest request data = ImmutableMap.of("request", request); } - rules.get().forEach(rule -> { + boolean enforceIsolation = false; + for (RoutingRule rule : requireNonNull(rules.get())) { if (rule.evaluateCondition(data, state)) { log.debug("%s evaluated to true on request: %s", rule, request); rule.evaluateAction(result, data, state); + enforceIsolation = rule.isEnforceIsolation(); } - }); - return new RoutingSelectorResponse(result.get(RESULTS_ROUTING_GROUP_KEY)); + } + return new RoutingSelectorResponse(result.get(RESULTS_ROUTING_GROUP_KEY), ImmutableMap.of(), enforceIsolation); } public List readRulesFromPath(Path rulesPath) 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..f89d1d8fc 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 @@ -37,6 +37,7 @@ public class MVELRoutingRule String name; String description; Integer priority; + boolean enforceIsolation; Serializable condition; List actions; ParserContext parserContext = new ParserContext(); @@ -46,6 +47,7 @@ public MVELRoutingRule( @JsonProperty("name") String name, @JsonProperty("description") String description, @JsonProperty("priority") Integer priority, + @JsonProperty("enforceIsolation") Boolean enforceIsolation, @JsonProperty("condition") Serializable condition, @JsonProperty("actions") List actions) { @@ -54,6 +56,7 @@ public MVELRoutingRule( this.name = requireNonNull(name, "name is null"); this.description = requireNonNullElse(description, ""); this.priority = requireNonNullElse(priority, 0); + this.enforceIsolation = requireNonNullElse(enforceIsolation, false); this.condition = requireNonNull( condition instanceof String stringCondition ? compileExpression(stringCondition, parserContext) : condition, "condition is null"); @@ -97,6 +100,12 @@ public Integer getPriority() return priority; } + @Override + public boolean isEnforceIsolation() + { + return enforceIsolation; + } + @Override public int compareTo(RoutingRule o) { @@ -132,6 +141,7 @@ public String toString() .add("name", name) .add("description", description) .add("priority", priority) + .add("enforceIsolation", enforceIsolation) .add("condition", decompile(condition)) .add("actions", String.join(",", actions.stream().map(DebugTools::decompile).toList())) .add("parserContext", parserContext) 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..7771bbda7 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 @@ -81,7 +81,8 @@ public interface RoutingManager * * @param routingGroup the routing group to use for backend selection * @param user the user requesting the backend + * @param enforceIsolation whether to enforce isolation * @return the backend configuration for the selected cluster */ - ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user); + ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user, Boolean enforceIsolation); } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRule.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRule.java index d9110aeff..8bb8c867a 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRule.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRule.java @@ -23,4 +23,6 @@ public interface RoutingRule void evaluateAction(Map result, Map data, Map state); Integer getPriority(); + + boolean isEnforceIsolation(); } 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..515a09845 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 @@ -28,10 +28,12 @@ public record ExternalRouterResponse( @Nullable String routingGroup, List errors, - @Nullable Map externalHeaders) + @Nullable Map externalHeaders, + @Nullable Boolean enforceIsolation) implements RoutingGroupResponse { public ExternalRouterResponse { externalHeaders = externalHeaders == null ? ImmutableMap.of() : ImmutableMap.copyOf(externalHeaders); + enforceIsolation = enforceIsolation != null && enforceIsolation; } } 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/RoutingGroupResponse.java index 73b667d3d..b8f39bf46 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/RoutingGroupResponse.java @@ -31,4 +31,6 @@ public interface RoutingGroupResponse @Nullable String routingGroup(); Map externalHeaders(); + + Boolean enforceIsolation(); } 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..861157a02 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 @@ -22,8 +22,10 @@ * Response from the routing service that includes: * - routingGroup: The target routing group for the request (Optional) * - externalHeaders: Headers that can be set in the request (Currently can only be set in ExternalRoutingGroupSelector) + * - enforceIsolation: If true, the handler must not fall back to default when target group has no available backend; + * instead, a 4xx should be returned. */ -public record RoutingSelectorResponse(@Nullable String routingGroup, Map externalHeaders) +public record RoutingSelectorResponse(@Nullable String routingGroup, Map externalHeaders, Boolean enforceIsolation) implements RoutingGroupResponse { public RoutingSelectorResponse { @@ -32,6 +34,6 @@ public record RoutingSelectorResponse(@Nullable String routingGroup, Map handler.resolveRouting(request)) @@ -293,7 +299,7 @@ void testPropagateErrorsTrueResponseGroupAndErrorsSet() { RoutingTargetHandler handler = createHandlerWithPropagateErrorsTrue(); - ExternalRouterResponse response = new ExternalRouterResponse("test-group", List.of("some-error"), ImmutableMap.of()); + ExternalRouterResponse response = new ExternalRouterResponse("test-group", List.of("some-error"), ImmutableMap.of(), null); when(httpClient.execute(any(), any())).thenReturn(response); assertThatThrownBy(() -> handler.resolveRouting(request)) 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/TestExternalRoutingGroupSelector.java index d5092f09a..713350cea 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/TestExternalRoutingGroupSelector.java @@ -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, ImmutableMap.of(), false); // Create ArgumentCaptor ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(Request.class); @@ -210,7 +210,8 @@ void testFindRoutingDestinationWithHeaderValues() ExternalRouterResponse mockResponse = new ExternalRouterResponse( "test-group", ImmutableList.of(), - ImmutableMap.of(headerKey, headerValue)); + ImmutableMap.of(headerKey, headerValue), + false); when(httpClient.execute(any(), any())).thenReturn(mockResponse); @@ -244,7 +245,8 @@ void testExcludedHeaders() ImmutableList.of(), ImmutableMap.of( allowedHeaderKey, allowedHeaderValue, - excludedHeaderKey, excludedHeaderValue)); + excludedHeaderKey, excludedHeaderValue), + false); when(httpClient.execute(any(), any())).thenReturn(mockResponse); @@ -272,7 +274,8 @@ void testHeaderModificationWithErrors() ExternalRouterResponse mockResponse = new ExternalRouterResponse( "test-group", ImmutableList.of("Error occurred"), - ImmutableMap.of(headerKey, headerValue)); + ImmutableMap.of(headerKey, headerValue), + false); when(httpClient.execute(any(), any())).thenReturn(mockResponse); @@ -292,7 +295,7 @@ void testHeaderModificationWithNoExternalHeaders() HttpServletRequest mockRequest = prepareMockRequest(); setMockHeaders(mockRequest); - ExternalRouterResponse mockResponse = new ExternalRouterResponse("test-group", ImmutableList.of(), null); + ExternalRouterResponse mockResponse = new ExternalRouterResponse("test-group", ImmutableList.of(), null, false); when(httpClient.execute(any(), any())).thenReturn(mockResponse); @@ -317,7 +320,8 @@ void testHeaderModificationWithEmptyRoutingGroup() ExternalRouterResponse mockResponse = new ExternalRouterResponse( "", ImmutableList.of(), - ImmutableMap.of(headerKey, headerValue)); + ImmutableMap.of(headerKey, headerValue), + false); when(httpClient.execute(any(), any())).thenReturn(mockResponse); @@ -340,7 +344,7 @@ void testPropagateErrorsFalseResponseWithErrors() setMockHeaders(mockRequest); ExternalRouterResponse mockResponse = new ExternalRouterResponse( - "test-group", List.of("some-error"), ImmutableMap.of()); + "test-group", List.of("some-error"), ImmutableMap.of(), false); when(httpClient.execute(any(), any())).thenReturn(mockResponse); @@ -363,7 +367,7 @@ void testPropagateErrorsTrueResponseWithErrors() setMockHeaders(mockRequest); ExternalRouterResponse mockResponse = new ExternalRouterResponse( - "test-group", List.of("some-error"), ImmutableMap.of()); + "test-group", List.of("some-error"), ImmutableMap.of(), false); 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..fa29157e9 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", "u1", false); 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", "u1", false); 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(), "u2", false); 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(), "u3", false); 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(), "u101", false); 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", "u1", false); 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", "u1", false); 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", "u1", false); String proxyTo = proxyConfig.getProxyTo(); assertThat(BACKEND_URL_5).isEqualTo(proxyTo); diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingAPI.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingAPI.java index ae9cfdac9..4c165c66b 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingAPI.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingAPI.java @@ -102,7 +102,7 @@ void testUpdateRoutingRulesAPI() throws Exception { //Update routing rules with a new rule - RoutingRule updatedRoutingRules = new RoutingRule("airflow", "if query from airflow, route to adhoc group", 0, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"JDBC\""); + RoutingRule updatedRoutingRules = new RoutingRule("airflow", "if query from airflow, route to adhoc group", 0, false, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"JDBC\""); RequestBody requestBody = RequestBody.create(OBJECT_MAPPER.writeValueAsString(updatedRoutingRules), MediaType.parse("application/json; charset=utf-8")); Request request = new Request.Builder() .url("http://localhost:" + routerPort + "/webapp/updateRoutingRules") @@ -136,7 +136,7 @@ void testUpdateRoutingRulesAPI() assertThat(routingRules[0].actions()).first().isEqualTo("result.put(\"routingGroup\", \"adhoc\")"); //Revert back to old routing rules to avoid any test failures - RoutingRule revertRoutingRules = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"etl\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\""); + RoutingRule revertRoutingRules = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, false, List.of("result.put(\"routingGroup\", \"etl\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\""); RequestBody requestBody3 = RequestBody.create(OBJECT_MAPPER.writeValueAsString(revertRoutingRules), MediaType.parse("application/json; charset=utf-8")); Request request3 = new Request.Builder() .url("http://localhost:" + routerPort + "/webapp/updateRoutingRules") 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..92deab330 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", "user", false)) .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..98170b782 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,6 +49,7 @@ void testGetRoutingRules() "airflow", "if query from airflow, route to etl group", null, + null, List.of("result.put(FileBasedRoutingGroupSelector.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( @@ -56,6 +57,7 @@ void testGetRoutingRules() "airflow special", "if query from airflow with special label, route to etl-special group", null, + null, List.of("result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\" && request.getHeader(\"X-Trino-Client-Tags\") contains \"label=special\"")); } @@ -84,13 +86,13 @@ void testUpdateRoutingRulesFile() configuration.setRoutingRules(routingRulesConfiguration); RoutingRulesManager routingRulesManager = new RoutingRulesManager(configuration); - RoutingRule routingRules = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"JDBC\""); + RoutingRule routingRules = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, false, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"JDBC\""); List updatedRoutingRules = routingRulesManager.updateRoutingRule(routingRules); assertThat(updatedRoutingRules.getFirst().actions().getFirst()).isEqualTo("result.put(\"routingGroup\", \"adhoc\")"); assertThat(updatedRoutingRules.getFirst().condition()).isEqualTo("request.getHeader(\"X-Trino-Source\") == \"JDBC\""); - RoutingRule originalRoutingRules = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"etl\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\""); + RoutingRule originalRoutingRules = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, false, List.of("result.put(\"routingGroup\", \"etl\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\""); List updateRoutingRules = routingRulesManager.updateRoutingRule(originalRoutingRules); assertThat(updateRoutingRules).hasSize(2); @@ -107,7 +109,7 @@ void testUpdateRoutingRulesNoSuchFileException() routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); configuration.setRoutingRules(routingRulesConfiguration); RoutingRulesManager routingRulesManager = new RoutingRulesManager(configuration); - RoutingRule routingRules = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"JDBC\""); + RoutingRule routingRules = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, false, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"JDBC\""); assertThatThrownBy(() -> routingRulesManager.updateRoutingRule(routingRules)).hasRootCauseInstanceOf(NoSuchFileException.class); } @@ -123,8 +125,8 @@ void testConcurrentUpdateRoutingRule() configuration.setRoutingRules(routingRulesConfiguration); RoutingRulesManager routingRulesManager = new RoutingRulesManager(configuration); - RoutingRule routingRule1 = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"etl\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\""); - RoutingRule routingRule2 = new RoutingRule("airflow", "if query from airflow, route to adhoc group", 0, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"datagrip\""); + RoutingRule routingRule1 = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, false, List.of("result.put(\"routingGroup\", \"etl\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\""); + RoutingRule routingRule2 = new RoutingRule("airflow", "if query from airflow, route to adhoc group", 0, false, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"datagrip\""); ExecutorService executorService = Executors.newFixedThreadPool(2); 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..afae8b467 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 @@ -17,6 +17,8 @@ import io.trino.gateway.ha.config.ProxyBackendConfiguration; import io.trino.gateway.ha.config.RoutingConfiguration; import io.trino.gateway.ha.persistence.JdbcConnectionManager; +import jakarta.ws.rs.WebApplicationException; +import jakarta.ws.rs.core.Response.Status; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.TestInstance; @@ -24,6 +26,7 @@ import static io.trino.gateway.ha.TestingJdbcConnectionManager.createTestingJdbcConnectionManager; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; @TestInstance(Lifecycle.PER_CLASS) final class TestStochasticRoutingManager @@ -50,12 +53,12 @@ void testAddMockBackends() String backend; for (int i = 0; i < numBackends; i++) { backend = groupName + i; - ProxyBackendConfiguration proxyBackend = new ProxyBackendConfiguration(); - proxyBackend.setActive(true); - proxyBackend.setRoutingGroup(groupName); - proxyBackend.setName(backend); - proxyBackend.setProxyTo(backend + ".trino.example.com"); - proxyBackend.setExternalUrl("trino.example.com"); + ProxyBackendConfiguration proxyBackend = createBackend( + backend, + groupName, + true, + backend + ".trino.example.com", + "trino.example.com"); backendManager.addBackend(proxyBackend); //set backend as healthy to start with haRoutingManager.updateBackEndHealth(backend, TrinoStatus.HEALTHY); @@ -69,7 +72,57 @@ void testAddMockBackends() haRoutingManager.updateBackEndHealth(backend, TrinoStatus.UNHEALTHY); } - assertThat(haRoutingManager.provideBackendConfiguration(groupName, "").getProxyTo()) + assertThat(haRoutingManager.provideBackendConfiguration(groupName, "", false).getProxyTo()) .isEqualTo("test_group0.trino.example.com"); } + + @Test + void testEnforceIsolationException() + { + ProxyBackendConfiguration inactiveBackend = createBackend( + "inactive-backend", + "inactive-group", + false, + "inactive.trino.example.com", + "https://inactive.example"); + backendManager.addBackend(inactiveBackend); + + assertNotFoundForIsolation("inactive-group"); + + ProxyBackendConfiguration unhealthyBackend = createBackend( + "unhealthy-backend", + "unhealthy-group", + true, + "unhealthy.trino.example.com", + "https://unhealthy.example"); + backendManager.addBackend(unhealthyBackend); + haRoutingManager.updateBackEndHealth(unhealthyBackend.getName(), TrinoStatus.UNHEALTHY); + + assertNotFoundForIsolation("unhealthy-group"); + + assertNotFoundForIsolation("missing-group"); + } + + private void assertNotFoundForIsolation(String routingGroup) + { + assertThatThrownBy(() -> haRoutingManager.provideBackendConfiguration(routingGroup, "user", true)) + .isInstanceOfSatisfying(WebApplicationException.class, exception -> + assertThat(exception.getResponse().getStatus()).isEqualTo(Status.NOT_FOUND.getStatusCode())); + } + + private static ProxyBackendConfiguration createBackend( + String name, + String routingGroup, + boolean active, + String proxyTo, + String externalUrl) + { + ProxyBackendConfiguration backend = new ProxyBackendConfiguration(); + backend.setName(name); + backend.setRoutingGroup(routingGroup); + backend.setActive(active); + backend.setProxyTo(proxyTo); + backend.setExternalUrl(externalUrl); + return backend; + } } diff --git a/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml b/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml index f492dfeaa..a22c44783 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml @@ -2,6 +2,7 @@ name: "airflow" description: "if query from airflow, route to adhoc group" priority: 0 +enforceIsolation: false actions: - "result.put(\"routingGroup\", \"adhoc\")" condition: "request.getHeader(\"X-Trino-Source\") == \"datagrip\"" diff --git a/gateway-ha/src/test/resources/rules/routing_rules_update.yml b/gateway-ha/src/test/resources/rules/routing_rules_update.yml index 3478bc45e..4723b3792 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_update.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_update.yml @@ -2,6 +2,7 @@ name: "airflow" description: "if query from airflow, route to etl group" priority: 0 +enforceIsolation: false actions: - "result.put(\"routingGroup\", \"etl\")" condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\"" @@ -9,6 +10,7 @@ condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\"" name: "airflow special" description: "if query from airflow with special label, route to etl-special group" priority: 1 +enforceIsolation: false actions: - "result.put(\"routingGroup\", \"etl-special\")" condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\" && request.getHeader(\"\ From 2810a25d02ab454ef1de72c9cb018a097973cfc2 Mon Sep 17 00:00:00 2001 From: Peiyingy Date: Thu, 4 Dec 2025 17:03:52 -0500 Subject: [PATCH 2/6] rename to strictRouting --- .../java/io/trino/gateway/ha/domain/RoutingRule.java | 4 ++-- .../gateway/ha/handler/RoutingTargetHandler.java | 2 +- .../trino/gateway/ha/router/BaseRoutingManager.java | 6 +++--- .../ha/router/ExternalRoutingGroupSelector.java | 2 +- .../ha/router/FileBasedRoutingGroupSelector.java | 6 +++--- .../io/trino/gateway/ha/router/MVELRoutingRule.java | 12 ++++++------ .../io/trino/gateway/ha/router/RoutingManager.java | 4 ++-- .../java/io/trino/gateway/ha/router/RoutingRule.java | 2 +- .../ha/router/schema/ExternalRouterResponse.java | 4 ++-- .../ha/router/schema/RoutingGroupResponse.java | 2 +- .../ha/router/schema/RoutingSelectorResponse.java | 4 ++-- .../ha/router/TestStochasticRoutingManager.java | 2 +- .../resources/rules/routing_rules_concurrent.yml | 2 +- .../test/resources/rules/routing_rules_update.yml | 4 ++-- 14 files changed, 28 insertions(+), 28 deletions(-) diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java b/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java index de208ba0c..7602f4421 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java @@ -33,7 +33,7 @@ public record RoutingRule( String name, String description, Integer priority, - Boolean enforceIsolation, + Boolean strictRouting, List actions, String condition) { @@ -41,7 +41,7 @@ public record RoutingRule( requireNonNull(name, "name is null"); description = requireNonNullElse(description, ""); priority = requireNonNullElse(priority, 0); - enforceIsolation = requireNonNullElse(enforceIsolation, false); + strictRouting = requireNonNullElse(strictRouting, false); actions = ImmutableList.copyOf(actions); requireNonNull(condition, "condition is null"); } 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 c9b5bbc04..d7d6a79a6 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 @@ -95,7 +95,7 @@ private RoutingTargetResponse getRoutingTargetResponse(HttpServletRequest reques String routingGroup = !isNullOrEmpty(routingDestination.routingGroup()) ? routingDestination.routingGroup() : defaultRoutingGroup; - ProxyBackendConfiguration backendConfiguration = routingManager.provideBackendConfiguration(routingGroup, user, routingDestination.enforceIsolation()); + ProxyBackendConfiguration backendConfiguration = routingManager.provideBackendConfiguration(routingGroup, user, routingDestination.strictRouting()); String clusterHost = backendConfiguration.getProxyTo(); String externalUrl = backendConfiguration.getExternalUrl(); // Apply headers from RoutingDestination if there are any 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 b93e58712..dfc4327d6 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 @@ -104,15 +104,15 @@ public ProxyBackendConfiguration provideDefaultBackendConfiguration(String user) /** * Performs routing to a given cluster group. This falls back to a default backend if the target group - * has no suitable backend unless {@code enforceIsolation} is true, in which case a 404 is returned. + * has no suitable backend unless {@code strictRouting} is true, in which case a 404 is returned. */ @Override - public ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user, Boolean enforceIsolation) + public ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user, Boolean strictRouting) { List backends = gatewayBackendManager.getActiveBackends(routingGroup).stream() .filter(backEnd -> isBackendHealthy(backEnd.getName())) .toList(); - if (backends.isEmpty() && enforceIsolation) { + if (backends.isEmpty() && strictRouting) { throw new WebApplicationException( Response.status(NOT_FOUND) .entity(String.format("No healthy backends available for routing group '%s' under enforced isolation for user '%s'", routingGroup, user)) 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/ExternalRoutingGroupSelector.java index 91ec819ba..32c39bb76 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/ExternalRoutingGroupSelector.java @@ -128,7 +128,7 @@ else if (response.errors() != null && !response.errors().isEmpty()) { log.info("External routing service modified headers to: %s", filteredHeaders); } } - return new RoutingSelectorResponse(response.routingGroup(), filteredHeaders, response.enforceIsolation()); + return new RoutingSelectorResponse(response.routingGroup(), filteredHeaders, response.strictRouting()); } catch (Exception e) { throwIfInstanceOf(e, WebApplicationException.class); 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/FileBasedRoutingGroupSelector.java index 4a014579c..573c3c5dc 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/FileBasedRoutingGroupSelector.java @@ -73,15 +73,15 @@ public RoutingSelectorResponse findRoutingDestination(HttpServletRequest request data = ImmutableMap.of("request", request); } - boolean enforceIsolation = false; + boolean strictRouting = false; for (RoutingRule rule : requireNonNull(rules.get())) { if (rule.evaluateCondition(data, state)) { log.debug("%s evaluated to true on request: %s", rule, request); rule.evaluateAction(result, data, state); - enforceIsolation = rule.isEnforceIsolation(); + strictRouting = rule.isStrictRouting(); } } - return new RoutingSelectorResponse(result.get(RESULTS_ROUTING_GROUP_KEY), ImmutableMap.of(), enforceIsolation); + return new RoutingSelectorResponse(result.get(RESULTS_ROUTING_GROUP_KEY), ImmutableMap.of(), strictRouting); } public List readRulesFromPath(Path rulesPath) 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 f89d1d8fc..4656d0fb6 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 @@ -37,7 +37,7 @@ public class MVELRoutingRule String name; String description; Integer priority; - boolean enforceIsolation; + boolean strictRouting; Serializable condition; List actions; ParserContext parserContext = new ParserContext(); @@ -47,7 +47,7 @@ public MVELRoutingRule( @JsonProperty("name") String name, @JsonProperty("description") String description, @JsonProperty("priority") Integer priority, - @JsonProperty("enforceIsolation") Boolean enforceIsolation, + @JsonProperty("strictRouting") Boolean strictRouting, @JsonProperty("condition") Serializable condition, @JsonProperty("actions") List actions) { @@ -56,7 +56,7 @@ public MVELRoutingRule( this.name = requireNonNull(name, "name is null"); this.description = requireNonNullElse(description, ""); this.priority = requireNonNullElse(priority, 0); - this.enforceIsolation = requireNonNullElse(enforceIsolation, false); + this.strictRouting = requireNonNullElse(strictRouting, false); this.condition = requireNonNull( condition instanceof String stringCondition ? compileExpression(stringCondition, parserContext) : condition, "condition is null"); @@ -101,9 +101,9 @@ public Integer getPriority() } @Override - public boolean isEnforceIsolation() + public boolean isStrictRouting() { - return enforceIsolation; + return strictRouting; } @Override @@ -141,7 +141,7 @@ public String toString() .add("name", name) .add("description", description) .add("priority", priority) - .add("enforceIsolation", enforceIsolation) + .add("strictRouting", strictRouting) .add("condition", decompile(condition)) .add("actions", String.join(",", actions.stream().map(DebugTools::decompile).toList())) .add("parserContext", parserContext) 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 7771bbda7..427eda061 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 @@ -81,8 +81,8 @@ public interface RoutingManager * * @param routingGroup the routing group to use for backend selection * @param user the user requesting the backend - * @param enforceIsolation whether to enforce isolation + * @param strictRouting whether to enforce isolation * @return the backend configuration for the selected cluster */ - ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user, Boolean enforceIsolation); + ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user, Boolean strictRouting); } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRule.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRule.java index 8bb8c867a..45890c7aa 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRule.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRule.java @@ -24,5 +24,5 @@ public interface RoutingRule Integer getPriority(); - boolean isEnforceIsolation(); + boolean isStrictRouting(); } 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 515a09845..dccca52dd 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 @@ -29,11 +29,11 @@ public record ExternalRouterResponse( @Nullable String routingGroup, List errors, @Nullable Map externalHeaders, - @Nullable Boolean enforceIsolation) + @Nullable Boolean strictRouting) implements RoutingGroupResponse { public ExternalRouterResponse { externalHeaders = externalHeaders == null ? ImmutableMap.of() : ImmutableMap.copyOf(externalHeaders); - enforceIsolation = enforceIsolation != null && enforceIsolation; + strictRouting = strictRouting != null && strictRouting; } } 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/RoutingGroupResponse.java index b8f39bf46..b4ce28673 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/RoutingGroupResponse.java @@ -32,5 +32,5 @@ public interface RoutingGroupResponse Map externalHeaders(); - Boolean enforceIsolation(); + Boolean strictRouting(); } 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 861157a02..1dd06bb27 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 @@ -22,10 +22,10 @@ * Response from the routing service that includes: * - routingGroup: The target routing group for the request (Optional) * - externalHeaders: Headers that can be set in the request (Currently can only be set in ExternalRoutingGroupSelector) - * - enforceIsolation: If true, the handler must not fall back to default when target group has no available backend; + * - strictRouting: If true, the handler must not fall back to default when target group has no available backend; * instead, a 4xx should be returned. */ -public record RoutingSelectorResponse(@Nullable String routingGroup, Map externalHeaders, Boolean enforceIsolation) +public record RoutingSelectorResponse(@Nullable String routingGroup, Map externalHeaders, Boolean strictRouting) implements RoutingGroupResponse { public RoutingSelectorResponse { 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 afae8b467..c340a4edf 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 @@ -77,7 +77,7 @@ void testAddMockBackends() } @Test - void testEnforceIsolationException() + void testStrictRoutingException() { ProxyBackendConfiguration inactiveBackend = createBackend( "inactive-backend", diff --git a/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml b/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml index a22c44783..2046c3079 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml @@ -2,7 +2,7 @@ name: "airflow" description: "if query from airflow, route to adhoc group" priority: 0 -enforceIsolation: false +strictRouting: false actions: - "result.put(\"routingGroup\", \"adhoc\")" condition: "request.getHeader(\"X-Trino-Source\") == \"datagrip\"" diff --git a/gateway-ha/src/test/resources/rules/routing_rules_update.yml b/gateway-ha/src/test/resources/rules/routing_rules_update.yml index 4723b3792..a5ca326e3 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_update.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_update.yml @@ -2,7 +2,7 @@ name: "airflow" description: "if query from airflow, route to etl group" priority: 0 -enforceIsolation: false +strictRouting: false actions: - "result.put(\"routingGroup\", \"etl\")" condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\"" @@ -10,7 +10,7 @@ condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\"" name: "airflow special" description: "if query from airflow with special label, route to etl-special group" priority: 1 -enforceIsolation: false +strictRouting: false actions: - "result.put(\"routingGroup\", \"etl-special\")" condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\" && request.getHeader(\"\ From ac28ac5a5067fc6a78e4c4a1aca31e801791035b Mon Sep 17 00:00:00 2001 From: Peiyingy Date: Thu, 4 Dec 2025 17:05:38 -0500 Subject: [PATCH 3/6] update error msg --- .../java/io/trino/gateway/ha/router/BaseRoutingManager.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 dfc4327d6..f40e080f9 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 @@ -115,7 +115,7 @@ public ProxyBackendConfiguration provideBackendConfiguration(String routingGroup if (backends.isEmpty() && strictRouting) { throw new WebApplicationException( Response.status(NOT_FOUND) - .entity(String.format("No healthy backends available for routing group '%s' under enforced isolation for user '%s'", routingGroup, user)) + .entity(String.format("No healthy backends available for routing group '%s' under strict routing for user '%s'", routingGroup, user)) .build()); } return selectBackend(backends, user).orElseGet(() -> provideDefaultBackendConfiguration(user)); From 54d3c378b31637955a8fb7fbf42a9dd5a76332c0 Mon Sep 17 00:00:00 2001 From: Peiyingy Date: Thu, 4 Dec 2025 17:17:44 -0500 Subject: [PATCH 4/6] address comment --- .../java/io/trino/gateway/ha/router/BaseRoutingManager.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 f40e080f9..04a270da7 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 @@ -112,7 +112,7 @@ public ProxyBackendConfiguration provideBackendConfiguration(String routingGroup List backends = gatewayBackendManager.getActiveBackends(routingGroup).stream() .filter(backEnd -> isBackendHealthy(backEnd.getName())) .toList(); - if (backends.isEmpty() && strictRouting) { + if (strictRouting && backends.isEmpty()) { throw new WebApplicationException( Response.status(NOT_FOUND) .entity(String.format("No healthy backends available for routing group '%s' under strict routing for user '%s'", routingGroup, user)) From a4159fcd135b988edad6ec92a97be5a088925b2b Mon Sep 17 00:00:00 2001 From: Peiyingy Date: Thu, 4 Dec 2025 17:30:37 -0500 Subject: [PATCH 5/6] small fixes --- .../java/io/trino/gateway/ha/router/BaseRoutingManager.java | 2 +- .../main/java/io/trino/gateway/ha/router/RoutingManager.java | 4 ++-- .../trino/gateway/ha/router/schema/RoutingGroupResponse.java | 3 ++- .../gateway/ha/router/schema/RoutingSelectorResponse.java | 4 ++-- 4 files changed, 7 insertions(+), 6 deletions(-) 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 04a270da7..74d374aaf 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 @@ -107,7 +107,7 @@ public ProxyBackendConfiguration provideDefaultBackendConfiguration(String user) * has no suitable backend unless {@code strictRouting} is true, in which case a 404 is returned. */ @Override - public ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user, Boolean strictRouting) + public ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user, boolean strictRouting) { List backends = gatewayBackendManager.getActiveBackends(routingGroup).stream() .filter(backEnd -> isBackendHealthy(backEnd.getName())) 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 427eda061..2f80c841d 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 @@ -81,8 +81,8 @@ public interface RoutingManager * * @param routingGroup the routing group to use for backend selection * @param user the user requesting the backend - * @param strictRouting whether to enforce isolation + * @param strictRouting whether to force strict routing * @return the backend configuration for the selected cluster */ - ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user, Boolean strictRouting); + ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user, boolean strictRouting); } 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/RoutingGroupResponse.java index b4ce28673..852e7b820 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/RoutingGroupResponse.java @@ -25,6 +25,7 @@ Implementations of this interface are used to: * Specify the target routing group for a request * Provide additional headers that should be added to the request + * Specify whether strict routing should be used */ public interface RoutingGroupResponse { @@ -32,5 +33,5 @@ public interface RoutingGroupResponse Map externalHeaders(); - Boolean strictRouting(); + @Nullable Boolean strictRouting(); } 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 1dd06bb27..d335846ab 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 @@ -22,10 +22,10 @@ * Response from the routing service that includes: * - routingGroup: The target routing group for the request (Optional) * - externalHeaders: Headers that can be set in the request (Currently can only be set in ExternalRoutingGroupSelector) - * - strictRouting: If true, the handler must not fall back to default when target group has no available backend; + * - strictRouting: If true, the handler must not fall back to default when target group has no available backend (Optional) * instead, a 4xx should be returned. */ -public record RoutingSelectorResponse(@Nullable String routingGroup, Map externalHeaders, Boolean strictRouting) +public record RoutingSelectorResponse(@Nullable String routingGroup, Map externalHeaders, @Nullable Boolean strictRouting) implements RoutingGroupResponse { public RoutingSelectorResponse { From eb736ee547d7cca82b78026a4cd275b718c49d07 Mon Sep 17 00:00:00 2001 From: Peiyingy Date: Fri, 5 Dec 2025 16:18:03 -0500 Subject: [PATCH 6/6] small fixes --- .../io/trino/gateway/ha/handler/RoutingTargetHandler.java | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) 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 d7d6a79a6..5439829c1 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 @@ -91,11 +91,14 @@ private RoutingTargetResponse getRoutingTargetResponse(HttpServletRequest reques RoutingSelectorResponse routingDestination = routingGroupSelector.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. + // When no cluster is found: + // - If strictRouting is false, fall back to the default routing group backend. + // - If strictRouting is true, return a 404 response. String routingGroup = !isNullOrEmpty(routingDestination.routingGroup()) ? routingDestination.routingGroup() : defaultRoutingGroup; - ProxyBackendConfiguration backendConfiguration = routingManager.provideBackendConfiguration(routingGroup, user, routingDestination.strictRouting()); + boolean strictRouting = Optional.ofNullable(routingDestination.strictRouting()).orElse(false); + ProxyBackendConfiguration backendConfiguration = routingManager.provideBackendConfiguration(routingGroup, user, strictRouting); String clusterHost = backendConfiguration.getProxyTo(); String externalUrl = backendConfiguration.getExternalUrl(); // Apply headers from RoutingDestination if there are any