Skip to content

Commit b662308

Browse files
seoyonieeseoyoniee
authored
Improve properties handling for DB JDBC connection
1 parent 7d4280f commit b662308

File tree

2 files changed

+168
-4
lines changed

2 files changed

+168
-4
lines changed

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

Lines changed: 37 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,17 @@
1313
*/
1414
package io.trino.gateway.ha.persistence;
1515

16+
import com.google.common.annotations.VisibleForTesting;
1617
import io.airlift.log.Logger;
1718
import io.trino.gateway.ha.config.DataStoreConfiguration;
1819
import io.trino.gateway.ha.persistence.dao.QueryHistoryDao;
1920
import jakarta.annotation.Nullable;
2021
import org.jdbi.v3.core.Jdbi;
2122
import org.jdbi.v3.sqlobject.SqlObjectPlugin;
2223

24+
import java.net.URI;
25+
import java.net.URISyntaxException;
26+
import java.nio.file.Path;
2327
import java.util.concurrent.Executors;
2428
import java.util.concurrent.ScheduledExecutorService;
2529
import java.util.concurrent.TimeUnit;
@@ -60,13 +64,42 @@ public Jdbi getJdbi(@Nullable String routingGroupDatabase)
6064
.registerRowMapper(new RecordAndAnnotatedConstructorMapper());
6165
}
6266

63-
private String buildJdbcUrl(@Nullable String routingGroupDatabase)
67+
@VisibleForTesting
68+
String buildJdbcUrl(@Nullable String routingGroupDatabase)
6469
{
6570
String jdbcUrl = configuration.getJdbcUrl();
66-
if (routingGroupDatabase != null) {
67-
jdbcUrl = jdbcUrl.substring(0, jdbcUrl.lastIndexOf('/') + 1) + routingGroupDatabase;
71+
if (jdbcUrl == null) {
72+
throw new IllegalArgumentException("JDBC URL cannot be null");
6873
}
69-
return jdbcUrl;
74+
if (routingGroupDatabase == null) {
75+
return jdbcUrl;
76+
}
77+
try {
78+
int index = jdbcUrl.indexOf("/") + 1;
79+
if (index == 0) {
80+
throw new IllegalArgumentException("Invalid JDBC URL: no '/' found in " + jdbcUrl);
81+
}
82+
83+
URI newUri = getUriWithRoutingGroupDatabase(routingGroupDatabase, index, jdbcUrl);
84+
return jdbcUrl.substring(0, index) + newUri;
85+
}
86+
catch (URISyntaxException e) {
87+
throw new RuntimeException(e);
88+
}
89+
}
90+
91+
private static URI getUriWithRoutingGroupDatabase(String routingGroupDatabase, int index, String jdbcUrl)
92+
throws URISyntaxException
93+
{
94+
URI uri = new URI(jdbcUrl.substring(index));
95+
return new URI(
96+
uri.getScheme(),
97+
uri.getUserInfo(),
98+
uri.getHost(),
99+
uri.getPort(),
100+
Path.of(uri.getPath()).resolveSibling(routingGroupDatabase).toString(),
101+
uri.getQuery(),
102+
uri.getFragment());
70103
}
71104

72105
private void startCleanUps()

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

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package io.trino.gateway.ha.persistence;
15+
16+
import io.trino.gateway.ha.config.DataStoreConfiguration;
17+
import org.jdbi.v3.core.Jdbi;
18+
import org.junit.jupiter.api.Test;
19+
import org.mockito.Mockito;
20+
21+
import static org.assertj.core.api.Assertions.assertThat;
22+
import static org.assertj.core.api.Assertions.assertThatThrownBy;
23+
import static org.mockito.Mockito.when;
24+
25+
final class TestJdbcConnectionManager
26+
{
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+
41+
@Test
42+
void testBuildJdbcUrlWithMySQLAndNoRoutingGroupDatabase()
43+
{
44+
JdbcConnectionManager connectionManager = createConnectionManager("jdbc:mysql://localhost:3306/mydb");
45+
assertThat(connectionManager.buildJdbcUrl(null)).isEqualTo("jdbc:mysql://localhost:3306/mydb");
46+
}
47+
48+
@Test
49+
void testBuildJdbcUrlWithMySQLAndRoutingGroupDatabase()
50+
{
51+
JdbcConnectionManager connectionManager = createConnectionManager("jdbc:mysql://localhost:3306/mydb");
52+
assertThat(connectionManager.buildJdbcUrl("newdb")).isEqualTo("jdbc:mysql://localhost:3306/newdb");
53+
}
54+
55+
@Test
56+
void testBuildJdbcUrlWithMySQLAndParametersAndRoutingGroupDatabase()
57+
{
58+
JdbcConnectionManager connectionManager = createConnectionManager("jdbc:mysql://localhost:3306/mydb?useSSL=false&serverTimezone=Asia/Seoul");
59+
assertThat(connectionManager.buildJdbcUrl("newdb")).isEqualTo("jdbc:mysql://localhost:3306/newdb?useSSL=false&serverTimezone=Asia/Seoul");
60+
}
61+
62+
@Test
63+
void testBuildJdbcUrlWithPostgreSQLAndNoRoutingGroupDatabase()
64+
{
65+
JdbcConnectionManager connectionManager = createConnectionManager("jdbc:postgresql://localhost:5432/mydb");
66+
assertThat(connectionManager.buildJdbcUrl(null)).isEqualTo("jdbc:postgresql://localhost:5432/mydb");
67+
}
68+
69+
@Test
70+
void testBuildJdbcUrlWithPostgreSQLAndRoutingGroupDatabase()
71+
{
72+
JdbcConnectionManager connectionManager = createConnectionManager("jdbc:postgresql://localhost:5432/mydb");
73+
assertThat(connectionManager.buildJdbcUrl("newdb")).isEqualTo("jdbc:postgresql://localhost:5432/newdb");
74+
}
75+
76+
@Test
77+
void testBuildJdbcUrlWithPostgreSQLAndParametersAndRoutingGroupDatabase()
78+
{
79+
JdbcConnectionManager connectionManager = createConnectionManager("jdbc:postgresql://localhost:5432/mydb?ssl=false&serverTimezone=Asia/Seoul");
80+
assertThat(connectionManager.buildJdbcUrl("newdb")).isEqualTo("jdbc:postgresql://localhost:5432/newdb?ssl=false&serverTimezone=Asia/Seoul");
81+
}
82+
83+
@Test
84+
void testBuildJdbcUrlWithOracleAndNoRoutingGroupDatabase()
85+
{
86+
JdbcConnectionManager connectionManager = createConnectionManager("jdbc:oracle:thin:@//localhost:1521/mydb");
87+
assertThat(connectionManager.buildJdbcUrl(null)).isEqualTo("jdbc:oracle:thin:@//localhost:1521/mydb");
88+
}
89+
90+
@Test
91+
void testBuildJdbcUrlWithOracleAndRoutingGroupDatabase()
92+
{
93+
JdbcConnectionManager connectionManager = createConnectionManager("jdbc:oracle:thin:@//localhost:1521/mydb");
94+
assertThat(connectionManager.buildJdbcUrl("newdb")).isEqualTo("jdbc:oracle:thin:@//localhost:1521/newdb");
95+
}
96+
97+
@Test
98+
void testBuildJdbcUrlWithOracleAndParametersAndRoutingGroupDatabase()
99+
{
100+
JdbcConnectionManager connectionManager = createConnectionManager("jdbc:oracle:thin:@//localhost:1521/mydb?sessionTimeZone=Asia/Seoul");
101+
assertThat(connectionManager.buildJdbcUrl("newdb")).isEqualTo("jdbc:oracle:thin:@//localhost:1521/newdb?sessionTimeZone=Asia/Seoul");
102+
}
103+
104+
@Test
105+
void testBuildJdbcUrlWithNullJdbcUrlThrowsException()
106+
{
107+
// Mock the behavior of DataStoreConfiguration.getJdbcUrl
108+
DataStoreConfiguration dataStoreConfiguration = Mockito.mock(DataStoreConfiguration.class);
109+
when(dataStoreConfiguration.getJdbcUrl()).thenReturn(null);
110+
111+
JdbcConnectionManager connectionManager = new JdbcConnectionManager(Jdbi.create("jdbc:h2:/mydb", "sa", "sa"), dataStoreConfiguration);
112+
assertThatThrownBy(() -> connectionManager.buildJdbcUrl(null))
113+
.isInstanceOf(IllegalArgumentException.class)
114+
.hasMessage("JDBC URL cannot be null");
115+
}
116+
117+
@Test
118+
void testBuildJdbcUrlWithNoSlashThrowsException()
119+
{
120+
JdbcConnectionManager connectionManager = createConnectionManager("jdbc:h2:mem:test");
121+
assertThatThrownBy(() -> connectionManager.buildJdbcUrl("newdb"))
122+
.isInstanceOf(IllegalArgumentException.class)
123+
.hasMessage("Invalid JDBC URL: no '/' found in jdbc:h2:mem:test");
124+
}
125+
126+
private static JdbcConnectionManager createConnectionManager(String jdbcUrl)
127+
{
128+
DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, "sa", "sa", "", 4, true);
129+
return new JdbcConnectionManager(Jdbi.create(jdbcUrl, "sa", "sa"), db);
130+
}
131+
}

0 commit comments

Comments
 (0)