View Javadoc
1   /*
2    * Licensed under the Apache License, Version 2.0 (the "License");
3    * you may not use this file except in compliance with the License.
4    * You may obtain a copy of the License at
5    *
6    * http://www.apache.org/licenses/LICENSE-2.0
7    *
8    * Unless required by applicable law or agreed to in writing, software
9    * distributed under the License is distributed on an "AS IS" BASIS,
10   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11   * See the License for the specific language governing permissions and
12   * limitations under the License.
13   */
14  
15  package de.softwareforge.testing.postgres.embedded;
16  
17  import static com.google.common.base.Preconditions.checkNotNull;
18  import static com.google.common.base.Preconditions.checkState;
19  import static de.softwareforge.testing.postgres.embedded.DatabaseInfo.PG_DEFAULT_USER;
20  import static java.lang.String.format;
21  
22  import jakarta.annotation.Nonnull;
23  import java.io.IOException;
24  import java.sql.Connection;
25  import java.sql.SQLException;
26  import java.sql.Statement;
27  import java.util.Set;
28  import java.util.concurrent.ExecutorService;
29  import java.util.concurrent.Executors;
30  import java.util.concurrent.SynchronousQueue;
31  import java.util.concurrent.atomic.AtomicBoolean;
32  import java.util.function.Supplier;
33  import javax.sql.DataSource;
34  
35  import com.google.common.collect.ImmutableSet;
36  import com.google.common.util.concurrent.ThreadFactoryBuilder;
37  import org.slf4j.Logger;
38  import org.slf4j.LoggerFactory;
39  
40  /**
41   * Controls database instances on a PostgreSQL server instances.
42   */
43  public final class DatabaseManager implements AutoCloseable {
44  
45      private static final String PG_DEFAULT_ENCODING = "utf8";
46  
47      private static final Logger LOG = LoggerFactory.getLogger(DatabaseManager.class);
48  
49      private final AtomicBoolean closed = new AtomicBoolean();
50      private final AtomicBoolean started = new AtomicBoolean();
51  
52      private final Set<EmbeddedPostgresPreparer<DataSource>> databasePreparers;
53      private final Set<EmbeddedPostgresPreparer<EmbeddedPostgres.Builder>> instancePreparers;
54      private final boolean multiMode;
55  
56      private volatile InstanceProvider instanceProvider = null;
57      private volatile EmbeddedPostgres pg = null;
58  
59      private DatabaseManager(Set<EmbeddedPostgresPreparer<DataSource>> databasePreparers,
60              Set<EmbeddedPostgresPreparer<EmbeddedPostgres.Builder>> instancePreparers,
61              boolean multiMode) {
62          this.databasePreparers = checkNotNull(databasePreparers, "databasePreparers is null");
63          this.instancePreparers = checkNotNull(instancePreparers, "instancePreparers is null");
64          this.multiMode = multiMode;
65      }
66  
67      /**
68       * Creates a new {@link Builder<DatabaseManager>} instance that will create a new database on each call to {@link DatabaseManager#getDatabaseInfo()}.
69       *
70       * @return A builder instance.
71       */
72      @Nonnull
73      public static Builder<DatabaseManager> multiDatabases() {
74          return new DatabaseManagerBuilder(true);
75      }
76  
77      /**
78       * Creates a new {@link Builder<DatabaseManager>} instance that will return a connection to the same database on each call to
79       * {@link DatabaseManager#getDatabaseInfo()}.
80       *
81       * @return A builder instance.
82       */
83      @Nonnull
84      public static Builder<DatabaseManager> singleDatabase() {
85          return new DatabaseManagerBuilder(false);
86      }
87  
88      /**
89       * Start the database server and the machinery that will provide new database instances.
90       *
91       * @return This object.
92       * @throws IOException  The server could not be started.
93       * @throws SQLException A SQL problem occured while trying to initialize the database.
94       */
95      @Nonnull
96      public DatabaseManager start() throws IOException, SQLException {
97          if (!started.getAndSet(true)) {
98  
99              // bring up the embedded postgres server and call all instance preparer instances on it.
100             EmbeddedPostgres.Builder builder = EmbeddedPostgres.builder();
101 
102             for (EmbeddedPostgresPreparer<EmbeddedPostgres.Builder> instancePreparer : instancePreparers) {
103                 instancePreparer.prepare(builder);
104             }
105 
106             this.pg = builder.build();
107 
108             final DataSource dataSource;
109 
110             if (multiMode) {
111                 // apply database setup to the template database.
112                 dataSource = pg.createTemplateDataSource();
113 
114                 // the provider pipeline will create new instances based on the template database.
115                 this.instanceProvider = new InstanceProviderPipeline();
116             } else {
117                 // apply database setup to the default database.
118                 dataSource = pg.createDefaultDataSource();
119 
120                 // always return a reference to the default database.
121                 this.instanceProvider = () -> pg.createDefaultDatabaseInfo();
122             }
123 
124             for (EmbeddedPostgresPreparer<DataSource> databasePreparer : databasePreparers) {
125                 databasePreparer.prepare(dataSource);
126             }
127 
128             this.instanceProvider.start();
129         }
130 
131         return this;
132     }
133 
134     @Override
135     public void close() throws Exception {
136         checkState(started.get(), "not yet started!");
137         if (!closed.getAndSet(true)) {
138 
139             if (instanceProvider != null) {
140                 instanceProvider.close();
141             }
142 
143             if (pg != null) {
144                 pg.close();
145             }
146         }
147     }
148 
149     /**
150      * 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
151      * new database instance. If it is in single mode, it will always return the same database instance.
152      *
153      * @return A {@link DatabaseInfo} instance. This is never null.
154      * @throws SQLException Any error that happened during the database creation is thrown here.
155      */
156     @Nonnull
157     public DatabaseInfo getDatabaseInfo() throws SQLException {
158         checkState(started.get(), "not yet started!");
159 
160         DatabaseInfo databaseInfo = instanceProvider.get();
161         if (databaseInfo.exception().isPresent()) {
162             throw databaseInfo.exception().get();
163         }
164 
165         return databaseInfo;
166     }
167 
168     /**
169      * Return the {@link EmbeddedPostgres} instance that manages the database server which holds all of the databases managed by this database manager.
170      *
171      * @return An {@link EmbeddedPostgres} instance. Never null.
172      */
173     @Nonnull
174     public EmbeddedPostgres getEmbeddedPostgres() {
175         checkState(started.get(), "not yet started!");
176         return pg;
177     }
178 
179     private interface InstanceProvider extends Supplier<DatabaseInfo>, AutoCloseable {
180 
181         default void start() {
182         }
183 
184         @Override
185         default void close() {
186         }
187 
188         @Override
189         DatabaseInfo get();
190     }
191 
192     private final class InstanceProviderPipeline implements InstanceProvider, Runnable {
193 
194         private final ExecutorService executor;
195         private final SynchronousQueue<DatabaseInfo> nextDatabase = new SynchronousQueue<>();
196 
197         private final AtomicBoolean closed = new AtomicBoolean();
198 
199         InstanceProviderPipeline() {
200             this.executor = Executors.newSingleThreadExecutor(
201                     new ThreadFactoryBuilder()
202                             .setDaemon(true)
203                             .setNameFormat("instance-creator-" + pg.instanceId() + "-%d")
204                             .build());
205 
206         }
207 
208         @Override
209         public void start() {
210             this.executor.submit(this);
211         }
212 
213         @Override
214         public void close() {
215             if (!this.closed.getAndSet(true)) {
216                 executor.shutdownNow();
217             }
218         }
219 
220         @Override
221         public void run() {
222             while (!closed.get()) {
223                 try {
224                     final String newDbName = EmbeddedUtil.randomLowercase(12);
225                     try {
226                         createDatabase(pg.createDefaultDataSource(), newDbName);
227                         nextDatabase.put(DatabaseInfo.builder()
228                                 .dbName(newDbName)
229                                 .port(pg.getPort())
230                                 .connectionProperties(pg.getConnectionProperties())
231                                 .build());
232                     } catch (SQLException e) {
233                         // https://www.postgresql.org/docs/13/errcodes-appendix.html - 57P01 admin_shutdown
234                         if (!e.getSQLState().equals("57P01")) {
235                             LOG.warn("Caught SQL Exception (" + e.getSQLState() + "):", e);
236                             nextDatabase.put(DatabaseInfo.forException(e));
237                         }
238                     }
239                 } catch (InterruptedException e) {
240                     Thread.currentThread().interrupt();
241                     return;
242                 } catch (Exception e) {
243                     LOG.warn("Caught exception in instance provider loop:", e);
244                 }
245             }
246         }
247 
248         @Override
249         public DatabaseInfo get() {
250             try {
251                 return nextDatabase.take();
252             } catch (final InterruptedException e) {
253                 Thread.currentThread().interrupt();
254                 throw new IllegalStateException(e);
255             }
256         }
257     }
258 
259     private static void createDatabase(final DataSource dataSource, final String databaseName) throws SQLException {
260         try (Connection c = dataSource.getConnection();
261                 Statement stmt = c.createStatement()) {
262             stmt.executeUpdate(format("CREATE DATABASE %s OWNER %s ENCODING = '%s'", databaseName, PG_DEFAULT_USER, PG_DEFAULT_ENCODING));
263         }
264     }
265 
266     /**
267      * Builder template.
268      *
269      * @param <T> Object to create.
270      */
271     public abstract static class Builder<T> {
272 
273         protected ImmutableSet.Builder<EmbeddedPostgresPreparer<DataSource>> databasePreparers = ImmutableSet.builder();
274         protected ImmutableSet.Builder<EmbeddedPostgresPreparer<EmbeddedPostgres.Builder>> instancePreparers = ImmutableSet.builder();
275         protected final boolean multiMode;
276 
277         /**
278          * Creates a new builder.
279          *
280          * @param multiMode True if the resulting object should be in multi mode (create multiple database instances) or single mode (use only one instance).
281          */
282         protected Builder(boolean multiMode) {
283             this.multiMode = multiMode;
284         }
285 
286         /**
287          * 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
288          * used to create tables, sequences etc. or preload the databases with information. In multi database mode, the template database is used and each
289          * created database will have this information cloned.
290          *
291          * @param databasePreparer A {@link EmbeddedPostgresPreparer<DataSource>} instance. Must not be null.
292          * @return This object instance.
293          * @since 3.0
294          */
295         @Nonnull
296         public Builder<T> withDatabasePreparer(@Nonnull EmbeddedPostgresPreparer<DataSource> databasePreparer) {
297             this.databasePreparers.add(checkNotNull(databasePreparer, "databasePreparer is null"));
298             return this;
299         }
300 
301         /**
302          * Add preparers for the template database. Each preparer is called once when the database manager starts to prepare the template database. This can be
303          * used to create tables, sequences etc. or preload the databases with information. In multi database mode, the template database is used and each
304          * created database will have this information cloned.
305          *
306          * @param databasePreparers A set of {@link EmbeddedPostgresPreparer<DataSource>} instances. Must not be null.
307          * @return This object instance.
308          * @since 3.0
309          */
310         @Nonnull
311         public Builder<T> withDatabasePreparers(@Nonnull Set<EmbeddedPostgresPreparer<DataSource>> databasePreparers) {
312             this.databasePreparers.addAll(checkNotNull(databasePreparers, "databasePreparers is null"));
313             return this;
314         }
315 
316         /**
317          * Add a preparer for the {@link EmbeddedPostgres.Builder} object. Each preparer is called once when the {@link EmbeddedPostgres} instance that manages
318          * the server is created.
319          *
320          * @param instancePreparer A {@link EmbeddedPostgresPreparer<EmbeddedPostgres.Builder>} instance. Must not be null.
321          * @return This object instance.
322          * @since 3.0
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          * @since 3.0
337          */
338         @Nonnull
339         public Builder<T> withInstancePreparers(@Nonnull Set<EmbeddedPostgresPreparer<EmbeddedPostgres.Builder>> instancePreparers) {
340             this.instancePreparers.addAll(checkNotNull(instancePreparers, "instancePreparers is null"));
341             return this;
342         }
343 
344         /**
345          * Creates a new instance.
346          *
347          * @return The instance to create.
348          */
349         @Nonnull
350         public abstract T build();
351     }
352 
353     /**
354      * Create new {@link DatabaseManager} instances.
355      */
356     public static final class DatabaseManagerBuilder extends Builder<DatabaseManager> {
357 
358         /**
359          * Creates a new builder for {@link DatabaseManager} instances.
360          *
361          * @param multiMode True if the database manager should return a new database instance for every {@link DatabaseManager#getDatabaseInfo()}} call, false
362          *                  if it should return the same database instance.
363          */
364         public DatabaseManagerBuilder(boolean multiMode) {
365             super(multiMode);
366         }
367 
368         /**
369          * Creates a new {@link DatabaseManager} instance from the builder.
370          *
371          * @return A database manager. Never null.
372          */
373         @Override
374         @Nonnull
375         public DatabaseManager build() {
376             return new DatabaseManager(databasePreparers.build(), instancePreparers.build(), multiMode);
377         }
378     }
379 }