Replace H2 use for testing with PostgreSQL

gateway-ha/config.yaml

@@ -10,9 +10,9 @@ routingRules:
   # rulesConfigPath: "src/main/resources/rules/routing_rules.yml"
 
 dataStore:
-  jdbcUrl: jdbc:postgresql://localhost:5432/trino_gateway_db
-  user: trino_gateway_db_admin
-  password: P0stG&es
+  jdbcUrl: jdbc:postgresql://${ENV:DB_HOSTNAME:localhost}:${ENV:DB_PORT:5432}/${ENV:DB_NAME:trino_gateway_db}
+  user: ${ENV:DB_USER:trino_gateway_db_admin}
+  password: ${ENV:DB_PASSWORD:P0stG&es}
   driver: org.postgresql.Driver
   queryHistoryHoursRetention: 24
 

gateway-ha/pom.xml

@@ -324,13 +324,6 @@
             <scope>test</scope>
         </dependency>
 
-        <dependency>
-            <groupId>com.h2database</groupId>
-            <artifactId>h2</artifactId>
-            <version>2.4.240</version>
-            <scope>test</scope>
-        </dependency>
-
         <dependency>
             <groupId>com.squareup.okhttp3</groupId>
             <artifactId>mockwebserver</artifactId>

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

@@ -14,14 +14,50 @@
 package io.trino.gateway.ha.persistence.dao;
 
 import io.trino.gateway.ha.router.ResourceGroupsManager;
+import org.jdbi.v3.core.mapper.ColumnMapper;
+import org.jdbi.v3.core.statement.StatementContext;
+import org.jdbi.v3.sqlobject.config.RegisterColumnMapper;
 import org.jdbi.v3.sqlobject.customizer.BindBean;
 import org.jdbi.v3.sqlobject.statement.SqlQuery;
 import org.jdbi.v3.sqlobject.statement.SqlUpdate;
 
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Timestamp;
+import java.time.LocalDateTime;
+import java.time.ZoneId;
+import java.time.format.DateTimeFormatter;
 import java.util.List;
 
+@RegisterColumnMapper(ExactMatchSourceSelectorsDao.TimestampColumnMapper.class)
 public interface ExactMatchSourceSelectorsDao
 {
+    class TimestampColumnMapper
+            implements ColumnMapper<String>
+    {
+        private static final DateTimeFormatter FORMATTER = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
+
+        @Override
+        public String map(ResultSet result, int columnNumber, StatementContext ctx)
+                throws SQLException
+        {
+            String columnName = result.getMetaData().getColumnName(columnNumber);
+            if ("update_time".equals(columnName)) {
+                try {
+                    Timestamp timestamp = result.getTimestamp(columnNumber);
+                    return timestamp != null ? timestamp.toLocalDateTime().format(FORMATTER) : null;
+                }
+                catch (SQLException e) {
+                    // Fallback to current timestamp when database-specific timestamp retrieval fails.
+                    // This can occur with certain database drivers or when the timestamp column format
+                    // is incompatible with JDBC's getTimestamp() method. The fallback ensures the mapper
+                    // always returns a valid timestamp string instead of throwing an exception.
+                    return LocalDateTime.now(ZoneId.systemDefault()).format(FORMATTER);
+                }
+            }
+            return result.getString(columnNumber);
+        }
+    }
     @SqlQuery("""
             SELECT * FROM exact_match_source_selectors
             """)
@@ -30,7 +66,20 @@ public interface ExactMatchSourceSelectorsDao
     @SqlUpdate("""
             INSERT INTO exact_match_source_selectors
             (resource_group_id, update_time, source, environment, query_type)
-            VALUES (:resourceGroupId, :updateTime, :source, :environment, :queryType)
+            VALUES (:resourceGroupId,
+                    CASE
+                        WHEN :updateTime IS NULL OR :updateTime = ''
+                        THEN CURRENT_TIMESTAMP
+                        ELSE CAST(:updateTime AS TIMESTAMP)
+                    END,
+                    :source, :environment, :queryType)
             """)
     void insert(@BindBean ResourceGroupsManager.ExactSelectorsDetail exactSelectors);
+
+    @SqlUpdate("""
+            INSERT INTO exact_match_source_selectors
+            (resource_group_id, update_time, source, environment, query_type)
+            VALUES (:resourceGroupId, CURRENT_TIMESTAMP, :source, :environment, :queryType)
+            """)
+    void insertWithCurrentTimestamp(@BindBean ResourceGroupsManager.ExactSelectorsDetail exactSelectors);
 }

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

@@ -26,6 +26,9 @@
 import io.trino.gateway.ha.persistence.dao.SelectorsDao;
 import jakarta.annotation.Nullable;
 
+import java.time.LocalDateTime;
+import java.time.ZoneId;
+import java.time.format.DateTimeFormatter;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -34,6 +37,8 @@
 public class HaResourceGroupsManager
         implements ResourceGroupsManager
 {
+    private static final DateTimeFormatter UPDATE_TIME_FORMATTER = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
+
     private final JdbcConnectionManager connectionManager;
     private final ExactMatchSourceSelectorsDao exactMatchSourceSelectorsDao;
 
@@ -237,11 +242,17 @@ public void deleteGlobalProperty(String name, @Nullable String routingGroupDatab
 
     /**
      * Creates exact match source selector for db.
+     * If no updateTime is provided, the current timestamp will be used automatically.
      */
     @Override
     public ExactSelectorsDetail createExactMatchSourceSelector(
             ExactSelectorsDetail exactSelectorDetail)
     {
+        // If updateTime is empty, set current timestamp
+        if (exactSelectorDetail.getUpdateTime().trim().isEmpty()) {
+            exactSelectorDetail.setUpdateTime(LocalDateTime.now(ZoneId.systemDefault()).format(UPDATE_TIME_FORMATTER));
+        }
+
         exactMatchSourceSelectorsDao.insert(exactSelectorDetail);
         return exactSelectorDetail;
     }

gateway-ha/src/test/java/io/trino/gateway/ha/HaGatewayTestUtils.java

@@ -69,12 +69,11 @@ public static void prepareMockBackend(
                 .setResponseCode(200));
     }
 
-    public static void seedRequiredData(String h2DbFilePath)
+    public static void seedRequiredData(PostgreSQLContainer container)
     {
-        String jdbcUrl = "jdbc:h2:" + h2DbFilePath + ";NON_KEYWORDS=NAME,VALUE";
-        Jdbi jdbi = Jdbi.create(jdbcUrl, "sa", "sa");
+        Jdbi jdbi = Jdbi.create(container.getJdbcUrl(), container.getUsername(), container.getPassword());
         try (Handle handle = jdbi.open()) {
-            handle.createUpdate(HaGatewayTestUtils.getResourceFileContent("gateway-ha-persistence-mysql.sql"))
+            handle.createUpdate(HaGatewayTestUtils.getResourceFileContent("gateway-ha-persistence-postgres.sql"))
                     .execute();
         }
     }

gateway-ha/src/test/java/io/trino/gateway/ha/TestingJdbcConnectionManager.java

@@ -17,26 +17,57 @@
 import io.trino.gateway.ha.module.HaGatewayProviderModule;
 import io.trino.gateway.ha.persistence.JdbcConnectionManager;
 import org.jdbi.v3.core.Jdbi;
-
-import java.io.File;
-import java.nio.file.Path;
+import org.testcontainers.containers.JdbcDatabaseContainer;
+import org.testcontainers.postgresql.PostgreSQLContainer;
 
 public final class TestingJdbcConnectionManager
 {
     private TestingJdbcConnectionManager() {}
 
+    public static JdbcConnectionManager createTestingJdbcConnectionManager()
+    {
+        PostgreSQLContainer postgres = new PostgreSQLContainer("postgres:14-alpine")
+                .withDatabaseName("testdb")
+                .withInitScript("gateway-ha-persistence-postgres.sql");
+        postgres.start();
+
+        DataStoreConfiguration db = new DataStoreConfiguration(
+                postgres.getJdbcUrl(),
+                postgres.getUsername(),
+                postgres.getPassword(),
+                postgres.getDriverClassName(),
+                4,
+                false);
+        Jdbi jdbi = HaGatewayProviderModule.createJdbi(db);
+
+        return new JdbcConnectionManager(jdbi, db);
+    }
+
     public static DataStoreConfiguration dataStoreConfig()
     {
-        File tempH2DbDir = Path.of(System.getProperty("java.io.tmpdir"), "h2db-" + System.currentTimeMillis()).toFile();
-        tempH2DbDir.deleteOnExit();
-        String jdbcUrl = "jdbc:h2:" + tempH2DbDir.getAbsolutePath() + ";NON_KEYWORDS=NAME,VALUE";
-        HaGatewayTestUtils.seedRequiredData(tempH2DbDir.getAbsolutePath());
-        return new DataStoreConfiguration(jdbcUrl, "sa", "sa", "org.h2.Driver", 4, false);
+        PostgreSQLContainer postgres = new PostgreSQLContainer("postgres:14-alpine")
+                .withDatabaseName("testdb")
+                .withInitScript("gateway-ha-persistence-postgres.sql");
+        postgres.start();
+
+        return new DataStoreConfiguration(
+                postgres.getJdbcUrl(),
+                postgres.getUsername(),
+                postgres.getPassword(),
+                postgres.getDriverClassName(),
+                4,
+                false);
     }
 
     public static JdbcConnectionManager createTestingJdbcConnectionManager(DataStoreConfiguration config)
     {
         Jdbi jdbi = HaGatewayProviderModule.createJdbi(config);
         return new JdbcConnectionManager(jdbi, config);
     }
+
+    public static JdbcConnectionManager createTestingJdbcConnectionManager(JdbcDatabaseContainer<?> container, DataStoreConfiguration config)
+    {
+        Jdbi jdbi = Jdbi.create(container.getJdbcUrl(), container.getUsername(), container.getPassword());
+        return new JdbcConnectionManager(jdbi, config);
+    }
 }

gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestJdbcConnectionManager.java

@@ -24,20 +24,6 @@
 
 final class TestJdbcConnectionManager
 {
-    @Test
-    void testBuildJdbcUrlWithH2AndNoRoutingGroupDatabase()
-    {
-        JdbcConnectionManager connectionManager = createConnectionManager("jdbc:h2:/mydb");
-        assertThat(connectionManager.buildJdbcUrl(null)).isEqualTo("jdbc:h2:/mydb");
-    }
-
-    @Test
-    void testBuildJdbcUrlWithH2AndRoutingGroupDatabase()
-    {
-        JdbcConnectionManager connectionManager = createConnectionManager("jdbc:h2:/mydb");
-        assertThat(connectionManager.buildJdbcUrl("newdb")).isEqualTo("jdbc:h2:/newdb");
-    }
-
     @Test
     void testBuildJdbcUrlWithMySQLAndNoRoutingGroupDatabase()
     {
@@ -108,7 +94,7 @@ void testBuildJdbcUrlWithNullJdbcUrlThrowsException()
         DataStoreConfiguration dataStoreConfiguration = Mockito.mock(DataStoreConfiguration.class);
         when(dataStoreConfiguration.getJdbcUrl()).thenReturn(null);
 
-        JdbcConnectionManager connectionManager = new JdbcConnectionManager(Jdbi.create("jdbc:h2:/mydb", "sa", "sa"), dataStoreConfiguration);
+        JdbcConnectionManager connectionManager = new JdbcConnectionManager(Jdbi.create("jdbc:postgresql://localhost:5432/mydb", "postgres", "postgres"), dataStoreConfiguration);
         assertThatThrownBy(() -> connectionManager.buildJdbcUrl(null))
                 .isInstanceOf(IllegalArgumentException.class)
                 .hasMessage("JDBC URL cannot be null");
@@ -117,10 +103,10 @@ void testBuildJdbcUrlWithNullJdbcUrlThrowsException()
     @Test
     void testBuildJdbcUrlWithNoSlashThrowsException()
     {
-        JdbcConnectionManager connectionManager = createConnectionManager("jdbc:h2:mem:test");
+        JdbcConnectionManager connectionManager = createConnectionManager("jdbc:postgresql:mydb");
         assertThatThrownBy(() -> connectionManager.buildJdbcUrl("newdb"))
                 .isInstanceOf(IllegalArgumentException.class)
-                .hasMessage("Invalid JDBC URL: no '/' found in jdbc:h2:mem:test");
+                .hasMessage("Invalid JDBC URL: no '/' found in jdbc:postgresql:mydb");
     }
 
     private static JdbcConnectionManager createConnectionManager(String jdbcUrl)

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

@@ -80,13 +80,13 @@ void testGatewayManager()
         assertThat(haGatewayManager.getActiveBackends("adhoc")).isEmpty();
         assertThat(haGatewayManager.getAllBackends())
                 .extracting(ProxyBackendConfiguration::getRoutingGroup)
-                .containsExactly("etl", "adhoc");
+                .containsExactlyInAnyOrder("etl", "adhoc");
 
         // Delete a backend
         haGatewayManager.deleteBackend("adhoc1");
         assertThat(haGatewayManager.getAllBackends())
                 .extracting(ProxyBackendConfiguration::getRoutingGroup)
-                .containsExactly("adhoc");
+                .containsExactlyInAnyOrder("adhoc");
 
         // Test default externalUrl to proxyUrl
         ProxyBackendConfiguration adhoc2 = new ProxyBackendConfiguration();