From 1dcc16e5831a98a4b430724780d06484dd9d3748 Mon Sep 17 00:00:00 2001
From: lucasward
* 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. *
* * @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); + } + } } diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/support/SingleColumnJdbcKeyCollector.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/support/SingleColumnJdbcKeyCollector.java index 7152e7513..3be56b020 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/support/SingleColumnJdbcKeyCollector.java +++ b/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); } /* diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/support/ColumnMapExecutionContextRowMapperTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/support/ColumnMapExecutionContextRowMapperTests.java index 391fcaf8b..52ddc267f 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/support/ColumnMapExecutionContextRowMapperTests.java +++ b/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(); } diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/support/MultipleColumnJdbcKeyGeneratorIntegrationTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/support/MultipleColumnJdbcKeyGeneratorIntegrationTests.java index 25a6b8160..68869ffd1 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/support/MultipleColumnJdbcKeyGeneratorIntegrationTests.java +++ b/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(){ diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/support/SingleColumnJdbcKeyGeneratorIntegrationTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/support/SingleColumnJdbcKeyGeneratorIntegrationTests.java index fe41c8d6a..061e420a5 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/support/SingleColumnJdbcKeyGeneratorIntegrationTests.java +++ b/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(){