BATCH-480: Quick refactoring to how the MultipleColumnJdbcKeyCollector works. It was nevery fully fixed from when the change to ExecutionContext was made from RestartData. Also some misc. documentation fixes based on feedback.

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/support/ColumnMapExecutionContextRowMapper.java

@@ -1,95 +0,0 @@
-/**
- * 
- */
-package org.springframework.batch.item.database.support;
-
-import java.sql.PreparedStatement;
-import java.sql.SQLException;
-import java.util.ArrayList;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Map.Entry;
-
-import org.springframework.batch.item.ExecutionContext;
-import org.springframework.jdbc.core.ColumnMapRowMapper;
-import org.springframework.jdbc.core.PreparedStatementSetter;
-import org.springframework.jdbc.core.SqlParameterValue;
-import org.springframework.jdbc.core.SqlTypeValue;
-import org.springframework.jdbc.core.StatementCreatorUtils;
-import org.springframework.util.Assert;
-import org.springframework.util.ClassUtils;
-
-/**
- * </p>Extension of the {@link ColumnMapRowMapper} that converts a column map to {@link ExecutionContext} and allows
- * {@link ExecutionContext} to be converted back as a PreparedStatementSetter.  This is useful in a restart 
- * scenario, as it allows for the standard functionality of the ColumnMapRowMapper to be used to 
- * create a map representing the columns returned by a query.  It should be noted that this column ordering
- * is preserved in the map using a link list version of Map.</p>
- * 
- * 
- * @author Lucas Ward
- * @author Dave Syer
- * @see ExecutionContextRowMapper
- * @see ColumnMapRowMapper
- */
-public class ColumnMapExecutionContextRowMapper extends ColumnMapRowMapper implements ExecutionContextRowMapper {
-	
-	public static final String KEY_PREFIX = ClassUtils.getQualifiedName(ColumnMapExecutionContextRowMapper.class) + ".KEY.";
-	
-	public PreparedStatementSetter createSetter(ExecutionContext executionContext) {
-		List columns = new ArrayList();
-		int count=0;
-		while(executionContext.containsKey(KEY_PREFIX+count)) {
-			Object column = executionContext.get(KEY_PREFIX+count);
-			columns.add(column);
-			count++;
-		}
-
-		return new ArgPreparedStatementSetter(columns.toArray());
-	}
-
-	public void mapKeys(Object key, ExecutionContext executionContext) {
-		Assert.isInstanceOf(Map.class, key, "Input to create ExecutionContext must be of type Map.");
-		Map keys = (Map) key;
-		int count = 0;
-		for (Iterator it = keys.entrySet().iterator(); it.hasNext();) {
-			Entry entry = (Entry)it.next();
-			executionContext.put(KEY_PREFIX+count, entry.getValue());
-			count++;
-		}
-	}
-
-	/*
-	 * Exact duplicate of Spring class of the same name, copied because it is
-	 * package private.
-	 */
-	private static class ArgPreparedStatementSetter implements PreparedStatementSetter {
-
-		private final Object[] args;
-
-		/**
-		 * Create a new ArgPreparedStatementSetter for the given arguments.
-		 * @param args the arguments to set
-		 */
-		public ArgPreparedStatementSetter(Object[] args) {
-			this.args = args;
-		}
-
-		public void setValues(PreparedStatement ps) throws SQLException {
-			if (this.args != null) {
-				for (int i = 0; i < this.args.length; i++) {
-					Object arg = this.args[i];
-					if (arg instanceof SqlParameterValue) {
-						SqlParameterValue paramValue = (SqlParameterValue) arg;
-						StatementCreatorUtils.setParameterValue(ps, i + 1, paramValue, paramValue.getValue());
-					}
-					else {
-						StatementCreatorUtils.setParameterValue(ps, i + 1, SqlTypeValue.TYPE_UNKNOWN, arg);
-					}
-				}
-			}
-		}
-	}
-
-}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/support/ColumnMapKeyMappingPreparedStatementSetter.java

@@ -0,0 +1,42 @@
+/**
+ * 
+ */
+package org.springframework.batch.item.database.support;
+
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Set;
+import java.util.Map.Entry;
+
+import org.springframework.jdbc.core.ColumnMapRowMapper;
+import org.springframework.jdbc.core.SqlTypeValue;
+import org.springframework.jdbc.core.StatementCreatorUtils;
+import org.springframework.util.Assert;
+
+/**
+ * </p>Implementation of the {@link KeyMappingPreparedStatementSetter} interface that assumes all
+ * keys are contained within a {@link Map} with the column name as the key.  It assumes nothing 
+ * about ordering, and assumes that the order the entry set can be iterated over is the same as
+ * the PreparedStatement should be set.</p>
+ * 
+ * @author Lucas Ward
+ * @author Dave Syer
+ * @see KeyMappingPreparedStatementSetter
+ * @see ColumnMapRowMapper
+ */
+public class ColumnMapKeyMappingPreparedStatementSetter implements KeyMappingPreparedStatementSetter {
+
+	public void setValues(PreparedStatement ps, Object key) throws SQLException {
+		Assert.isInstanceOf(Map.class, key, "Input to map PreparedStatement parameters must be of type Map.");
+		Set keySet = ((Map)key).entrySet();
+		int counter = 1;
+		for(Iterator it = keySet.iterator(); it.hasNext();){
+			Entry entry = (Entry)it.next();
+			StatementCreatorUtils.setParameterValue(ps, counter, SqlTypeValue.TYPE_UNKNOWN, entry.getValue());
+			counter++;
+		}
+	}
+
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/support/KeyMappingPreparedStatementSetter.java

@@ -15,6 +15,9 @@
  */
 package org.springframework.batch.item.database.support;
 
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
+
 import org.springframework.batch.item.ExecutionContext;
 import org.springframework.batch.item.database.DrivingQueryItemReader;
 import org.springframework.batch.item.database.KeyCollector;
@@ -45,7 +48,7 @@ import org.springframework.jdbc.core.RowMapper;
  * {@link MultipleColumnJdbcKeyCollector}. In general, it's much simpler to use
  * one key as a column, and if the data itself allows this, then it should be
  * used.  However, in certain cases there is no choice and multiple columns must
- * be used. Using a {@link ExecutionContextRowMapper}, developers can create
+ * be used. Using a {@link KeyMappingPreparedStatementSetter}, developers can create
  * each unique key to suite their specific needs, and also describe how such a
  * key would be converted to {@link ExecutionContext}, so that it can be
  * serialized and stored.
@@ -56,18 +59,7 @@ import org.springframework.jdbc.core.RowMapper;
  * @see KeyCollector
  * @since 1.0
  */
-public interface ExecutionContextRowMapper extends RowMapper {
-
-	/**
-	 * Given the provided composite key, return a {@link ExecutionContext}
-	 * representation.
-	 * 
-	 * @param key
-	 * @return ExecutionContext representing the composite key.
-	 * @throws IllegalArgumentException
-	 *             if key is null or of an unsupported type.
-	 */
-	public void mapKeys(Object key, ExecutionContext executionContext);
+public interface KeyMappingPreparedStatementSetter {
 
 	/**
 	 * Given the provided restart data, return a PreparedStatementSeter that can
@@ -77,6 +69,5 @@ public interface ExecutionContextRowMapper extends RowMapper {
 	 * @return an array of objects that can be used as arguments to a
 	 *         JdbcTemplate.
 	 */
-	public PreparedStatementSetter createSetter(
-			ExecutionContext executionContext);
+	void setValues(PreparedStatement ps, Object key) throws SQLException;
 }

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/support/MultipleColumnJdbcKeyCollector.java

@@ -15,38 +15,50 @@
  */
 package org.springframework.batch.item.database.support;
 
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
 import java.util.List;
 
 import org.springframework.batch.item.ExecutionContext;
+import org.springframework.batch.item.ExecutionContextUserSupport;
 import org.springframework.batch.item.database.DrivingQueryItemReader;
 import org.springframework.batch.item.database.KeyCollector;
+import org.springframework.jdbc.core.ColumnMapRowMapper;
 import org.springframework.jdbc.core.JdbcTemplate;
+import org.springframework.jdbc.core.PreparedStatementSetter;
+import org.springframework.jdbc.core.RowMapper;
 import org.springframework.util.Assert;
+import org.springframework.util.ClassUtils;
 import org.springframework.util.StringUtils;
 
 /**
  * <p>
  * Jdbc implementation of the {@link KeyCollector} interface that works for
  * composite keys. (i.e. keys represented by multiple columns) A sql query to be
- * used to return the keys and a {@link ExecutionContextRowMapper} to map each
+ * used to return the keys and a {@link KeyMappingPreparedStatementSetter} to map each
  * row in the resultset to an Object must be set in order to work correctly.
  * </p>
  * 
  * @author Lucas Ward
  * @see DrivingQueryItemReader
- * @see ExecutionContextRowMapper
+ * @see KeyMappingPreparedStatementSetter
  */
-public class MultipleColumnJdbcKeyCollector implements KeyCollector {
+public class MultipleColumnJdbcKeyCollector extends ExecutionContextUserSupport implements KeyCollector {
 
+	private static final String CURRENT_KEY = "current.key";
+	
 	private JdbcTemplate jdbcTemplate;
 
-	private ExecutionContextRowMapper keyMapper = new ColumnMapExecutionContextRowMapper();
+	private RowMapper keyMapper = new ColumnMapRowMapper();
+	
+	private KeyMappingPreparedStatementSetter keyMappingSetter = new ColumnMapKeyMappingPreparedStatementSetter();
 
 	private String sql;
 
 	private String restartSql;
 
 	public MultipleColumnJdbcKeyCollector() {
+		setName(ClassUtils.getShortName(MultipleColumnJdbcKeyCollector.class));
 	}
 
 	/**
@@ -74,9 +86,10 @@ public class MultipleColumnJdbcKeyCollector implements KeyCollector {
 		Assert.state(keyMapper != null, "KeyMapper must not be null.");
 		Assert.state(StringUtils.hasText(restartSql), "The RestartQuery must not be null or empty"
 				+ " in order to restart.");
-
+		
 		if (executionContext.size() > 0) {
-			return jdbcTemplate.query(restartSql, keyMapper.createSetter(executionContext), keyMapper);
+			Object key = executionContext.get(getKey(CURRENT_KEY));
+			return jdbcTemplate.query(restartSql, new PreparedStatementSetterKeyWrapper(key, keyMappingSetter), keyMapper);
 		}
 		else {
 			return jdbcTemplate.query(sql, keyMapper);
@@ -88,9 +101,9 @@ public class MultipleColumnJdbcKeyCollector implements KeyCollector {
 	 * @see org.springframework.batch.io.driving.KeyGenerator#getKeyAsExecutionContext(java.lang.Object)
 	 */
 	public void updateContext(Object key, ExecutionContext executionContext) {
-		Assert.state(keyMapper != null, "Key mapper must not be null.");
 		Assert.notNull(key, "The key must not be null");
-		keyMapper.mapKeys(key, executionContext);
+		Assert.notNull(executionContext, "The ExecutionContext must not be null");
+		executionContext.put(getKey(CURRENT_KEY), key);
 	}
 
 	/**
@@ -114,12 +127,12 @@ public class MultipleColumnJdbcKeyCollector implements KeyCollector {
 	}
 
 	/**
-	 * Set the {@link ExecutionContextRowMapper} to be used to map a resultset
+	 * Set the {@link KeyMappingPreparedStatementSetter} to be used to map a resultset
 	 * to keys.
 	 * 
 	 * @param keyMapper
 	 */
-	public void setKeyMapper(ExecutionContextRowMapper keyMapper) {
+	public void setKeyMapper(RowMapper keyMapper) {
 		this.keyMapper = keyMapper;
 	}
 
@@ -135,4 +148,24 @@ public class MultipleColumnJdbcKeyCollector implements KeyCollector {
 	public void setJdbcTemplate(JdbcTemplate jdbcTemplate) {
 		this.jdbcTemplate = jdbcTemplate;
 	}
+	
+	public void setKeyMappingSetter(
+			KeyMappingPreparedStatementSetter keyMappingSetter) {
+		this.keyMappingSetter = keyMappingSetter;
+	}
+	
+	private class PreparedStatementSetterKeyWrapper implements PreparedStatementSetter{
+		
+		private Object key;
+		private KeyMappingPreparedStatementSetter pss;
+		
+		public PreparedStatementSetterKeyWrapper(Object key, KeyMappingPreparedStatementSetter pss) {
+			this.key = key;
+			this.pss = pss;
+		}
+
+		public void setValues(PreparedStatement ps) throws SQLException {
+			pss.setValues(ps, key);
+		}
+	}
 }

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/support/SingleColumnJdbcKeyCollector.java

@@ -89,10 +89,10 @@ public class SingleColumnJdbcKeyCollector extends ExecutionContextUserSupport im
 
 		Assert.notNull(executionContext, "The ExecutionContext must not be null");
 
-		if (executionContext.containsKey(RESTART_KEY)) {
+		if (executionContext.containsKey(getKey(RESTART_KEY))) {
 			Assert.state(StringUtils.hasText(restartSql), "The restart sql query must not be null or empty"
 			        + " in order to restart.");
-			return jdbcTemplate.query(restartSql, new Object[] { executionContext.get(RESTART_KEY) }, keyMapper);
+			return jdbcTemplate.query(restartSql, new Object[] { executionContext.get(getKey(RESTART_KEY)) }, keyMapper);
 		} else {
 			return jdbcTemplate.query(sql, keyMapper);
 		}
@@ -107,7 +107,7 @@ public class SingleColumnJdbcKeyCollector extends ExecutionContextUserSupport im
 	public void updateContext(Object key, ExecutionContext executionContext) {
 		Assert.notNull(key, "The key must not be null.");
 		Assert.notNull(executionContext, "The ExecutionContext must not be null");
-		executionContext.put(RESTART_KEY, key);
+		executionContext.put(getKey(RESTART_KEY), key);
 	}
 
 	/*

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/support/ColumnMapExecutionContextRowMapperTests.java

@@ -6,88 +6,68 @@ package org.springframework.batch.item.database.support;
 import java.sql.PreparedStatement;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.Properties;
 
 import junit.framework.TestCase;
 
 import org.easymock.MockControl;
-import org.springframework.batch.item.ExecutionContext;
 import org.springframework.core.CollectionFactory;
-import org.springframework.jdbc.core.PreparedStatementSetter;
 
 /**
  * @author Lucas Ward
  */
 public class ColumnMapExecutionContextRowMapperTests extends TestCase {
 
-	private ColumnMapExecutionContextRowMapper mapper;
+	private ColumnMapKeyMappingPreparedStatementSetter mapper;
 	
 	private Map key;
 	
 	private MockControl psControl = MockControl.createControl(PreparedStatement.class);
 	private PreparedStatement ps;
-	
-	private ExecutionContext executionContext;
-	
+		
 	protected void setUp() throws Exception {
 		super.setUp();
 	
-		mapper = new ColumnMapExecutionContextRowMapper();
+		ps = (PreparedStatement)psControl.getMock();
+		mapper = new ColumnMapKeyMappingPreparedStatementSetter();
 		
 		key = CollectionFactory.createLinkedCaseInsensitiveMapIfPossible(2);
 		key.put("1", new Integer(1));
 		key.put("2", new Integer(2));
-		
-		executionContext = new ExecutionContext();
 	}
 	
-	public void testCreateExecutionContextWithInvalidType() throws Exception {
+	public void testSetValuesWithInvalidType() throws Exception {
 		
 		try{
-			mapper.mapKeys(new Object(), executionContext);
+			mapper.setValues(ps, new Object());
 			fail();
 		}catch(IllegalArgumentException ex){
 			//expected
 		}
 	}
 	
-	public void testCreateExecutionContextWithNull(){
+	public void testCreateExecutionContextWithNull() throws Exception{
 		
 		try{
-			mapper.mapKeys(null, null);
+			mapper.setValues(ps, null);
 			fail();
 		}catch(IllegalArgumentException ex){
 			//expected
 		}
 	}
 	
-	public void testCreateExecutionContext() throws Exception {
-		mapper.mapKeys(key, executionContext);
-		Properties props = executionContext.getProperties();
-		assertEquals("1", props.getProperty(ColumnMapExecutionContextRowMapper.KEY_PREFIX+"0"));
-		assertEquals("2", props.getProperty(ColumnMapExecutionContextRowMapper.KEY_PREFIX+"1"));
-	}
-	
 	public void testCreateExecutionContextFromEmptyKeys() throws Exception {
 		
-		mapper.mapKeys(new HashMap(), executionContext);
-		assertEquals(0, executionContext.size());
+		psControl.replay();
+		mapper.setValues(ps, new HashMap());
+		psControl.verify();
 	}
 	
 	public void testCreateSetter() throws Exception {
 		
-		ExecutionContext streamContext = new ExecutionContext();
-		streamContext.putString(ColumnMapExecutionContextRowMapper.KEY_PREFIX+"1", "1");
-		streamContext.putString(ColumnMapExecutionContextRowMapper.KEY_PREFIX+"0", "2");
-		PreparedStatementSetter setter = mapper.createSetter(streamContext);
-		ps = (PreparedStatement)psControl.getMock();
-		
-		ps.setString(1, "2");
-		ps.setString(2, "1");
+		ps.setObject(1, new Integer(1));
+		ps.setObject(2, new Integer(2));
 		psControl.replay();
-		
-		setter.setValues(ps);
-		
+		mapper.setValues(ps, key);	
 		psControl.verify();
 	}
 	

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/support/MultipleColumnJdbcKeyGeneratorIntegrationTests.java

@@ -3,18 +3,13 @@
  */
 package org.springframework.batch.item.database.support;
 
-import java.sql.ResultSet;
-import java.sql.SQLException;
-import java.util.Iterator;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Properties;
-import java.util.Map.Entry;
 
 import org.springframework.batch.item.ExecutionContext;
-import org.springframework.core.CollectionFactory;
-import org.springframework.jdbc.core.PreparedStatementSetter;
 import org.springframework.test.AbstractTransactionalDataSourceSpringContextTests;
+import org.springframework.util.ClassUtils;
 
 /**
  * @author Lucas Ward
@@ -54,8 +49,10 @@ public class MultipleColumnJdbcKeyGeneratorIntegrationTests extends AbstractTran
 	
 	public void testRestoreKeys(){
 		
-		executionContext.putString(ColumnMapExecutionContextRowMapper.KEY_PREFIX + "0", "3");
-		executionContext.putString(ColumnMapExecutionContextRowMapper.KEY_PREFIX + "1", "3");
+		Map keyMap = new LinkedHashMap();
+		keyMap.put("ID", "3");
+		keyMap.put("VALUE", "3");
+		executionContext.put(ClassUtils.getShortName(MultipleColumnJdbcKeyCollector.class)+ ".current.key", keyMap);
 		
 		List keys = keyStrategy.retrieveKeys(executionContext);
 		
@@ -68,36 +65,36 @@ public class MultipleColumnJdbcKeyGeneratorIntegrationTests extends AbstractTran
 		assertEquals(new Integer(5), key.get("VALUE"));
 	}
 	
-	public void testGetKeyAsExecutionContext(){
-		
-		Map key = CollectionFactory.createLinkedCaseInsensitiveMapIfPossible(1);
-		key.put("ID", new Long(3));
-		key.put("VALUE", new Integer(4));
-		
-		keyStrategy.setKeyMapper(new ExecutionContextRowMapper() {
-			public PreparedStatementSetter createSetter(ExecutionContext executionContext) {
-				return null;
-			}
-			public void mapKeys(Object key, ExecutionContext executionContext) {
-				// Just slap the key as a map into the context
-				Map keys = (Map) key;
-				for (Iterator it = keys.entrySet().iterator(); it.hasNext();) {
-					Entry entry = (Entry)it.next();
-					executionContext.put(entry.getKey().toString(), entry.getValue());
-				}
-			}
-			public Object mapRow(ResultSet rs, int rowNum) throws SQLException {
-				return null;
-			}
-		});
-		keyStrategy.updateContext(key, executionContext);
-		Properties props = executionContext.getProperties();
-		
-		assertEquals(2, props.size());
-		System.err.println(props);
-		assertEquals("3", props.get("ID"));
-		assertEquals("4", props.get("VALUE"));
-	}
+//	public void testGetKeyAsExecutionContext(){
+//		
+//		Map key = CollectionFactory.createLinkedCaseInsensitiveMapIfPossible(1);
+//		key.put("ID", new Long(3));
+//		key.put("VALUE", new Integer(4));
+//		
+//		keyStrategy.setKeyMapper(new KeyMappingPreparedStatementSetter() {
+//			public PreparedStatementSetter createSetter(ExecutionContext executionContext) {
+//				return null;
+//			}
+//			public void mapKeys(Object key, ExecutionContext executionContext) {
+//				// Just slap the key as a map into the context
+//				Map keys = (Map) key;
+//				for (Iterator it = keys.entrySet().iterator(); it.hasNext();) {
+//					Entry entry = (Entry)it.next();
+//					executionContext.put(entry.getKey().toString(), entry.getValue());
+//				}
+//			}
+//			public Object mapRow(ResultSet rs, int rowNum) throws SQLException {
+//				return null;
+//			}
+//		});
+//		keyStrategy.updateContext(key, executionContext);
+//		Properties props = executionContext.getProperties();
+//		
+//		assertEquals(2, props.size());
+//		System.err.println(props);
+//		assertEquals("3", props.get("ID"));
+//		assertEquals("4", props.get("VALUE"));
+//	}
 	
 	public void testGetNullKeyAsStreamContext(){
 		

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/support/SingleColumnJdbcKeyGeneratorIntegrationTests.java

@@ -4,6 +4,7 @@ import java.util.List;
 
 import org.springframework.batch.item.ExecutionContext;
 import org.springframework.test.AbstractTransactionalDataSourceSpringContextTests;
+import org.springframework.util.ClassUtils;
 
 /**
  * 
@@ -66,7 +67,7 @@ public class SingleColumnJdbcKeyGeneratorIntegrationTests extends AbstractTransa
 		keyStrategy.updateContext(new Long(3), executionContext);
 		
 		assertEquals(1, executionContext.size());
-		assertEquals(new Long(3), executionContext.get("key"));
+		assertEquals(new Long(3), executionContext.get(ClassUtils.getShortName(SingleColumnJdbcKeyCollector.class) + ".key"));
 	}
 	
 	public void testGetNullKeyAsStreamContext(){