From 06e74bfde7417c1b77307b0a65fc963e1cf41b9c Mon Sep 17 00:00:00 2001
From: Kazuki Shimizu <kazuki43zoo@gmail.com>
Date: Sun, 9 Jan 2022 06:48:47 +0900
Subject: [PATCH] Add SqlSessionFactoryBeanCustomizer See gh-617

---
 .../MybatisAutoConfiguration.java             | 16 +++++++-
 .../SqlSessionFactoryBeanCustomizer.java      | 37 +++++++++++++++++++
 .../src/site/markdown/index.md                | 25 ++++++++++++-
 .../MybatisAutoConfigurationTest.java         | 25 +++++++++++++
 4 files changed, 100 insertions(+), 3 deletions(-)
 create mode 100644 mybatis-spring-boot-autoconfigure/src/main/java/org/mybatis/spring/boot/autoconfigure/SqlSessionFactoryBeanCustomizer.java

diff --git a/mybatis-spring-boot-autoconfigure/src/main/java/org/mybatis/spring/boot/autoconfigure/MybatisAutoConfiguration.java b/mybatis-spring-boot-autoconfigure/src/main/java/org/mybatis/spring/boot/autoconfigure/MybatisAutoConfiguration.java
index 60ff53629..7dca1a74c 100644
--- a/mybatis-spring-boot-autoconfigure/src/main/java/org/mybatis/spring/boot/autoconfigure/MybatisAutoConfiguration.java
+++ b/mybatis-spring-boot-autoconfigure/src/main/java/org/mybatis/spring/boot/autoconfigure/MybatisAutoConfiguration.java
@@ -105,10 +105,13 @@ public class MybatisAutoConfiguration implements InitializingBean {
 
   private final List<ConfigurationCustomizer> configurationCustomizers;
 
+  private final List<SqlSessionFactoryBeanCustomizer> sqlSessionFactoryBeanCustomizers;
+
   public MybatisAutoConfiguration(MybatisProperties properties, ObjectProvider<Interceptor[]> interceptorsProvider,
       ObjectProvider<TypeHandler[]> typeHandlersProvider, ObjectProvider<LanguageDriver[]> languageDriversProvider,
       ResourceLoader resourceLoader, ObjectProvider<DatabaseIdProvider> databaseIdProvider,
-      ObjectProvider<List<ConfigurationCustomizer>> configurationCustomizersProvider) {
+      ObjectProvider<List<ConfigurationCustomizer>> configurationCustomizersProvider,
+      ObjectProvider<List<SqlSessionFactoryBeanCustomizer>> sqlSessionFactoryBeanCustomizers) {
     this.properties = properties;
     this.interceptors = interceptorsProvider.getIfAvailable();
     this.typeHandlers = typeHandlersProvider.getIfAvailable();
@@ -116,6 +119,7 @@ public MybatisAutoConfiguration(MybatisProperties properties, ObjectProvider<Int
     this.resourceLoader = resourceLoader;
     this.databaseIdProvider = databaseIdProvider.getIfAvailable();
     this.configurationCustomizers = configurationCustomizersProvider.getIfAvailable();
+    this.sqlSessionFactoryBeanCustomizers = sqlSessionFactoryBeanCustomizers.getIfAvailable();
   }
 
   @Override
@@ -180,7 +184,7 @@ public SqlSessionFactory sqlSessionFactory(DataSource dataSource) throws Excepti
       // Need to mybatis-spring 2.0.2+
       factory.setDefaultScriptingLanguageDriver(defaultLanguageDriver);
     }
-
+    applySqlSessionFactoryBeanCustomizers(factory);
     return factory.getObject();
   }
 
@@ -197,6 +201,14 @@ private void applyConfiguration(SqlSessionFactoryBean factory) {
     factory.setConfiguration(configuration);
   }
 
+  private void applySqlSessionFactoryBeanCustomizers(SqlSessionFactoryBean factory) {
+    if (!CollectionUtils.isEmpty(this.sqlSessionFactoryBeanCustomizers)) {
+      for (SqlSessionFactoryBeanCustomizer customizer : this.sqlSessionFactoryBeanCustomizers) {
+        customizer.customize(factory);
+      }
+    }
+  }
+
   @Bean
   @ConditionalOnMissingBean
   public SqlSessionTemplate sqlSessionTemplate(SqlSessionFactory sqlSessionFactory) {
diff --git a/mybatis-spring-boot-autoconfigure/src/main/java/org/mybatis/spring/boot/autoconfigure/SqlSessionFactoryBeanCustomizer.java b/mybatis-spring-boot-autoconfigure/src/main/java/org/mybatis/spring/boot/autoconfigure/SqlSessionFactoryBeanCustomizer.java
new file mode 100644
index 000000000..8cb7bf519
--- /dev/null
+++ b/mybatis-spring-boot-autoconfigure/src/main/java/org/mybatis/spring/boot/autoconfigure/SqlSessionFactoryBeanCustomizer.java
@@ -0,0 +1,37 @@
+/*
+ *    Copyright 2015-2022 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 org.mybatis.spring.boot.autoconfigure;
+
+import org.mybatis.spring.SqlSessionFactoryBean;
+
+/**
+ * Callback interface that can be customized a {@link SqlSessionFactoryBean} object generated on auto-configuration.
+ *
+ * @author Kazuki Shimizu
+ * @since 2.2.2
+ */
+@FunctionalInterface
+public interface SqlSessionFactoryBeanCustomizer {
+
+  /**
+   * Customize the given a {@link SqlSessionFactoryBean} object.
+   *
+   * @param factoryBean
+   *          the factory bean object to customize
+   */
+  void customize(SqlSessionFactoryBean factoryBean);
+
+}
diff --git a/mybatis-spring-boot-autoconfigure/src/site/markdown/index.md b/mybatis-spring-boot-autoconfigure/src/site/markdown/index.md
index 47bdad904..85ad03a6b 100644
--- a/mybatis-spring-boot-autoconfigure/src/site/markdown/index.md
+++ b/mybatis-spring-boot-autoconfigure/src/site/markdown/index.md
@@ -183,7 +183,7 @@ mybatis:
 ## Using a ConfigurationCustomizer
 
 The MyBatis-Spring-Boot-Starter provide opportunity to customize a MyBatis configuration generated by auto-configuration using Java Config.
-The MyBatis-Spring-Boot-Starter will search beans that implements the `ConfigurationCustomizer` interface by automatically,
+The MyBatis-Spring-Boot-Starter will search beans that implement the `ConfigurationCustomizer` interface by automatically,
 and call a method that customize a MyBatis configuration. (Available since 1.2.1 or above)
 
 For example:
@@ -203,6 +203,29 @@ public class MyBatisConfig {
 }
 ```
 
+## Using a SqlSessionFactoryBeanCustomizer
+
+The MyBatis-Spring-Boot-Starter provide opportunity to customize a `SqlSessionFactoryBean` generated by auto-configuration using Java Config.
+The MyBatis-Spring-Boot-Starter will search beans that implement the `SqlSessionFactoryBeanCustomizer` interface by automatically,
+and call a method that customize a `SqlSessionFactoryBean`. (Available since 2.2.2 or above)
+
+For example:
+
+```java
+@Configuration
+public class MyBatisConfig {
+  @Bean
+  SqlSessionFactoryBeanCustomizer sqlSessionFactoryBeanCustomizer() {
+    return new SqlSessionFactoryBeanCustomizer() {
+      @Override
+      public void customize(SqlSessionFactoryBean factoryBean) {
+        // customize ...
+      }
+    };
+  }
+}
+```
+
 ## Using the SpringBootVFS
 
 The MyBatis-Spring-Boot-Starter provides the `SpringBootVFS` as an implementation class of `VFS`.
diff --git a/mybatis-spring-boot-autoconfigure/src/test/java/org/mybatis/spring/boot/autoconfigure/MybatisAutoConfigurationTest.java b/mybatis-spring-boot-autoconfigure/src/test/java/org/mybatis/spring/boot/autoconfigure/MybatisAutoConfigurationTest.java
index ab40bd0f7..f12e835b3 100644
--- a/mybatis-spring-boot-autoconfigure/src/test/java/org/mybatis/spring/boot/autoconfigure/MybatisAutoConfigurationTest.java
+++ b/mybatis-spring-boot-autoconfigure/src/test/java/org/mybatis/spring/boot/autoconfigure/MybatisAutoConfigurationTest.java
@@ -563,6 +563,17 @@ void testWithMyBatisConfigurationCustomizer() {
     assertThat(sqlSessionFactory.getConfiguration().getCache("test")).isNotNull();
   }
 
+  @Test
+  void testWithSqlSessionFactoryBeanCustomizer() {
+    this.context.register(EmbeddedDataSourceConfiguration.class, MybatisAutoConfiguration.class,
+        SqlSessionFactoryBeanCustomizerConfiguration.class);
+    this.context.refresh();
+    SqlSessionFactory sqlSessionFactory = this.context.getBean(SqlSessionFactory.class);
+    assertThat(sqlSessionFactory.getConfiguration().getTypeHandlerRegistry().getTypeHandler(BigInteger.class))
+        .isInstanceOf(DummyTypeHandler.class);
+    assertThat(sqlSessionFactory.getConfiguration().getCache("test")).isNotNull();
+  }
+
   @Test
   void testConfigFileAndConfigurationWithTogether() {
     TestPropertyValues
@@ -952,6 +963,20 @@ ConfigurationCustomizer cacheConfigurationCustomizer() {
     }
   }
 
+  @Configuration
+  @EnableAutoConfiguration
+  static class SqlSessionFactoryBeanCustomizerConfiguration {
+    @Bean
+    SqlSessionFactoryBeanCustomizer typeHandlerSqlSessionFactoryBeanCustomizer() {
+      return factoryBean -> factoryBean.setTypeHandlers(new DummyTypeHandler());
+    }
+
+    @Bean
+    SqlSessionFactoryBeanCustomizer cacheSqlSessionFactoryBeanCustomizer() {
+      return factoryBean -> factoryBean.setCache(new PerpetualCache("test"));
+    }
+  }
+
   @Configuration
   static class SingleLanguageDriverConfiguration {
     @Bean