diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ExactMatchSourceSelectors.java b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ExactMatchSourceSelectors.java index 772626797..cee365b1c 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ExactMatchSourceSelectors.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ExactMatchSourceSelectors.java @@ -13,66 +13,21 @@ */ package io.trino.gateway.ha.persistence.dao; -import org.javalite.activejdbc.Model; -import org.javalite.activejdbc.annotations.Cached; -import org.javalite.activejdbc.annotations.CompositePK; -import org.javalite.activejdbc.annotations.Table; +import org.jdbi.v3.core.mapper.reflect.ColumnName; -import java.util.ArrayList; -import java.util.List; +import static java.util.Objects.requireNonNull; -import static io.trino.gateway.ha.router.ResourceGroupsManager.ExactSelectorsDetail; - -@CompositePK({"environment", "source", "query_type"}) -@Table("exact_match_source_selectors") // located in gateway-ha-persistence-*.sql -@Cached -public class ExactMatchSourceSelectors - extends Model +public record ExactMatchSourceSelectors( + @ColumnName("resource_group_id") String resourceGroupId, + @ColumnName("update_time") String updateTime, + @ColumnName("source") String source, + @ColumnName("environment") String environment, + @ColumnName("query_type") String queryType) { - private static final String resourceGroupId = "resource_group_id"; - private static final String updateTime = "update_time"; - - private static final String source = "source"; - private static final String environment = "environment"; - private static final String queryType = "query_type"; - - /** - * Returns the most specific exact-match selector for a given environment, source and query type. - * NULL values in the environment and query type fields signify wildcards. - * - * @return List of ExactMatchSourceSelectors - */ - public static List upcast( - List exactMatchSourceSelectorsList) - { - List exactSelectors = new ArrayList<>(); - for (ExactMatchSourceSelectors dao : exactMatchSourceSelectorsList) { - ExactSelectorsDetail exactSelectorDetail = new ExactSelectorsDetail(); - exactSelectorDetail.setResourceGroupId(dao.getString(resourceGroupId)); - exactSelectorDetail.setUpdateTime(dao.getString(updateTime)); - - exactSelectorDetail.setSource(dao.getString(source)); - exactSelectorDetail.setEnvironment(dao.getString(environment)); - exactSelectorDetail.setQueryType(dao.getString(queryType)); - - exactSelectors.add(exactSelectorDetail); - } - return exactSelectors; - } - - /** - * Create a new exactMatchSourceSelector. - */ - public static void create( - ExactMatchSourceSelectors model, ExactSelectorsDetail exactSelectorsDetail) + public ExactMatchSourceSelectors { - model.set(resourceGroupId, exactSelectorsDetail.getResourceGroupId()); - model.set(updateTime, exactSelectorsDetail.getUpdateTime()); - - model.set(source, exactSelectorsDetail.getSource()); - model.set(environment, exactSelectorsDetail.getEnvironment()); - model.set(queryType, exactSelectorsDetail.getQueryType()); - - model.insert(); + requireNonNull(resourceGroupId); + requireNonNull(updateTime); + requireNonNull(source); } } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ExactMatchSourceSelectorsDao.java b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ExactMatchSourceSelectorsDao.java new file mode 100644 index 000000000..145a62168 --- /dev/null +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ExactMatchSourceSelectorsDao.java @@ -0,0 +1,48 @@ +/* + * 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 io.trino.gateway.ha.router.ResourceGroupsManager; +import org.jdbi.v3.sqlobject.customizer.BindBean; +import org.jdbi.v3.sqlobject.statement.SqlQuery; +import org.jdbi.v3.sqlobject.statement.SqlUpdate; + +import java.util.List; + +public interface ExactMatchSourceSelectorsDao +{ + @SqlQuery(""" + SELECT * FROM exact_match_source_selectors + """) + List findAll(); + + @SqlQuery(""" + SELECT * FROM exact_match_source_selectors + WHERE + resource_group_id = :resourceGroupId + AND update_time = :updateTime + AND source = :source + AND environment = :environment + AND query_type = :queryType + LIMIT 1 + """) + ExactMatchSourceSelectors findFirst(@BindBean ResourceGroupsManager.ExactSelectorsDetail exactSelectors); + + @SqlUpdate(""" + INSERT INTO exact_match_source_selectors + (resource_group_id, update_time, source, environment, query_type) + VALUES (:resourceGroupId, :updateTime, :source, :environment, :queryType) + """) + void insert(@BindBean ResourceGroupsManager.ExactSelectorsDetail exactSelectors); +} 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 9e71217f2..0d6727d74 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 @@ -15,6 +15,7 @@ import io.trino.gateway.ha.persistence.JdbcConnectionManager; import io.trino.gateway.ha.persistence.dao.ExactMatchSourceSelectors; +import io.trino.gateway.ha.persistence.dao.ExactMatchSourceSelectorsDao; import io.trino.gateway.ha.persistence.dao.ResourceGroups; import io.trino.gateway.ha.persistence.dao.ResourceGroupsGlobalProperties; import io.trino.gateway.ha.persistence.dao.ResourceGroupsGlobalPropertiesDao; @@ -24,16 +25,19 @@ import java.util.ArrayList; import java.util.List; +import static com.google.common.collect.ImmutableList.toImmutableList; import static java.lang.String.format; public class HaResourceGroupsManager implements ResourceGroupsManager { private final JdbcConnectionManager connectionManager; + private final ExactMatchSourceSelectorsDao exactMatchSourceSelectorsDao; public HaResourceGroupsManager(JdbcConnectionManager connectionManager) { this.connectionManager = connectionManager; + this.exactMatchSourceSelectorsDao = connectionManager.getJdbi().onDemand(ExactMatchSourceSelectorsDao.class); } /** @@ -324,13 +328,7 @@ public void deleteGlobalProperty(String name, @Nullable String routingGroupDatab public ExactSelectorsDetail createExactMatchSourceSelector( ExactSelectorsDetail exactSelectorDetail) { - try { - connectionManager.open(); - ExactMatchSourceSelectors.create(new ExactMatchSourceSelectors(), exactSelectorDetail); - } - finally { - connectionManager.close(); - } + exactMatchSourceSelectorsDao.insert(exactSelectorDetail); return exactSelectorDetail; } @@ -340,15 +338,10 @@ public ExactSelectorsDetail createExactMatchSourceSelector( @Override public List readExactMatchSourceSelector() { - try { - connectionManager.open(); - List exactMatchSourceSelectorList = - ExactMatchSourceSelectors.findAll(); - return ExactMatchSourceSelectors.upcast(exactMatchSourceSelectorList); - } - finally { - connectionManager.close(); - } + List exactMatchSourceSelectors = exactMatchSourceSelectorsDao.findAll(); + return exactMatchSourceSelectors.stream() + .map(HaResourceGroupsManager::upcastExactSelectors) + .collect(toImmutableList()); } /** @@ -358,32 +351,8 @@ public List readExactMatchSourceSelector() public ExactSelectorsDetail getExactMatchSourceSelector( ExactSelectorsDetail exactSelectorDetail) { - try { - connectionManager.open(); - ExactMatchSourceSelectors model = - ExactMatchSourceSelectors.findFirst( - "resource_group_id = ? and update_time = ? " - + "and source = ? and environment = ? and query_type = ?", - exactSelectorDetail.getResourceGroupId(), - exactSelectorDetail.getUpdateTime(), - exactSelectorDetail.getSource(), - exactSelectorDetail.getEnvironment(), - exactSelectorDetail.getQueryType()); - - List exactMatchSourceSelectorList = new ArrayList(); - exactMatchSourceSelectorList.add(model); - - if (model == null) { - return null; - } - else { - ExactMatchSourceSelectors.upcast(exactMatchSourceSelectorList); - } - } - finally { - connectionManager.close(); - } - return exactSelectorDetail; + ExactMatchSourceSelectors exactSelector = exactMatchSourceSelectorsDao.findFirst(exactSelectorDetail); + return upcastExactSelectors(exactSelector); } public String getMatchingString(Object detail) @@ -414,4 +383,15 @@ private static List upcast(List