Support HttpClient options for external routing group selector

Author: Chaho12

docs/routing-rules.md

@@ -23,7 +23,7 @@ To enable the routing rules engine, find the following lines in
 
 * Set `rulesEngineEnabled` to `true`, then `rulesType` as `FILE` or `EXTERNAL`.
 * If you set `rulesType: FILE`, then set `rulesConfigPath` to the path to your
-  rules config file. 
+  rules config file.
 * If you set `rulesType: EXTERNAL`, set `rulesExternalConfiguration` to the URL
   of an external service for routing rules processing.
 * `rulesType` is by default `FILE` unless specified.
@@ -50,15 +50,30 @@ If there is error parsing the routing rules configuration file, an error is
 logged, and requests are routed using the routing group header
 `X-Trino-Routing-Group` as default.
 
+### Configuring API Requests with HttpClient Config
+
+Trino Gateway uses the Airlift HttpClient library when making HTTP requests.
+You can configure the HttpClient by adding the following configuration to
+the `serverConfig:` section with the `router` prefix.
+
+```yaml
+serverConfig:
+    router.http-client.request-timeout: 1s
+```
+
+Please refer to the [Trino HTTP Client properties](
+https://trino.io/docs/current/admin/properties-http-client.html)
+documentation for more.
+
 ### Use an external service for routing rules
 
 You can use an external service for processing your routing by setting the
 `rulesType` to `EXTERNAL` and configuring the `rulesExternalConfiguration`.
 
-Trino Gateway then sends all headers, other than those specified in 
+Trino Gateway then sends all headers, other than those specified in
 `excludeHeaders`, as a map in the body of a POST request to the external
 service. If `requestAnalyzerConfig.analyzeRequest` is set to `true`,
-`TrinoRequestUser` and `TrinoQueryProperties` are also included. 
+`TrinoRequestUser` and `TrinoQueryProperties` are also included.
 
 Additionally, the following HTTP information is included:
 
@@ -77,7 +92,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, then query would route to default routing group adhoc
 
 ```json
 {
@@ -120,8 +135,8 @@ object called `request`. Rules may also utilize
 [trinoQueryProperties](#trinoqueryproperties)
 objects, 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 
+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
 whole routing rule is redundant.
 
 The condition and actions are written in [MVEL](http://mvel.documentnode.com/),
@@ -146,7 +161,7 @@ group.
 ### TrinoStatus
 
 The `TrinoStatus` class attempts to track the current state of the configured
-Trino clusters.  The three possible states of these cluster are updated with
+Trino clusters. The three possible states of these cluster are updated with
 every healthcheck:
 
 - `PENDING`: A Trino cluster shows this state when it is still starting up. It
@@ -169,7 +184,7 @@ request, in the following order:
 3. `Authorization: Bearer` header. Requires configuring an OAuth2 User Info URL.
 4. `Trino-UI-Token` or `__Secure-Trino-ID-Token` cookie.
 
-Kerberos and Certificate authentication are not currently supported. If the 
+Kerberos and Certificate authentication are not currently supported. If the
 request contains the `Authorization: Bearer` header, an attempt is made to treat
 the token as a JWT and deserialize it. If this is successful, the value of the
 claim named in `requestAnalyzerConfig.tokenUserField` is used as the username.
@@ -201,7 +216,7 @@ syntactic analysis is performed.
 
 If a query references a view, then that view is not expanded, and tables
 referenced by the view are not recognized. Views and materialized views are
-treated as tables and  added to the list of tables in all contexts, including
+treated as tables and added to the list of tables in all contexts, including
 statements such as `CREATE VIEW`.
 
 A routing rule can call the following methods on the `trinoQueryProperties`
@@ -210,7 +225,7 @@ object:
 * `String errorMessage()`: the error message, only if there was any error while
   creating the `trinoQueryProperties` object.
 * `boolean isNewQuerySubmission()`: boolean flag to indicate if the
-  request is  a POST to the `v1/statement` query endpoint.
+  request is a POST to the `v1/statement` query endpoint.
 * `String getQueryType()`: the class name of the `Statement`, e.g. `ShowCreate`.
   Note that these are not mapped to the `ResourceGroup` query types. For a full
   list of potential query types, see the classes in
@@ -220,7 +235,7 @@ object:
   the Trino documentation](https://trino.io/docs/current/admin/resource-groups.html#selector-rules)
 * `String getDefaultCatalog()`: the default catalog, if set. It may or may not
   be referenced in the actual SQL
-* `String getDefaultSchema()`: the default schema,  if set. It may or may not
+* `String getDefaultSchema()`: the default schema, if set. It may or may not
   be referenced in the actual SQL
 * `Set<String> getCatalogs()`: the set of catalogs used in the query. Includes
   the default catalog if used by a non-fully qualified table reference
@@ -248,7 +263,7 @@ Set to `True` to make `trinoQueryProperties` and `trinoRequestUser` available.
 `maxBodySize`:
 
 By default, the max body size is 1,000,000 characters. This can be modified by
-configuring `maxBodySize`. If the request body is greater or equal to this 
+configuring `maxBodySize`. If the request body is greater or equal to this
 limit, Trino Gateway does not process the query. A buffer of length
 `maxBodySize` is allocated per query. Reduce this value if you observe
 excessive garbage collection at runtime. `maxBodySize` cannot be set to values
@@ -265,8 +280,8 @@ configuration.
 `tokenUserField`:
 
 When extracting the user from a JWT token, this field is used as the username.
-By default, the `email` claim is used. 
- 
+By default, the `email` claim is used.
+
 `oauthTokenInfoUrl`:
 
 If configured, Trino Gateway attempts to retrieve the user info by exchanging
@@ -275,7 +290,7 @@ minutes to avoid triggering rate limits.
 
 ### Execution of rules
 
-All rules whose conditions are satisfied fire. For example, in the "airflow" 
+All rules whose conditions are satisfied fire. For example, in the "airflow"
 and "airflow special" example rules from the following rule priority section, a
 query with source `airflow` and label `special` satisfies both rules. The
 `routingGroup` is set to `etl` and then to `etl-special` because of the order in

gateway-ha/src/main/java/io/trino/gateway/baseapp/BaseApp.java

@@ -31,6 +31,7 @@
 import io.trino.gateway.ha.resource.LoginResource;
 import io.trino.gateway.ha.resource.PublicResource;
 import io.trino.gateway.ha.resource.TrinoResource;
+import io.trino.gateway.ha.router.ForRouter;
 import io.trino.gateway.ha.security.AuthorizedExceptionMapper;
 import io.trino.gateway.proxyserver.ForProxy;
 import io.trino.gateway.proxyserver.ProxyRequestHandler;
@@ -170,5 +171,6 @@ private static void registerProxyResources(Binder binder)
         jaxrsBinder(binder).bind(ProxyRequestHandler.class);
         httpClientBinder(binder).bindHttpClient("proxy", ForProxy.class);
         httpClientBinder(binder).bindHttpClient("monitor", ForMonitor.class);
+        httpClientBinder(binder).bindHttpClient("router", ForRouter.class);
     }
 }

gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java

@@ -17,6 +17,7 @@
 import com.google.inject.AbstractModule;
 import com.google.inject.Provides;
 import com.google.inject.Singleton;
+import io.airlift.http.client.HttpClient;
 import io.trino.gateway.ha.config.AuthenticationConfiguration;
 import io.trino.gateway.ha.config.AuthorizationConfiguration;
 import io.trino.gateway.ha.config.GatewayCookieConfigurationPropertiesProvider;
@@ -26,6 +27,7 @@
 import io.trino.gateway.ha.config.RulesExternalConfiguration;
 import io.trino.gateway.ha.config.UserConfiguration;
 import io.trino.gateway.ha.router.BackendStateManager;
+import io.trino.gateway.ha.router.ForRouter;
 import io.trino.gateway.ha.router.RoutingGroupSelector;
 import io.trino.gateway.ha.security.ApiAuthenticator;
 import io.trino.gateway.ha.security.AuthorizationManager;
@@ -176,7 +178,7 @@ public BackendStateManager getBackendStateConnectionManager()
 
     @Provides
     @Singleton
-    public RoutingGroupSelector getRoutingGroupSelector()
+    public RoutingGroupSelector getRoutingGroupSelector(@ForRouter HttpClient httpClient)
     {
         RoutingRulesConfiguration routingRulesConfig = configuration.getRoutingRules();
         if (routingRulesConfig.isRulesEngineEnabled()) {
@@ -188,7 +190,7 @@ public RoutingGroupSelector getRoutingGroupSelector()
                     }
                     case EXTERNAL -> {
                         RulesExternalConfiguration rulesExternalConfiguration = routingRulesConfig.getRulesExternalConfiguration();
-                        yield RoutingGroupSelector.byRoutingExternal(rulesExternalConfiguration, configuration.getRequestAnalyzerConfig());
+                        yield RoutingGroupSelector.byRoutingExternal(rulesExternalConfiguration, configuration.getRequestAnalyzerConfig(), httpClient);
                     }
                 };
             }

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

@@ -18,11 +18,9 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Multimap;
 import io.airlift.http.client.HttpClient;
-import io.airlift.http.client.HttpClientConfig;
 import io.airlift.http.client.JsonBodyGenerator;
 import io.airlift.http.client.JsonResponseHandler;
 import io.airlift.http.client.Request;
-import io.airlift.http.client.jetty.JettyHttpClient;
 import io.airlift.json.JsonCodec;
 import io.airlift.log.Logger;
 import io.trino.gateway.ha.config.RequestAnalyzerConfig;
@@ -59,7 +57,7 @@ public class ExternalRoutingGroupSelector
             createJsonResponseHandler(jsonCodec(RoutingGroupExternalResponse.class));
 
     @VisibleForTesting
-    ExternalRoutingGroupSelector(RulesExternalConfiguration rulesExternalConfiguration, RequestAnalyzerConfig requestAnalyzerConfig)
+    ExternalRoutingGroupSelector(RulesExternalConfiguration rulesExternalConfiguration, RequestAnalyzerConfig requestAnalyzerConfig, HttpClient httpClient)
     {
         this.excludeHeaders = ImmutableSet.<String>builder()
                 .add("Content-Length")
@@ -76,7 +74,7 @@ public class ExternalRoutingGroupSelector
             throw new RuntimeException("Invalid URL provided, using "
                     + "routing group header as default.", e);
         }
-        httpClient = new JettyHttpClient(new HttpClientConfig());
+        this.httpClient = httpClient;
     }
 
     @Override
@@ -87,12 +85,13 @@ public String findRoutingGroup(HttpServletRequest servletRequest)
         try {
             RoutingGroupExternalBody requestBody = createRequestBody(servletRequest);
             requestBodyGenerator = jsonBodyGenerator(ROUTING_GROUP_EXTERNAL_BODY_JSON_CODEC, requestBody);
-            request = preparePost()
+            Request.Builder requestBuilder = preparePost()
                     .addHeader(CONTENT_TYPE, JSON_UTF_8.toString())
                     .addHeaders(getValidHeaders(servletRequest))
                     .setUri(uri)
-                    .setBodyGenerator(requestBodyGenerator)
-                    .build();
+                    .setBodyGenerator(requestBodyGenerator);
+
+            request = requestBuilder.build();
 
             // Execute the request and get the response
             RoutingGroupExternalResponse response = httpClient.execute(request, ROUTING_GROUP_EXTERNAL_RESPONSE_JSON_RESPONSE_HANDLER);

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

@@ -0,0 +1,31 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.gateway.ha.router;
+
+import com.google.inject.BindingAnnotation;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.ElementType.FIELD;
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.ElementType.PARAMETER;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+@Retention(RUNTIME)
+@Target({FIELD, PARAMETER, METHOD})
+@BindingAnnotation
+public @interface ForRouter
+{
+}

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

@@ -13,6 +13,7 @@
  */
 package io.trino.gateway.ha.router;
 
+import io.airlift.http.client.HttpClient;
 import io.trino.gateway.ha.config.RequestAnalyzerConfig;
 import io.trino.gateway.ha.config.RulesExternalConfiguration;
 import jakarta.servlet.http.HttpServletRequest;
@@ -46,9 +47,10 @@ static RoutingGroupSelector byRoutingRulesEngine(String rulesConfigPath, Request
      * Routing group selector that uses RESTful API
      * to determine the right routing group.
      */
-    static RoutingGroupSelector byRoutingExternal(RulesExternalConfiguration rulesExternalConfiguration, RequestAnalyzerConfig requestAnalyzerConfig)
+    static RoutingGroupSelector byRoutingExternal(RulesExternalConfiguration rulesExternalConfiguration,
+            RequestAnalyzerConfig requestAnalyzerConfig, HttpClient httpClient)
     {
-        return new ExternalRoutingGroupSelector(rulesExternalConfiguration, requestAnalyzerConfig);
+        return new ExternalRoutingGroupSelector(rulesExternalConfiguration, requestAnalyzerConfig, httpClient);
     }
 
     /**

gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingGroupSelectorExternal.java

@@ -133,7 +133,7 @@ void testApiFailure()
     {
         RulesExternalConfiguration rulesExternalConfiguration = provideRoutingRuleExternalConfig();
         RoutingGroupSelector routingGroupSelector =
-                RoutingGroupSelector.byRoutingExternal(rulesExternalConfiguration, requestAnalyzerConfig);
+                RoutingGroupSelector.byRoutingExternal(rulesExternalConfiguration, requestAnalyzerConfig, httpClient);
 
         HttpServletRequest mockRequest = prepareMockRequest();
         setMockHeaders(mockRequest);
@@ -162,7 +162,7 @@ void testNullUri()
         rulesExternalConfiguration.setUrlPath(null);
 
         // Assert that a RuntimeException is thrown with message
-        assertThatThrownBy(() -> RoutingGroupSelector.byRoutingExternal(rulesExternalConfiguration, requestAnalyzerConfig))
+        assertThatThrownBy(() -> RoutingGroupSelector.byRoutingExternal(rulesExternalConfiguration, requestAnalyzerConfig, httpClient))
                 .isInstanceOf(RuntimeException.class)
                 .hasMessage("Invalid URL provided, using routing group header as default.");
     }
@@ -175,7 +175,7 @@ void testExcludeHeader()
         rulesExternalConfiguration.setExcludeHeaders(List.of("test-exclude-header"));
 
         RoutingGroupSelector routingGroupSelector =
-                RoutingGroupSelector.byRoutingExternal(rulesExternalConfiguration, requestAnalyzerConfig);
+                RoutingGroupSelector.byRoutingExternal(rulesExternalConfiguration, requestAnalyzerConfig, httpClient);
 
         // Mock headers to be read by mockRequest
         HttpServletRequest mockRequest = mock(HttpServletRequest.class);