Use JDBI for gateway_backend

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

@@ -83,7 +83,7 @@ public HaGatewayProviderModule(HaGatewayConfiguration configuration, Environment
         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(connectionManager);
+        gatewayBackendManager = new HaGatewayManager(jdbi);
         queryHistoryManager = new HaQueryHistoryManager(jdbi);
         routingManager =
                 new HaRoutingManager(gatewayBackendManager, queryHistoryManager);

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

@@ -13,66 +13,25 @@
  */
 package io.trino.gateway.ha.persistence.dao;
 
-import io.trino.gateway.ha.config.ProxyBackendConfiguration;
-import org.javalite.activejdbc.Model;
-import org.javalite.activejdbc.annotations.Cached;
-import org.javalite.activejdbc.annotations.IdName;
-import org.javalite.activejdbc.annotations.Table;
+import org.jdbi.v3.core.mapper.reflect.ColumnName;
+import org.jdbi.v3.core.mapper.reflect.JdbiConstructor;
 
-import java.util.ArrayList;
-import java.util.List;
+import static java.util.Objects.requireNonNull;
 
-@Table("gateway_backend")
-@IdName("name")
-@Cached
-public class GatewayBackend
-        extends Model
+public record GatewayBackend(String name, String routingGroup, String backendUrl, String externalUrl, boolean active)
 {
-    private static final String name = "name";
-    private static final String routingGroup = "routing_group";
-    private static final String backendUrl = "backend_url";
-    private static final String externalUrl = "external_url";
-    private static final String active = "active";
-
-    public static List<ProxyBackendConfiguration> upcast(List<GatewayBackend> gatewayBackendList)
-    {
-        List<ProxyBackendConfiguration> proxyBackendConfigurations = new ArrayList<>();
-        for (GatewayBackend model : gatewayBackendList) {
-            ProxyBackendConfiguration backendConfig = new ProxyBackendConfiguration();
-            backendConfig.setActive(model.getBoolean(active));
-            backendConfig.setRoutingGroup(model.getString(routingGroup));
-            backendConfig.setProxyTo(model.getString(backendUrl));
-            backendConfig.setExternalUrl(model.getString(externalUrl));
-            backendConfig.setName(model.getString(name));
-            proxyBackendConfigurations.add(backendConfig);
-        }
-        return proxyBackendConfigurations;
-    }
-
-    public static void update(GatewayBackend model, ProxyBackendConfiguration backend)
-    {
-        model
-                .set(name, backend.getName())
-                .set(routingGroup, backend.getRoutingGroup())
-                .set(backendUrl, backend.getProxyTo())
-                .set(externalUrl, backend.getExternalUrl())
-                .set(active, backend.isActive())
-                .saveIt();
-    }
-
-    public static void create(GatewayBackend model, ProxyBackendConfiguration backend)
+    @JdbiConstructor
+    public GatewayBackend(
+            @ColumnName("name") String name,
+            @ColumnName("routing_group") String routingGroup,
+            @ColumnName("backend_url") String backendUrl,
+            @ColumnName("external_url") String externalUrl,
+            @ColumnName("active") boolean active)
     {
-        create(
-                        name,
-                        backend.getName(),
-                        routingGroup,
-                        backend.getRoutingGroup(),
-                        backendUrl,
-                        backend.getProxyTo(),
-                        externalUrl,
-                        backend.getExternalUrl(),
-                        active,
-                        backend.isActive())
-                .insert();
+        this.name = requireNonNull(name, "name is null");
+        this.routingGroup = requireNonNull(routingGroup, "routingGroup is null");
+        this.backendUrl = requireNonNull(backendUrl, "backendUrl is null");
+        this.externalUrl = requireNonNull(externalUrl, "externalUrl is null");
+        this.active = active;
     }
 }

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

@@ -0,0 +1,89 @@
+/*
+ * 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.persistence.dao;
+
+import org.jdbi.v3.sqlobject.statement.SqlQuery;
+import org.jdbi.v3.sqlobject.statement.SqlUpdate;
+
+import java.util.List;
+
+public interface GatewayBackendDao
+{
+    @SqlQuery("SELECT * FROM gateway_backend")
+    List<GatewayBackend> findAll();
+
+    @SqlQuery("""
+            SELECT * FROM gateway_backend
+            WHERE active = true
+            """)
+    List<GatewayBackend> findActiveBackend();
+
+    @SqlQuery("""
+            SELECT * FROM gateway_backend
+            WHERE active = true AND routing_group = 'adhoc'
+            """)
+    List<GatewayBackend> findActiveAdhocBackend();
+
+    @SqlQuery("""
+            SELECT * FROM gateway_backend
+            WHERE active = true AND routing_group = :routingGroup
+            """)
+    List<GatewayBackend> findActiveBackendByRoutingGroup(String routingGroup);
+
+    @SqlQuery("""
+            SELECT * FROM gateway_backend
+            WHERE name = :name
+            """)
+    List<GatewayBackend> findByName(String name);
+
+    @SqlQuery("""
+            SELECT * FROM gateway_backend
+            WHERE name = :name
+            LIMIT 1
+            """)
+    GatewayBackend findFirstByName(String name);
+
+    @SqlUpdate("""
+            INSERT INTO gateway_backend (name, routing_group, backend_url, external_url, active)
+            VALUES (:name, :routingGroup, :backendUrl, :externalUrl, :active)
+            """)
+    void create(String name, String routingGroup, String backendUrl, String externalUrl, boolean active);
+
+    @SqlUpdate("""
+            UPDATE gateway_backend
+            SET routing_group = :routingGroup, backend_url = :backendUrl, external_url = :externalUrl, active = :active
+            WHERE name = :name
+            """)
+    void update(String name, String routingGroup, String backendUrl, String externalUrl, boolean active);
+
+    @SqlUpdate("""
+            UPDATE gateway_backend
+            SET active = false
+            WHERE name = :name
+            """)
+    void deactivate(String name);
+
+    @SqlUpdate("""
+            UPDATE gateway_backend
+            SET active = true
+            WHERE name = :name
+            """)
+    void activate(String name);
+
+    @SqlUpdate("""
+            DELETE FROM gateway_backend
+            WHERE name = :name
+            """)
+    void deleteByName(String name);
+}

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

@@ -15,161 +15,120 @@
 
 import com.google.common.collect.ImmutableList;
 import io.trino.gateway.ha.config.ProxyBackendConfiguration;
-import io.trino.gateway.ha.persistence.JdbcConnectionManager;
 import io.trino.gateway.ha.persistence.dao.GatewayBackend;
+import io.trino.gateway.ha.persistence.dao.GatewayBackendDao;
+import org.jdbi.v3.core.Jdbi;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 
+import static java.util.Objects.requireNonNull;
+
 public class HaGatewayManager
         implements GatewayBackendManager
 {
     private static final Logger log = LoggerFactory.getLogger(HaGatewayManager.class);
-    private final JdbcConnectionManager connectionManager;
 
-    public HaGatewayManager(JdbcConnectionManager connectionManager)
+    private final GatewayBackendDao dao;
+
+    public HaGatewayManager(Jdbi jdbi)
     {
-        this.connectionManager = connectionManager;
+        dao = requireNonNull(jdbi, "jdbi is null").onDemand(GatewayBackendDao.class);
     }
 
     @Override
     public List<ProxyBackendConfiguration> getAllBackends()
     {
-        try {
-            connectionManager.open();
-            List<GatewayBackend> proxyBackendList = GatewayBackend.findAll();
-            return GatewayBackend.upcast(proxyBackendList);
-        }
-        finally {
-            connectionManager.close();
-        }
+        List<GatewayBackend> proxyBackendList = dao.findAll();
+        return upcast(proxyBackendList);
     }
 
     @Override
     public List<ProxyBackendConfiguration> getAllActiveBackends()
     {
-        try {
-            connectionManager.open();
-            List<GatewayBackend> proxyBackendList = GatewayBackend.where("active = ?", true);
-            return GatewayBackend.upcast(proxyBackendList);
-        }
-        finally {
-            connectionManager.close();
-        }
+        List<GatewayBackend> proxyBackendList = dao.findActiveBackend();
+        return upcast(proxyBackendList);
     }
 
     @Override
     public List<ProxyBackendConfiguration> getActiveAdhocBackends()
     {
         try {
-            connectionManager.open();
-            List<GatewayBackend> proxyBackendList =
-                    GatewayBackend.where("active = ? and routing_group = ?", true, "adhoc");
-            return GatewayBackend.upcast(proxyBackendList);
+            List<GatewayBackend> proxyBackendList = dao.findActiveAdhocBackend();
+            return upcast(proxyBackendList);
         }
         catch (Exception e) {
             log.info("Error fetching all backends: {}", e.getLocalizedMessage());
         }
-        finally {
-            connectionManager.close();
-        }
         return ImmutableList.of();
     }
 
     @Override
     public List<ProxyBackendConfiguration> getActiveBackends(String routingGroup)
     {
-        try {
-            connectionManager.open();
-            List<GatewayBackend> proxyBackendList =
-                    GatewayBackend.where("active = ? and routing_group = ?", true, routingGroup);
-            return GatewayBackend.upcast(proxyBackendList);
-        }
-        finally {
-            connectionManager.close();
-        }
+        List<GatewayBackend> proxyBackendList = dao.findActiveBackendByRoutingGroup(routingGroup);
+        return upcast(proxyBackendList);
     }
 
     @Override
     public Optional<ProxyBackendConfiguration> getBackendByName(String name)
     {
-        try {
-            connectionManager.open();
-            List<GatewayBackend> proxyBackendList =
-                    GatewayBackend.where("name = ?", name);
-            return GatewayBackend.upcast(proxyBackendList).stream().findAny();
-        }
-        finally {
-            connectionManager.close();
-        }
+        List<GatewayBackend> proxyBackendList = dao.findByName(name);
+        return upcast(proxyBackendList).stream().findAny();
     }
 
     @Override
     public void deactivateBackend(String backendName)
     {
-        try {
-            connectionManager.open();
-            GatewayBackend.findFirst("name = ?", backendName).set("active", false).saveIt();
-        }
-        finally {
-            connectionManager.close();
-        }
+        dao.deactivate(backendName);
     }
 
     @Override
     public void activateBackend(String backendName)
     {
-        try {
-            connectionManager.open();
-            GatewayBackend.findFirst("name = ?", backendName).set("active", true).saveIt();
-        }
-        finally {
-            connectionManager.close();
-        }
+        dao.activate(backendName);
     }
 
     @Override
     public ProxyBackendConfiguration addBackend(ProxyBackendConfiguration backend)
     {
-        try {
-            connectionManager.open();
-            GatewayBackend.create(new GatewayBackend(), backend);
-        }
-        finally {
-            connectionManager.close();
-        }
+        dao.create(backend.getName(), backend.getRoutingGroup(), backend.getProxyTo(), backend.getExternalUrl(), backend.isActive());
         return backend;
     }
 
     @Override
     public ProxyBackendConfiguration updateBackend(ProxyBackendConfiguration backend)
     {
-        try {
-            connectionManager.open();
-            GatewayBackend model = GatewayBackend.findFirst("name = ?", backend.getName());
-            if (model == null) {
-                GatewayBackend.create(model, backend);
-            }
-            else {
-                GatewayBackend.update(model, backend);
-            }
+        GatewayBackend model = dao.findFirstByName(backend.getName());
+        if (model == null) {
+            dao.create(backend.getName(), backend.getRoutingGroup(), backend.getProxyTo(), backend.getExternalUrl(), backend.isActive());
         }
-        finally {
-            connectionManager.close();
+        else {
+            dao.update(backend.getName(), backend.getRoutingGroup(), backend.getProxyTo(), backend.getExternalUrl(), backend.isActive());
         }
         return backend;
     }
 
     public void deleteBackend(String name)
     {
-        try {
-            connectionManager.open();
-            GatewayBackend.delete("name = ?", name);
-        }
-        finally {
-            connectionManager.close();
-        }
+        dao.deleteByName(name);
+    }
+
+    private static List<ProxyBackendConfiguration> upcast(List<GatewayBackend> gatewayBackendList)
+    {
+        List<ProxyBackendConfiguration> proxyBackendConfigurations = new ArrayList<>();
+        for (GatewayBackend model : gatewayBackendList) {
+            ProxyBackendConfiguration backendConfig = new ProxyBackendConfiguration();
+            backendConfig.setActive(model.active());
+            backendConfig.setRoutingGroup(model.routingGroup());
+            backendConfig.setProxyTo(model.backendUrl());
+            backendConfig.setExternalUrl(model.externalUrl());
+            backendConfig.setName(model.name());
+            proxyBackendConfigurations.add(backendConfig);
+        }
+        return proxyBackendConfigurations;
     }
 }

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

@@ -32,7 +32,7 @@ public class TestHaGatewayManager
     public void setUp()
     {
         JdbcConnectionManager connectionManager = createTestingJdbcConnectionManager();
-        haGatewayManager = new HaGatewayManager(connectionManager);
+        haGatewayManager = new HaGatewayManager(connectionManager.getJdbi());
     }
 
     @Test