com.opentable.db.postgres.embedded.PreparedDbProvider Maven / Gradle / Ivy
/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.opentable.db.postgres.embedded;
import java.io.IOException;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Locale;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.SynchronousQueue;
import javax.annotation.Nonnull;
import javax.annotation.concurrent.GuardedBy;
import javax.sql.DataSource;
import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableMap;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.apache.commons.lang3.RandomStringUtils;
import org.postgresql.ds.PGSimpleDataSource;
public class PreparedDbProvider
{
private static final String JDBC_FORMAT = "jdbc:postgresql://localhost:%d/%s";
/**
* Each database cluster's template1 database has a unique set of schema
* loaded so that the databases may be cloned.
*/
@GuardedBy("PreparedDbProvider.class")
private static final Map CLUSTERS = new HashMap<>();
private final PrepPipeline dbPreparer;
public static PreparedDbProvider forPreparer(DatabasePreparer preparer) {
return new PreparedDbProvider(preparer);
}
private PreparedDbProvider(DatabasePreparer preparer)
{
try {
dbPreparer = createOrFindPreparer(preparer);
} catch (final IOException | SQLException e) {
throw Throwables.propagate(e);
}
}
/**
* Each schema set has its own database cluster. The template1 database has the schema preloaded so that
* each test case need only create a new database and not re-invoke your preparer.
*/
private synchronized PrepPipeline createOrFindPreparer(DatabasePreparer preparer) throws IOException, SQLException
{
PrepPipeline result = CLUSTERS.get(preparer);
if (result != null) {
return result;
}
final EmbeddedPostgres pg = EmbeddedPostgres.start();
preparer.prepare(pg.getTemplateDatabase());
result = new PrepPipeline(pg).start();
CLUSTERS.put(preparer, result);
return result;
}
/**
* Create a new database, and return it as a JDBC connection string.
* No two invocations will return the same database.
*/
public String createDatabase() throws SQLException
{
final DbInfo db = dbPreparer.getNextDb();
return getJdbcUri(db);
}
/**
* Create a new database, and return it as a DataSource.
* No two invocations will return the same database.
*/
public DataSource createDataSource() throws SQLException
{
final DbInfo db = dbPreparer.getNextDb();
final PGSimpleDataSource ds = new PGSimpleDataSource();
ds.setPortNumber(db.port);
ds.setDatabaseName(db.dbName);
ds.setUser(db.user);
return ds;
}
private String getJdbcUri(DbInfo db)
{
return String.format(JDBC_FORMAT, db.port, db.dbName);
}
/**
* Return configuration tweaks in a format appropriate for otj-jdbc DatabaseModule.
*/
public ImmutableMap getConfigurationTweak(String dbModuleName) throws SQLException
{
final DbInfo db = dbPreparer.getNextDb();
return ImmutableMap.of("ot.db." + dbModuleName + ".uri", getJdbcUri(db),
"ot.db." + dbModuleName + ".ds.user", db.user);
}
/**
* Spawns a background thread that prepares databases ahead of time for speed, and then uses a
* synchronous queue to hand the prepared databases off to test cases.
*/
private static class PrepPipeline implements Runnable
{
private final EmbeddedPostgres pg;
private final SynchronousQueue nextDatabase = new SynchronousQueue();
PrepPipeline(EmbeddedPostgres pg)
{
this.pg = pg;
}
PrepPipeline start()
{
final ExecutorService service = Executors.newSingleThreadExecutor(new ThreadFactoryBuilder()
.setDaemon(true).setNameFormat("cluster-" + pg + "-preparer").build());
service.submit(this);
service.shutdown();
return this;
}
DbInfo getNextDb() throws SQLException
{
try {
final DbInfo next = nextDatabase.take();
if (next.ex != null) {
throw next.ex;
}
return next;
} catch (final InterruptedException e) {
Thread.currentThread().interrupt();
throw new IllegalStateException(e);
}
}
@Override
public void run()
{
while (true) {
final String newDbName = RandomStringUtils.randomAlphabetic(12).toLowerCase(Locale.ENGLISH);
SQLException failure = null;
try {
create(pg.getPostgresDatabase(), newDbName, "postgres");
} catch (SQLException e) {
failure = e;
}
try {
if (failure == null) {
nextDatabase.put(new DbInfo(newDbName, pg.getPort(), "postgres"));
} else {
nextDatabase.put(new DbInfo(failure));
}
} catch (final InterruptedException e) {
Thread.currentThread().interrupt();
return;
}
}
}
}
private static void create(final DataSource connectDb, @Nonnull final String dbName, @Nonnull final String userName) throws SQLException
{
Preconditions.checkArgument(dbName != null, "the database name must not be null!");
Preconditions.checkArgument(userName != null, "the user name must not be null!");
try (Connection c = connectDb.getConnection();
PreparedStatement stmt = c.prepareStatement(String.format("CREATE DATABASE %s OWNER %s ENCODING = 'utf8'", dbName, userName))) {
stmt.execute();
}
}
private static class DbInfo
{
private final String dbName;
private final int port;
private final String user;
private final SQLException ex;
DbInfo(String dbName, int port, String user) {
this.dbName = dbName;
this.port = port;
this.user = user;
this.ex = null;
}
DbInfo(SQLException e) {
this.dbName = null;
this.port = -1;
this.user = null;
this.ex = e;
}
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy