diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/JdbcConnectionManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/JdbcConnectionManager.java index d014e1e6b..fe55123be 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/JdbcConnectionManager.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/JdbcConnectionManager.java @@ -51,6 +51,17 @@ public Jdbi getJdbi() return jdbi; } + public Jdbi getJdbi(@Nullable String routingGroupDatabase) + { + if (routingGroupDatabase == null) { + return jdbi; + } + + return Jdbi.create(buildJdbcUrl(routingGroupDatabase), configuration.getUser(), configuration.getPassword()) + .installPlugin(new SqlObjectPlugin()) + .registerRowMapper(new RecordAndAnnotatedConstructorMapper()); + } + public void open() { this.open(null); @@ -58,10 +69,7 @@ public void open() public void open(@Nullable String routingGroupDatabase) { - String jdbcUrl = configuration.getJdbcUrl(); - if (routingGroupDatabase != null) { - jdbcUrl = jdbcUrl.substring(0, jdbcUrl.lastIndexOf('/') + 1) + routingGroupDatabase; - } + String jdbcUrl = buildJdbcUrl(routingGroupDatabase); log.debug("Jdbc url is " + jdbcUrl); Base.open( configuration.getDriver(), @@ -71,6 +79,15 @@ public void open(@Nullable String routingGroupDatabase) log.debug("Connection opened"); } + private String buildJdbcUrl(@Nullable String routingGroupDatabase) + { + String jdbcUrl = configuration.getJdbcUrl(); + if (routingGroupDatabase != null) { + jdbcUrl = jdbcUrl.substring(0, jdbcUrl.lastIndexOf('/') + 1) + routingGroupDatabase; + } + return jdbcUrl; + } + public void close() { Base.close(); diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ResourceGroupsGlobalProperties.java b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ResourceGroupsGlobalProperties.java index e52e39af8..3710ce5f3 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ResourceGroupsGlobalProperties.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ResourceGroupsGlobalProperties.java @@ -13,64 +13,19 @@ */ package io.trino.gateway.ha.persistence.dao; -import io.trino.gateway.ha.router.ResourceGroupsManager.GlobalPropertiesDetail; -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; -@IdName("name") -@Table("resource_groups_global_properties") // located in gateway-ha-persistence.sql -@Cached -public class ResourceGroupsGlobalProperties - extends Model +public record ResourceGroupsGlobalProperties(String name, String value) { - private static final String name = "name"; - private static final String value = "value"; - - /** - * Reads all existing global properties and returns them in a List. - * - * @return List of ResourceGroupGlobalProperties - */ - public static List upcast( - List globalPropertiesList) - { - List globalProperties = new ArrayList<>(); - for (ResourceGroupsGlobalProperties dao : globalPropertiesList) { - GlobalPropertiesDetail globalPropertiesDetail = new GlobalPropertiesDetail(); - globalPropertiesDetail.setName(dao.getString(name)); - globalPropertiesDetail.setValue(dao.getString(value)); - - globalProperties.add(globalPropertiesDetail); - } - return globalProperties; - } - - /** - * Creates a new global property. - */ - public static void create( - ResourceGroupsGlobalProperties model, GlobalPropertiesDetail globalPropertiesDetail) + @JdbiConstructor + public ResourceGroupsGlobalProperties( + @ColumnName("name") String name, + @ColumnName("value") String value) { - model.set(name, globalPropertiesDetail.getName()); - model.set(value, globalPropertiesDetail.getValue()); - - model.insert(); - } - - /** - * Updates existing global property. - */ - public static void update( - ResourceGroupsGlobalProperties model, GlobalPropertiesDetail globalPropertiesDetail) - { - model.set(name, globalPropertiesDetail.getName()); - model.set(value, globalPropertiesDetail.getValue()); - - model.saveIt(); + this.name = requireNonNull(name, "name is null"); + this.value = requireNonNull(value, "value is null"); } } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ResourceGroupsGlobalPropertiesDao.java b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ResourceGroupsGlobalPropertiesDao.java new file mode 100644 index 000000000..0cec6d598 --- /dev/null +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ResourceGroupsGlobalPropertiesDao.java @@ -0,0 +1,59 @@ +/* + * 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 ResourceGroupsGlobalPropertiesDao +{ + @SqlQuery(""" + SELECT * FROM resource_groups_global_properties + """) + List findAll(); + + @SqlQuery(""" + SELECT * FROM resource_groups_global_properties + WHERE name = :name + """) + List findByName(String name); + + @SqlQuery(""" + SELECT * FROM resource_groups_global_properties + WHERE name = :name + LIMIT 1 + """) + ResourceGroupsGlobalProperties findFirstByName(String name); + + @SqlUpdate(""" + INSERT INTO resource_groups_global_properties (name, value) + VALUES (:name, :value) + """) + void insert(String name, String value); + + @SqlUpdate(""" + UPDATE resource_groups_global_properties + SET value = :value + WHERE name = :name + """) + void update(String name, String value); + + @SqlUpdate(""" + DELETE FROM resource_groups_global_properties + WHERE name = :name + """) + void deleteByName(String name); +} diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaResourceGroupsManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaResourceGroupsManager.java index 78711f49c..f9091271b 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaResourceGroupsManager.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaResourceGroupsManager.java @@ -17,6 +17,7 @@ import io.trino.gateway.ha.persistence.dao.ExactMatchSourceSelectors; import io.trino.gateway.ha.persistence.dao.ResourceGroups; import io.trino.gateway.ha.persistence.dao.ResourceGroupsGlobalProperties; +import io.trino.gateway.ha.persistence.dao.ResourceGroupsGlobalPropertiesDao; import io.trino.gateway.ha.persistence.dao.Selectors; import jakarta.annotation.Nullable; @@ -260,14 +261,7 @@ public void deleteSelector(SelectorsDetail selector, @Nullable String routingGro public GlobalPropertiesDetail createGlobalProperty(GlobalPropertiesDetail globalPropertyDetail, @Nullable String routingGroupDatabase) { - try { - connectionManager.open(routingGroupDatabase); - ResourceGroupsGlobalProperties.create( - new ResourceGroupsGlobalProperties(), globalPropertyDetail); - } - finally { - connectionManager.close(); - } + getDao(routingGroupDatabase).insert(globalPropertyDetail.getName(), globalPropertyDetail.getValue()); return globalPropertyDetail; } @@ -278,15 +272,8 @@ public GlobalPropertiesDetail createGlobalProperty(GlobalPropertiesDetail global public List readAllGlobalProperties( @Nullable String routingGroupDatabase) { - try { - connectionManager.open(routingGroupDatabase); - List globalPropertyList = - ResourceGroupsGlobalProperties.findAll(); - return ResourceGroupsGlobalProperties.upcast(globalPropertyList); - } - finally { - connectionManager.close(); - } + List globalPropertyList = getDao(routingGroupDatabase).findAll(); + return upcast(globalPropertyList); } /** @@ -296,15 +283,8 @@ public List readAllGlobalProperties( public List readGlobalProperty(String name, @Nullable String routingGroupDatabase) { - try { - connectionManager.open(routingGroupDatabase); - List globalPropertyList = - ResourceGroupsGlobalProperties.where("name = ?", name); - return ResourceGroupsGlobalProperties.upcast(globalPropertyList); - } - finally { - connectionManager.close(); - } + List globalPropertyList = getDao(routingGroupDatabase).findByName(name); + return upcast(globalPropertyList); } /** @@ -314,20 +294,14 @@ public List readGlobalProperty(String name, public GlobalPropertiesDetail updateGlobalProperty(GlobalPropertiesDetail globalProperty, @Nullable String routingGroupDatabase) { - try { - connectionManager.open(routingGroupDatabase); - ResourceGroupsGlobalProperties model = - ResourceGroupsGlobalProperties.findFirst("name = ?", globalProperty.getName()); + ResourceGroupsGlobalPropertiesDao dao = getDao(routingGroupDatabase); + ResourceGroupsGlobalProperties model = dao.findFirstByName(globalProperty.getName()); - if (model == null) { - ResourceGroupsGlobalProperties.create(new ResourceGroupsGlobalProperties(), globalProperty); - } - else { - ResourceGroupsGlobalProperties.update(model, globalProperty); - } + if (model == null) { + dao.insert(globalProperty.getName(), globalProperty.getValue()); } - finally { - connectionManager.close(); + else { + dao.update(globalProperty.getName(), globalProperty.getValue()); } return globalProperty; } @@ -338,13 +312,7 @@ public GlobalPropertiesDetail updateGlobalProperty(GlobalPropertiesDetail global @Override public void deleteGlobalProperty(String name, @Nullable String routingGroupDatabase) { - try { - connectionManager.open(routingGroupDatabase); - ResourceGroupsGlobalProperties.delete("name = ?", name); - } - finally { - connectionManager.close(); - } + getDao(routingGroupDatabase).deleteByName(name); } /** @@ -426,4 +394,22 @@ else if (detail.getClass().equals(String.class)) { } return "= " + detail; } + + private ResourceGroupsGlobalPropertiesDao getDao(@Nullable String routingGroupDatabase) + { + return connectionManager.getJdbi(routingGroupDatabase).onDemand(ResourceGroupsGlobalPropertiesDao.class); + } + + private static List upcast(List globalPropertiesList) + { + List globalProperties = new ArrayList<>(); + for (ResourceGroupsGlobalProperties dao : globalPropertiesList) { + GlobalPropertiesDetail globalPropertiesDetail = new GlobalPropertiesDetail(); + globalPropertiesDetail.setName(dao.name()); + globalPropertiesDetail.setValue(dao.value()); + + globalProperties.add(globalPropertiesDetail); + } + return globalProperties; + } } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestSpecificDbResourceGroupsManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestSpecificDbResourceGroupsManager.java index aa89acc39..04bef86c9 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestSpecificDbResourceGroupsManager.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestSpecificDbResourceGroupsManager.java @@ -48,7 +48,7 @@ public void setUp() new HaGatewayTestUtils.TestConfig("", tempH2DbDir.getAbsolutePath())); DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, "sa", "sa", "org.h2.Driver", 4); - Jdbi jdbi = Jdbi.create(jdbcUrl); + Jdbi jdbi = Jdbi.create(jdbcUrl, "sa", "sa"); JdbcConnectionManager connectionManager = new JdbcConnectionManager(jdbi, db); super.resourceGroupManager = new HaResourceGroupsManager(connectionManager); }