BATCH-1743: extract target from proxy before making assumptions about type

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<I, O> 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<I, O> 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<I, O> fb) {
 		if (name != null) {
 			fb.setBeanName(name);

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)

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<Object, Object> fb = new StepParserStepFactoryBean<Object, Object>();
+		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<Object, Object> fb = new StepParserStepFactoryBean<Object, Object>();