Replace H2 use for testing with PostgreSQL

Author: Akanksha-kedia

gateway-ha/config.yaml

@@ -7,9 +7,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

@@ -322,12 +322,6 @@
         </dependency>
 
         <!-- Test deps -->
-        <dependency>
-            <groupId>com.h2database</groupId>
-            <artifactId>h2</artifactId>
-            <version>1.4.192</version>
-            <scope>test</scope>
-        </dependency>
 
         <dependency>
             <groupId>com.squareup.okhttp3</groupId>

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

@@ -14,14 +14,46 @@
 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) {
+                    return LocalDateTime.now(ZoneId.systemDefault()).format(FORMATTER);
+                }
+            }
+            return result.getString(columnNumber);
+        }
+    }
     @SqlQuery("""
             SELECT * FROM exact_match_source_selectors
             """)
@@ -30,7 +62,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

@@ -25,6 +25,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;
 
@@ -235,11 +238,18 @@ 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 null or empty, set current timestamp
+        if (exactSelectorDetail.getUpdateTime() == null || exactSelectorDetail.getUpdateTime().trim().isEmpty()) {
+            DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
+            exactSelectorDetail.setUpdateTime(LocalDateTime.now(ZoneId.systemDefault()).format(formatter));
+        }
+
         exactMatchSourceSelectorsDao.insert(exactSelectorDetail);
         return exactSelectorDetail;
     }

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

@@ -70,12 +70,11 @@ public static void prepareMockBackend(
                 .setResponseCode(200));
     }
 
-    public static void seedRequiredData(String h2DbFilePath)
+    public static void seedRequiredData(PostgreSQLContainer<?> container)
     {
-        String jdbcUrl = "jdbc:h2:" + h2DbFilePath;
-        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,22 +17,26 @@
 import io.trino.gateway.ha.persistence.JdbcConnectionManager;
 import org.jdbi.v3.core.Jdbi;
 import org.testcontainers.containers.JdbcDatabaseContainer;
-
-import java.io.File;
-import java.nio.file.Path;
+import org.testcontainers.containers.PostgreSQLContainer;
 
 public final class TestingJdbcConnectionManager
 {
     private TestingJdbcConnectionManager() {}
 
     public static JdbcConnectionManager createTestingJdbcConnectionManager()
     {
-        File tempH2DbDir = Path.of(System.getProperty("java.io.tmpdir"), "h2db-" + System.currentTimeMillis()).toFile();
-        tempH2DbDir.deleteOnExit();
-        String jdbcUrl = "jdbc:h2:" + tempH2DbDir.getAbsolutePath();
-        HaGatewayTestUtils.seedRequiredData(tempH2DbDir.getAbsolutePath());
-        DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, "sa", "sa", "org.h2.Driver", 4, false);
-        Jdbi jdbi = Jdbi.create(jdbcUrl, "sa", "sa");
+        PostgreSQLContainer<?> postgres = new PostgreSQLContainer<>("postgres:14-alpine")
+                .withDatabaseName("testdb")
+                .withInitScript("gateway-ha-persistence-postgres.sql");
+        postgres.start();
+
+        String jdbcUrl = postgres.getJdbcUrl();
+        String username = postgres.getUsername();
+        String password = postgres.getPassword();
+
+        DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, username, password, "org.postgresql.Driver", 4, false);
+        Jdbi jdbi = Jdbi.create(jdbcUrl, username, password);
+
         return new JdbcConnectionManager(jdbi, db);
     }
 

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)