Use JDBI for ExactMatchSourceSelectors

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<ExactSelectorsDetail> upcast(
-            List<ExactMatchSourceSelectors> exactMatchSourceSelectorsList)
-    {
-        List<ExactSelectorsDetail> 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);
     }
 }

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<ExactMatchSourceSelectors> 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);
+}

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<ExactSelectorsDetail> readExactMatchSourceSelector()
     {
-        try {
-            connectionManager.open();
-            List<ExactMatchSourceSelectors> exactMatchSourceSelectorList =
-                    ExactMatchSourceSelectors.findAll();
-            return ExactMatchSourceSelectors.upcast(exactMatchSourceSelectorList);
-        }
-        finally {
-            connectionManager.close();
-        }
+        List<ExactMatchSourceSelectors> exactMatchSourceSelectors = exactMatchSourceSelectorsDao.findAll();
+        return exactMatchSourceSelectors.stream()
+                .map(HaResourceGroupsManager::upcastExactSelectors)
+                .collect(toImmutableList());
     }
 
     /**
@@ -358,32 +351,8 @@ public List<ExactSelectorsDetail> 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<ExactMatchSourceSelectors> 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<GlobalPropertiesDetail> upcast(List<ResourceGroupsGlobalProp
         }
         return globalProperties;
     }
+
+    private static ExactSelectorsDetail upcastExactSelectors(ExactMatchSourceSelectors exactMatchSourceSelector)
+    {
+        ExactSelectorsDetail exactSelectorDetail = new ExactSelectorsDetail();
+        exactSelectorDetail.setResourceGroupId(exactMatchSourceSelector.resourceGroupId());
+        exactSelectorDetail.setUpdateTime(exactMatchSourceSelector.updateTime());
+        exactSelectorDetail.setSource(exactMatchSourceSelector.source());
+        exactSelectorDetail.setEnvironment(exactMatchSourceSelector.environment());
+        exactSelectorDetail.setQueryType(exactMatchSourceSelector.queryType());
+        return exactSelectorDetail;
+    }
 }