Allow custom mapper methods to initiate dynamic injection

Author: qmdx

src/main/java/org/apache/ibatis/builder/SqlInjector.java

@@ -0,0 +1,32 @@
+/*
+ *    Copyright 2009-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
+ *
+ *       https://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.apache.ibatis.builder;
+
+/**
+ * SQL Injector
+ *
+ * @author hubin
+ */
+public interface SqlInjector {
+
+  /**
+   * inject custom XML SQL scripts
+   *
+   * @param mapperBuilderAssistant {@link MapperBuilderAssistant}
+   * @param type                   Mapper type
+   */
+  void inject(MapperBuilderAssistant mapperBuilderAssistant, Class<?> type);
+}

src/main/java/org/apache/ibatis/builder/annotation/MapperAnnotationBuilder.java

@@ -63,6 +63,7 @@
 import org.apache.ibatis.builder.CacheRefResolver;
 import org.apache.ibatis.builder.IncompleteElementException;
 import org.apache.ibatis.builder.MapperBuilderAssistant;
+import org.apache.ibatis.builder.SqlInjector;
 import org.apache.ibatis.builder.xml.XMLMapperBuilder;
 import org.apache.ibatis.cursor.Cursor;
 import org.apache.ibatis.executor.keygen.Jdbc3KeyGenerator;
@@ -106,7 +107,7 @@ public class MapperAnnotationBuilder {
 
   public MapperAnnotationBuilder(Configuration configuration, Class<?> type) {
     String resource = type.getName().replace('.', '/') + ".java (best guess)";
-    this.assistant = new MapperBuilderAssistant(configuration, resource);
+    this.assistant = configuration.getMapperBuilderAssistant(resource);
     this.configuration = configuration;
     this.type = type;
   }
@@ -133,6 +134,12 @@ public void parse() {
           configuration.addIncompleteMethod(new MethodResolver(this, method));
         }
       }
+
+      // inject custom XML SQL scripts
+      SqlInjector sqlInjector = configuration.getSqlInjector();
+      if (null != sqlInjector) {
+        sqlInjector.inject(assistant, type);
+      }
     }
     configuration.parsePendingMethods(false);
   }

src/main/java/org/apache/ibatis/session/Configuration.java

@@ -32,7 +32,9 @@
 import org.apache.ibatis.binding.MapperRegistry;
 import org.apache.ibatis.builder.CacheRefResolver;
 import org.apache.ibatis.builder.IncompleteElementException;
+import org.apache.ibatis.builder.MapperBuilderAssistant;
 import org.apache.ibatis.builder.ResultMapResolver;
+import org.apache.ibatis.builder.SqlInjector;
 import org.apache.ibatis.builder.annotation.MethodResolver;
 import org.apache.ibatis.builder.xml.XMLStatementBuilder;
 import org.apache.ibatis.cache.Cache;
@@ -276,6 +278,23 @@ public void setDefaultSqlProviderType(Class<?> defaultSqlProviderType) {
     this.defaultSqlProviderType = defaultSqlProviderType;
   }
 
+  /**
+   * SqlInjector to provide subclass rewriting.
+   */
+  public SqlInjector getSqlInjector() {
+    return null;
+  }
+
+  /**
+   * Building MapperBuilderAssistant Cooperate with SQL Injector
+   *
+   * @param resource the resource path
+   * @return {@link MapperBuilderAssistant}
+   */
+  public MapperBuilderAssistant getMapperBuilderAssistant(String resource) {
+    return new MapperBuilderAssistant(this, resource);
+  }
+
   public boolean isCallSettersOnNulls() {
     return callSettersOnNulls;
   }

src/test/java/org/apache/ibatis/builder/SqlInjectorTest.java

@@ -0,0 +1,44 @@
+package org.apache.ibatis.builder;
+
+import org.apache.ibatis.builder.annotation.MapperAnnotationBuilder;
+import org.apache.ibatis.mapping.SqlCommandType;
+import org.apache.ibatis.mapping.SqlSource;
+import org.apache.ibatis.mapping.StatementType;
+import org.apache.ibatis.scripting.LanguageDriver;
+import org.apache.ibatis.session.Configuration;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class SqlInjectorTest {
+  private final String customMapperMethod = "selectCount";
+  private static Configuration configuration;
+  private static MapperAnnotationBuilder mapperAnnotationBuilder;
+
+  @BeforeEach
+  void setUp() {
+    // custom configuration
+    configuration = new Configuration() {
+
+      @Override
+      public SqlInjector getSqlInjector() {
+
+        // custom sql injector
+        return (mapperBuilderAssistant, type) -> {
+          LanguageDriver languageDriver = configuration.getDefaultScriptingLanguageInstance();
+          SqlSource sqlSource = languageDriver.createSqlSource(configuration, "select count(*) from user", null);
+          mapperBuilderAssistant.addMappedStatement(customMapperMethod, sqlSource, StatementType.PREPARED, SqlCommandType.SELECT, null, null,
+            null, null, null, Long.class, null, false, true, false,
+            null, null, null, configuration.getDatabaseId(), languageDriver, null);
+        };
+      }
+    };
+    mapperAnnotationBuilder = new MapperAnnotationBuilder(configuration, AnnotationMapperBuilderTest.Mapper.class);
+  }
+
+  @Test
+  void testSqlInjector() {
+    mapperAnnotationBuilder.parse();
+    Assertions.assertNotNull(configuration.getMappedStatement(customMapperMethod));
+  }
+}