XD-1622: Add tasklet and step type to Step's ExecutionContext for XD usage

mminella · mminella · commit bf525de0e299 · 2014-05-30T11:31:20.000-05:00
diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/Step.java b/spring-batch-core/src/main/java/org/springframework/batch/core/Step.java
@@ -24,6 +24,7 @@
  */
 public interface Step {
 
+	static final String STEP_TYPE_KEY = "batch.stepType";
 	/**
 	 * @return the name of this step.
 	 */
diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/job/flow/FlowStep.java b/spring-batch-core/src/main/java/org/springframework/batch/core/job/flow/FlowStep.java
@@ -84,6 +84,7 @@ public void afterPropertiesSet() throws Exception {
 	@Override
 	protected void doExecute(StepExecution stepExecution) throws Exception {
 		try {
+			stepExecution.getExecutionContext().put(STEP_TYPE_KEY, this.getClass().getName());
 			StepHandler stepHandler = new SimpleStepHandler(getJobRepository(), stepExecution.getExecutionContext());
 			FlowExecutor executor = new JobFlowExecutor(getJobRepository(), stepHandler, stepExecution.getJobExecution());
 			executor.updateJobExecutionStatus(flow.start(executor).getStatus());
diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/PartitionStep.java b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/PartitionStep.java
@@ -16,8 +16,6 @@
 
 package org.springframework.batch.core.partition.support;
 
-import java.util.Collection;
-
 import org.springframework.batch.core.BatchStatus;
 import org.springframework.batch.core.JobExecutionException;
 import org.springframework.batch.core.Step;
@@ -28,6 +26,8 @@
 import org.springframework.batch.item.ExecutionContext;
 import org.springframework.util.Assert;
 
+import java.util.Collection;
+
 /**
  * Implementation of {@link Step} which partitions the execution and spreads the
  * load using a {@link PartitionHandler}.
@@ -100,6 +100,7 @@ public void afterPropertiesSet() throws Exception {
 	 */
 	@Override
 	protected void doExecute(StepExecution stepExecution) throws Exception {
+		stepExecution.getExecutionContext().put(STEP_TYPE_KEY, this.getClass().getName());
 
 		// Wait for task completion and then aggregate the results
 		Collection<StepExecution> executions = partitionHandler.handle(stepExecutionSplitter, stepExecution);
diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/step/job/JobStep.java b/spring-batch-core/src/main/java/org/springframework/batch/core/step/job/JobStep.java
@@ -103,6 +103,8 @@ protected void doExecute(StepExecution stepExecution) throws Exception {
 
 		ExecutionContext executionContext = stepExecution.getExecutionContext();
 
+		executionContext.put(STEP_TYPE_KEY, this.getClass().getName());
+
 		JobParameters jobParameters;
 		if (executionContext.containsKey(JOB_PARAMETERS_KEY)) {
 			jobParameters = (JobParameters) executionContext.get(JOB_PARAMETERS_KEY);
diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/step/tasklet/TaskletStep.java b/spring-batch-core/src/main/java/org/springframework/batch/core/step/tasklet/TaskletStep.java
@@ -100,6 +100,8 @@ public boolean rollbackOn(Throwable ex) {
 
 	private Tasklet tasklet;
 
+	public static final String TASKLET_TYPE_KEY = "batch.taskletType";
+
 	/**
 	 * Default constructor.
 	 */
@@ -242,6 +244,8 @@ public void setInterruptionPolicy(StepInterruptionPolicy interruptionPolicy) {
 	 */
 	@Override
 	protected void doExecute(StepExecution stepExecution) throws Exception {
+		stepExecution.getExecutionContext().put(TASKLET_TYPE_KEY, tasklet.getClass().getName());
+		stepExecution.getExecutionContext().put(STEP_TYPE_KEY, this.getClass().getName());
 
 		stream.update(stepExecution.getExecutionContext());
 		getJobRepository().updateExecutionContext(stepExecution);
diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/step/item/TaskletStepExceptionTests.java b/spring-batch-core/src/test/java/org/springframework/batch/core/step/item/TaskletStepExceptionTests.java
@@ -15,23 +15,14 @@
  */
 package org.springframework.batch.core.step.item;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-import static org.springframework.batch.core.BatchStatus.COMPLETED;
-import static org.springframework.batch.core.BatchStatus.FAILED;
-import static org.springframework.batch.core.BatchStatus.STOPPED;
-import static org.springframework.batch.core.BatchStatus.UNKNOWN;
-
-import java.util.Collection;
-
 import org.junit.Before;
 import org.junit.Test;
 import org.springframework.batch.core.ExitStatus;
 import org.springframework.batch.core.JobExecution;
 import org.springframework.batch.core.JobInstance;
 import org.springframework.batch.core.JobInterruptedException;
 import org.springframework.batch.core.JobParameters;
+import org.springframework.batch.core.Step;
 import org.springframework.batch.core.StepContribution;
 import org.springframework.batch.core.StepExecution;
 import org.springframework.batch.core.StepExecutionListener;
@@ -54,6 +45,16 @@
 import org.springframework.transaction.support.DefaultTransactionStatus;
 import org.springframework.transaction.support.TransactionSynchronizationManager;
 
+import java.util.Collection;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.springframework.batch.core.BatchStatus.COMPLETED;
+import static org.springframework.batch.core.BatchStatus.FAILED;
+import static org.springframework.batch.core.BatchStatus.STOPPED;
+import static org.springframework.batch.core.BatchStatus.UNKNOWN;
+
 /**
  * Tests for the behavior of TaskletStep in a failure scenario.
  *
@@ -245,7 +246,9 @@ public RepeatStatus execute(StepContribution contribution, ChunkContext attribut
 		assertEquals(1, stepExecution.getRollbackCount()); // Failed transaction
 		// counts as
 		// rollback
-		assertEquals(0, stepExecution.getExecutionContext().size());
+		assertEquals(2, stepExecution.getExecutionContext().size());
+		assertTrue(stepExecution.getExecutionContext().containsKey(Step.STEP_TYPE_KEY));
+		assertTrue(stepExecution.getExecutionContext().containsKey(TaskletStep.TASKLET_TYPE_KEY));
 	}
 
 	@SuppressWarnings("serial")
@@ -278,7 +281,9 @@ public RepeatStatus execute(StepContribution contribution, ChunkContext attribut
 		assertEquals(1, stepExecution.getRollbackCount()); // Failed transaction
 		// counts as
 		// rollback
-		assertEquals(0, stepExecution.getExecutionContext().size());
+		assertEquals(2, stepExecution.getExecutionContext().size());
+		assertTrue(stepExecution.getExecutionContext().containsKey(Step.STEP_TYPE_KEY));
+		assertTrue(stepExecution.getExecutionContext().containsKey(TaskletStep.TASKLET_TYPE_KEY));
 	}
 
 	@Test
