Fix step execution context is not persisted and restored

1. Step execution context is not persisted in `SimpleStepExecutionSplitter::split`
2. Step execution context is not restored in `SimpleJobRepository::getStepExecution`

Closes GH-5138

Signed-off-by: Yanming Zhou <[email protected]>

Author: quaff

spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/SimpleStepExecutionSplitter.java

@@ -45,6 +45,7 @@
  *
  * @author Dave Syer
  * @author Mahmoud Ben Hassine
+ * @author Yanming Zhou
  * @since 2.0
  */
 public class SimpleStepExecutionSplitter implements StepExecutionSplitter {
@@ -138,13 +139,15 @@ public Set<StepExecution> split(StepExecution stepExecution, int gridSize) throw
 			if (lastStepExecution == null) { // fresh start
 				StepExecution currentStepExecution = jobRepository.createStepExecution(stepName, jobExecution);
 				currentStepExecution.setExecutionContext(context.getValue());
+				jobRepository.updateExecutionContext(currentStepExecution);
 				set.add(currentStepExecution);
 			}
 			else { // restart
 				if (lastStepExecution.getStatus() != BatchStatus.COMPLETED
 						&& shouldStart(allowStartIfComplete, stepExecution, lastStepExecution)) {
 					StepExecution currentStepExecution = jobRepository.createStepExecution(stepName, jobExecution);
 					currentStepExecution.setExecutionContext(lastStepExecution.getExecutionContext());
+					jobRepository.updateExecutionContext(currentStepExecution);
 					set.add(currentStepExecution);
 				}
 			}

spring-batch-core/src/main/java/org/springframework/batch/core/repository/explore/support/SimpleJobExplorer.java

@@ -44,6 +44,7 @@
  * @author Mahmoud Ben Hassine
  * @author Parikshit Dutta
  * @author Glenn Renfro
+ * @author Yanming Zhou
  * @see JobExplorer
  * @see JobInstanceDao
  * @see JobExecutionDao
@@ -287,7 +288,7 @@ public long getStepExecutionCount(JobInstance jobInstance, String stepName) thro
 		return stepExecutionDao.countStepExecutions(jobInstance, stepName);
 	}
 
-	private void getStepExecutionDependencies(StepExecution stepExecution) {
+	protected void getStepExecutionDependencies(StepExecution stepExecution) {
 		if (stepExecution != null) {
 			stepExecution.setExecutionContext(ecDao.getExecutionContext(stepExecution));
 		}

spring-batch-core/src/main/java/org/springframework/batch/core/repository/support/SimpleJobRepository.java

@@ -52,6 +52,7 @@
  * @author Baris Cubukcuoglu
  * @author Parikshit Dutta
  * @author Mark John Moreno
+ * @author Yanming Zhou
  * @see JobRepository
  * @see JobInstanceDao
  * @see JobExecutionDao
@@ -82,7 +83,9 @@ public List<JobInstance> findJobInstances(String jobName) {
 	@Nullable
 	@Override
 	public StepExecution getStepExecution(long executionId) {
-		return this.stepExecutionDao.getStepExecution(executionId);
+		StepExecution stepExecution = this.stepExecutionDao.getStepExecution(executionId);
+		getStepExecutionDependencies(stepExecution);
+		return stepExecution;
 	}
 
 	/**

spring-batch-core/src/test/java/org/springframework/batch/core/partition/PartitionStepTests.java

@@ -16,9 +16,7 @@
 package org.springframework.batch.core.partition;
 
 import java.time.LocalDateTime;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Set;
+import java.util.*;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.junit.jupiter.api.BeforeEach;
@@ -41,10 +39,12 @@
 import org.springframework.jdbc.datasource.embedded.EmbeddedDatabaseBuilder;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 
 /**
  * @author Dave Syer
  * @author Mahmoud Ben Hassine
+ * @author Yanming Zhou
  *
  */
 class PartitionStepTests {
@@ -71,12 +71,24 @@ void setUp() throws Exception {
 	@Test
 	void testVanillaStepExecution() throws Exception {
 		SimpleStepExecutionSplitter stepExecutionSplitter = new SimpleStepExecutionSplitter(jobRepository,
-				step.getName(), new SimplePartitioner());
+				step.getName(), gridSize -> {
+					Map<String, ExecutionContext> map = new HashMap<>(gridSize);
+					for (int i = 0; i < gridSize; i++) {
+						ExecutionContext context = new ExecutionContext();
+						context.putString("foo", "foo" + i);
+						map.put("partition" + i, context);
+					}
+					return map;
+				});
 		stepExecutionSplitter.setAllowStartIfComplete(true);
 		step.setStepExecutionSplitter(stepExecutionSplitter);
 		step.setPartitionHandler((stepSplitter, stepExecution) -> {
 			Set<StepExecution> executions = stepSplitter.split(stepExecution, 2);
 			for (StepExecution execution : executions) {
+				// Query from repository to ensure it's persisted
+				ExecutionContext context = jobRepository.getStepExecution(execution.getId()).getExecutionContext();
+				assertNotNull(context.getString("foo"));
+
 				execution.setStatus(BatchStatus.COMPLETED);
 				execution.setExitStatus(ExitStatus.COMPLETED);
 				jobRepository.update(execution);
@@ -144,7 +156,9 @@ void testRestartStepExecution() throws Exception {
 			else {
 				for (StepExecution execution : executions) {
 					// On restart the execution context should have been restored
-					assertEquals(execution.getStepName(), execution.getExecutionContext().getString("foo"));
+					// Query from repository to ensure it's persisted
+					ExecutionContext context = jobRepository.getStepExecution(execution.getId()).getExecutionContext();
+					assertEquals(execution.getStepName(), context.getString("foo"));
 				}
 			}
 			for (StepExecution execution : executions) {