Java code to launch postgres instances. Useful for testing etc.

PostgresInstance.java

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) {
            }
        }));

    }
}