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 b3c35ce38..8f7edba34 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 @@ -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); diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/GatewayBackend.java b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/GatewayBackend.java index 0fe8b5d1f..a33521a2d 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/GatewayBackend.java +++ b/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 upcast(List gatewayBackendList) - { - List 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; } } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/GatewayBackendDao.java b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/GatewayBackendDao.java new file mode 100644 index 000000000..87009a03b --- /dev/null +++ b/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 findAll(); + + @SqlQuery(""" + SELECT * FROM gateway_backend + WHERE active = true + """) + List findActiveBackend(); + + @SqlQuery(""" + SELECT * FROM gateway_backend + WHERE active = true AND routing_group = 'adhoc' + """) + List findActiveAdhocBackend(); + + @SqlQuery(""" + SELECT * FROM gateway_backend + WHERE active = true AND routing_group = :routingGroup + """) + List findActiveBackendByRoutingGroup(String routingGroup); + + @SqlQuery(""" + SELECT * FROM gateway_backend + WHERE name = :name + """) + List 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); +} diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaGatewayManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaGatewayManager.java index 3047d9f4d..846ee94a1 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaGatewayManager.java +++ b/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 getAllBackends() { - try { - connectionManager.open(); - List proxyBackendList = GatewayBackend.findAll(); - return GatewayBackend.upcast(proxyBackendList); - } - finally { - connectionManager.close(); - } + List proxyBackendList = dao.findAll(); + return upcast(proxyBackendList); } @Override public List getAllActiveBackends() { - try { - connectionManager.open(); - List proxyBackendList = GatewayBackend.where("active = ?", true); - return GatewayBackend.upcast(proxyBackendList); - } - finally { - connectionManager.close(); - } + List proxyBackendList = dao.findActiveBackend(); + return upcast(proxyBackendList); } @Override public List getActiveAdhocBackends() { try { - connectionManager.open(); - List proxyBackendList = - GatewayBackend.where("active = ? and routing_group = ?", true, "adhoc"); - return GatewayBackend.upcast(proxyBackendList); + List 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 getActiveBackends(String routingGroup) { - try { - connectionManager.open(); - List proxyBackendList = - GatewayBackend.where("active = ? and routing_group = ?", true, routingGroup); - return GatewayBackend.upcast(proxyBackendList); - } - finally { - connectionManager.close(); - } + List proxyBackendList = dao.findActiveBackendByRoutingGroup(routingGroup); + return upcast(proxyBackendList); } @Override public Optional getBackendByName(String name) { - try { - connectionManager.open(); - List proxyBackendList = - GatewayBackend.where("name = ?", name); - return GatewayBackend.upcast(proxyBackendList).stream().findAny(); - } - finally { - connectionManager.close(); - } + List 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 upcast(List gatewayBackendList) + { + List 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; } } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestHaGatewayManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestHaGatewayManager.java index fc1521c3c..dbf4e0eea 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestHaGatewayManager.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestHaGatewayManager.java @@ -38,7 +38,7 @@ public class TestHaGatewayManager public void setUp() { JdbcConnectionManager connectionManager = createTestingJdbcConnectionManager(); - haGatewayManager = new HaGatewayManager(connectionManager); + haGatewayManager = new HaGatewayManager(connectionManager.getJdbi()); } @Test diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestHaRoutingManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestHaRoutingManager.java index 8b2f5d831..93a6e4855 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestHaRoutingManager.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestHaRoutingManager.java @@ -34,7 +34,7 @@ public class TestHaRoutingManager public void setUp() { JdbcConnectionManager connectionManager = createTestingJdbcConnectionManager(); - backendManager = new HaGatewayManager(connectionManager); + backendManager = new HaGatewayManager(connectionManager.getJdbi()); historyManager = new HaQueryHistoryManager(connectionManager.getJdbi()); haRoutingManager = new HaRoutingManager(backendManager, historyManager); } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestTrinoQueueLengthRoutingTable.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestTrinoQueueLengthRoutingTable.java index 2fb902c2f..2b4a1a1f9 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestTrinoQueueLengthRoutingTable.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestTrinoQueueLengthRoutingTable.java @@ -56,7 +56,7 @@ public class TestTrinoQueueLengthRoutingTable public void setUp() { JdbcConnectionManager connectionManager = createTestingJdbcConnectionManager(); - backendManager = new HaGatewayManager(connectionManager); + backendManager = new HaGatewayManager(connectionManager.getJdbi()); historyManager = new HaQueryHistoryManager(connectionManager.getJdbi()); routingTable = new TrinoQueueLengthRoutingTable(backendManager, historyManager);