001/* 002 * Licensed under the Apache License, Version 2.0 (the "License"); 003 * you may not use this file except in compliance with the License. 004 * You may obtain a copy of the License at 005 * 006 * http://www.apache.org/licenses/LICENSE-2.0 007 * 008 * Unless required by applicable law or agreed to in writing, software 009 * distributed under the License is distributed on an "AS IS" BASIS, 010 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 011 * See the License for the specific language governing permissions and 012 * limitations under the License. 013 */ 014package de.softwareforge.testing.postgres.embedded; 015 016import static com.google.common.base.Preconditions.checkNotNull; 017import static com.google.common.base.Preconditions.checkState; 018import static de.softwareforge.testing.postgres.embedded.DatabaseInfo.PG_DEFAULT_USER; 019import static java.lang.String.format; 020 021import java.io.IOException; 022import java.sql.Connection; 023import java.sql.SQLException; 024import java.sql.Statement; 025import java.util.Set; 026import java.util.concurrent.ExecutorService; 027import java.util.concurrent.Executors; 028import java.util.concurrent.SynchronousQueue; 029import java.util.concurrent.atomic.AtomicBoolean; 030import java.util.function.Supplier; 031import javax.sql.DataSource; 032 033import com.google.common.collect.ImmutableSet; 034import com.google.common.util.concurrent.ThreadFactoryBuilder; 035import edu.umd.cs.findbugs.annotations.NonNull; 036import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; 037import org.slf4j.Logger; 038import org.slf4j.LoggerFactory; 039 040/** 041 * Controls database instances on a PostgreSQL server instances. 042 */ 043public final class DatabaseManager implements AutoCloseable { 044 045 private static final String PG_DEFAULT_ENCODING = "utf8"; 046 047 public static final Logger LOG = LoggerFactory.getLogger(DatabaseManager.class); 048 049 private final AtomicBoolean closed = new AtomicBoolean(); 050 private final AtomicBoolean started = new AtomicBoolean(); 051 052 private final Set<EmbeddedPostgresPreparer<DataSource>> databasePreparers; 053 private final Set<EmbeddedPostgresPreparer<EmbeddedPostgres.Builder>> instancePreparers; 054 private final boolean multiMode; 055 056 private volatile InstanceProvider instanceProvider = null; 057 private volatile EmbeddedPostgres pg = null; 058 059 private DatabaseManager(Set<EmbeddedPostgresPreparer<DataSource>> databasePreparers, 060 Set<EmbeddedPostgresPreparer<EmbeddedPostgres.Builder>> instancePreparers, 061 boolean multiMode) { 062 this.databasePreparers = checkNotNull(databasePreparers, "databasePreparers is null"); 063 this.instancePreparers = checkNotNull(instancePreparers, "instancePreparers is null"); 064 this.multiMode = multiMode; 065 } 066 067 /** 068 * Creates a new {@link Builder<DatabaseManager>} instance that will create a new database on each call to {@link #getDatabaseInfo()}. 069 * 070 * @return A builder instance. 071 */ 072 @NonNull 073 public static Builder<DatabaseManager> multiDatabases() { 074 return new DatabaseManagerBuilder(true); 075 } 076 077 /** 078 * Creates a new {@link Builder<DatabaseManager>} instance that will return a connection to the same database on each call to {@link #getDatabaseInfo()}. 079 * 080 * @return A builder instance. 081 */ 082 @NonNull 083 public static Builder<DatabaseManager> singleDatabase() { 084 return new DatabaseManagerBuilder(false); 085 } 086 087 /** 088 * Start the database server and the machinery that will provide new database instances. 089 * 090 * @return This object. 091 * @throws IOException The server could not be started. 092 * @throws SQLException A SQL problem occured while trying to initialize the database. 093 */ 094 @NonNull 095 public DatabaseManager start() throws IOException, SQLException { 096 if (!started.getAndSet(true)) { 097 098 // bring up the embedded postgres server and call all instance preparer instances on it. 099 EmbeddedPostgres.Builder builder = EmbeddedPostgres.builder(); 100 101 for (EmbeddedPostgresPreparer<EmbeddedPostgres.Builder> instancePreparer : instancePreparers) { 102 instancePreparer.prepare(builder); 103 } 104 105 this.pg = builder.build(); 106 107 final DataSource dataSource; 108 109 if (multiMode) { 110 // apply database setup to the template database. 111 dataSource = pg.createTemplateDataSource(); 112 113 // the provider pipeline will create new instances based on the template database. 114 this.instanceProvider = new InstanceProviderPipeline(); 115 } else { 116 // apply database setup to the default database. 117 dataSource = pg.createDefaultDataSource(); 118 119 // always return a reference to the default database. 120 this.instanceProvider = () -> pg.createDefaultDatabaseInfo(); 121 } 122 123 for (EmbeddedPostgresPreparer<DataSource> databasePreparer : databasePreparers) { 124 databasePreparer.prepare(dataSource); 125 } 126 127 this.instanceProvider.start(); 128 } 129 130 return this; 131 } 132 133 @Override 134 public void close() throws Exception { 135 checkState(started.get(), "not yet started!"); 136 if (!closed.getAndSet(true)) { 137 instanceProvider.close(); 138 pg.close(); 139 } 140 } 141 142 /** 143 * Returns a {@link DatabaseInfo} instance that describes a database. If this database provider is in multi mode, every call to this method will return a 144 * new database instance. If it is in single mode, it will always return the same database instance. 145 * 146 * @return A {@link DatabaseInfo} instance. This is never null. 147 * @throws SQLException Any error that happened during the database creation is thrown here. 148 */ 149 @NonNull 150 public DatabaseInfo getDatabaseInfo() throws SQLException { 151 checkState(started.get(), "not yet started!"); 152 153 DatabaseInfo databaseInfo = instanceProvider.get(); 154 if (databaseInfo.exception().isPresent()) { 155 throw databaseInfo.exception().get(); 156 } 157 158 return databaseInfo; 159 } 160 161 /** 162 * Return the {@link EmbeddedPostgres} instance that manages the database server which holds all of the databases managed by this database manager. 163 * 164 * @return An {@link EmbeddedPostgres} instance. Never null. 165 */ 166 @NonNull 167 public EmbeddedPostgres getEmbeddedPostgres() { 168 checkState(started.get(), "not yet started!"); 169 return pg; 170 } 171 172 private interface InstanceProvider extends Supplier<DatabaseInfo>, AutoCloseable { 173 174 default void start() { 175 } 176 177 @Override 178 default void close() { 179 } 180 181 DatabaseInfo get(); 182 } 183 184 private final class InstanceProviderPipeline implements InstanceProvider, Runnable { 185 186 private final ExecutorService executor; 187 private final SynchronousQueue<DatabaseInfo> nextDatabase = new SynchronousQueue<>(); 188 189 private final AtomicBoolean closed = new AtomicBoolean(); 190 191 InstanceProviderPipeline() { 192 this.executor = Executors.newSingleThreadExecutor( 193 new ThreadFactoryBuilder() 194 .setDaemon(true) 195 .setNameFormat("instance-creator-" + pg.instanceId() + "-%d") 196 .build()); 197 198 } 199 200 @Override 201 public void start() { 202 this.executor.submit(this); 203 } 204 205 @Override 206 public void close() { 207 if (!this.closed.getAndSet(true)) { 208 executor.shutdownNow(); 209 } 210 } 211 212 @Override 213 public void run() { 214 while (!closed.get()) { 215 try { 216 final String newDbName = EmbeddedUtil.randomLowercase(12); 217 try { 218 createDatabase(pg.createDefaultDataSource(), newDbName); 219 nextDatabase.put(DatabaseInfo.builder() 220 .dbName(newDbName) 221 .port(pg.getPort()) 222 .connectionProperties(pg.getConnectionProperties()) 223 .build()); 224 } catch (SQLException e) { 225 // https://www.postgresql.org/docs/13/errcodes-appendix.html - 57P01 admin_shutdown 226 if (!e.getSQLState().equals("57P01")) { 227 LOG.warn("Caught SQL Exception (" + e.getSQLState() + "):", e); 228 nextDatabase.put(DatabaseInfo.forException(e)); 229 } 230 } 231 } catch (InterruptedException e) { 232 Thread.currentThread().interrupt(); 233 return; 234 } catch (Throwable t) { 235 LOG.warn("Caught Throwable in instance provider loop:", t); 236 } 237 } 238 } 239 240 @Override 241 public DatabaseInfo get() { 242 try { 243 return nextDatabase.take(); 244 } catch (final InterruptedException e) { 245 Thread.currentThread().interrupt(); 246 throw new IllegalStateException(e); 247 } 248 } 249 } 250 251 @SuppressFBWarnings("RCN_REDUNDANT_NULLCHECK_WOULD_HAVE_BEEN_A_NPE") 252 private static void createDatabase(final DataSource dataSource, final String databaseName) throws SQLException { 253 try (Connection c = dataSource.getConnection(); 254 Statement stmt = c.createStatement()) { 255 stmt.executeUpdate(format("CREATE DATABASE %s OWNER %s ENCODING = '%s'", databaseName, PG_DEFAULT_USER, PG_DEFAULT_ENCODING)); 256 } 257 } 258 259 /** 260 * Builder template. 261 * 262 * @param <T> Object to create. 263 */ 264 public abstract static class Builder<T> { 265 266 protected ImmutableSet.Builder<EmbeddedPostgresPreparer<DataSource>> databasePreparers = ImmutableSet.builder(); 267 protected ImmutableSet.Builder<EmbeddedPostgresPreparer<EmbeddedPostgres.Builder>> instancePreparers = ImmutableSet.builder(); 268 protected final boolean multiMode; 269 270 /** 271 * Creates a new builder. 272 * 273 * @param multiMode True if the resulting object should be in multi mode (create multiple database instances) or single mode (use only one instance). 274 */ 275 protected Builder(boolean multiMode) { 276 this.multiMode = multiMode; 277 } 278 279 /** 280 * @deprecated Use {@link #withDatabasePreparer(EmbeddedPostgresPreparer)}. 281 */ 282 @Deprecated 283 @NonNull 284 public Builder<T> withPreparer(@NonNull DatabasePreparer databasePreparer) { 285 checkNotNull(databasePreparer, "databasePreparer is null"); 286 return withDatabasePreparer(databasePreparer::prepare); 287 } 288 289 /** 290 * Add a preparer for the template database. Each preparer is called once when the database manager starts to prepare the template database. This can be 291 * used to create tables, sequences etc. or preload the databases with information. In multi database mode, the template database is used and each 292 * created database will have this information cloned. 293 * 294 * @param databasePreparer A {@link EmbeddedPostgresPreparer<DataSource>} instance. Must not be null. 295 * @return This object instance. 296 */ 297 @NonNull 298 public Builder<T> withDatabasePreparer(@NonNull EmbeddedPostgresPreparer<DataSource> databasePreparer) { 299 this.databasePreparers.add(checkNotNull(databasePreparer, "databasePreparer is null")); 300 return this; 301 } 302 303 /** 304 * Add preparers for the template database. Each preparer is called once when the database manager starts to prepare the template database. This can be 305 * used to create tables, sequences etc. or preload the databases with information. In multi database mode, the template database is used and each 306 * created database will have this information cloned. 307 * 308 * @param databasePreparers A set of {@link EmbeddedPostgresPreparer<DataSource>} instances. Must not be null. 309 * @return This object instance. 310 */ 311 @NonNull 312 public Builder<T> withDatabasePreparers(@NonNull Set<EmbeddedPostgresPreparer<DataSource>> databasePreparers) { 313 this.databasePreparers.addAll(checkNotNull(databasePreparers, "databasePreparers is null")); 314 return this; 315 } 316 317 /** 318 * Add a preparer for the {@link EmbeddedPostgres.Builder} object. Each preparer is called once when the {@link EmbeddedPostgres} instance that manages 319 * the server is created. 320 * 321 * @param instancePreparer A {@link EmbeddedPostgresPreparer<EmbeddedPostgres.Builder>} instance. Must not be null. 322 * @return This object instance. 323 */ 324 @NonNull 325 public Builder<T> withInstancePreparer(@NonNull EmbeddedPostgresPreparer<EmbeddedPostgres.Builder> instancePreparer) { 326 this.instancePreparers.add(checkNotNull(instancePreparer, "instancePreparer is null")); 327 return this; 328 } 329 330 /** 331 * Add preparers for the {@link EmbeddedPostgres.Builder} object. Each preparer is called once when the {@link EmbeddedPostgres} instance that manages 332 * the server is created. 333 * 334 * @param instancePreparers A set of {@link EmbeddedPostgresPreparer<EmbeddedPostgres.Builder>} instances. Must not be null. 335 * @return This object instance. 336 */ 337 @NonNull 338 public Builder<T> withInstancePreparers(@NonNull Set<EmbeddedPostgresPreparer<EmbeddedPostgres.Builder>> instancePreparers) { 339 this.instancePreparers.addAll(checkNotNull(instancePreparers, "instancePreparers is null")); 340 return this; 341 } 342 343 /** 344 * @deprecated Use {@link #withInstancePreparer(EmbeddedPostgresPreparer)}. 345 */ 346 @Deprecated 347 @NonNull 348 public Builder<T> withCustomizer(@NonNull EmbeddedPostgres.BuilderCustomizer customizer) { 349 checkNotNull(customizer, "customizer is null"); 350 this.instancePreparers.add(customizer::customize); 351 return this; 352 } 353 354 /** 355 * Creates a new instance. 356 * 357 * @return The instance to create. 358 */ 359 @NonNull 360 public abstract T build(); 361 } 362 363 /** 364 * Create new {@link DatabaseManager} instances. 365 */ 366 public static final class DatabaseManagerBuilder extends Builder<DatabaseManager> { 367 368 /** 369 * Creates a new builder for {@link DatabaseManager} instances. 370 * 371 * @param multiMode True if the database manager should return a new database instance for every {@link DatabaseManager#getDatabaseInfo()}} call, false 372 * if it should return the same database instance. 373 */ 374 public DatabaseManagerBuilder(boolean multiMode) { 375 super(multiMode); 376 } 377 378 /** 379 * Creates a new {@link DatabaseManager} instance from the builder. 380 * 381 * @return A database manager. Never null. 382 */ 383 @Override 384 @NonNull 385 public DatabaseManager build() { 386 return new DatabaseManager(databasePreparers.build(), instancePreparers.build(), multiMode); 387 } 388 } 389}