Skip to content

Commit 43840de

Browse files
tomix26tomix26
committed
fix propagation of connection parameters
1 parent def2770 commit 43840de

File tree

4 files changed

+116
-19
lines changed

4 files changed

+116
-19
lines changed

‎src/main/java/io/zonky/test/db/postgres/embedded/ConnectionInfo.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,15 +13,28 @@
1313
*/
1414
package io.zonky.test.db.postgres.embedded;
1515

16+
import java.util.HashMap;
17+
import java.util.Map;
18+
19+
import static java.util.Collections.emptyMap;
20+
import static java.util.Collections.unmodifiableMap;
21+
1622
public class ConnectionInfo {
23+
1724
private final String dbName;
1825
private final int port;
1926
private final String user;
27+
private final Map<String, String> properties;
2028

2129
public ConnectionInfo(final String dbName, final int port, final String user) {
30+
this(dbName, port, user, emptyMap());
31+
}
32+
33+
public ConnectionInfo(final String dbName, final int port, final String user, final Map<String, String> properties) {
2234
this.dbName = dbName;
2335
this.port = port;
2436
this.user = user;
37+
this.properties = new HashMap<>(properties);
2538
}
2639

2740
public String getUser() {
@@ -35,4 +48,8 @@ public String getDbName() {
3548
public int getPort() {
3649
return port;
3750
}
51+
52+
public Map<String, String> getProperties() {
53+
return unmodifiableMap(properties);
54+
}
3855
}

‎src/main/java/io/zonky/test/db/postgres/embedded/EmbeddedPostgres.java‎

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,6 @@
4545
import java.time.Duration;
4646
import java.util.ArrayList;
4747
import java.util.Arrays;
48-
import java.util.Collections;
4948
import java.util.HashMap;
5049
import java.util.List;
5150
import java.util.Map;
@@ -76,6 +75,7 @@
7675

7776
import static java.nio.file.StandardOpenOption.CREATE_NEW;
7877
import static java.nio.file.StandardOpenOption.WRITE;
78+
import static java.util.Collections.unmodifiableMap;
7979

8080
@SuppressWarnings("PMD.AvoidDuplicateLiterals") // "postgres"
8181
public class EmbeddedPostgres implements Closeable
@@ -100,6 +100,7 @@ public class EmbeddedPostgres implements Closeable
100100

101101
private final Map<String, String> postgresConfig;
102102
private final Map<String, String> localeConfig;
103+
private final Map<String, String> connectConfig;
103104

104105
private volatile FileOutputStream lockStream;
105106
private volatile FileLock lock;
@@ -122,6 +123,7 @@ public class EmbeddedPostgres implements Closeable
122123
this.cleanDataDirectory = cleanDataDirectory;
123124
this.postgresConfig = new HashMap<>(postgresConfig);
124125
this.localeConfig = new HashMap<>(localeConfig);
126+
this.connectConfig = new HashMap<>(connectConfig);
125127
this.port = port;
126128
this.pgDir = prepareBinaries(pgBinaryResolver);
127129
this.errorRedirector = errorRedirector;
@@ -153,7 +155,7 @@ public class EmbeddedPostgres implements Closeable
153155
}
154156

155157
lock();
156-
startPostmaster(connectConfig);
158+
startPostmaster();
157159
}
158160

159161
public DataSource getTemplateDatabase()
@@ -176,7 +178,7 @@ public DataSource getPostgresDatabase(Map<String, String> properties)
176178
}
177179

178180
public DataSource getDatabase(String userName, String dbName) {
179-
return getDatabase(userName, dbName, Collections.emptyMap());
181+
return getDatabase(userName, dbName, connectConfig);
180182
}
181183

182184
public DataSource getDatabase(String userName, String dbName, Map<String, String> properties)
@@ -207,6 +209,11 @@ public int getPort()
207209
return port;
208210
}
209211

212+
Map<String, String> getConnectConfig()
213+
{
214+
return unmodifiableMap(connectConfig);
215+
}
216+
210217
private static int detectPort() throws IOException
211218
{
212219
try (ServerSocket socket = new ServerSocket(0)) {
@@ -235,7 +242,7 @@ private void initdb()
235242
LOG.info("{} initdb completed in {}", instanceId, watch);
236243
}
237244

238-
private void startPostmaster(Map<String, String> connectConfig) throws IOException
245+
private void startPostmaster() throws IOException
239246
{
240247
final StopWatch watch = new StopWatch();
241248
watch.start();
@@ -266,7 +273,7 @@ private void startPostmaster(Map<String, String> connectConfig) throws IOExcepti
266273

267274
Runtime.getRuntime().addShutdownHook(newCloserThread());
268275

269-
waitForServerStartup(watch, connectConfig);
276+
waitForServerStartup(watch);
270277
}
271278

272279
private List<String> createInitOptions()
@@ -294,14 +301,14 @@ private List<String> createLocaleOptions()
294301
return localeOptions;
295302
}
296303

297-
private void waitForServerStartup(StopWatch watch, Map<String, String> connectConfig) throws IOException
304+
private void waitForServerStartup(StopWatch watch) throws IOException
298305
{
299306
Throwable lastCause = null;
300307
final long start = System.nanoTime();
301308
final long maxWaitNs = TimeUnit.NANOSECONDS.convert(pgStartupWait.toMillis(), TimeUnit.MILLISECONDS);
302309
while (System.nanoTime() - start < maxWaitNs) {
303310
try {
304-
verifyReady(connectConfig);
311+
verifyReady();
305312
LOG.info("{} postmaster startup finished in {}", instanceId, watch);
306313
return;
307314
} catch (final SQLException e) {
@@ -319,7 +326,7 @@ private void waitForServerStartup(StopWatch watch, Map<String, String> connectCo
319326
throw new IOException("Gave up waiting for server to start after " + pgStartupWait.toMillis() + "ms", lastCause);
320327
}
321328

322-
private void verifyReady(Map<String, String> connectConfig) throws SQLException
329+
private void verifyReady() throws SQLException
323330
{
324331
final InetAddress localhost;
325332
try {
@@ -344,7 +351,7 @@ private void verifyReady(Map<String, String> connectConfig) throws SQLException
344351
LOG.trace("i/o exception closing test socket", e);
345352
}
346353
}
347-
try (Connection c = getPostgresDatabase(connectConfig).getConnection() ;
354+
try (Connection c = getPostgresDatabase().getConnection() ;
348355
Statement s = c.createStatement() ;
349356
ResultSet rs = s.executeQuery("SELECT 1"))
350357
{

‎src/main/java/io/zonky/test/db/postgres/embedded/PreparedDbProvider.java‎

Lines changed: 30 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,11 @@
1313
*/
1414
package io.zonky.test.db.postgres.embedded;
1515

16+
import io.zonky.test.db.postgres.embedded.EmbeddedPostgres.Builder;
17+
import org.apache.commons.lang3.RandomStringUtils;
18+
import org.postgresql.ds.PGSimpleDataSource;
19+
20+
import javax.sql.DataSource;
1621
import java.io.IOException;
1722
import java.sql.Connection;
1823
import java.sql.PreparedStatement;
@@ -21,17 +26,16 @@
2126
import java.util.HashMap;
2227
import java.util.Locale;
2328
import java.util.Map;
29+
import java.util.Map.Entry;
2430
import java.util.Objects;
31+
import java.util.Set;
2532
import java.util.concurrent.ExecutorService;
2633
import java.util.concurrent.Executors;
2734
import java.util.concurrent.SynchronousQueue;
2835
import java.util.function.Consumer;
2936

30-
import javax.sql.DataSource;
31-
32-
import io.zonky.test.db.postgres.embedded.EmbeddedPostgres.Builder;
33-
import org.apache.commons.lang3.RandomStringUtils;
34-
import org.postgresql.ds.PGSimpleDataSource;
37+
import static java.util.Collections.emptyMap;
38+
import static java.util.Collections.unmodifiableMap;
3539

3640
public class PreparedDbProvider
3741
{
@@ -108,7 +112,7 @@ private DbInfo createNewDB() throws SQLException
108112
public ConnectionInfo createNewDatabase() throws SQLException
109113
{
110114
final DbInfo dbInfo = createNewDB();
111-
return dbInfo == null || !dbInfo.isSuccess() ? null : new ConnectionInfo(dbInfo.getDbName(), dbInfo.getPort(), dbInfo.getUser());
115+
return dbInfo == null || !dbInfo.isSuccess() ? null : new ConnectionInfo(dbInfo.getDbName(), dbInfo.getPort(), dbInfo.getUser(), dbInfo.getProperties());
112116
}
113117

114118
/**
@@ -121,6 +125,12 @@ public DataSource createDataSourceFromConnectionInfo(final ConnectionInfo connec
121125
ds.setPortNumber(connectionInfo.getPort());
122126
ds.setDatabaseName(connectionInfo.getDbName());
123127
ds.setUser(connectionInfo.getUser());
128+
129+
Set<Entry<String, String>> properties = connectionInfo.getProperties().entrySet();
130+
for (Entry<String, String> property : properties) {
131+
ds.setProperty(property.getKey(), property.getValue());
132+
}
133+
124134
return ds;
125135
}
126136

@@ -204,7 +214,7 @@ public void run()
204214
}
205215
try {
206216
if (failure == null) {
207-
nextDatabase.put(DbInfo.ok(newDbName, pg.getPort(), "postgres"));
217+
nextDatabase.put(DbInfo.ok(newDbName, pg.getPort(), "postgres", pg.getConnectConfig()));
208218
} else {
209219
nextDatabase.put(DbInfo.error(failure));
210220
}
@@ -264,22 +274,28 @@ public int hashCode() {
264274
public static class DbInfo
265275
{
266276
public static DbInfo ok(final String dbName, final int port, final String user) {
267-
return new DbInfo(dbName, port, user, null);
277+
return ok(dbName, port, user, emptyMap());
278+
}
279+
280+
private static DbInfo ok(final String dbName, final int port, final String user, final Map<String, String> properties) {
281+
return new DbInfo(dbName, port, user, properties, null);
268282
}
269283

270284
public static DbInfo error(SQLException e) {
271-
return new DbInfo(null, -1, null, e);
285+
return new DbInfo(null, -1, null, emptyMap(), e);
272286
}
273287

274288
private final String dbName;
275289
private final int port;
276290
private final String user;
291+
private final Map<String, String> properties;
277292
private final SQLException ex;
278293

279-
private DbInfo(final String dbName, final int port, final String user, final SQLException e) {
294+
private DbInfo(final String dbName, final int port, final String user, final Map<String, String> properties, final SQLException e) {
280295
this.dbName = dbName;
281296
this.port = port;
282297
this.user = user;
298+
this.properties = properties;
283299
this.ex = e;
284300
}
285301

@@ -295,6 +311,10 @@ public String getUser() {
295311
return user;
296312
}
297313

314+
public Map<String, String> getProperties() {
315+
return unmodifiableMap(properties);
316+
}
317+
298318
public SQLException getException() {
299319
return ex;
300320
}

‎src/test/java/io/zonky/test/db/postgres/embedded/ConnectConfigTest.java‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
package io.zonky.test.db.postgres.embedded;
2+
3+
import io.zonky.test.db.postgres.junit.EmbeddedPostgresRules;
4+
import io.zonky.test.db.postgres.junit.PreparedDbRule;
5+
import org.junit.Rule;
6+
import org.junit.Test;
7+
import org.postgresql.ds.common.BaseDataSource;
8+
9+
import javax.sql.DataSource;
10+
import java.sql.SQLException;
11+
import java.util.Map;
12+
13+
import static org.junit.Assert.assertEquals;
14+
15+
public class ConnectConfigTest {
16+
17+
private final CapturingDatabasePreparer preparer = new CapturingDatabasePreparer();
18+
19+
@Rule
20+
public PreparedDbRule db = EmbeddedPostgresRules.preparedDatabase(preparer)
21+
.customize(builder -> builder.setConnectConfig("connectTimeout", "20"));
22+
23+
@Test
24+
public void test() throws SQLException {
25+
ConnectionInfo connectionInfo = db.getConnectionInfo();
26+
27+
Map<String, String> properties = connectionInfo.getProperties();
28+
assertEquals(1, properties.size());
29+
assertEquals("20", properties.get("connectTimeout"));
30+
31+
BaseDataSource testDatabase = (BaseDataSource) db.getTestDatabase();
32+
assertEquals("20", testDatabase.getProperty("connectTimeout"));
33+
34+
BaseDataSource preparerDataSource = (BaseDataSource) preparer.getDataSource();
35+
assertEquals("20", preparerDataSource.getProperty("connectTimeout"));
36+
}
37+
38+
private class CapturingDatabasePreparer implements DatabasePreparer {
39+
40+
private DataSource dataSource;
41+
42+
@Override
43+
public void prepare(DataSource ds) {
44+
if (dataSource != null)
45+
throw new IllegalStateException("database preparer has been called multiple times");
46+
dataSource = ds;
47+
}
48+
49+
public DataSource getDataSource() {
50+
return dataSource;
51+
}
52+
}
53+
}

0 commit comments

Comments
 (0)