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..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,6 +33,7 @@ public record RoutingRule( String name, String description, Integer priority, + Boolean strictRouting, List actions, String condition) { @@ -40,6 +41,7 @@ public record RoutingRule( requireNonNull(name, "name is null"); description = requireNonNullElse(description, ""); priority = requireNonNullElse(priority, 0); + 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 16521c44e..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); + 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 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..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 @@ -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 strictRouting} 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 strictRouting) { List backends = gatewayBackendManager.getActiveBackends(routingGroup).stream() .filter(backEnd -> isBackendHealthy(backEnd.getName())) .toList(); + 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)) + .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..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); + 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 96106b889..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 @@ -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 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); + strictRouting = rule.isStrictRouting(); } - }); - return new RoutingSelectorResponse(result.get(RESULTS_ROUTING_GROUP_KEY)); + } + 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 8b2700460..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,6 +37,7 @@ public class MVELRoutingRule String name; String description; Integer priority; + boolean strictRouting; 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("strictRouting") Boolean strictRouting, @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.strictRouting = requireNonNullElse(strictRouting, 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 isStrictRouting() + { + return strictRouting; + } + @Override public int compareTo(RoutingRule o) { @@ -132,6 +141,7 @@ public String toString() .add("name", name) .add("description", description) .add("priority", priority) + .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 2f112f6dd..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,7 +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 force strict routing * @return the backend configuration for the selected cluster */ - ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user); + 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 d9110aeff..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 @@ -23,4 +23,6 @@ public interface RoutingRule void evaluateAction(Map result, Map data, Map state); Integer getPriority(); + + 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 88b6db164..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 @@ -28,10 +28,12 @@ public record ExternalRouterResponse( @Nullable String routingGroup, List errors, - @Nullable Map externalHeaders) + @Nullable Map externalHeaders, + @Nullable Boolean strictRouting) implements RoutingGroupResponse { public ExternalRouterResponse { externalHeaders = externalHeaders == null ? ImmutableMap.of() : ImmutableMap.copyOf(externalHeaders); + 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 73b667d3d..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,10 +25,13 @@ 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 { @Nullable String routingGroup(); Map externalHeaders(); + + @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 60b44900a..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,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) + * - 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) +public record RoutingSelectorResponse(@Nullable String routingGroup, Map externalHeaders, @Nullable Boolean strictRouting) 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..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 @@ -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 testStrictRoutingException() + { + 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..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,6 +2,7 @@ name: "airflow" description: "if query from airflow, route to adhoc group" priority: 0 +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 3478bc45e..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,6 +2,7 @@ name: "airflow" description: "if query from airflow, route to etl group" priority: 0 +strictRouting: 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 +strictRouting: false actions: - "result.put(\"routingGroup\", \"etl-special\")" condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\" && request.getHeader(\"\