diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/configuration/xml/StepParserStepFactoryBean.java b/spring-batch-core/src/main/java/org/springframework/batch/core/configuration/xml/StepParserStepFactoryBean.java index c5ec72977..90d489cef 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/configuration/xml/StepParserStepFactoryBean.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/configuration/xml/StepParserStepFactoryBean.java @@ -24,6 +24,7 @@ import java.util.Map; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; +import org.springframework.aop.framework.Advised; import org.springframework.batch.classify.BinaryExceptionClassifier; import org.springframework.batch.core.ChunkListener; import org.springframework.batch.core.Job; @@ -301,13 +302,14 @@ class StepParserStepFactoryBean implements FactoryBean, BeanNameAware { handler = partitionHandler; } + PartitionHandler targetHandler = (PartitionHandler) extractTarget(handler, PartitionHandler.class); // BATCH-1659 - if (handler instanceof TaskExecutorPartitionHandler) { + if (targetHandler instanceof TaskExecutorPartitionHandler) { // Only for a local partition handler is the step required Assert.state(step != null, "A Step must be provided for a partition step with a TaskExecutorPartitionHandler"); try { - TaskExecutorPartitionHandler taskExecutorPartitionHandler = (TaskExecutorPartitionHandler) handler; + TaskExecutorPartitionHandler taskExecutorPartitionHandler = (TaskExecutorPartitionHandler) targetHandler; taskExecutorPartitionHandler.setStep(step); taskExecutorPartitionHandler.afterPropertiesSet(); } @@ -342,6 +344,25 @@ class StepParserStepFactoryBean implements FactoryBean, BeanNameAware { } } + private Object extractTarget(Object target, Class type) { + if (target instanceof Advised) { + Object source; + try { + source = ((Advised) target).getTargetSource().getTarget(); + } + catch (Exception e) { + throw new IllegalStateException("Could not extract target from proxy", e); + } + if (source instanceof Advised) { + source = extractTarget(source, type); + } + if (type.isAssignableFrom(source.getClass())) { + target = source; + } + } + return target; + } + private void configureSimple(SimpleStepFactoryBean fb) { if (name != null) { fb.setBeanName(name); diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/TaskExecutorPartitionHandler.java b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/TaskExecutorPartitionHandler.java index 3eacc1b71..7731b934b 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/TaskExecutorPartitionHandler.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/TaskExecutorPartitionHandler.java @@ -30,6 +30,7 @@ import org.springframework.batch.core.Step; import org.springframework.batch.core.StepExecution; import org.springframework.batch.core.partition.PartitionHandler; import org.springframework.batch.core.partition.StepExecutionSplitter; +import org.springframework.batch.core.step.StepHolder; import org.springframework.beans.factory.InitializingBean; import org.springframework.core.task.SyncTaskExecutor; import org.springframework.core.task.TaskExecutor; @@ -45,7 +46,7 @@ import org.springframework.util.Assert; * @author Dave Syer * @since 2.0 */ -public class TaskExecutorPartitionHandler implements PartitionHandler, InitializingBean { +public class TaskExecutorPartitionHandler implements PartitionHandler, StepHolder, InitializingBean { private int gridSize = 1; @@ -89,6 +90,16 @@ public class TaskExecutorPartitionHandler implements PartitionHandler, Initializ public void setStep(Step step) { this.step = step; } + + /** + * The step instance that will be executed in parallel by this handler. + * + * @return the step instance that will be used + * @see StepHolder#getStep() + */ + public Step getStep() { + return this.step; + } /** * @see PartitionHandler#handle(StepExecutionSplitter, StepExecution) diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/configuration/xml/StepParserStepFactoryBeanTests.java b/spring-batch-core/src/test/java/org/springframework/batch/core/configuration/xml/StepParserStepFactoryBeanTests.java index c874f11fd..c5f14a64c 100644 --- a/spring-batch-core/src/test/java/org/springframework/batch/core/configuration/xml/StepParserStepFactoryBeanTests.java +++ b/spring-batch-core/src/test/java/org/springframework/batch/core/configuration/xml/StepParserStepFactoryBeanTests.java @@ -22,10 +22,13 @@ import static org.junit.Assert.assertTrue; import java.util.HashMap; import org.junit.Test; +import org.springframework.aop.framework.Advised; +import org.springframework.aop.framework.ProxyFactory; import org.springframework.batch.core.StepListener; import org.springframework.batch.core.job.flow.FlowStep; import org.springframework.batch.core.job.flow.support.SimpleFlow; import org.springframework.batch.core.listener.StepExecutionListenerSupport; +import org.springframework.batch.core.partition.PartitionHandler; import org.springframework.batch.core.partition.support.PartitionStep; import org.springframework.batch.core.partition.support.SimplePartitioner; import org.springframework.batch.core.partition.support.TaskExecutorPartitionHandler; @@ -254,6 +257,28 @@ public class StepParserStepFactoryBeanTests { assertTrue(handler instanceof TaskExecutorPartitionHandler); } + @Test + public void testPartitionStepWithProxyHandler() throws Exception { + StepParserStepFactoryBean fb = new StepParserStepFactoryBean(); + fb.setBeanName("step1"); + fb.setAllowStartIfComplete(true); + fb.setJobRepository(new JobRepositorySupport()); + fb.setStartLimit(5); + fb.setListeners(new StepListener[] { new StepExecutionListenerSupport() }); + fb.setTaskExecutor(new SyncTaskExecutor()); + + SimplePartitioner partitioner = new SimplePartitioner(); + fb.setPartitioner(partitioner); + fb.setStep(new StepSupport("foo")); + ProxyFactory factory = new ProxyFactory(new TaskExecutorPartitionHandler()); + fb.setPartitionHandler((PartitionHandler) factory.getProxy()); + + Object step = fb.getObject(); + assertTrue(step instanceof PartitionStep); + Object handler = ReflectionTestUtils.getField(step, "partitionHandler"); + assertTrue(handler instanceof Advised); + } + @Test public void testFlowStep() throws Exception { StepParserStepFactoryBean fb = new StepParserStepFactoryBean();