Add support for DataSourceScriptDatabaseInitializer interception

Intercept Spring Boot's DataSourceScriptDatabaseInitializer via
BeanPostProcessor to register SQL init scripts as DatabasePreparer,
enabling optimized template-based database initialization and
parallel prefetching using ThreadLocal DataSource approach.

Author: tomix26

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() {
+        return new DataSourceScriptDatabaseExtension();
+    }
+
     @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,73 @@
+/*
+ * 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 javax.sql.DataSource;
+
+public class DataSourceScriptDatabaseExtension implements BeanPostProcessor, Ordered {
+
+    @Override
+    public int getOrder() {
+        return Ordered.HIGHEST_PRECEDENCE + 1;
+    }
+
+    @Override
+    public Object postProcessBeforeInitialization(Object bean, String beanName) {
+        if (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 SuppressedDataSourceInitializer(initializer);
+            }
+        }
+
+        return bean;
+    }
+
+    @Override
+    public Object postProcessAfterInitialization(Object bean, String beanName) {
+        return bean;
+    }
+
+    public static class SuppressedDataSourceInitializer {
+
+        private final DataSourceScriptDatabaseInitializer originalInitializer;
+
+        public SuppressedDataSourceInitializer(DataSourceScriptDatabaseInitializer originalInitializer) {
+            this.originalInitializer = originalInitializer;
+        }
+
+        public DataSourceScriptDatabaseInitializer getOriginalInitializer() {
+            return originalInitializer;
+        }
+
+        @Override
+        public String toString() {
+            return "SuppressedDataSourceInitializer{originalType=" + originalInitializer.getClass().getName() + "}";
+        }
+    }
+}

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

@@ -0,0 +1,70 @@
+/*
+ * 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.boot.sql.init.DatabaseInitializationSettings;
+
+import javax.sql.DataSource;
+
+public class DataSourceScriptDatabasePreparer implements DatabasePreparer {
+
+    private final DataSourceScriptDatabaseInitializer initializer;
+    private final ThreadLocalDataSource threadLocalDataSource;
+    private final DatabaseInitializationSettings settings;
+
+    public DataSourceScriptDatabasePreparer(DataSourceScriptDatabaseInitializer initializer) {
+        this.initializer = initializer;
+        this.settings = ReflectionUtils.getField(initializer, "settings");
+        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;
+        return initializer.getClass() == that.initializer.getClass()
+                && DeepEquals.deepEquals(settings, that.settings);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = initializer.getClass().hashCode();
+        result = 31 * result + DeepEquals.deepHashCode(settings);
+        return result;
+    }
+}

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

@@ -0,0 +1,82 @@
+/*
+ * 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 javax.sql.DataSource;
+import java.io.PrintWriter;
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.SQLFeatureNotSupportedException;
+import java.util.logging.Logger;
+
+class ThreadLocalDataSource implements DataSource {
+
+    private final ThreadLocal<DataSource> current = new ThreadLocal<>();
+
+    void set(DataSource dataSource) {
+        current.set(dataSource);
+    }
+
+    void clear() {
+        current.remove();
+    }
+
+    @Override
+    public Connection getConnection() throws SQLException {
+        return current.get().getConnection();
+    }
+
+    @Override
+    public Connection getConnection(String username, String password) throws SQLException {
+        return current.get().getConnection(username, password);
+    }
+
+    @Override
+    public PrintWriter getLogWriter() throws SQLException {
+        return current.get().getLogWriter();
+    }
+
+    @Override
+    public void setLogWriter(PrintWriter out) throws SQLException {
+        current.get().setLogWriter(out);
+    }
+
+    @Override
+    public void setLoginTimeout(int seconds) throws SQLException {
+        current.get().setLoginTimeout(seconds);
+    }
+
+    @Override
+    public int getLoginTimeout() throws SQLException {
+        return current.get().getLoginTimeout();
+    }
+
+    @Override
+    public Logger getParentLogger() throws SQLFeatureNotSupportedException {
+        return current.get().getParentLogger();
+    }
+
+    @Override
+    public <T> T unwrap(Class<T> iface) throws SQLException {
+        return current.get().unwrap(iface);
+    }
+
+    @Override
+    public boolean isWrapperFor(Class<?> iface) throws SQLException {
+        return current.get().isWrapperFor(iface);
+    }
+}

embedded-database-spring-test/src/test/java/io/zonky/test/db/SpringBootSqlInitIntegrationTest.java

@@ -0,0 +1,93 @@
+/*
+ * 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;
+
+import io.zonky.test.category.FlywayTestSuite;
+import io.zonky.test.support.ConditionalTestRule;
+import io.zonky.test.support.TestAssumptions;
+import org.junit.ClassRule;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.autoconfigure.jdbc.JdbcTest;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.jdbc.core.JdbcTemplate;
+import org.springframework.test.context.TestPropertySource;
+import org.springframework.test.context.junit4.SpringRunner;
+
+import javax.sql.DataSource;
+import java.util.List;
+import java.util.Map;
+
+import static io.zonky.test.db.AutoConfigureEmbeddedDatabase.DatabaseType.POSTGRES;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.entry;
+
+@RunWith(SpringRunner.class)
+@Category(FlywayTestSuite.class)
+@AutoConfigureEmbeddedDatabase(type = POSTGRES)
+@JdbcTest
+@TestPropertySource(properties = {
+        "spring.sql.init.mode=always",
+        "spring.sql.init.schema-locations=" +
+                "classpath:/db/schema/init-schema.sql," +
+                "classpath:/db/migration/V0001_1__create_person_table.sql," +
+                "classpath:/db/migration/V0002_1__rename_surname_column.sql",
+
+        "flyway.enabled=false",
+        "spring.flyway.enabled=false",
+        "liquibase.enabled=false",
+        "spring.liquibase.enabled=false"
+})
+public class SpringBootSqlInitIntegrationTest {
+
+    @ClassRule
+    public static ConditionalTestRule conditionalTestRule = new ConditionalTestRule(() -> {
+        TestAssumptions.assumeSpringBootSupportsJdbcTestAnnotation();
+        TestAssumptions.assumeSpringBootSupportsSqlInit();
+    });
+
+    private static final String SQL_SELECT_PERSONS = "select * from test.person";
+
+    @Configuration
+    static class Config {
+
+        @Bean
+        public JdbcTemplate jdbcTemplate(DataSource dataSource) {
+            return new JdbcTemplate(dataSource);
+        }
+    }
+
+    @Autowired
+    private JdbcTemplate jdbcTemplate;
+
+    @Test
+    public void testSqlInitScripts() {
+        assertThat(jdbcTemplate).isNotNull();
+
+        List<Map<String, Object>> persons = jdbcTemplate.queryForList(SQL_SELECT_PERSONS);
+        assertThat(persons).isNotNull().hasSize(1);
+
+        Map<String, Object> person = persons.get(0);
+        assertThat(person).containsExactly(
+                entry("id", 1L),
+                entry("first_name", "Dave"),
+                entry("last_name", "Syer"));
+    }
+}