Skip to content

Commit 3c528b2

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

File tree

10 files changed

+509
-177
lines changed

10 files changed

+509
-177
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: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,46 @@
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+
return LocalDateTime.now(ZoneId.systemDefault()).format(FORMATTER);
52+
}
53+
}
54+
return result.getString(columnNumber);
55+
}
56+
}
2557
@SqlQuery("""
2658
SELECT * FROM exact_match_source_selectors
2759
""")
@@ -30,7 +62,20 @@ public interface ExactMatchSourceSelectorsDao
3062
@SqlUpdate("""
3163
INSERT INTO exact_match_source_selectors
3264
(resource_group_id, update_time, source, environment, query_type)
33-
VALUES (:resourceGroupId, :updateTime, :source, :environment, :queryType)
65+
VALUES (:resourceGroupId,
66+
CASE
67+
WHEN :updateTime IS NULL OR :updateTime = ''
68+
THEN CURRENT_TIMESTAMP
69+
ELSE CAST(:updateTime AS TIMESTAMP)
70+
END,
71+
:source, :environment, :queryType)
3472
""")
3573
void insert(@BindBean ResourceGroupsManager.ExactSelectorsDetail exactSelectors);
74+
75+
@SqlUpdate("""
76+
INSERT INTO exact_match_source_selectors
77+
(resource_group_id, update_time, source, environment, query_type)
78+
VALUES (:resourceGroupId, CURRENT_TIMESTAMP, :source, :environment, :queryType)
79+
""")
80+
void insertWithCurrentTimestamp(@BindBean ResourceGroupsManager.ExactSelectorsDetail exactSelectors);
3681
}

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

Lines changed: 10 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

@@ -237,11 +240,18 @@ public void deleteGlobalProperty(String name, @Nullable String routingGroupDatab
237240

238241
/**
239242
* Creates exact match source selector for db.
243+
* If no updateTime is provided, the current timestamp will be used automatically.
240244
*/
241245
@Override
242246
public ExactSelectorsDetail createExactMatchSourceSelector(
243247
ExactSelectorsDetail exactSelectorDetail)
244248
{
249+
// If updateTime is null or empty, set current timestamp
250+
if (exactSelectorDetail.getUpdateTime() == null || exactSelectorDetail.getUpdateTime().trim().isEmpty()) {
251+
DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
252+
exactSelectorDetail.setUpdateTime(LocalDateTime.now(ZoneId.systemDefault()).format(formatter));
253+
}
254+
245255
exactMatchSourceSelectorsDao.insert(exactSelectorDetail);
246256
return exactSelectorDetail;
247257
}

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

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
import static java.nio.charset.StandardCharsets.UTF_8;
5151
import static java.util.Objects.requireNonNull;
5252
import static org.assertj.core.api.Assertions.assertThat;
53+
import static org.junit.jupiter.api.Assumptions.assumeTrue;
5354

5455
public class HaGatewayTestUtils
5556
{
@@ -69,12 +70,11 @@ public static void prepareMockBackend(
6970
.setResponseCode(200));
7071
}
7172

72-
public static void seedRequiredData(String h2DbFilePath)
73+
public static void seedRequiredData(PostgreSQLContainer<?> container)
7374
{
74-
String jdbcUrl = "jdbc:h2:" + h2DbFilePath + ";NON_KEYWORDS=NAME,VALUE";
75-
Jdbi jdbi = Jdbi.create(jdbcUrl, "sa", "sa");
75+
Jdbi jdbi = Jdbi.create(container.getJdbcUrl(), container.getUsername(), container.getPassword());
7676
try (Handle handle = jdbi.open()) {
77-
handle.createUpdate(HaGatewayTestUtils.getResourceFileContent("gateway-ha-persistence-mysql.sql"))
77+
handle.createUpdate(HaGatewayTestUtils.getResourceFileContent("gateway-ha-persistence-postgres.sql"))
7878
.execute();
7979
}
8080
}
@@ -166,7 +166,16 @@ public static void setUpBackend(
166166

167167
public static OracleContainer getOracleContainer()
168168
{
169-
return new OracleContainer("gvenzl/oracle-free:23.9-slim");
169+
// gvenzl/oracle-xe:18.4.0-slim is x86 only, and tests using this image will most likely timeout if
170+
// run on other CPU architectures. This test should run in three circumstances:
171+
// * in CI, defined by the presence of GitHub environment variables
172+
// * if the CPU architecture is x86_64
173+
// * if they are explicitly enabled by setting TG_RUN_ORACLE_TESTS=true in the test environment
174+
// reference: https://docs.github.com/en/actions/writing-workflows/choosing-what-your-workflow-does/store-information-in-variables#default-environment-variables
175+
assumeTrue("true".equals(System.getenv("GITHUB_ACTIONS"))
176+
|| "x86_64".equalsIgnoreCase(System.getProperty("os.arch"))
177+
|| "true".equals(System.getenv("TG_RUN_ORACLE_TESTS")));
178+
return new OracleContainer("gvenzl/oracle-xe:18.4.0-slim");
170179
}
171180

172181
private static void verifyTrinoStatus(int port, String name)

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

Lines changed: 37 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,26 +17,55 @@
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.containers.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+
String jdbcUrl = postgres.getJdbcUrl();
35+
String username = postgres.getUsername();
36+
String password = postgres.getPassword();
37+
38+
DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, username, password, "org.postgresql.Driver", 4, false);
39+
Jdbi jdbi = HaGatewayProviderModule.createJdbi(db);
40+
41+
return new JdbcConnectionManager(jdbi, db);
42+
}
43+
2844
public static DataStoreConfiguration dataStoreConfig()
2945
{
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);
46+
PostgreSQLContainer<?> postgres = new PostgreSQLContainer<>("postgres:14-alpine")
47+
.withDatabaseName("testdb")
48+
.withInitScript("gateway-ha-persistence-postgres.sql");
49+
postgres.start();
50+
51+
return new DataStoreConfiguration(
52+
postgres.getJdbcUrl(),
53+
postgres.getUsername(),
54+
postgres.getPassword(),
55+
"org.postgresql.Driver",
56+
4,
57+
false);
3558
}
3659

3760
public static JdbcConnectionManager createTestingJdbcConnectionManager(DataStoreConfiguration config)
3861
{
3962
Jdbi jdbi = HaGatewayProviderModule.createJdbi(config);
4063
return new JdbcConnectionManager(jdbi, config);
4164
}
65+
66+
public static JdbcConnectionManager createTestingJdbcConnectionManager(JdbcDatabaseContainer<?> container, DataStoreConfiguration config)
67+
{
68+
Jdbi jdbi = Jdbi.create(container.getJdbcUrl(), container.getUsername(), container.getPassword());
69+
return new JdbcConnectionManager(jdbi, config);
70+
}
4271
}

‎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: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ void testGatewayManager()
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)