From 1762e02c1497fcbcc9b84aff528d9e4582700b95 Mon Sep 17 00:00:00 2001
From: Will Morrison <will.morrison@starburstdata.com>
Date: Thu, 26 Sep 2024 17:58:25 -0400
Subject: [PATCH] Inject RoutingTargetHandler

---
 .../java/io/trino/gateway/baseapp/BaseApp.java    |  2 ++
 .../gateway/ha/handler/RoutingTargetHandler.java  | 10 ++++++----
 .../ha/module/HaGatewayProviderModule.java        | 15 ---------------
 3 files changed, 8 insertions(+), 19 deletions(-)

diff --git a/gateway-ha/src/main/java/io/trino/gateway/baseapp/BaseApp.java b/gateway-ha/src/main/java/io/trino/gateway/baseapp/BaseApp.java
index 7c169a531..250e7bc5d 100644
--- a/gateway-ha/src/main/java/io/trino/gateway/baseapp/BaseApp.java
+++ b/gateway-ha/src/main/java/io/trino/gateway/baseapp/BaseApp.java
@@ -21,6 +21,7 @@
 import io.trino.gateway.ha.clustermonitor.ForMonitor;
 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;
@@ -119,6 +120,7 @@ public void configure(Binder binder)
         registerAuthFilters(binder);
         registerResources(binder);
         registerProxyResources(binder);
+        jaxrsBinder(binder).bind(RoutingTargetHandler.class);
         addManagedApps(this.haGatewayConfiguration, binder);
         jaxrsBinder(binder).bind(AuthorizedExceptionMapper.class);
         binder.bind(ProxyHandlerStats.class).in(Scopes.SINGLETON);
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
index 6b8e5d621..676782602 100644
--- 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
@@ -13,8 +13,10 @@
  */
 package io.trino.gateway.ha.handler;
 
+import com.google.inject.Inject;
 import io.airlift.log.Logger;
 import io.trino.gateway.ha.config.GatewayCookieConfigurationPropertiesProvider;
+import io.trino.gateway.ha.config.HaGatewayConfiguration;
 import io.trino.gateway.ha.router.GatewayCookie;
 import io.trino.gateway.ha.router.RoutingGroupSelector;
 import io.trino.gateway.ha.router.RoutingManager;
@@ -47,16 +49,16 @@ public class RoutingTargetHandler
     private final List<Pattern> extraWhitelistPaths;
     private final boolean cookiesEnabled;
 
+    @Inject
     public RoutingTargetHandler(
             RoutingManager routingManager,
             RoutingGroupSelector routingGroupSelector,
-            List<String> statementPaths,
-            List<String> extraWhitelistPaths)
+            HaGatewayConfiguration haGatewayConfiguration)
     {
         this.routingManager = requireNonNull(routingManager);
         this.routingGroupSelector = requireNonNull(routingGroupSelector);
-        this.statementPaths = requireNonNull(statementPaths);
-        this.extraWhitelistPaths = extraWhitelistPaths.stream().map(Pattern::compile).collect(toImmutableList());
+        this.statementPaths = requireNonNull(haGatewayConfiguration.getStatementPaths());
+        this.extraWhitelistPaths = requireNonNull(haGatewayConfiguration.getExtraWhitelistPaths()).stream().map(Pattern::compile).collect(toImmutableList());
         cookiesEnabled = GatewayCookieConfigurationPropertiesProvider.getInstance().isEnabled();
     }
 
diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java b/gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java
index c9c925140..550c98daf 100644
--- a/gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java
+++ b/gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java
@@ -25,10 +25,8 @@
 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.handler.RoutingTargetHandler;
 import io.trino.gateway.ha.router.BackendStateManager;
 import io.trino.gateway.ha.router.RoutingGroupSelector;
-import io.trino.gateway.ha.router.RoutingManager;
 import io.trino.gateway.ha.security.ApiAuthenticator;
 import io.trino.gateway.ha.security.AuthorizationManager;
 import io.trino.gateway.ha.security.BasicAuthFilter;
@@ -200,17 +198,4 @@ public RoutingGroupSelector getRoutingGroupSelector()
         }
         return RoutingGroupSelector.byRoutingGroupHeader();
     }
-
-    @Provides
-    @Singleton
-    public RoutingTargetHandler getRoutingTargetHandler(
-            RoutingManager routingManager,
-            RoutingGroupSelector routingGroupSelector)
-    {
-        return new RoutingTargetHandler(
-                routingManager,
-                routingGroupSelector,
-                configuration.getStatementPaths(),
-                configuration.getExtraWhitelistPaths());
-    }
 }