add configurable default routing group, add indicative errors, use ex…

gateway-ha/src/main/java/io/trino/gateway/ha/config/RoutingConfiguration.java

@@ -23,6 +23,8 @@ public class RoutingConfiguration
 
     private boolean addXForwardedHeaders = true;
 
+    private String defaultRoutingGroup = "adhoc";
+
     public Duration getAsyncTimeout()
     {
         return asyncTimeout;
@@ -42,4 +44,14 @@ public void setAddXForwardedHeaders(boolean addXForwardedHeaders)
     {
         this.addXForwardedHeaders = addXForwardedHeaders;
     }
+
+    public String getDefaultRoutingGroup()
+    {
+        return defaultRoutingGroup;
+    }
+
+    public void setDefaultRoutingGroup(String defaultRoutingGroup)
+    {
+        this.defaultRoutingGroup = defaultRoutingGroup;
+    }
 }

gateway-ha/src/main/java/io/trino/gateway/ha/handler/RoutingTargetHandler.java

@@ -21,6 +21,9 @@
 import io.trino.gateway.ha.router.RoutingGroupSelector;
 import io.trino.gateway.ha.router.RoutingManager;
 import jakarta.servlet.http.HttpServletRequest;
+import jakarta.ws.rs.NotFoundException;
+import jakarta.ws.rs.WebApplicationException;
+import jakarta.ws.rs.core.Response;
 
 import java.util.Arrays;
 import java.util.List;
@@ -68,11 +71,23 @@ public RoutingTargetHandler(
 
     public String getRoutingDestination(HttpServletRequest request)
     {
-        Optional<String> previousBackend = getPreviousBackend(request);
-        String clusterHost = previousBackend.orElseGet(() -> getBackendFromRoutingGroup(request));
-        logRewrite(clusterHost, request);
+        try {
+            Optional<String> previousBackend = getPreviousBackend(request);
+            String clusterHost = previousBackend.orElseGet(() -> getBackendFromRoutingGroup(request));
+            logRewrite(clusterHost, request);
 
-        return buildUriWithNewBackend(clusterHost, request);
+            return buildUriWithNewBackend(clusterHost, request);
+        }
+        catch (NotFoundException e) {
+            throw new WebApplicationException(
+                    Response.status(Response.Status.NOT_FOUND)
+                            .entity(e.getMessage())
+                            .build());
+        }
+        catch (WebApplicationException e) {
+            // Re-throw other WebApplicationExceptions (like BAD_REQUEST from query validation)
+            throw e;
+        }
     }
 
     public boolean isPathWhiteListed(String path)
@@ -92,7 +107,7 @@ private String getBackendFromRoutingGroup(HttpServletRequest request)
         String routingGroup = routingGroupSelector.findRoutingGroup(request);
         String user = request.getHeader(USER_HEADER);
         if (!isNullOrEmpty(routingGroup)) {
-            // This falls back on adhoc backend if there is no cluster found for the routing group.
+            // This falls back on default routing group backend if there is no cluster found for the routing group.
             return routingManager.provideBackendForRoutingGroup(routingGroup, user);
         }
         return routingManager.provideAdhocBackend(user);

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

@@ -26,7 +26,7 @@ public class QueryCountBasedRouterProvider
     public QueryCountBasedRouterProvider(HaGatewayConfiguration configuration)
     {
         super(configuration);
-        routingManager = new QueryCountBasedRouter(gatewayBackendManager, queryHistoryManager);
+        routingManager = new QueryCountBasedRouter(gatewayBackendManager, queryHistoryManager, configuration.getRouting());
     }
 
     @Provides

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

@@ -38,7 +38,7 @@ public RouterBaseModule(HaGatewayConfiguration configuration)
         Jdbi jdbi = Jdbi.create(configuration.getDataStore().getJdbcUrl(), configuration.getDataStore().getUser(), configuration.getDataStore().getPassword());
         connectionManager = new JdbcConnectionManager(jdbi, configuration.getDataStore());
         resourceGroupsManager = new HaResourceGroupsManager(connectionManager);
-        gatewayBackendManager = new HaGatewayManager(jdbi);
+        gatewayBackendManager = new HaGatewayManager(jdbi, configuration.getRouting());
         queryHistoryManager = new HaQueryHistoryManager(jdbi, configuration.getDataStore().getJdbcUrl().startsWith("jdbc:oracle"));
     }
 

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

@@ -26,7 +26,7 @@ public class StochasticRoutingManagerProvider
     public StochasticRoutingManagerProvider(HaGatewayConfiguration configuration)
     {
         super(configuration);
-        routingManager = new StochasticRoutingManager(gatewayBackendManager, queryHistoryManager);
+        routingManager = new StochasticRoutingManager(gatewayBackendManager, queryHistoryManager, configuration.getRouting());
     }
 
     @Provides

gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/GatewayBackendDao.java

@@ -29,6 +29,10 @@ public interface GatewayBackendDao
             """)
     List<GatewayBackend> findActiveBackend();
 
+    /**
+     * @deprecated Use {@link #findActiveBackendByRoutingGroup(String)} with the configured default routing group
+     */
+    @Deprecated
     @SqlQuery("""
             SELECT * FROM gateway_backend
             WHERE active = true AND routing_group = 'adhoc'

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

@@ -28,6 +28,8 @@
 import io.trino.gateway.ha.router.schema.RoutingGroupExternalBody;
 import io.trino.gateway.ha.router.schema.RoutingGroupExternalResponse;
 import jakarta.servlet.http.HttpServletRequest;
+import jakarta.ws.rs.WebApplicationException;
+import jakarta.ws.rs.core.Response;
 
 import java.net.URI;
 import java.net.URISyntaxException;
@@ -98,10 +100,18 @@ public String findRoutingGroup(HttpServletRequest servletRequest)
                 throw new RuntimeException("Unexpected response: null");
             }
             else if (response.errors() != null && !response.errors().isEmpty()) {
-                throw new RuntimeException("Response with error: " + String.join(", ", response.errors()));
+                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());
             }
             return response.routingGroup();
         }
+        catch (WebApplicationException e) {
+            // Re-throw WebApplicationException to preserve status and entity
+            throw e;
+        }
         catch (Exception e) {
             log.error(e, "Error occurred while retrieving routing group "
                     + "from external routing rules processing at " + uri);

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

@@ -24,6 +24,10 @@ public interface GatewayBackendManager
 
     List<ProxyBackendConfiguration> getAllActiveBackends();
 
+    /**
+     * @deprecated Use {@link #getActiveBackends(String)} with the configured default routing group
+     */
+    @Deprecated
     List<ProxyBackendConfiguration> getActiveAdhocBackends();
 
     List<ProxyBackendConfiguration> getActiveBackends(String routingGroup);

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

@@ -16,6 +16,7 @@
 import com.google.common.collect.ImmutableList;
 import io.airlift.log.Logger;
 import io.trino.gateway.ha.config.ProxyBackendConfiguration;
+import io.trino.gateway.ha.config.RoutingConfiguration;
 import io.trino.gateway.ha.persistence.dao.GatewayBackend;
 import io.trino.gateway.ha.persistence.dao.GatewayBackendDao;
 import org.jdbi.v3.core.Jdbi;
@@ -32,10 +33,12 @@ public class HaGatewayManager
     private static final Logger log = Logger.get(HaGatewayManager.class);
 
     private final GatewayBackendDao dao;
+    private final String defaultRoutingGroup;
 
-    public HaGatewayManager(Jdbi jdbi)
+    public HaGatewayManager(Jdbi jdbi, RoutingConfiguration routingConfiguration)
     {
         dao = requireNonNull(jdbi, "jdbi is null").onDemand(GatewayBackendDao.class);
+        this.defaultRoutingGroup = requireNonNull(routingConfiguration, "routingConfiguration is null").getDefaultRoutingGroup();
     }
 
     @Override
@@ -56,11 +59,10 @@ public List<ProxyBackendConfiguration> getAllActiveBackends()
     public List<ProxyBackendConfiguration> getActiveAdhocBackends()
     {
         try {
-            List<GatewayBackend> proxyBackendList = dao.findActiveAdhocBackend();
-            return upcast(proxyBackendList);
+            return getActiveBackends(defaultRoutingGroup);
         }
         catch (Exception e) {
-            log.info("Error fetching all backends: %s", e.getLocalizedMessage());
+            log.info("Error fetching backends for default routing group: %s", e.getLocalizedMessage());
         }
         return ImmutableList.of();
     }

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

@@ -20,6 +20,7 @@
 import io.airlift.log.Logger;
 import io.trino.gateway.ha.clustermonitor.ClusterStats;
 import io.trino.gateway.ha.clustermonitor.TrinoStatus;
+import io.trino.gateway.ha.config.RoutingConfiguration;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -35,6 +36,7 @@ public class QueryCountBasedRouter
     private static final Logger log = Logger.get(QueryCountBasedRouter.class);
     @GuardedBy("this")
     private List<LocalStats> clusterStats;
+    private final String defaultRoutingGroup;
 
     @VisibleForTesting
     synchronized List<LocalStats> clusterStats()
@@ -136,9 +138,11 @@ public void userQueuedCount(Map<String, Integer> userQueuedCount)
 
     public QueryCountBasedRouter(
             GatewayBackendManager gatewayBackendManager,
-            QueryHistoryManager queryHistoryManager)
+            QueryHistoryManager queryHistoryManager,
+            RoutingConfiguration routingConfiguration)
     {
-        super(gatewayBackendManager, queryHistoryManager);
+        super(gatewayBackendManager, queryHistoryManager, routingConfiguration);
+        this.defaultRoutingGroup = routingConfiguration.getDefaultRoutingGroup();
         clusterStats = new ArrayList<>();
     }
 
@@ -224,7 +228,7 @@ private synchronized Optional<String> getBackendForRoutingGroup(String routingGr
     @Override
     public String provideAdhocBackend(String user)
     {
-        return getBackendForRoutingGroup("adhoc", user).orElseThrow(() -> new RouterException("did not find any cluster for the adhoc routing group"));
+        return getBackendForRoutingGroup(defaultRoutingGroup, user).orElseThrow(() -> new RouterException("did not find any cluster for the default routing group: " + defaultRoutingGroup));
     }
 
     @Override

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

@@ -20,7 +20,9 @@
 import io.trino.gateway.ha.clustermonitor.ClusterStats;
 import io.trino.gateway.ha.clustermonitor.TrinoStatus;
 import io.trino.gateway.ha.config.ProxyBackendConfiguration;
+import io.trino.gateway.ha.config.RoutingConfiguration;
 import jakarta.ws.rs.HttpMethod;
+import jakarta.ws.rs.NotFoundException;
 
 import java.net.HttpURLConnection;
 import java.net.URL;
@@ -47,10 +49,12 @@ public abstract class RoutingManager
     private final ExecutorService executorService = Executors.newFixedThreadPool(5);
     private final GatewayBackendManager gatewayBackendManager;
     private final ConcurrentHashMap<String, TrinoStatus> backendToStatus;
+    private final String defaultRoutingGroup;
 
-    public RoutingManager(GatewayBackendManager gatewayBackendManager)
+    public RoutingManager(GatewayBackendManager gatewayBackendManager, RoutingConfiguration routingConfiguration)
     {
         this.gatewayBackendManager = gatewayBackendManager;
+        this.defaultRoutingGroup = routingConfiguration.getDefaultRoutingGroup();
         queryIdBackendCache =
                 CacheBuilder.newBuilder()
                         .maximumSize(10000)
@@ -83,7 +87,7 @@ public void setBackendForQueryId(String queryId, String backend)
      */
     public String provideAdhocBackend(String user)
     {
-        List<ProxyBackendConfiguration> backends = this.gatewayBackendManager.getActiveAdhocBackends();
+        List<ProxyBackendConfiguration> backends = this.gatewayBackendManager.getActiveBackends(defaultRoutingGroup);
         backends.removeIf(backend -> isBackendNotHealthy(backend.getName()));
         if (backends.size() == 0) {
             throw new IllegalStateException("Number of active backends found zero");
@@ -100,6 +104,10 @@ public String provideBackendForRoutingGroup(String routingGroup, String user)
     {
         List<ProxyBackendConfiguration> backends =
                 gatewayBackendManager.getActiveBackends(routingGroup);
+        // Check if any backends exist for the routing group (even before filtering unhealthy ones)
+        if (backends.isEmpty()) {
+            throw new NotFoundException("Routing group does not exist: " + routingGroup);
+        }
         backends.removeIf(backend -> isBackendNotHealthy(backend.getName()));
         if (backends.isEmpty()) {
             return provideAdhocBackend(user);
@@ -177,7 +185,7 @@ protected String findBackendForUnknownQueryId(String queryId)
             log.warn("Query id [%s] not found", queryId);
         }
         // Fallback on first active backend if queryId mapping not found.
-        return gatewayBackendManager.getActiveAdhocBackends().get(0).getProxyTo();
+        return gatewayBackendManager.getActiveBackends(defaultRoutingGroup).get(0).getProxyTo();
     }
 
     // Predicate helper function to remove the backends from the list

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

@@ -15,6 +15,7 @@
 
 import com.google.common.base.Strings;
 import io.airlift.log.Logger;
+import io.trino.gateway.ha.config.RoutingConfiguration;
 
 public class StochasticRoutingManager
         extends RoutingManager
@@ -25,7 +26,15 @@ public class StochasticRoutingManager
     public StochasticRoutingManager(
             GatewayBackendManager gatewayBackendManager, QueryHistoryManager queryHistoryManager)
     {
-        super(gatewayBackendManager);
+        this(gatewayBackendManager, queryHistoryManager, null);
+    }
+
+    public StochasticRoutingManager(
+            GatewayBackendManager gatewayBackendManager,
+            QueryHistoryManager queryHistoryManager,
+            RoutingConfiguration routingConfiguration)
+    {
+        super(gatewayBackendManager, routingConfiguration);
         this.queryHistoryManager = queryHistoryManager;
     }