Implement Strict Routing Semantics for Routing Rules

docs/routing-rules.md

@@ -149,6 +149,26 @@ actions:
   - 'result.put("routingGroup", "etl-special")'
 ```
 
+#### Strict routing
+
+By default, if no healthy backend is available for the resolved routing group,
+Trino Gateway falls back to the default routing group. You can override this
+behavior by setting `strictRouting` to `true` on a rule. When strict routing is
+enabled and no healthy backend is available for the routing group, Trino Gateway
+returns a `503 Service Unavailable` error to the client instead of falling back.
+
+```yaml
+---
+name: "airflow"
+description: "if query from airflow, route to etl group with strict routing"
+strictRouting: true
+condition: 'request.getHeader("X-Trino-Source") == "airflow"'
+actions:
+  - 'result.put("routingGroup", "etl")'
+```
+
+The `strictRouting` property defaults to `false` if not specified.
+
 Three objects are available by default. They are
 * `request`, the incoming request as an [HttpServletRequest](https://docs.oracle.com/javaee/6/api/javax/servlet/http/HttpServletRequest.html)
 * `state`, a `HashMap<String, Object>` that allows passing arbitrary state from one rule evaluation to the next

gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java

@@ -33,13 +33,15 @@ public record RoutingRule(
         String name,
         String description,
         Integer priority,
+        Boolean strictRouting,
         List<String> actions,
         String condition)
 {
     public 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");
     }

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

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.SERVICE_UNAVAILABLE;
+
 /**
  * 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 503 is returned.
      */
     @Override
-    public ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user)
+    public ProxyBackendConfiguration provideBackendConfiguration(String routingGroup, String user, boolean strictRouting)
     {
         List<ProxyBackendConfiguration> backends = gatewayBackendManager.getActiveBackends(routingGroup).stream()
                 .filter(backEnd -> isBackendHealthy(backEnd.getName()))
                 .toList();
+        if (strictRouting && backends.isEmpty()) {
+            throw new WebApplicationException(
+                             Response.status(SERVICE_UNAVAILABLE)
+                            .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));
     }
 

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);

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<RoutingRule> readRulesFromPath(Path rulesPath)

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<Serializable> 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<Serializable> 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)

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);
 }

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

@@ -23,4 +23,6 @@ public interface RoutingRule
     void evaluateAction(Map<String, String> result, Map<String, Object> data, Map<String, Object> state);
 
     Integer getPriority();
+
+    boolean isStrictRouting();
 }

gateway-ha/src/main/java/io/trino/gateway/ha/router/schema/ExternalRouterResponse.java

@@ -28,10 +28,12 @@
 public record ExternalRouterResponse(
         @Nullable String routingGroup,
         List<String> errors,
-        @Nullable Map<String, String> externalHeaders)
+        @Nullable Map<String, String> externalHeaders,
+        @Nullable Boolean strictRouting)
         implements RoutingGroupResponse
 {
     public ExternalRouterResponse {
         externalHeaders = externalHeaders == null ? ImmutableMap.of() : ImmutableMap.copyOf(externalHeaders);
+        strictRouting = strictRouting != null && strictRouting;
     }
 }

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<String, String> externalHeaders();
+
+    @Nullable Boolean strictRouting();
 }

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<String, String> externalHeaders)
+public record RoutingSelectorResponse(@Nullable String routingGroup, Map<String, String> externalHeaders, @Nullable Boolean strictRouting)
         implements RoutingGroupResponse
 {
     public RoutingSelectorResponse {
@@ -32,6 +34,6 @@ public record RoutingSelectorResponse(@Nullable String routingGroup, Map<String,
 
     public RoutingSelectorResponse(String routingGroup)
     {
-        this(routingGroup, ImmutableMap.of());
+        this(routingGroup, ImmutableMap.of(), false);
     }
 }

gateway-ha/src/main/resources/routing_rules.yml

@@ -0,0 +1,7 @@
+---
+name: "airflow"
+description: "if query from airflow, route to adhoc"
+priority: 0
+actions:
+- "result.put(\"routingGroup\", \"adhoc\")"
+condition: "request.getHeader(\"X-Trino-Source\") == \"datagrip\""

gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestRoutingTargetHandler.java

@@ -43,6 +43,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.Mockito.when;
 
 @ExtendWith(MockitoExtension.class)
@@ -108,7 +109,7 @@ 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(), anyBoolean())).thenReturn(new ProxyBackendConfiguration());
         request = prepareMockRequest();
 
         // Initialize the handler with the configuration
@@ -128,7 +129,8 @@ void testBasicHeaderModification()
         ExternalRouterResponse mockResponse = new ExternalRouterResponse(
                 "test-group",
                 Collections.emptyList(),
-                modifiedHeaders);
+                modifiedHeaders,
+                false);
         when(httpClient.execute(any(), any())).thenReturn(mockResponse);
 
         // Execute
@@ -152,7 +154,8 @@ void testExcludedHeaders()
         ExternalRouterResponse mockResponse = new ExternalRouterResponse(
                 "test-group",
                 Collections.emptyList(),
-                modifiedHeaders);
+                modifiedHeaders,
+                false);
         when(httpClient.execute(any(), any())).thenReturn(mockResponse);
 
         // Execute
@@ -173,7 +176,8 @@ void testNoHeaderModification()
         ExternalRouterResponse mockResponse = new ExternalRouterResponse(
                 "test-group",
                 Collections.emptyList(),
-                ImmutableMap.of());
+                ImmutableMap.of(),
+                false);
         when(httpClient.execute(any(), any())).thenReturn(mockResponse);
 
         // Execute
@@ -195,7 +199,8 @@ void testEmptyHeader()
         ExternalRouterResponse mockResponse = new ExternalRouterResponse(
                 "test-group",
                 Collections.emptyList(),
-                modifiedHeaders);
+                modifiedHeaders,
+                false);
         when(httpClient.execute(any(), any())).thenReturn(mockResponse);
 
         // Execute
@@ -218,7 +223,8 @@ void testEmptyRoutingGroup()
         ExternalRouterResponse mockResponse = new ExternalRouterResponse(
                 "",
                 Collections.emptyList(),
-                modifiedHeaders);
+                modifiedHeaders,
+                false);
         when(httpClient.execute(any(), any())).thenReturn(mockResponse);
 
         // Execute
@@ -233,7 +239,7 @@ void testEmptyRoutingGroup()
     @Test
     void testResponsePropertiesNull()
     {
-        ExternalRouterResponse mockResponse = new ExternalRouterResponse(null, null, ImmutableMap.of());
+        ExternalRouterResponse mockResponse = new ExternalRouterResponse(null, null, ImmutableMap.of(), null);
         when(httpClient.execute(any(), any())).thenReturn(mockResponse);
 
         RoutingTargetResponse result = handler.resolveRouting(request);
@@ -245,7 +251,7 @@ void testResponsePropertiesNull()
     void testResponseGroupSetResponseErrorsNull()
     {
         ExternalRouterResponse mockResponse = new ExternalRouterResponse(
-                "test-group", null, ImmutableMap.of());
+                "test-group", null, ImmutableMap.of(), null);
         when(httpClient.execute(any(), any())).thenReturn(mockResponse);
 
         RoutingTargetResponse result = handler.resolveRouting(request);
@@ -256,7 +262,7 @@ void testResponseGroupSetResponseErrorsNull()
     @Test
     void testPropagateErrorsFalseResponseGroupNullResponseErrorsSet()
     {
-        ExternalRouterResponse mockResponse = new ExternalRouterResponse(null, List.of("some-error"), ImmutableMap.of());
+        ExternalRouterResponse mockResponse = new ExternalRouterResponse(null, List.of("some-error"), ImmutableMap.of(), null);
         when(httpClient.execute(any(), any())).thenReturn(mockResponse);
 
         RoutingTargetResponse result = handler.resolveRouting(request);
@@ -267,7 +273,7 @@ void testPropagateErrorsFalseResponseGroupNullResponseErrorsSet()
     @Test
     void testPropagateErrorsFalseResponseGroupAndErrorsSet()
     {
-        ExternalRouterResponse mockResponse = new ExternalRouterResponse("test-group", List.of("some-error"), ImmutableMap.of());
+        ExternalRouterResponse mockResponse = new ExternalRouterResponse("test-group", List.of("some-error"), ImmutableMap.of(), null);
         when(httpClient.execute(any(), any())).thenReturn(mockResponse);
 
         RoutingTargetResponse result = handler.resolveRouting(request);
@@ -281,7 +287,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, List.of("some-error"), ImmutableMap.of(), null);
         when(httpClient.execute(any(), any())).thenReturn(mockResponse);
 
         assertThatThrownBy(() -> 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))