testcontainers#26 Support init scripts for test containers

zzzLobster · Oct 9, 2023 · 7d07ce1 · 7d07ce1
1 parent a83a0eb
commit 7d07ce1
Show file tree

Hide file tree

Showing 11 changed files with 91 additions and 21 deletions.
diff --git a/README.md b/README.md
@@ -24,6 +24,7 @@ To configure a target database, you need to specify at least database `type` pro
 | username       |          | Provided from database container if not specified                          | Database username for container                                |
 | password       |          | Provided from database container if not specified                          | Database password for container                                |
 | databaseName   |          | Provided from database container if not specified                          | Database name for container                                    |
+| initScript     |          |                                                                            | Path to SQL script to run after container creation                     |
 
 #### `database` block configuration
 
@@ -35,6 +36,7 @@ To configure a target database, you need to specify at least database `type` pro
     <username>test</username>
     <password>test</password>
     <databaseName>test</databaseName>
+    <initScript>filesystem:src/test/resources/db/init.sql</initScript>
 </database>
 ```
 

diff --git a/src/main/java/org/testcontainers/jooq/codegen/Plugin.java b/src/main/java/org/testcontainers/jooq/codegen/Plugin.java
@@ -8,6 +8,7 @@
 import java.net.URL;
 import java.net.URLClassLoader;
 import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 import javax.inject.Inject;
 import org.apache.maven.plugin.AbstractMojo;
@@ -62,15 +63,12 @@ public void execute() throws MojoExecutionException {
             throw new MojoExecutionException("Property 'type' should be specified inside 'database' block");
         }
 
-        final var oldCL = Thread.currentThread().getContextClassLoader();
-        final var mavenClassloader = getMavenClassloader();
-
-        try (var targetDatasource = TargetDatasource.createOrJoinExisting(jooq, database)) {
-            doExecute(mavenClassloader, targetDatasource);
-        } catch (Exception ex) {
-            throw new MojoExecutionException("Error running jOOQ code generation tool", ex);
-        } finally {
-            closeClassloader(oldCL, mavenClassloader);
+        try (var closableContextClassLoader = new ClosableContextClassLoader(getMavenClassloader())) {
+            try (var targetDatasource = TargetDatasource.createOrJoinExisting(jooq, database)) {
+                doExecute(closableContextClassLoader.getClassLoader(), targetDatasource);
+            } catch (Exception ex) {
+                throw new MojoExecutionException("Error running jOOQ code generation tool", ex);
+            }
         }
     }
 
@@ -81,7 +79,6 @@ private void doExecute(URLClassLoader mavenClassloader, TargetDatasource targetD
                 .log(getLog())
                 .mavenClassloader(mavenClassloader)
                 .build();
-        Thread.currentThread().setContextClassLoader(mavenClassloader);
 
         final var oFlyway = Optional.<MigrationRunner>ofNullable(flyway);
         final var oLiquibase = Optional.<MigrationRunner>ofNullable(liquibase);
@@ -104,15 +101,6 @@ private void doExecute(URLClassLoader mavenClassloader, TargetDatasource targetD
         jooqGenerator.generateSources(properties, jooq);
     }
 
-    private void closeClassloader(ClassLoader oldCL, URLClassLoader mavenClassloader) {
-        Thread.currentThread().setContextClassLoader(oldCL);
-        try {
-            mavenClassloader.close();
-        } catch (Throwable e) {
-            getLog().error("Couldn't close the classloader.", e);
-        }
-    }
-
     private URLClassLoader getMavenClassloader() throws MojoExecutionException {
         try {
             List<String> classpathElements = project.getRuntimeClasspathElements();
@@ -130,4 +118,33 @@ private URLClassLoader getMavenClassloader() throws MojoExecutionException {
             throw new MojoExecutionException("Couldn't create a classloader.", e);
         }
     }
+
+    private class ClosableContextClassLoader implements AutoCloseable {
+        private URLClassLoader newClassLoader;
+        private ClassLoader oldClassLoader;
+
+        public ClosableContextClassLoader(URLClassLoader cl) {
+            newClassLoader = Objects.requireNonNull(cl);
+            oldClassLoader = Thread.currentThread().getContextClassLoader();
+            Thread.currentThread().setContextClassLoader(newClassLoader);
+        }
+
+        public URLClassLoader getClassLoader() {
+            return newClassLoader;
+        }
+
+        @Override
+        public void close() {
+            if (newClassLoader == null) {
+                return;
+            }
+            Thread.currentThread().setContextClassLoader(oldClassLoader);
+            try {
+                newClassLoader.close();
+                newClassLoader = null;
+            } catch (Throwable e) {
+                getLog().error("Couldn't close the classloader.", e);
+            }
+        }
+    }
 }
diff --git a/src/main/java/org/testcontainers/jooq/codegen/database/DatabaseProps.java b/src/main/java/org/testcontainers/jooq/codegen/database/DatabaseProps.java
@@ -34,4 +34,9 @@ public class DatabaseProps {
      */
     @Parameter
     private String databaseName;
+    /**
+     * Optional
+     */
+    @Parameter
+    private String initScript;
 }
diff --git a/src/main/java/org/testcontainers/jooq/codegen/datasource/ContainerTargetDatasource.java b/src/main/java/org/testcontainers/jooq/codegen/datasource/ContainerTargetDatasource.java
@@ -1,26 +1,38 @@
 package org.testcontainers.jooq.codegen.datasource;
 
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Path;
 import java.sql.Driver;
 import java.util.Objects;
+import javax.script.ScriptException;
 import lombok.experimental.Delegate;
 import org.testcontainers.containers.JdbcDatabaseContainer;
 import org.testcontainers.containers.wait.strategy.HostPortWaitStrategy;
+import org.testcontainers.ext.ScriptUtils;
+import org.testcontainers.jdbc.JdbcDatabaseDelegate;
+import org.testcontainers.shaded.org.apache.commons.io.FileUtils;
+import org.testcontainers.shaded.org.apache.commons.lang3.StringUtils;
 
 /**
  * Containerized target datasource
  */
 public final class ContainerTargetDatasource implements TargetDatasource {
+    private static final String FILESYSTEM_PREFIX = "filesystem:";
+    private static final String CLASSPATH_PREFIX = "classpath:";
+
     /**
      * Getting datasource properties from container, auto stopping container <br/>
      * {@link AutoCloseable} is implemented by container and {@code close()} delegated to {@code container.stop()}
      */
     @Delegate
     private final JdbcDatabaseContainer<?> container;
 
-    public ContainerTargetDatasource(JdbcDatabaseContainer<?> container) {
+    public ContainerTargetDatasource(JdbcDatabaseContainer<?> container, String initScript) {
         this.container = Objects.requireNonNull(container);
         this.container.setWaitStrategy(new HostPortWaitStrategy());
         this.container.start();
+        runInitScript(initScript);
     }
 
     @Override
@@ -32,4 +44,30 @@ public String getUrl() {
     public Driver getDriverInstance() {
         return container.getJdbcDriverInstance();
     }
+
+    private void runInitScript(String initScript) {
+        if (StringUtils.isEmpty(initScript)) {
+            return;
+        }
+        try (var jdbcDelegate = new JdbcDatabaseDelegate(container, "")) {
+            if (initScript.startsWith(FILESYSTEM_PREFIX)) {
+                var file = Path.of(initScript.substring(FILESYSTEM_PREFIX.length()))
+                        .toAbsolutePath()
+                        .toFile();
+                try {
+                    var scriptBody = FileUtils.readFileToString(file, StandardCharsets.UTF_8);
+                    ScriptUtils.executeDatabaseScript(jdbcDelegate, initScript, scriptBody);
+                    return;
+                } catch (IOException e) {
+                    throw new RuntimeException("Failed to load " + file.getAbsolutePath(), e);
+                } catch (ScriptException e) {
+                    throw new RuntimeException("Failed to execute " + file.getAbsolutePath(), e);
+                }
+            }
+            var scriptClassPath = initScript.startsWith(CLASSPATH_PREFIX)
+                    ? initScript.substring(CLASSPATH_PREFIX.length())
+                    : initScript;
+            ScriptUtils.runInitScript(jdbcDelegate, scriptClassPath);
+        }
+    }
 }
diff --git a/src/main/java/org/testcontainers/jooq/codegen/datasource/TargetDatasource.java b/src/main/java/org/testcontainers/jooq/codegen/datasource/TargetDatasource.java
@@ -16,7 +16,7 @@ public interface TargetDatasource extends AutoCloseable {
     static TargetDatasource createOrJoinExisting(JooqProps jooq, DatabaseProps database) {
         if (needSpinContainer(jooq)) {
             var databaseContainer = DatabaseProvider.getDatabaseContainer(database);
-            return new ContainerTargetDatasource(databaseContainer);
+            return new ContainerTargetDatasource(databaseContainer, database.getInitScript());
         }
 
         return new ExistingTargetDatasource(jooq.getJdbc());

diff --git a/src/test/resources/db/postgres/init.sql b/src/test/resources/db/postgres/init.sql
@@ -0,0 +1 @@
+CREATE USER read_user WITH PASSWORD 'test123';
diff --git a/src/test/resources/pom/postgis-flyway/pom.xml b/src/test/resources/pom/postgis-flyway/pom.xml
@@ -67,6 +67,7 @@
                         <username>test</username>
                         <password>test</password>
                         <databaseName>test</databaseName>
+                        <initScript>classpath:db/postgres/init.sql</initScript>
                     </database>
                     <flyway>
                         <locations>

diff --git a/...sources/pom/postgis-flyway/src/main/resources/db/migration/postgres/V1__create_tables.sql b/...sources/pom/postgis-flyway/src/main/resources/db/migration/postgres/V1__create_tables.sql
@@ -14,3 +14,5 @@ create table users
     primary key (id),
     CONSTRAINT user_email_unique UNIQUE (email)
 );
+
+grant select on users to read_user;
diff --git a/src/test/resources/pom/postgres-flyway/pom.xml b/src/test/resources/pom/postgres-flyway/pom.xml
@@ -67,6 +67,7 @@
                         <username>test</username>
                         <password>test</password>
                         <databaseName>test</databaseName>
+                        <initScript>filesystem:${project.basedir}/src/test/resources/db/init.sql</initScript>
                     </database>
                     <flyway>
                         <locations>

diff --git a/...ources/pom/postgres-flyway/src/main/resources/db/migration/postgres/V1__create_tables.sql b/...ources/pom/postgres-flyway/src/main/resources/db/migration/postgres/V1__create_tables.sql
@@ -14,3 +14,5 @@ create table users
     primary key (id),
     CONSTRAINT user_email_unique UNIQUE (email)
 );
+
+grant select on users to read_user;
diff --git a/src/test/resources/pom/postgres-flyway/src/test/resources/db/init.sql b/src/test/resources/pom/postgres-flyway/src/test/resources/db/init.sql
@@ -0,0 +1 @@
+CREATE USER read_user WITH PASSWORD 'test123';