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

@@ -26,20 +26,23 @@
  * @param name name of the routing rule
  * @param description description of the routing rule. Defaults to an empty string if not provided, indicating the user intends it to be blank.
  * @param priority priority of the routing rule. Higher number represents higher priority. If two rules have same priority then order of execution is not guaranteed.
+ * @param strictRouting if true, return 503 instead of falling back to the default backend when no healthy backend is available for the routing group
  * @param actions actions of the routing rule
  * @param condition condition of the routing rule
  */
 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

@@ -40,6 +40,7 @@
 import static io.trino.gateway.ha.handler.ProxyUtils.buildUriWithNewCluster;
 import static io.trino.gateway.ha.handler.ProxyUtils.extractQueryIdIfPresent;
 import static java.util.Objects.requireNonNull;
+import static java.util.Objects.requireNonNullElse;
 
 public class RoutingTargetHandler
 {
@@ -91,11 +92,11 @@ 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.
         String routingGroup = !isNullOrEmpty(routingDestination.routingGroup())
                 ? routingDestination.routingGroup()
                 : defaultRoutingGroup;
-        ProxyBackendConfiguration backendConfiguration = routingManager.provideBackendConfiguration(routingGroup, user);
+        boolean strictRouting = requireNonNullElse(routingDestination.strictRouting(), 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

@@ -31,12 +31,14 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 import static com.google.common.base.Suppliers.memoizeWithExpiration;
 import static io.trino.gateway.ha.handler.HttpUtils.TRINO_QUERY_PROPERTIES;
 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 +74,18 @@ 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);
+                String previousRoutingGroup = result.get(RESULTS_ROUTING_GROUP_KEY);
                 rule.evaluateAction(result, data, state);
+                if (!Objects.equals(previousRoutingGroup, result.get(RESULTS_ROUTING_GROUP_KEY))) {
+                    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

@@ -24,14 +24,17 @@
  * - routingGroup: The target routing group for the request (optional)
  * - errors: Any errors that occurred during routing
  * - externalHeaders: Headers that can be set in the request
+ * - strictRouting: If true, a 503 error is returned instead of falling back when no backend is available for the routing group
  */
 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\""