Modify external routing to configurably pass errors to client

EdenKik · Chaho12 · commit b40f08723894 · 2025-07-13T16:16:50.000+09:00
diff --git a/docs/routing-rules.md b/docs/routing-rules.md
@@ -3,8 +3,14 @@
 Trino Gateway includes a routing rules engine.
 
 By default, Trino Gateway reads the `X-Trino-Routing-Group` request header to
-route requests. If this header is not specified, requests are sent to the
-default routing group called `adhoc`.
+route requests. Use the `defaultRoutingGroup` setting to specify a fallback group, 
+defaults to the `adhoc` group.
+
+
+```yaml
+routing:
+    defaultRoutingGroup: "test-group"
+```
 
 The routing rules engine feature enables you to either write custom logic to
 route requests based on the request info such as any of the [request
@@ -40,13 +46,16 @@ routingRules:
         excludeHeaders:
             - 'Authorization'
             - 'Accept-Encoding'
+        propagateErrors: false
 ```
 
 * Redirect URLs are not supported.
 * Optionally, add headers to the `excludeHeaders` list to exclude requests with
   corresponding header values from being sent in the POST request.
 * Check headers to exclude when making API requests, specifics depend on the
   network configuration.
+* Set `propagateErrors` to `true` to forward routing service errors to  
+  clients if present in the response.
 
 If there is error parsing the routing rules configuration file, an error is
 logged, and requests are routed using the routing group header
@@ -93,7 +102,7 @@ return a result with the following criteria:
 * Response status code of OK (200)
 * Message in JSON format
 * Only one group can be returned
-* If errors is not null, then query would route to default routing group adhoc
+* If `errors` is not null, the query is routed to the configured default group
 
 #### Request headers modification
 
@@ -151,7 +160,7 @@ In addition to the default objects, rules may optionally utilize
 , which provide information about the user and query respectively.
 You must include an action of the form `result.put(\"routingGroup\", \"foo\")`
 to trigger routing of a request that satisfies the condition to the specific
-routing group. Without this action, the default adhoc group is used and the
+routing group. Without this action, the configured default group is used and the
 whole routing rule is redundant.
 
 The condition and actions are written in [MVEL](http://mvel.documentnode.com/),
@@ -184,8 +193,7 @@ You can use the `contains` operator
 condition: 'request.getHeader("X-Trino-Client-Tags") contains "label=foo"'
 ```
 
-If no rules match, then the request is routed to the default `adhoc` routing
-group.
+If no rules match, then the request is routed to the configured default routing group.
 
 ### TrinoStatus
 
diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/config/RulesExternalConfiguration.java b/gateway-ha/src/main/java/io/trino/gateway/ha/config/RulesExternalConfiguration.java
@@ -19,6 +19,7 @@ public class RulesExternalConfiguration
 {
     private String urlPath;
     private List<String> excludeHeaders;
+    private boolean propagateErrors;
 
     public String getUrlPath()
     {
@@ -39,4 +40,14 @@ public void setExcludeHeaders(List<String> excludeHeaders)
     {
         this.excludeHeaders = excludeHeaders;
     }
+
+    public boolean isPropagateErrors()
+    {
+        return this.propagateErrors;
+    }
+
+    public void setPropagateErrors(boolean propagateErrors)
+    {
+        this.propagateErrors = propagateErrors;
+    }
 }
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
@@ -54,6 +54,7 @@ public class RoutingTargetHandler
     private static final Logger log = Logger.get(RoutingTargetHandler.class);
     private final RoutingManager routingManager;
     private final RoutingGroupSelector routingGroupSelector;
+    private final String defaultRoutingGroup;
     private final List<String> statementPaths;
     private final List<Pattern> extraWhitelistPaths;
     private final boolean requestAnalyserClientsUseV2Format;
@@ -68,6 +69,7 @@ public RoutingTargetHandler(
     {
         this.routingManager = requireNonNull(routingManager);
         this.routingGroupSelector = requireNonNull(routingGroupSelector);
+        this.defaultRoutingGroup = haGatewayConfiguration.getRouting().getDefaultRoutingGroup();
         statementPaths = requireNonNull(haGatewayConfiguration.getStatementPaths());
         extraWhitelistPaths = requireNonNull(haGatewayConfiguration.getExtraWhitelistPaths()).stream().map(Pattern::compile).collect(toImmutableList());
         requestAnalyserClientsUseV2Format = haGatewayConfiguration.getRequestAnalyzerConfig().isClientsUseV2Format();
@@ -79,15 +81,17 @@ public RoutingTargetResponse resolveRouting(HttpServletRequest request)
     {
         Optional<String> queryId = extractQueryIdIfPresent(request, statementPaths, requestAnalyserClientsUseV2Format, requestAnalyserMaxBodySize);
         Optional<String> previousCluster = getPreviousCluster(queryId, request);
+
         RoutingTargetResponse routingTargetResponse = previousCluster.map(cluster -> {
             String routingGroup = queryId.map(routingManager::findRoutingGroupForQueryId)
-                    .orElse("adhoc");
+                    .orElse(defaultRoutingGroup);
             String externalUrl = queryId.map(routingManager::findExternalUrlForQueryId)
                     .orElse(cluster);
             return new RoutingTargetResponse(
                     new RoutingDestination(routingGroup, cluster, buildUriWithNewCluster(cluster, request), externalUrl),
                     request);
         }).orElse(getRoutingTargetResponse(request));
+
         logRewrite(routingTargetResponse.routingDestination().clusterHost(), request);
         return routingTargetResponse;
     }
@@ -96,10 +100,11 @@ private RoutingTargetResponse getRoutingTargetResponse(HttpServletRequest reques
     {
         RoutingSelectorResponse routingDestination = routingGroupSelector.findRoutingDestination(request);
         String user = request.getHeader(USER_HEADER);
-        // This falls back on adhoc routing group if there is no cluster found (or value is empty) for the routing group.
-        String routingGroup = (routingDestination.routingGroup() != null && !routingDestination.routingGroup().isEmpty())
+
+        // 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()
-                : "adhoc";
+                : defaultRoutingGroup;
         ProxyBackendConfiguration backendConfiguration = routingManager.provideBackendConfiguration(routingGroup, user);
         String clusterHost = backendConfiguration.getProxyTo();
         String externalUrl = backendConfiguration.getExternalUrl();
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
@@ -29,6 +29,8 @@
 import io.trino.gateway.ha.router.schema.RoutingGroupExternalBody;
 import io.trino.gateway.ha.router.schema.RoutingSelectorResponse;
 import jakarta.servlet.http.HttpServletRequest;
+import jakarta.ws.rs.WebApplicationException;
+import jakarta.ws.rs.core.Response;
 
 import java.net.URI;
 import java.net.URISyntaxException;
@@ -37,6 +39,7 @@
 import java.util.Optional;
 import java.util.Set;
 
+import static com.google.common.base.Throwables.throwIfInstanceOf;
 import static com.google.common.net.HttpHeaders.CONTENT_TYPE;
 import static com.google.common.net.MediaType.JSON_UTF_8;
 import static io.airlift.http.client.JsonBodyGenerator.jsonBodyGenerator;
@@ -52,6 +55,7 @@ public class ExternalRoutingGroupSelector
     private static final Logger log = Logger.get(ExternalRoutingGroupSelector.class);
     private final Set<String> excludeHeaders;
     private final URI uri;
+    private final boolean propagateErrors;
     private final HttpClient httpClient;
     private final RequestAnalyzerConfig requestAnalyzerConfig;
     private final TrinoRequestUser.TrinoRequestUserProvider trinoRequestUserProvider;
@@ -67,6 +71,7 @@ public class ExternalRoutingGroupSelector
                 .add("Content-Length")
                 .addAll(rulesExternalConfiguration.getExcludeHeaders())
                 .build();
+        propagateErrors = rulesExternalConfiguration.isPropagateErrors();
 
         this.requestAnalyzerConfig = requestAnalyzerConfig;
         trinoRequestUserProvider = new TrinoRequestUser.TrinoRequestUserProvider(requestAnalyzerConfig);
@@ -101,7 +106,13 @@ public RoutingSelectorResponse findRoutingDestination(HttpServletRequest servlet
                 throw new RuntimeException("Unexpected response: null");
             }
             else if (response.errors() != null && !response.errors().isEmpty()) {
-                throw new RuntimeException("Response with error: " + String.join(", ", response.errors()));
+                if (propagateErrors) {
+                    log.warn("Query validation failed with errors: %s", String.join(", ", response.errors()));
+                    throw new WebApplicationException(
+                            Response.status(Response.Status.BAD_REQUEST)
+                                    .entity(response.errors())
+                                    .build());
+                }
             }
 
             // Filter out excluded headers and null values
@@ -120,6 +131,7 @@ else if (response.errors() != null && !response.errors().isEmpty()) {
             return new RoutingSelectorResponse(response.routingGroup(), filteredHeaders);
         }
         catch (Exception e) {
+            throwIfInstanceOf(e, WebApplicationException.class);
             log.error(e, "Error occurred while retrieving routing group "
                     + "from external routing rules processing at " + uri);
         }
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
@@ -125,7 +125,7 @@ public void setRoutingGroupForQueryId(String queryId, String routingGroup)
     /**
      * Performs routing to a default backend.
      */
-    private ProxyBackendConfiguration provideDefaultCluster()
+    private ProxyBackendConfiguration provideDefaultBackendConfiguration()
     {
         List<ProxyBackendConfiguration> backends = gatewayBackendManager.getActiveDefaultBackends();
         backends.removeIf(backend -> isBackendNotHealthy(backend.getName()));
@@ -136,6 +136,11 @@ private ProxyBackendConfiguration provideDefaultCluster()
         return backends.get(backendId);
     }
 
+    public String provideDefaultCluster(String user)
+    {
+        return provideDefaultBackendConfiguration().getProxyTo();
+    }
+
     /**
      * Performs routing to a given cluster group. This falls back to a default backend, if no scheduled
      * backend is found.
@@ -146,7 +151,7 @@ public ProxyBackendConfiguration provideBackendConfiguration(String routingGroup
                 gatewayBackendManager.getActiveBackends(routingGroup);
         backends.removeIf(backend -> isBackendNotHealthy(backend.getName()));
         if (backends.isEmpty()) {
-            return provideDefaultCluster();
+            return provideDefaultBackendConfiguration();
         }
         int backendId = Math.abs(RANDOM.nextInt()) % backends.size();
         return backends.get(backendId);
diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestRoutingTargetHandler.java b/gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestRoutingTargetHandler.java
@@ -27,20 +27,21 @@
 import io.trino.gateway.ha.router.schema.ExternalRouterResponse;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.ws.rs.HttpMethod;
+import jakarta.ws.rs.WebApplicationException;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mockito;
 import org.mockito.junit.jupiter.MockitoExtension;
 
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
 import static io.trino.gateway.ha.handler.HttpUtils.USER_HEADER;
 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.Mockito.when;
 
@@ -60,6 +61,8 @@ static HaGatewayConfiguration provideGatewayConfiguration()
         HaGatewayConfiguration config = new HaGatewayConfiguration();
         config.setRequestAnalyzerConfig(new RequestAnalyzerConfig());
 
+        config.getRouting().setDefaultRoutingGroup("default-group");
+
         // Configure excluded headers
         RulesExternalConfiguration rulesExternalConfig = new RulesExternalConfiguration();
         rulesExternalConfig.setExcludeHeaders(List.of("Authorization", "Cookie"));
@@ -109,7 +112,9 @@ void setUp()
         request = prepareMockRequest();
 
         // Initialize the handler with the configuration
-        handler = new RoutingTargetHandler(routingManager, RoutingGroupSelector.byRoutingExternal(httpClient, config.getRoutingRules().getRulesExternalConfiguration(), config.getRequestAnalyzerConfig()), config);
+        handler = new RoutingTargetHandler(
+                routingManager,
+                RoutingGroupSelector.byRoutingExternal(httpClient, config.getRoutingRules().getRulesExternalConfiguration(), config.getRequestAnalyzerConfig()), config);
     }
 
     @Test
@@ -220,29 +225,86 @@ void testEmptyRoutingGroup()
         RoutingTargetResponse response = handler.resolveRouting(request);
 
         // Verify that when no routing group header is set, we default to "adhoc"
-        assertThat(response.routingDestination().routingGroup()).isEqualTo("adhoc");
+        assertThat(response.routingDestination().routingGroup()).isEqualTo("default-group");
         assertThat(response.modifiedRequest().getHeader("X-Empty-Group-Header"))
                 .isEqualTo("should-be-set");
     }
 
     @Test
-    void testErrorsGiven()
-            throws Exception
+    void testResponsePropertiesNull()
+    {
+        ExternalRouterResponse mockResponse = new ExternalRouterResponse(null, null, ImmutableMap.of());
+        when(httpClient.execute(any(), any())).thenReturn(mockResponse);
+
+        RoutingTargetResponse result = handler.resolveRouting(request);
+
+        assertThat(result.routingDestination().routingGroup()).isEqualTo("default-group");
+    }
+
+    @Test
+    void testResponseGroupSetResponseErrorsNull()
     {
-        // Setup routing group selector response with errors
-        Map<String, String> modifiedHeaders = ImmutableMap.of(
-                "X-New-Header", "new-value");
-        List<String> someErrors = Arrays.asList("ErrorA", "ErrorB");
         ExternalRouterResponse mockResponse = new ExternalRouterResponse(
-                "test-group",  // This value should be ignored due to errors
-                someErrors,
-                modifiedHeaders);
+                "test-group", null, ImmutableMap.of());
         when(httpClient.execute(any(), any())).thenReturn(mockResponse);
 
-        // Execute
-        RoutingTargetResponse response = handler.resolveRouting(request);
+        RoutingTargetResponse result = handler.resolveRouting(request);
+
+        assertThat(result.routingDestination().routingGroup()).isEqualTo("test-group");
+    }
+
+    @Test
+    void testPropagateErrorsFalseResponseGroupNullResponseErrorsSet()
+    {
+        ExternalRouterResponse mockResponse = new ExternalRouterResponse(null, List.of("some-error"), ImmutableMap.of());
+        when(httpClient.execute(any(), any())).thenReturn(mockResponse);
+
+        RoutingTargetResponse result = handler.resolveRouting(request);
+
+        assertThat(result.routingDestination().routingGroup()).isEqualTo("default-group");
+    }
+
+    @Test
+    void testPropagateErrorsFalseResponseGroupAndErrorsSet()
+    {
+        ExternalRouterResponse mockResponse = new ExternalRouterResponse("test-group", List.of("some-error"), ImmutableMap.of());
+        when(httpClient.execute(any(), any())).thenReturn(mockResponse);
+
+        RoutingTargetResponse result = handler.resolveRouting(request);
+
+        assertThat(result.routingDestination().routingGroup()).isEqualTo("test-group");
+    }
+
+    @Test
+    void testPropagateErrorsTrueResponseGroupNullResponseErrorsSet()
+    {
+        RoutingTargetHandler handler = createHandlerWithPropagateErrorsTrue();
 
-        // Verify that when errors are present, we default to "adhoc" group
-        assertThat(response.routingDestination().routingGroup()).isEqualTo("adhoc");
+        config.getRoutingRules().getRulesExternalConfiguration().setPropagateErrors(true);
+        ExternalRouterResponse mockResponse = new ExternalRouterResponse(null, List.of("some-error"), ImmutableMap.of());
+        when(httpClient.execute(any(), any())).thenReturn(mockResponse);
+
+        assertThatThrownBy(() -> handler.resolveRouting(request))
+                .isInstanceOf(WebApplicationException.class);
+    }
+
+    @Test
+    void testPropagateErrorsTrueResponseGroupAndErrorsSet()
+    {
+        RoutingTargetHandler handler = createHandlerWithPropagateErrorsTrue();
+
+        ExternalRouterResponse response = new ExternalRouterResponse("test-group", List.of("some-error"), ImmutableMap.of());
+        when(httpClient.execute(any(), any())).thenReturn(response);
+
+        assertThatThrownBy(() -> handler.resolveRouting(request))
+                .isInstanceOf(WebApplicationException.class);
+    }
+
+    private RoutingTargetHandler createHandlerWithPropagateErrorsTrue()
+    {
+        config.getRoutingRules().getRulesExternalConfiguration().setPropagateErrors(true);
+        return new RoutingTargetHandler(
+                routingManager,
+                RoutingGroupSelector.byRoutingExternal(httpClient, config.getRoutingRules().getRulesExternalConfiguration(), config.getRequestAnalyzerConfig()), config);
     }
 }
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
diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestHaGatewayManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestHaGatewayManager.java
diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingManagerExternalUrlCache.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingManagerExternalUrlCache.java

Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,7 @@ public class RulesExternalConfiguration`
`19`	`19`	`{`
`20`	`20`	`private String urlPath;`
`21`	`21`	`private List<String> excludeHeaders;`
	`22`	`+ private boolean propagateErrors;`
`22`	`23`
`23`	`24`	`public String getUrlPath()`
`24`	`25`	`{`
`@@ -39,4 +40,14 @@ public void setExcludeHeaders(List<String> excludeHeaders)`
`39`	`40`	`{`
`40`	`41`	`this.excludeHeaders = excludeHeaders;`
`41`	`42`	`}`
	`43`	`+`
	`44`	`+ public boolean isPropagateErrors()`
	`45`	`+ {`
	`46`	`+ return this.propagateErrors;`
	`47`	`+ }`
	`48`	`+`
	`49`	`+ public void setPropagateErrors(boolean propagateErrors)`
	`50`	`+ {`
	`51`	`+ this.propagateErrors = propagateErrors;`
	`52`	`+ }`
`42`	`53`	`}`