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}