Update support for Spring Boot's DataSourceScriptDatabaseInitializer

Author: tomix26

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

@@ -18,9 +18,9 @@
 
 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.init.SpringScriptDatabaseExtension;
 import io.zonky.test.db.liquibase.LiquibaseDatabaseExtension;
 import io.zonky.test.db.liquibase.LiquibasePropertiesPostProcessor;
 import io.zonky.test.db.provider.DatabaseProvider;
@@ -292,10 +292,10 @@ public BeanPostProcessor 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);
+    @ConditionalOnClass(name = "org.springframework.boot.sql.init.AbstractScriptDatabaseInitializer")
+    @ConditionalOnMissingBean(name = "springScriptDatabaseExtension")
+    public SpringScriptDatabaseExtension springScriptDatabaseExtension(Environment environment) {
+        return new SpringScriptDatabaseExtension(environment);
     }
 
     @Bean

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

@@ -0,0 +1,115 @@
+/*
+ * 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.aopalliance.aop.Advice;
+import org.aopalliance.intercept.MethodInterceptor;
+import org.aopalliance.intercept.MethodInvocation;
+import org.springframework.aop.Advisor;
+import org.springframework.aop.framework.Advised;
+import org.springframework.aop.framework.AopInfrastructureBean;
+import org.springframework.aop.framework.ProxyFactory;
+import org.springframework.aop.support.NameMatchMethodPointcutAdvisor;
+import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.boot.sql.init.AbstractScriptDatabaseInitializer;
+import org.springframework.core.Ordered;
+import org.springframework.core.env.Environment;
+
+import javax.sql.DataSource;
+
+import static com.google.common.base.Preconditions.checkState;
+
+public class SpringScriptDatabaseExtension implements BeanPostProcessor, Ordered {
+
+    private final boolean enabled;
+
+    public SpringScriptDatabaseExtension(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 AopInfrastructureBean) {
+            return bean;
+        }
+
+        if (bean instanceof AbstractScriptDatabaseInitializer) {
+            AbstractScriptDatabaseInitializer initializer = (AbstractScriptDatabaseInitializer) bean;
+            DataSource dataSource = ReflectionUtils.getField(initializer, "dataSource");
+            DatabaseContext context = AopProxyUtils.getDatabaseContext(dataSource);
+
+            if (context != null) {
+                if (bean instanceof Advised && !((Advised) bean).isFrozen()) {
+                    ((Advised) bean).addAdvisor(0, createAdvisor(initializer));
+                    return bean;
+                } else {
+                    ProxyFactory proxyFactory = new ProxyFactory(bean);
+                    proxyFactory.addAdvisor(createAdvisor(initializer));
+                    proxyFactory.setProxyTargetClass(true);
+                    return proxyFactory.getProxy();
+                }
+            }
+        }
+
+        return bean;
+    }
+
+    @Override
+    public Object postProcessAfterInitialization(Object bean, String beanName) {
+        return bean;
+    }
+
+    protected Advisor createAdvisor(AbstractScriptDatabaseInitializer initializer) {
+        Advice advice = new SpringScriptDatabaseExtensionInterceptor(initializer);
+        NameMatchMethodPointcutAdvisor advisor = new NameMatchMethodPointcutAdvisor(advice);
+        advisor.setMappedNames("afterPropertiesSet", "initializeDatabase");
+        return advisor;
+    }
+
+    protected static class SpringScriptDatabaseExtensionInterceptor implements MethodInterceptor {
+
+        private final AbstractScriptDatabaseInitializer initializer;
+
+        protected SpringScriptDatabaseExtensionInterceptor(AbstractScriptDatabaseInitializer initializer) {
+            this.initializer = initializer;
+        }
+
+        @Override
+        public Object invoke(MethodInvocation invocation) throws Throwable {
+            String methodName = invocation.getMethod().getName();
+            if (!"afterPropertiesSet".equals(methodName) && !"initializeDatabase".equals(methodName)) {
+                return invocation.proceed();
+            }
+
+            DataSource dataSource = ReflectionUtils.getField(initializer, "dataSource");
+            DatabaseContext context = AopProxyUtils.getDatabaseContext(dataSource);
+            checkState(context != null, "Data source context cannot be resolved");
+
+            context.apply(new SpringScriptDatabasePreparer(initializer));
+
+            return null;
+        }
+    }
+}

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

@@ -19,8 +19,7 @@
 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.AbstractScriptDatabaseInitializer;
 import org.springframework.util.ReflectionUtils.FieldFilter;
 
 import javax.sql.DataSource;
@@ -33,17 +32,17 @@
 
 import static org.springframework.util.ReflectionUtils.makeAccessible;
 
-public class DataSourceScriptDatabasePreparer implements DatabasePreparer {
+public class SpringScriptDatabasePreparer 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 AbstractScriptDatabaseInitializer initializer;
     private final ThreadLocalDataSource threadLocalDataSource;
 
-    public DataSourceScriptDatabasePreparer(DataSourceScriptDatabaseInitializer initializer) {
+    public SpringScriptDatabasePreparer(AbstractScriptDatabaseInitializer initializer) {
         this.initializer = initializer;
         this.threadLocalDataSource = new ThreadLocalDataSource();
         ReflectionUtils.setField(initializer, "dataSource", threadLocalDataSource);
@@ -68,7 +67,7 @@ public void prepare(DataSource dataSource) {
     public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
-        DataSourceScriptDatabasePreparer that = (DataSourceScriptDatabasePreparer) o;
+        SpringScriptDatabasePreparer that = (SpringScriptDatabasePreparer) o;
         if (initializer.getClass() != that.initializer.getClass()) return false;
         AtomicBoolean equal = new AtomicBoolean(true);
         org.springframework.util.ReflectionUtils.doWithFields(initializer.getClass(),

…it/DataSourceScriptDatabasePreparer.java …b/init/SpringScriptDatabasePreparer.javaembedded-database-spring-test/src/main/java/io/zonky/test/db/init/DataSourceScriptDatabasePreparer.java renamed to embedded-database-spring-test/src/main/java/io/zonky/test/db/init/SpringScriptDatabasePreparer.java embedded-database-spring-test/src/main/java/io/zonky/test/db/init/DataSourceScriptDatabasePreparer.java renamed to embedded-database-spring-test/src/main/java/io/zonky/test/db/init/SpringScriptDatabasePreparer.java

@@ -7,9 +7,9 @@
  * https://github.com/spring-projects/spring-boot
  */
 
-package org.springframework.boot.jdbc.init;
+package org.springframework.boot.sql.init;
 
-public class DataSourceScriptDatabaseInitializer {
+public class AbstractScriptDatabaseInitializer {
 
     public boolean initializeDatabase() {
         return false;