Source code

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 */
014
015package de.softwareforge.testing.postgres.embedded;
016
017import static com.google.common.base.Preconditions.checkNotNull;
018import static com.google.common.base.Preconditions.checkState;
019import static de.softwareforge.testing.postgres.embedded.DatabaseInfo.PG_DEFAULT_USER;
020import static java.lang.String.format;
021
022import jakarta.annotation.Nonnull;
023import java.io.IOException;
024import java.sql.Connection;
025import java.sql.SQLException;
026import java.sql.Statement;
027import java.util.Set;
028import java.util.concurrent.ExecutorService;
029import java.util.concurrent.Executors;
030import java.util.concurrent.SynchronousQueue;
031import java.util.concurrent.atomic.AtomicBoolean;
032import java.util.function.Supplier;
033import javax.sql.DataSource;
034
035import com.google.common.collect.ImmutableSet;
036import com.google.common.util.concurrent.ThreadFactoryBuilder;
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    private 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 DatabaseManager#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
079     * {@link DatabaseManager#getDatabaseInfo()}.
080     *
081     * @return A builder instance.
082     */
083    @Nonnull
084    public static Builder<DatabaseManager> singleDatabase() {
085        return new DatabaseManagerBuilder(false);
086    }
087
088    /**
089     * Start the database server and the machinery that will provide new database instances.
090     *
091     * @return This object.
092     * @throws IOException  The server could not be started.
093     * @throws SQLException A SQL problem occured while trying to initialize the database.
094     */
095    @Nonnull
096    public DatabaseManager start() throws IOException, SQLException {
097        if (!started.getAndSet(true)) {
098
099            // 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}