Skip to content

Instantly share code, notes, and snippets.

@danielflower
Created February 28, 2020 02:46
Show Gist options
  • Save danielflower/fa03b29d499105339c026714803d863b to your computer and use it in GitHub Desktop.
Save danielflower/fa03b29d499105339c026714803d863b to your computer and use it in GitHub Desktop.
Java code to launch postgres instances. Useful for testing etc.
package scaffolding;
import com.zaxxer.hikari.HikariDataSource;
import io.muserver.Mutils;
import org.flywaydb.core.Flyway;
import org.junit.jupiter.api.Assumptions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import javax.sql.DataSource;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Random;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import static io.muserver.Mutils.fullPath;
public class PostgresInstance {
private static final Logger log = LoggerFactory.getLogger(PostgresInstance.class);
private final Process pgProcess;
private final int port;
private final String dbName;
private final ExecutorService executor = Executors.newSingleThreadExecutor();
private final CountDownLatch readyLatch = new CountDownLatch(1);
private final HikariDataSource dataSource;
private static PostgresInstance singleton;
public DataSource dataSource() {
return dataSource;
}
public synchronized static PostgresInstance testInstance() {
if (singleton == null) {
String pgHomePath = System.getenv("PG_HOME");
if (Mutils.nullOrEmpty(pgHomePath)) {
log.info("PG_HOME not set so will not test the database");
return null;
}
File pgHome = new File(new File(pgHomePath), "bin");
File instanceDir = new File("db/instances/test");
int port = new Random().nextInt(40000) + 20000;
try {
String dbName = "test" + System.currentTimeMillis();
PostgresInstance pg = PostgresInstance.start(instanceDir, pgHome, port, dbName);
Flyway flyway = Flyway.configure().dataSource(pg.dataSource()).load();
flyway.migrate();
singleton = pg;
Runtime.getRuntime().addShutdownHook(new Thread(() -> {
pg.dataSource.close();
try (Connection con = DriverManager.getConnection(jdbcUrl(port, "postgres"));
Statement statement = con.createStatement()) {
statement.execute("DROP DATABASE " + dbName);
} catch (SQLException e) {
log.info("Could not delete test DB after tests", e);
}
try {
pg.stop();
} catch (InterruptedException e) {
log.info("Exception stopping PG instance", e);
}
}));
} catch (Exception e) {
throw new RuntimeException("Error starting PG instance", e);
}
}
return singleton;
}
private PostgresInstance(Process pgProcess, int port, String dbName) {
this.pgProcess = pgProcess;
this.port = port;
this.dbName = dbName;
this.dataSource = new HikariDataSource();
this.dataSource.setJdbcUrl(jdbcUrl());
executor.submit(() -> {
InputStreamReader outputStream = new InputStreamReader(pgProcess.getInputStream(), StandardCharsets.UTF_8);
StringBuilder allOutput = new StringBuilder();
char[] buffer = new char[8192];
int read;
try {
while ((read = outputStream.read(buffer)) > -1) {
if (read > 0) {
String s = new String(buffer, 0, read);
System.out.print(s);
if (allOutput != null) {
allOutput.append(s);
if (allOutput.toString().contains("database system is ready to accept connections")) {
allOutput = null;
readyLatch.countDown();
}
}
}
}
} catch (IOException e) {
System.out.println("Got exception " + e);
}
System.out.println("PG has been shut down");
});
}
public static PostgresInstance assumeAvailable() {
PostgresInstance pg = PostgresInstance.testInstance();
Assumptions.assumeTrue(pg != null, "No local PG to test against");
return pg;
}
public void stop() throws InterruptedException {
dataSource.close();
pgProcess.destroy();
pgProcess.waitFor(1, TimeUnit.MINUTES);
executor.shutdownNow();
}
public long pid() {
return pgProcess.pid();
}
public String jdbcUrl() {
return jdbcUrl(port, dbName);
}
public static String jdbcUrl(int port, String dbName) {
return "jdbc:postgresql://localhost:" + port + "/" + dbName + "?user=crickamapp&password=testpassword";
}
public static PostgresInstance start(File instanceDir, File pgHome, int port, String dbName) throws IOException, InterruptedException {
if (!pgHome.isDirectory()) {
throw new RuntimeException("Did not find the pg binaries at " + Mutils.fullPath(pgHome));
}
boolean createDB = false;
if (!new File(instanceDir, "postgresql.conf").isFile()) {
createDB = true;
log.info("Going to initialise " + fullPath(instanceDir) + " as a postgres db data dir");
instanceDir.mkdirs();
Process initDBProcess = new ProcessBuilder()
.directory(instanceDir)
.inheritIO()
.command(new File(pgHome, "initdb").getCanonicalPath(), "-D", ".", "-E", "UTF-8", "--lc-collate=en_US.UTF-8", "--lc-ctype=en_US.UTF-8")
.start();
boolean completed = initDBProcess.waitFor(1, TimeUnit.MINUTES);
log.info("Completed initdb? " + completed);
}
Process startDBProcess = new ProcessBuilder()
.directory(instanceDir)
.redirectErrorStream(true)
.redirectOutput(ProcessBuilder.Redirect.PIPE)
.command(new File(pgHome, "postgres").getCanonicalPath(), /*"-s",*/ "-p", String.valueOf(port), "-D", ".")
.start();
PostgresInstance pg = new PostgresInstance(startDBProcess, port, dbName);
if (!pg.readyLatch.await(1, TimeUnit.MINUTES)) {
pg.stop();
throw new RuntimeException("Could not start up pg");
}
try (Connection connection = DriverManager.getConnection("jdbc:postgresql://localhost:" + port + "/postgres")) {
try (Statement statement = connection.createStatement()) {
statement.execute("CREATE USER crickamapp WITH PASSWORD 'testpassword'");
} catch (SQLException se) {
// expected if the user exists already
}
try (Statement statement = connection.createStatement()) {
statement.execute("CREATE DATABASE " + dbName + " " +
" WITH OWNER crickamapp " +
" TABLESPACE pg_default " +
" ENCODING 'UTF8' " +
" LC_COLLATE = 'en_US.UTF-8' " +
" LC_CTYPE = 'en_US.UTF-8'");
}
} catch (SQLException e) {
if (!e.getMessage().contains("already exists")) {
throw new RuntimeException("Error setting up user and DB", e);
}
}
if (createDB) {
File creationScript = new File("src/test/resources/create-db.sql");
if (!creationScript.isFile()) {
throw new RuntimeException("Could not find script at " + fullPath(creationScript));
}
Process initDBProcess = new ProcessBuilder()
.directory(instanceDir)
.inheritIO()
.command(new File(pgHome, "psql").getCanonicalPath(), "-h", "localhost", "-p", String.valueOf(port), "-d", "postgres", "-f", fullPath(creationScript))
.start();
boolean completed = initDBProcess.waitFor(1, TimeUnit.MINUTES);
log.info("Completed db creation? " + completed);
}
return pg;
}
}
class RunLocalDB {
private static final Logger log = LoggerFactory.getLogger(RunLocalDB.class);
static final String LOCAL_DB_NAME = "crickam";
static final int LOCAL_DB_PORT = 4444;
public static void main(String[] args) throws Exception {
String pgHomePath = System.getenv("PG_HOME");
if (Mutils.nullOrEmpty(pgHomePath)) {
throw new RuntimeException("Please set the PG_HOME env var to point at your postgres install directory");
}
File instanceDir = new File("db/instances/local");
File pgHome = new File(new File(pgHomePath), "bin");
if (!pgHome.isDirectory()) {
throw new RuntimeException("Did not find the pg binaries at " + Mutils.fullPath(pgHome));
}
PostgresInstance pg = PostgresInstance.start(instanceDir, pgHome, LOCAL_DB_PORT, LOCAL_DB_NAME);
log.info("Started postgres on PID " + pg.pid() + " at " + pg.jdbcUrl());
Runtime.getRuntime().addShutdownHook(new Thread(() -> {
log.info("Shutting down");
try {
pg.stop();
} catch (InterruptedException ignored) {
}
}));
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment