make MyBatisCursorItemReader & MyBatisPagingItemReader support lazy/dynamic query parameter

Author: happysca

src/main/java/org/mybatis/spring/batch/MyBatisCursorItemReader.java

@@ -21,6 +21,7 @@
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.function.Supplier;
 
 import org.apache.ibatis.cursor.Cursor;
 import org.apache.ibatis.session.ExecutorType;
@@ -41,6 +42,7 @@ public class MyBatisCursorItemReader<T> extends AbstractItemCountingItemStreamIt
   private SqlSession sqlSession;
 
   private Map<String, Object> parameterValues;
+  private Supplier<Map<String, Object>> parameterSupplier;
 
   private Cursor<T> cursor;
   private Iterator<T> cursorIterator;
@@ -65,6 +67,11 @@ protected void doOpen() throws Exception {
       parameters.putAll(parameterValues);
     }
 
+    Map<String, Object> dynamicParameters;
+    if (parameterSupplier != null && (dynamicParameters = parameterSupplier.get()) != null) {
+      parameters.putAll(dynamicParameters);
+    }
+
     sqlSession = sqlSessionFactory.openSession(ExecutorType.SIMPLE);
     cursor = sqlSession.selectCursor(queryId, parameters);
     cursorIterator = cursor.iterator();
@@ -121,4 +128,14 @@ public void setQueryId(String queryId) {
   public void setParameterValues(Map<String, Object> parameterValues) {
     this.parameterValues = parameterValues;
   }
+
+  /**
+   * The parameter supplier used to get parameter values for the query execution.
+   *
+   * @param parameterSupplier
+   *          the supplier used to get values keyed by the parameter named used in the query string.
+   */
+  public void setParameterSupplier(Supplier<Map<String, Object>> parameterSupplier) {
+    this.parameterSupplier = parameterSupplier;
+  }
 }

src/main/java/org/mybatis/spring/batch/MyBatisPagingItemReader.java

@@ -21,6 +21,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.function.Supplier;
 
 import org.apache.ibatis.session.ExecutorType;
 import org.apache.ibatis.session.SqlSession;
@@ -47,6 +48,8 @@ public class MyBatisPagingItemReader<T> extends AbstractPagingItemReader<T> {
 
   private Map<String, Object> parameterValues;
 
+  private Supplier<Map<String, Object>> parameterSupplier;
+
   public MyBatisPagingItemReader() {
     setName(getShortName(MyBatisPagingItemReader.class));
   }
@@ -81,6 +84,16 @@ public void setParameterValues(Map<String, Object> parameterValues) {
     this.parameterValues = parameterValues;
   }
 
+  /**
+   * The parameter supplier used to get parameter values for the query execution.
+   *
+   * @param parameterSupplier
+   *          the supplier used to get values keyed by the parameter named used in the query string.
+   */
+  public void setParameterSupplier(Supplier<Map<String, Object>> parameterSupplier) {
+    this.parameterSupplier = parameterSupplier;
+  }
+
   /**
    * Check mandatory properties.
    *
@@ -102,6 +115,10 @@ protected void doReadPage() {
     if (parameterValues != null) {
       parameters.putAll(parameterValues);
     }
+    Map<String, Object> dynamicParameters;
+    if (parameterSupplier != null && (dynamicParameters = parameterSupplier.get()) != null) {
+      parameters.putAll(dynamicParameters);
+    }
     parameters.put("_page", getPage());
     parameters.put("_pagesize", getPageSize());
     parameters.put("_skiprows", getPage() * getPageSize());

src/main/java/org/mybatis/spring/batch/builder/MyBatisCursorItemReaderBuilder.java

@@ -17,6 +17,7 @@
 
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Supplier;
 
 import org.apache.ibatis.session.SqlSessionFactory;
 import org.mybatis.spring.batch.MyBatisCursorItemReader;
@@ -35,6 +36,7 @@ public class MyBatisCursorItemReaderBuilder<T> {
   private SqlSessionFactory sqlSessionFactory;
   private String queryId;
   private Map<String, Object> parameterValues;
+  private Supplier<Map<String, Object>> parameterSupplier;
   private Boolean saveState;
   private Integer maxItemCount;
 
@@ -83,6 +85,21 @@ public MyBatisCursorItemReaderBuilder<T> parameterValues(Map<String, Object> par
     return this;
   }
 
+  /**
+   * Set the parameter supplier to be used to get parameters for the query execution.
+   *
+   * @param parameterSupplier
+   *          the parameter supplier to be used to get parameters for the query execution
+   *
+   * @return this instance for method chaining
+   *
+   * @see MyBatisCursorItemReader#setParameterSupplier(Supplier)
+   */
+  public MyBatisCursorItemReaderBuilder<T> parameterSupplier(Supplier<Map<String, Object>> parameterSupplier) {
+    this.parameterSupplier = parameterSupplier;
+    return this;
+  }
+
   /**
    * Configure if the state of the {@link org.springframework.batch.item.ItemStreamSupport} should be persisted within
    * the {@link org.springframework.batch.item.ExecutionContext} for restart purposes.
@@ -124,6 +141,7 @@ public MyBatisCursorItemReader<T> build() {
     reader.setSqlSessionFactory(this.sqlSessionFactory);
     reader.setQueryId(this.queryId);
     reader.setParameterValues(this.parameterValues);
+    reader.setParameterSupplier(this.parameterSupplier);
     Optional.ofNullable(this.saveState).ifPresent(reader::setSaveState);
     Optional.ofNullable(this.maxItemCount).ifPresent(reader::setMaxItemCount);
     return reader;

src/main/java/org/mybatis/spring/batch/builder/MyBatisPagingItemReaderBuilder.java

@@ -17,6 +17,7 @@
 
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Supplier;
 
 import org.apache.ibatis.session.SqlSessionFactory;
 import org.mybatis.spring.batch.MyBatisPagingItemReader;
@@ -35,6 +36,7 @@ public class MyBatisPagingItemReaderBuilder<T> {
   private SqlSessionFactory sqlSessionFactory;
   private String queryId;
   private Map<String, Object> parameterValues;
+  private Supplier<Map<String, Object>> parameterSupplier;
   private Integer pageSize;
   private Boolean saveState;
   private Integer maxItemCount;
@@ -84,6 +86,21 @@ public MyBatisPagingItemReaderBuilder<T> parameterValues(Map<String, Object> par
     return this;
   }
 
+  /**
+   * Set the parameter supplier to be used to get parameters for the query execution.
+   *
+   * @param parameterSupplier
+   *          the parameter supplier to be used to get parameters for the query execution
+   *
+   * @return this instance for method chaining
+   *
+   * @see MyBatisPagingItemReader#setParameterSupplier(Supplier)
+   */
+  public MyBatisPagingItemReaderBuilder<T> parameterSupplier(Supplier<Map<String, Object>> parameterSupplier) {
+    this.parameterSupplier = parameterSupplier;
+    return this;
+  }
+
   /**
    * The number of records to request per page/query. Defaults to 10. Must be greater than zero.
    *
@@ -140,6 +157,7 @@ public MyBatisPagingItemReader<T> build() {
     reader.setSqlSessionFactory(this.sqlSessionFactory);
     reader.setQueryId(this.queryId);
     reader.setParameterValues(this.parameterValues);
+    reader.setParameterSupplier(this.parameterSupplier);
     Optional.ofNullable(this.pageSize).ifPresent(reader::setPageSize);
     Optional.ofNullable(this.saveState).ifPresent(reader::setSaveState);
     Optional.ofNullable(this.maxItemCount).ifPresent(reader::setMaxItemCount);

src/test/java/org/mybatis/spring/batch/builder/MyBatisCursorItemReaderBuilderTest.java

@@ -17,7 +17,9 @@
 
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import org.apache.ibatis.cursor.Cursor;
 import org.apache.ibatis.session.ExecutorType;
@@ -56,7 +58,10 @@ void setUp() {
 
     Mockito.when(this.sqlSessionFactory.openSession(ExecutorType.SIMPLE)).thenReturn(this.sqlSession);
     Mockito.when(this.cursor.iterator()).thenReturn(getFoos().iterator());
-    Mockito.when(this.sqlSession.selectCursor("selectFoo", Collections.singletonMap("id", 1))).thenReturn(this.cursor);
+    Map<String, Object> parameters = new HashMap<>();
+    parameters.put("id", 1);
+    parameters.put("name", "Doe");
+    Mockito.when(this.sqlSession.selectCursor("selectFoo", parameters)).thenReturn(this.cursor);
   }
 
   @Test
@@ -67,6 +72,7 @@ void testConfiguration() throws Exception {
             .sqlSessionFactory(this.sqlSessionFactory)
             .queryId("selectFoo")
             .parameterValues(Collections.singletonMap("id", 1))
+            .parameterSupplier(() -> Collections.singletonMap("name", "Doe"))
             .build();
     // @formatter:on
     itemReader.afterPropertiesSet();
@@ -93,6 +99,7 @@ void testConfigurationSaveStateIsFalse() throws Exception {
             .sqlSessionFactory(this.sqlSessionFactory)
             .queryId("selectFoo")
             .parameterValues(Collections.singletonMap("id", 1))
+            .parameterSupplier(() -> Collections.singletonMap("name", "Doe"))
             .saveState(false)
             .build();
     // @formatter:on
@@ -118,6 +125,7 @@ void testConfigurationMaxItemCount() throws Exception {
             .sqlSessionFactory(this.sqlSessionFactory)
             .queryId("selectFoo")
             .parameterValues(Collections.singletonMap("id", 1))
+            .parameterSupplier(() -> Collections.singletonMap("name", "Doe"))
             .maxItemCount(2)
             .build();
     // @formatter:on

src/test/java/org/mybatis/spring/batch/builder/MyBatisPagingItemReaderBuilderTest.java

@@ -67,6 +67,7 @@ void setUp() {
     Mockito.when(this.sqlSessionFactory.openSession(ExecutorType.BATCH)).thenReturn(this.sqlSession);
     Map<String, Object> parameters = new HashMap<>();
     parameters.put("id", 1);
+    parameters.put("name", "Doe");
     parameters.put("_page", 0);
     parameters.put("_pagesize", 10);
     parameters.put("_skiprows", 0);
@@ -80,6 +81,7 @@ void testConfiguration() throws Exception {
             .sqlSessionFactory(this.sqlSessionFactory)
             .queryId("selectFoo")
             .parameterValues(Collections.singletonMap("id", 1))
+            .parameterSupplier(() -> Collections.singletonMap("name", "Doe"))
             .build();
     // @formatter:on
     itemReader.afterPropertiesSet();
@@ -105,6 +107,7 @@ void testConfigurationSaveStateIsFalse() throws Exception {
             .sqlSessionFactory(this.sqlSessionFactory)
             .queryId("selectFoo")
             .parameterValues(Collections.singletonMap("id", 1))
+            .parameterSupplier(() -> Collections.singletonMap("name", "Doe"))
             .saveState(false)
             .build();
     // @formatter:on
@@ -128,6 +131,7 @@ void testConfigurationMaxItemCount() throws Exception {
             .sqlSessionFactory(this.sqlSessionFactory)
             .queryId("selectFoo")
             .parameterValues(Collections.singletonMap("id", 1))
+            .parameterSupplier(() -> Collections.singletonMap("name", "Doe"))
             .maxItemCount(2)
             .build();
     // @formatter:on
@@ -152,6 +156,7 @@ void testConfigurationPageSize() throws Exception {
             .sqlSessionFactory(this.sqlSessionFactory)
             .queryId("selectFoo")
             .parameterValues(Collections.singletonMap("id", 1))
+            .parameterSupplier(() -> Collections.singletonMap("name", "Doe"))
             .pageSize(2)
             .build();
     // @formatter:on
@@ -160,6 +165,7 @@ void testConfigurationPageSize() throws Exception {
     Map<String, Object> parameters = new HashMap<>();
     parameters.put("id", 1);
     parameters.put("_page", 0);
+    parameters.put("name", "Doe");
     parameters.put("_pagesize", 2);
     parameters.put("_skiprows", 0);
     Mockito.when(this.sqlSession.selectList("selectFoo", parameters)).thenReturn(getFoos());