From 88f73791fbf79d198df7f491e85a31ce96f93154 Mon Sep 17 00:00:00 2001 From: Michael Minella Date: Tue, 20 May 2014 10:14:05 -0500 Subject: [PATCH] XD-1622: Add tasklet and step type to Step's ExecutionContext for XD usage --- .../org/springframework/batch/core/Step.java | 1 + .../batch/core/job/flow/FlowStep.java | 1 + .../core/partition/support/PartitionStep.java | 5 ++-- .../batch/core/step/job/JobStep.java | 2 ++ .../batch/core/step/tasklet/TaskletStep.java | 4 +++ .../step/item/TaskletStepExceptionTests.java | 29 +++++++++++-------- 6 files changed, 28 insertions(+), 14 deletions(-) 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 index c5e83e4b8..5ae3ed7fe 100644 --- 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 @@ package org.springframework.batch.core; */ 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 index dd2b7a1fb..67ae795ab 100644 --- 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 class FlowStep extends AbstractStep { @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 index fc29e4576..b6a4383ca 100644 --- 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.core.step.AbstractStep; 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 class PartitionStep extends AbstractStep { */ @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 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 index 77a931bb1..93ecb313d 100644 --- 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 @@ public class JobStep extends AbstractStep { 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 index 096570916..784ba4a39 100644 --- 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 class TaskletStep extends AbstractStep { private Tasklet tasklet; + public static final String TASKLET_TYPE_KEY = "batch.taskletType"; + /** * Default constructor. */ @@ -242,6 +244,8 @@ public class TaskletStep extends AbstractStep { */ @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 index 82f454dad..c66cfa8f9 100644 --- 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,16 +15,6 @@ */ 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; @@ -32,6 +22,7 @@ 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.UnexpectedRollbackException; 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 class TaskletStepExceptionTests { 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 class TaskletStepExceptionTests { 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