Merge pull request #328 from zonkyio/spring-database-initializer

Add support for Spring Boot's DataSourceScriptDatabaseInitializer

Author: tomix26

README.md

@@ -434,6 +434,14 @@ Before you use Liquibase, you have to add the following Maven dependency:
 Given that Liquibase does not offer an analogy to the `@FlywayTest` annotation,
 you may consider using the [refresh mode](#refreshing-the-database-during-tests) to refresh an embedded database during tests.
 
+### Spring Boot SQL Init
+
+The library also includes optimized processing of Spring Boot's `spring.sql.init` scripts. If this optimization causes any issues, it can be disabled:
+
+```properties
+zonky.test.database.spring.optimized-sql-init.enabled=false
+```
+
 ## Database Providers
 
 The library can be combined with different database providers.

embedded-database-spring-test/src/main/java/io/zonky/test/db/config/EmbeddedDatabaseAutoConfiguration.java

@@ -18,6 +18,7 @@
 
 import io.zonky.test.db.flyway.FlywayDatabaseExtension;
 import io.zonky.test.db.flyway.FlywayPropertiesPostProcessor;
+import io.zonky.test.db.init.DataSourceScriptDatabaseExtension;
 import io.zonky.test.db.init.EmbeddedDatabaseInitializer;
 import io.zonky.test.db.init.ScriptDatabasePreparer;
 import io.zonky.test.db.liquibase.LiquibaseDatabaseExtension;
@@ -289,6 +290,14 @@ public BeanPostProcessor liquibasePropertiesPostProcessor() {
         return new LiquibasePropertiesPostProcessor();
     }
 
+    @Bean
+    @Role(BeanDefinition.ROLE_INFRASTRUCTURE)
+    @ConditionalOnClass(name = "org.springframework.boot.jdbc.init.DataSourceScriptDatabaseInitializer")
+    @ConditionalOnMissingBean(name = "dataSourceScriptDatabaseExtension")
+    public DataSourceScriptDatabaseExtension dataSourceScriptDatabaseExtension(Environment environment) {
+        return new DataSourceScriptDatabaseExtension(environment);
+    }
+
     @Bean
     @Role(BeanDefinition.ROLE_INFRASTRUCTURE)
     @ConditionalOnMissingBean(name = "embeddedDatabaseInitializer")

embedded-database-spring-test/src/main/java/io/zonky/test/db/init/DataSourceScriptDatabaseExtension.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2025 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.zonky.test.db.init;
+
+import io.zonky.test.db.context.DatabaseContext;
+import io.zonky.test.db.util.AopProxyUtils;
+import io.zonky.test.db.util.ReflectionUtils;
+import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.boot.jdbc.init.DataSourceScriptDatabaseInitializer;
+import org.springframework.core.Ordered;
+import org.springframework.core.env.Environment;
+
+import javax.sql.DataSource;
+
+public class DataSourceScriptDatabaseExtension implements BeanPostProcessor, Ordered {
+
+    private final boolean enabled;
+
+    public DataSourceScriptDatabaseExtension(Environment environment) {
+        this.enabled = environment.getProperty("zonky.test.database.spring.optimized-sql-init.enabled", boolean.class, true);
+    }
+
+    @Override
+    public int getOrder() {
+        return Ordered.HIGHEST_PRECEDENCE + 1;
+    }
+
+    @Override
+    public Object postProcessBeforeInitialization(Object bean, String beanName) {
+        if (enabled && bean instanceof DataSourceScriptDatabaseInitializer) {
+            DataSourceScriptDatabaseInitializer initializer = (DataSourceScriptDatabaseInitializer) bean;
+            DataSource dataSource = ReflectionUtils.getField(initializer, "dataSource");
+            DatabaseContext context = AopProxyUtils.getDatabaseContext(dataSource);
+
+            if (context != null) {
+                context.apply(new DataSourceScriptDatabasePreparer(initializer));
+                return new SuppressedInitializerWrapper(initializer);
+            }
+        }
+
+        return bean;
+    }
+
+    @Override
+    public Object postProcessAfterInitialization(Object bean, String beanName) {
+        return bean;
+    }
+
+    public static class SuppressedInitializerWrapper {
+
+        private final DataSourceScriptDatabaseInitializer initializer;
+
+        public SuppressedInitializerWrapper(DataSourceScriptDatabaseInitializer initializer) {
+            this.initializer = initializer;
+        }
+
+        @Override
+        public String toString() {
+            return "SuppressedInitializerWrapper{initializer=" + initializer.getClass().getName() + "}";
+        }
+    }
+}

embedded-database-spring-test/src/main/java/io/zonky/test/db/init/DataSourceScriptDatabasePreparer.java

@@ -0,0 +1,97 @@
+/*
+ * Copyright 2025 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.zonky.test.db.init;
+
+import com.cedarsoftware.util.DeepEquals;
+import io.zonky.test.db.preparer.DatabasePreparer;
+import io.zonky.test.db.util.ReflectionUtils;
+import org.springframework.boot.jdbc.init.DataSourceScriptDatabaseInitializer;
+
+import org.springframework.util.ReflectionUtils.FieldFilter;
+
+import javax.sql.DataSource;
+import java.lang.reflect.Modifier;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.springframework.util.ReflectionUtils.makeAccessible;
+
+public class DataSourceScriptDatabasePreparer implements DatabasePreparer {
+
+    private static final Set<String> EXCLUDED_FIELDS = new HashSet<>(Arrays.asList("dataSource", "resourceLoader"));
+
+    private static final FieldFilter FIELD_FILTER =
+            field -> !Modifier.isStatic(field.getModifiers()) && !EXCLUDED_FIELDS.contains(field.getName());
+
+    private final DataSourceScriptDatabaseInitializer initializer;
+    private final ThreadLocalDataSource threadLocalDataSource;
+
+    public DataSourceScriptDatabasePreparer(DataSourceScriptDatabaseInitializer initializer) {
+        this.initializer = initializer;
+        this.threadLocalDataSource = new ThreadLocalDataSource();
+        ReflectionUtils.setField(initializer, "dataSource", threadLocalDataSource);
+    }
+
+    @Override
+    public long estimatedDuration() {
+        return 10;
+    }
+
+    @Override
+    public void prepare(DataSource dataSource) {
+        threadLocalDataSource.set(dataSource);
+        try {
+            initializer.initializeDatabase();
+        } finally {
+            threadLocalDataSource.clear();
+        }
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        DataSourceScriptDatabasePreparer that = (DataSourceScriptDatabasePreparer) o;
+        if (initializer.getClass() != that.initializer.getClass()) return false;
+        AtomicBoolean equal = new AtomicBoolean(true);
+        org.springframework.util.ReflectionUtils.doWithFields(initializer.getClass(),
+                field -> {
+                    if (!equal.get()) return;
+                    makeAccessible(field);
+                    if (!DeepEquals.deepEquals(field.get(initializer), field.get(that.initializer))) {
+                        equal.set(false);
+                    }
+                },
+                FIELD_FILTER);
+        return equal.get();
+    }
+
+    @Override
+    public int hashCode() {
+        AtomicInteger result = new AtomicInteger(initializer.getClass().hashCode());
+        org.springframework.util.ReflectionUtils.doWithFields(initializer.getClass(),
+                field -> {
+                    makeAccessible(field);
+                    result.set(31 * result.get() + DeepEquals.deepHashCode(field.get(initializer)));
+                },
+                FIELD_FILTER);
+        return result.get();
+    }
+}

embedded-database-spring-test/src/main/java/io/zonky/test/db/init/ThreadLocalDataSource.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2025 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.zonky.test.db.init;
+
+import io.zonky.test.db.provider.support.AbstractDelegatingDataSource;
+
+import javax.sql.DataSource;
+
+class ThreadLocalDataSource extends AbstractDelegatingDataSource {
+
+    private final ThreadLocal<DataSource> current = new ThreadLocal<>();
+
+    void set(DataSource dataSource) {
+        current.set(dataSource);
+    }
+
+    void clear() {
+        current.remove();
+    }
+
+    @Override
+    protected DataSource getDataSource() {
+        return current.get();
+    }
+}

embedded-database-spring-test/src/main/java/io/zonky/test/db/provider/support/AbstractDelegatingDataSource.java

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2025 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.zonky.test.db.provider.support;
+
+import javax.sql.DataSource;
+import java.io.PrintWriter;
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.SQLFeatureNotSupportedException;
+import java.util.logging.Logger;
+
+public abstract class AbstractDelegatingDataSource implements DataSource {
+
+    protected abstract DataSource getDataSource();
+
+    @Override
+    public Connection getConnection() throws SQLException {
+        return getDataSource().getConnection();
+    }
+
+    @Override
+    public Connection getConnection(String username, String password) throws SQLException {
+        return getDataSource().getConnection(username, password);
+    }
+
+    @Override
+    public PrintWriter getLogWriter() throws SQLException {
+        return getDataSource().getLogWriter();
+    }
+
+    @Override
+    public void setLogWriter(PrintWriter out) throws SQLException {
+        getDataSource().setLogWriter(out);
+    }
+
+    @Override
+    public int getLoginTimeout() throws SQLException {
+        return getDataSource().getLoginTimeout();
+    }
+
+    @Override
+    public void setLoginTimeout(int seconds) throws SQLException {
+        getDataSource().setLoginTimeout(seconds);
+    }
+
+    @Override
+    public <T> T unwrap(Class<T> iface) throws SQLException {
+        if (iface.isAssignableFrom(getClass())) {
+            return iface.cast(this);
+        }
+        if (iface.isAssignableFrom(getDataSource().getClass())) {
+            return iface.cast(getDataSource());
+        }
+        return getDataSource().unwrap(iface);
+    }
+
+    @Override
+    public boolean isWrapperFor(Class<?> iface) throws SQLException {
+        if (iface.isAssignableFrom(getClass())) {
+            return true;
+        }
+        if (iface.isAssignableFrom(getDataSource().getClass())) {
+            return true;
+        }
+        return getDataSource().isWrapperFor(iface);
+    }
+
+    @Override
+    public Logger getParentLogger() throws SQLFeatureNotSupportedException {
+        return getDataSource().getParentLogger();
+    }
+}