Skip to content

Commit 146e761

Browse files
Akanksha-kediaAkanksha-kedia
committed
Replace H2 use for testing with PostgreSQL
1 parent 70359dc commit 146e761

File tree

10 files changed

+489
-164
lines changed

10 files changed

+489
-164
lines changed

‎gateway-ha/config.yaml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,9 @@ routingRules:
1010
# rulesConfigPath: "src/main/resources/rules/routing_rules.yml"
1111

1212
dataStore:
13-
jdbcUrl: jdbc:postgresql://localhost:5432/trino_gateway_db
14-
user: trino_gateway_db_admin
15-
password: P0stG&es
13+
jdbcUrl: jdbc:postgresql://${ENV:DB_HOSTNAME:localhost}:${ENV:DB_PORT:5432}/${ENV:DB_NAME:trino_gateway_db}
14+
user: ${ENV:DB_USER:trino_gateway_db_admin}
15+
password: ${ENV:DB_PASSWORD:P0stG&es}
1616
driver: org.postgresql.Driver
1717
queryHistoryHoursRetention: 24
1818

‎gateway-ha/pom.xml‎

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -324,13 +324,6 @@
324324
<scope>test</scope>
325325
</dependency>
326326

327-
<dependency>
328-
<groupId>com.h2database</groupId>
329-
<artifactId>h2</artifactId>
330-
<version>2.4.240</version>
331-
<scope>test</scope>
332-
</dependency>
333-
334327
<dependency>
335328
<groupId>com.squareup.okhttp3</groupId>
336329
<artifactId>mockwebserver</artifactId>

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

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,50 @@
1414
package io.trino.gateway.ha.persistence.dao;
1515

1616
import io.trino.gateway.ha.router.ResourceGroupsManager;
17+
import org.jdbi.v3.core.mapper.ColumnMapper;
18+
import org.jdbi.v3.core.statement.StatementContext;
19+
import org.jdbi.v3.sqlobject.config.RegisterColumnMapper;
1720
import org.jdbi.v3.sqlobject.customizer.BindBean;
1821
import org.jdbi.v3.sqlobject.statement.SqlQuery;
1922
import org.jdbi.v3.sqlobject.statement.SqlUpdate;
2023

24+
import java.sql.ResultSet;
25+
import java.sql.SQLException;
26+
import java.sql.Timestamp;
27+
import java.time.LocalDateTime;
28+
import java.time.ZoneId;
29+
import java.time.format.DateTimeFormatter;
2130
import java.util.List;
2231

32+
@RegisterColumnMapper(ExactMatchSourceSelectorsDao.TimestampColumnMapper.class)
2333
public interface ExactMatchSourceSelectorsDao
2434
{
35+
class TimestampColumnMapper
36+
implements ColumnMapper<String>
37+
{
38+
private static final DateTimeFormatter FORMATTER = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
39+
40+
@Override
41+
public String map(ResultSet result, int columnNumber, StatementContext ctx)
42+
throws SQLException
43+
{
44+
String columnName = result.getMetaData().getColumnName(columnNumber);
45+
if ("update_time".equals(columnName)) {
46+
try {
47+
Timestamp timestamp = result.getTimestamp(columnNumber);
48+
return timestamp != null ? timestamp.toLocalDateTime().format(FORMATTER) : null;
49+
}
50+
catch (SQLException e) {
51+
// Fallback to current timestamp when database-specific timestamp retrieval fails.
52+
// This can occur with certain database drivers or when the timestamp column format
53+
// is incompatible with JDBC's getTimestamp() method. The fallback ensures the mapper
54+
// always returns a valid timestamp string instead of throwing an exception.
55+
return LocalDateTime.now(ZoneId.systemDefault()).format(FORMATTER);
56+
}
57+
}
58+
return result.getString(columnNumber);
59+
}
60+
}
2561
@SqlQuery("""
2662
SELECT * FROM exact_match_source_selectors
2763
""")
@@ -30,7 +66,20 @@ public interface ExactMatchSourceSelectorsDao
3066
@SqlUpdate("""
3167
INSERT INTO exact_match_source_selectors
3268
(resource_group_id, update_time, source, environment, query_type)
33-
VALUES (:resourceGroupId, :updateTime, :source, :environment, :queryType)
69+
VALUES (:resourceGroupId,
70+
CASE
71+
WHEN :updateTime IS NULL OR :updateTime = ''
72+
THEN CURRENT_TIMESTAMP
73+
ELSE CAST(:updateTime AS TIMESTAMP)
74+
END,
75+
:source, :environment, :queryType)
3476
""")
3577
void insert(@BindBean ResourceGroupsManager.ExactSelectorsDetail exactSelectors);
78+
79+
@SqlUpdate("""
80+
INSERT INTO exact_match_source_selectors
81+
(resource_group_id, update_time, source, environment, query_type)
82+
VALUES (:resourceGroupId, CURRENT_TIMESTAMP, :source, :environment, :queryType)
83+
""")
84+
void insertWithCurrentTimestamp(@BindBean ResourceGroupsManager.ExactSelectorsDetail exactSelectors);
3685
}

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@
2626
import io.trino.gateway.ha.persistence.dao.SelectorsDao;
2727
import jakarta.annotation.Nullable;
2828

29+
import java.time.LocalDateTime;
30+
import java.time.ZoneId;
31+
import java.time.format.DateTimeFormatter;
2932
import java.util.ArrayList;
3033
import java.util.List;
3134

@@ -34,6 +37,8 @@
3437
public class HaResourceGroupsManager
3538
implements ResourceGroupsManager
3639
{
40+
private static final DateTimeFormatter UPDATE_TIME_FORMATTER = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
41+
3742
private final JdbcConnectionManager connectionManager;
3843
private final ExactMatchSourceSelectorsDao exactMatchSourceSelectorsDao;
3944

@@ -237,11 +242,17 @@ public void deleteGlobalProperty(String name, @Nullable String routingGroupDatab
237242

238243
/**
239244
* Creates exact match source selector for db.
245+
* If no updateTime is provided, the current timestamp will be used automatically.
240246
*/
241247
@Override
242248
public ExactSelectorsDetail createExactMatchSourceSelector(
243249
ExactSelectorsDetail exactSelectorDetail)
244250
{
251+
// If updateTime is empty, set current timestamp
252+
if (exactSelectorDetail.getUpdateTime().trim().isEmpty()) {
253+
exactSelectorDetail.setUpdateTime(LocalDateTime.now(ZoneId.systemDefault()).format(UPDATE_TIME_FORMATTER));
254+
}
255+
245256
exactMatchSourceSelectorsDao.insert(exactSelectorDetail);
246257
return exactSelectorDetail;
247258
}

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -69,12 +69,11 @@ public static void prepareMockBackend(
6969
.setResponseCode(200));
7070
}
7171

72-
public static void seedRequiredData(String h2DbFilePath)
72+
public static void seedRequiredData(PostgreSQLContainer container)
7373
{
74-
String jdbcUrl = "jdbc:h2:" + h2DbFilePath + ";NON_KEYWORDS=NAME,VALUE";
75-
Jdbi jdbi = Jdbi.create(jdbcUrl, "sa", "sa");
74+
Jdbi jdbi = Jdbi.create(container.getJdbcUrl(), container.getUsername(), container.getPassword());
7675
try (Handle handle = jdbi.open()) {
77-
handle.createUpdate(HaGatewayTestUtils.getResourceFileContent("gateway-ha-persistence-mysql.sql"))
76+
handle.createUpdate(HaGatewayTestUtils.getResourceFileContent("gateway-ha-persistence-postgres.sql"))
7877
.execute();
7978
}
8079
}

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

Lines changed: 39 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,26 +17,57 @@
1717
import io.trino.gateway.ha.module.HaGatewayProviderModule;
1818
import io.trino.gateway.ha.persistence.JdbcConnectionManager;
1919
import org.jdbi.v3.core.Jdbi;
20-
21-
import java.io.File;
22-
import java.nio.file.Path;
20+
import org.testcontainers.containers.JdbcDatabaseContainer;
21+
import org.testcontainers.postgresql.PostgreSQLContainer;
2322

2423
public final class TestingJdbcConnectionManager
2524
{
2625
private TestingJdbcConnectionManager() {}
2726

27+
public static JdbcConnectionManager createTestingJdbcConnectionManager()
28+
{
29+
PostgreSQLContainer<?> postgres = new PostgreSQLContainer<>("postgres:14-alpine")
30+
.withDatabaseName("testdb")
31+
.withInitScript("gateway-ha-persistence-postgres.sql");
32+
postgres.start();
33+
34+
DataStoreConfiguration db = new DataStoreConfiguration(
35+
postgres.getJdbcUrl(),
36+
postgres.getUsername(),
37+
postgres.getPassword(),
38+
postgres.getDriverClassName(),
39+
4,
40+
false);
41+
Jdbi jdbi = HaGatewayProviderModule.createJdbi(db);
42+
43+
return new JdbcConnectionManager(jdbi, db);
44+
}
45+
2846
public static DataStoreConfiguration dataStoreConfig()
2947
{
30-
File tempH2DbDir = Path.of(System.getProperty("java.io.tmpdir"), "h2db-" + System.currentTimeMillis()).toFile();
31-
tempH2DbDir.deleteOnExit();
32-
String jdbcUrl = "jdbc:h2:" + tempH2DbDir.getAbsolutePath() + ";NON_KEYWORDS=NAME,VALUE";
33-
HaGatewayTestUtils.seedRequiredData(tempH2DbDir.getAbsolutePath());
34-
return new DataStoreConfiguration(jdbcUrl, "sa", "sa", "org.h2.Driver", 4, false);
48+
PostgreSQLContainer<?> postgres = new PostgreSQLContainer<>("postgres:14-alpine")
49+
.withDatabaseName("testdb")
50+
.withInitScript("gateway-ha-persistence-postgres.sql");
51+
postgres.start();
52+
53+
return new DataStoreConfiguration(
54+
postgres.getJdbcUrl(),
55+
postgres.getUsername(),
56+
postgres.getPassword(),
57+
postgres.getDriverClassName(),
58+
4,
59+
false);
3560
}
3661

3762
public static JdbcConnectionManager createTestingJdbcConnectionManager(DataStoreConfiguration config)
3863
{
3964
Jdbi jdbi = HaGatewayProviderModule.createJdbi(config);
4065
return new JdbcConnectionManager(jdbi, config);
4166
}
67+
68+
public static JdbcConnectionManager createTestingJdbcConnectionManager(JdbcDatabaseContainer<?> container, DataStoreConfiguration config)
69+
{
70+
Jdbi jdbi = Jdbi.create(container.getJdbcUrl(), container.getUsername(), container.getPassword());
71+
return new JdbcConnectionManager(jdbi, config);
72+
}
4273
}

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

Lines changed: 3 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -24,20 +24,6 @@
2424

2525
final class TestJdbcConnectionManager
2626
{
27-
@Test
28-
void testBuildJdbcUrlWithH2AndNoRoutingGroupDatabase()
29-
{
30-
JdbcConnectionManager connectionManager = createConnectionManager("jdbc:h2:/mydb");
31-
assertThat(connectionManager.buildJdbcUrl(null)).isEqualTo("jdbc:h2:/mydb");
32-
}
33-
34-
@Test
35-
void testBuildJdbcUrlWithH2AndRoutingGroupDatabase()
36-
{
37-
JdbcConnectionManager connectionManager = createConnectionManager("jdbc:h2:/mydb");
38-
assertThat(connectionManager.buildJdbcUrl("newdb")).isEqualTo("jdbc:h2:/newdb");
39-
}
40-
4127
@Test
4228
void testBuildJdbcUrlWithMySQLAndNoRoutingGroupDatabase()
4329
{
@@ -108,7 +94,7 @@ void testBuildJdbcUrlWithNullJdbcUrlThrowsException()
10894
DataStoreConfiguration dataStoreConfiguration = Mockito.mock(DataStoreConfiguration.class);
10995
when(dataStoreConfiguration.getJdbcUrl()).thenReturn(null);
11096

111-
JdbcConnectionManager connectionManager = new JdbcConnectionManager(Jdbi.create("jdbc:h2:/mydb", "sa", "sa"), dataStoreConfiguration);
97+
JdbcConnectionManager connectionManager = new JdbcConnectionManager(Jdbi.create("jdbc:postgresql://localhost:5432/mydb", "postgres", "postgres"), dataStoreConfiguration);
11298
assertThatThrownBy(() -> connectionManager.buildJdbcUrl(null))
11399
.isInstanceOf(IllegalArgumentException.class)
114100
.hasMessage("JDBC URL cannot be null");
@@ -117,10 +103,10 @@ void testBuildJdbcUrlWithNullJdbcUrlThrowsException()
117103
@Test
118104
void testBuildJdbcUrlWithNoSlashThrowsException()
119105
{
120-
JdbcConnectionManager connectionManager = createConnectionManager("jdbc:h2:mem:test");
106+
JdbcConnectionManager connectionManager = createConnectionManager("jdbc:postgresql:mydb");
121107
assertThatThrownBy(() -> connectionManager.buildJdbcUrl("newdb"))
122108
.isInstanceOf(IllegalArgumentException.class)
123-
.hasMessage("Invalid JDBC URL: no '/' found in jdbc:h2:mem:test");
109+
.hasMessage("Invalid JDBC URL: no '/' found in jdbc:postgresql:mydb");
124110
}
125111

126112
private static JdbcConnectionManager createConnectionManager(String jdbcUrl)

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -80,13 +80,13 @@ void testGatewayManager()
8080
assertThat(haGatewayManager.getActiveBackends("adhoc")).isEmpty();
8181
assertThat(haGatewayManager.getAllBackends())
8282
.extracting(ProxyBackendConfiguration::getRoutingGroup)
83-
.containsExactly("etl", "adhoc");
83+
.containsExactlyInAnyOrder("etl", "adhoc");
8484

8585
// Delete a backend
8686
haGatewayManager.deleteBackend("adhoc1");
8787
assertThat(haGatewayManager.getAllBackends())
8888
.extracting(ProxyBackendConfiguration::getRoutingGroup)
89-
.containsExactly("adhoc");
89+
.containsExactlyInAnyOrder("adhoc");
9090

9191
// Test default externalUrl to proxyUrl
9292
ProxyBackendConfiguration adhoc2 = new ProxyBackendConfiguration();

0 commit comments

Comments
 (0)