Use OptionalBinder for RoutingManager

Author: andythsu

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

@@ -24,8 +24,6 @@
 import io.trino.gateway.ha.config.HaGatewayConfiguration;
 import io.trino.gateway.ha.handler.ProxyHandlerStats;
 import io.trino.gateway.ha.handler.RoutingTargetHandler;
-import io.trino.gateway.ha.module.RouterBaseModule;
-import io.trino.gateway.ha.module.StochasticRoutingManagerProvider;
 import io.trino.gateway.ha.resource.EntityEditorResource;
 import io.trino.gateway.ha.resource.GatewayHealthCheckResource;
 import io.trino.gateway.ha.resource.GatewayResource;
@@ -36,7 +34,9 @@
 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.router.RoutingManager;
 import io.trino.gateway.ha.router.RoutingRulesManager;
+import io.trino.gateway.ha.router.StochasticRoutingManager;
 import io.trino.gateway.ha.security.AuthorizedExceptionMapper;
 import io.trino.gateway.proxyserver.ForProxy;
 import io.trino.gateway.proxyserver.ProxyRequestHandler;
@@ -47,9 +47,8 @@
 import java.lang.reflect.Constructor;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Optional;
 
-import static com.google.common.collect.MoreCollectors.toOptional;
+import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
 import static io.airlift.http.client.HttpClientBinder.httpClientBinder;
 import static io.airlift.jaxrs.JaxrsBinder.jaxrsBinder;
 import static java.lang.String.format;
@@ -91,18 +90,6 @@ private static Module newModule(String clazz, HaGatewayConfiguration configurati
         return null;
     }
 
-    private static void addDefaultRouterProviderModules(List<Module> modules, HaGatewayConfiguration configuration)
-    {
-        Optional<Module> routerProvider = modules.stream()
-                .filter(module -> module instanceof RouterBaseModule)
-                .collect(toOptional());
-        if (routerProvider.isEmpty()) {
-            logger.warn("Router provider doesn't exist in the config, using the StochasticRoutingManagerProvider");
-            String clazz = StochasticRoutingManagerProvider.class.getCanonicalName();
-            modules.add(newModule(clazz, configuration));
-        }
-    }
-
     public static List<Module> addModules(HaGatewayConfiguration configuration)
     {
         List<Module> modules = new ArrayList<>();
@@ -115,8 +102,6 @@ public static List<Module> addModules(HaGatewayConfiguration configuration)
                 modules.add(newModule(clazz, configuration));
             }
         }
-        addDefaultRouterProviderModules(modules, configuration);
-
         return modules;
     }
 
@@ -148,6 +133,10 @@ public void configure(Binder binder)
         newExporter(binder).export(ProxyHandlerStats.class).withGeneratedName();
         binder.bind(RoutingRulesManager.class);
         binder.bind(ClusterMetricsStatsExporter.class).in(Scopes.SINGLETON);
+        newOptionalBinder(binder, RoutingManager.class)
+                .setDefault()
+                .to(StochasticRoutingManager.class)
+                .in(Scopes.SINGLETON);
     }
 
     private static void addManagedApps(HaGatewayConfiguration configuration, Binder binder)

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

@@ -39,8 +39,15 @@
 import io.trino.gateway.ha.config.RoutingRulesConfiguration;
 import io.trino.gateway.ha.config.RulesExternalConfiguration;
 import io.trino.gateway.ha.config.UserConfiguration;
+import io.trino.gateway.ha.persistence.JdbcConnectionManager;
 import io.trino.gateway.ha.router.BackendStateManager;
 import io.trino.gateway.ha.router.ForRouter;
+import io.trino.gateway.ha.router.GatewayBackendManager;
+import io.trino.gateway.ha.router.HaGatewayManager;
+import io.trino.gateway.ha.router.HaQueryHistoryManager;
+import io.trino.gateway.ha.router.HaResourceGroupsManager;
+import io.trino.gateway.ha.router.QueryHistoryManager;
+import io.trino.gateway.ha.router.ResourceGroupsManager;
 import io.trino.gateway.ha.router.RoutingGroupSelector;
 import io.trino.gateway.ha.router.RoutingManager;
 import io.trino.gateway.ha.security.ApiAuthenticator;
@@ -59,6 +66,7 @@
 import io.trino.gateway.ha.security.util.Authorizer;
 import io.trino.gateway.ha.security.util.ChainedAuthFilter;
 import jakarta.ws.rs.container.ContainerRequestFilter;
+import org.jdbi.v3.core.Jdbi;
 
 import java.util.List;
 import java.util.Map;
@@ -77,11 +85,17 @@ public class HaGatewayProviderModule
     private final BackendStateManager backendStateConnectionManager;
     private final ResourceSecurityDynamicFeature resourceSecurityDynamicFeature;
     private final HaGatewayConfiguration configuration;
+    private final ResourceGroupsManager resourceGroupsManager;
+    private final GatewayBackendManager gatewayBackendManager;
+    private final QueryHistoryManager queryHistoryManager;
 
     @Override
     protected void configure()
     {
         jaxrsBinder(binder()).bindInstance(resourceSecurityDynamicFeature);
+        binder().bind(ResourceGroupsManager.class).toInstance(resourceGroupsManager);
+        binder().bind(GatewayBackendManager.class).toInstance(gatewayBackendManager);
+        binder().bind(QueryHistoryManager.class).toInstance(queryHistoryManager);
     }
 
     public HaGatewayProviderModule(HaGatewayConfiguration configuration)
@@ -101,6 +115,12 @@ public HaGatewayProviderModule(HaGatewayConfiguration configuration)
 
         OAuth2GatewayCookieConfigurationPropertiesProvider oAuth2GatewayCookieConfigurationPropertiesProvider = OAuth2GatewayCookieConfigurationPropertiesProvider.getInstance();
         oAuth2GatewayCookieConfigurationPropertiesProvider.initialize(configuration.getOauth2GatewayCookieConfiguration());
+
+        Jdbi jdbi = Jdbi.create(configuration.getDataStore().getJdbcUrl(), configuration.getDataStore().getUser(), configuration.getDataStore().getPassword());
+        JdbcConnectionManager connectionManager = new JdbcConnectionManager(jdbi, configuration.getDataStore());
+        resourceGroupsManager = new HaResourceGroupsManager(connectionManager);
+        gatewayBackendManager = new HaGatewayManager(jdbi);
+        queryHistoryManager = new HaQueryHistoryManager(jdbi, configuration.getDataStore().getJdbcUrl().startsWith("jdbc:oracle"));
     }
 
     private LbOAuthManager getOAuthManager(HaGatewayConfiguration configuration)
@@ -253,6 +273,7 @@ public List<TrinoClusterStatsObserver> getClusterStatsObservers(
     }
 
     @Provides
+    @Singleton
     public MonitorConfiguration getMonitorConfiguration()
     {
         return configuration.getMonitor();

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

@@ -13,22 +13,33 @@
  */
 package io.trino.gateway.ha.module;
 
+import com.google.inject.AbstractModule;
 import com.google.inject.Scopes;
+import io.airlift.log.Logger;
 import io.trino.gateway.ha.config.HaGatewayConfiguration;
 import io.trino.gateway.ha.router.QueryCountBasedRouter;
 import io.trino.gateway.ha.router.RoutingManager;
 
+import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
+
 public class QueryCountBasedRouterProvider
-          extends RouterBaseModule
+        extends AbstractModule
 {
+    private static final Logger logger = Logger.get(QueryCountBasedRouterProvider.class);
+
+    // We require all modules to take HaGatewayConfiguration as the only parameter
     public QueryCountBasedRouterProvider(HaGatewayConfiguration configuration)
     {
-        super(configuration);
+        // no-op
     }
 
     @Override
     public void configure()
     {
-        bind(RoutingManager.class).to(QueryCountBasedRouter.class).in(Scopes.SINGLETON);
+        logger.info("Using QueryCountBasedRouterProvider instead of default");
+        newOptionalBinder(binder(), RoutingManager.class)
+                .setBinding()
+                .to(QueryCountBasedRouter.class)
+                .in(Scopes.SINGLETON);
     }
 }