#331 Refactor Spring SQL init support to use AOP proxying

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;
@@ -232,7 +232,7 @@ public DatabaseProviderFactory defaultDatabaseProviderFactory(AutowireCapableBea
         int concurrency = environment.getProperty("zonky.test.database.prefetching.concurrency", int.class, 3);
         int pipelineCacheSize = environment.getProperty("zonky.test.database.prefetching.pipeline-cache-size", int.class, 5);
         int maxPreparedTemplates = environment.getProperty("zonky.test.database.prefetching.max-prepared-templates", int.class, 10);
-        int maxPreparedDatabases = (maxPreparedTemplates * 2/3 * 2) + pipelineCacheSize;
+        int maxPreparedDatabases = (maxPreparedTemplates * 2 / 3 * 2) + pipelineCacheSize;
 
         return new DatabaseProviderFactory(beanFactory)
                 .customizeTemplating(builder -> builder
@@ -293,9 +293,9 @@ 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);
+    @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/SpringScriptDatabaseExtension.java

@@ -0,0 +1,114 @@
+/*
+ * 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.jdbc.init.DataSourceScriptDatabaseInitializer;
+import org.springframework.core.Ordered;
+import org.springframework.core.env.Environment;
+
+import javax.sql.DataSource;
+
+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 DataSourceScriptDatabaseInitializer) {
+            DataSourceScriptDatabaseInitializer initializer = (DataSourceScriptDatabaseInitializer) 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, context));
+                    return bean;
+                } else {
+                    ProxyFactory proxyFactory = new ProxyFactory(bean);
+                    proxyFactory.addAdvisor(createAdvisor(initializer, context));
+                    proxyFactory.setProxyTargetClass(true);
+                    return proxyFactory.getProxy();
+                }
+            }
+        }
+
+        return bean;
+    }
+
+    @Override
+    public Object postProcessAfterInitialization(Object bean, String beanName) {
+        return bean;
+    }
+
+    protected Advisor createAdvisor(DataSourceScriptDatabaseInitializer initializer, DatabaseContext context) {
+        Advice advice = new SpringScriptDatabaseExtensionInterceptor(initializer, context);
+        NameMatchMethodPointcutAdvisor advisor = new NameMatchMethodPointcutAdvisor(advice);
+        advisor.setMappedNames("afterPropertiesSet", "initializeDatabase");
+        return advisor;
+    }
+
+    protected static class SpringScriptDatabaseExtensionInterceptor implements MethodInterceptor {
+
+        private final SpringScriptDatabasePreparer preparer;
+        private final DatabaseContext context;
+
+        protected SpringScriptDatabaseExtensionInterceptor(DataSourceScriptDatabaseInitializer initializer,
+                                                           DatabaseContext context) {
+            this.preparer = new SpringScriptDatabasePreparer(initializer);
+            this.context = context;
+        }
+
+        @Override
+        public Object invoke(MethodInvocation invocation) throws Throwable {
+            switch (invocation.getMethod().getName()) {
+                case "afterPropertiesSet":
+                    context.apply(preparer);
+                    return null;
+                case "initializeDatabase":
+                    context.apply(preparer);
+                    return true;
+                default:
+                    return invocation.proceed();
+            }
+        }
+    }
+}

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

@@ -20,7 +20,6 @@
 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;
@@ -33,19 +32,19 @@
 
 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 static final ThreadLocalDataSource threadLocalDataSource = new ThreadLocalDataSource();
+
     private final DataSourceScriptDatabaseInitializer initializer;
-    private final ThreadLocalDataSource threadLocalDataSource;
 
-    public DataSourceScriptDatabasePreparer(DataSourceScriptDatabaseInitializer initializer) {
+    public SpringScriptDatabasePreparer(DataSourceScriptDatabaseInitializer 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(),