/*
* Copyright 2013-2014 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.batch.core.jsr.step;
import java.util.Collection;
import javax.batch.api.partition.PartitionReducer;
import javax.batch.api.partition.PartitionReducer.PartitionStatus;
import org.springframework.batch.core.BatchStatus;
import org.springframework.batch.core.JobExecutionException;
import org.springframework.batch.core.Step;
import org.springframework.batch.core.StepExecution;
import org.springframework.batch.core.jsr.partition.JsrPartitionHandler;
import org.springframework.batch.core.jsr.partition.support.JsrStepExecutionAggregator;
import org.springframework.batch.core.partition.PartitionHandler;
import org.springframework.batch.core.partition.StepExecutionSplitter;
import org.springframework.batch.core.partition.support.StepExecutionAggregator;
import org.springframework.batch.core.step.NoSuchStepException;
import org.springframework.batch.core.step.StepLocator;
import org.springframework.batch.item.ExecutionContext;
/**
* An extension of the {@link PartitionStep} that provides additional semantics
* required by JSR-352. Specifically, this implementation adds the required
* lifecycle calls to the {@link PartitionReducer} if it is used.
*
* @author Michael Minella
* @since 3.0
*/
public class PartitionStep extends org.springframework.batch.core.partition.support.PartitionStep implements StepLocator {
private PartitionReducer reducer;
private boolean hasReducer = false;
private StepExecutionAggregator stepExecutionAggregator = new JsrStepExecutionAggregator();
public void setPartitionReducer(PartitionReducer reducer) {
this.reducer = reducer;
hasReducer = reducer != null;
}
/**
* Delegate execution to the {@link PartitionHandler} provided. The
* {@link StepExecution} passed in here becomes the parent or master
* execution for the partition, summarizing the status on exit of the
* logical grouping of work carried out by the {@link PartitionHandler}. The
* individual step executions and their input parameters (through
* {@link ExecutionContext}) for the partition elements are provided by the
* {@link StepExecutionSplitter}.
*
* @param stepExecution the master step execution for the partition
*
* @see Step#execute(StepExecution)
*/
@Override
protected void doExecute(StepExecution stepExecution) throws Exception {
if(hasReducer) {
reducer.beginPartitionedStep();
}
// Wait for task completion and then aggregate the results
Collection<StepExecution> stepExecutions = getPartitionHandler().handle(null, stepExecution);
stepExecution.upgradeStatus(BatchStatus.COMPLETED);
stepExecutionAggregator.aggregate(stepExecution, stepExecutions);
if (stepExecution.getStatus().isUnsuccessful()) {
if (hasReducer) {
reducer.rollbackPartitionedStep();
reducer.afterPartitionedStepCompletion(PartitionStatus.ROLLBACK);
}
throw new JobExecutionException("Partition handler returned an unsuccessful step");
}
if (hasReducer) {
reducer.beforePartitionedStepCompletion();
reducer.afterPartitionedStepCompletion(PartitionStatus.COMMIT);
}
}
/* (non-Javadoc)
* @see org.springframework.batch.core.step.StepLocator#getStepNames()
*/
@Override
public Collection<String> getStepNames() {
return ((JsrPartitionHandler) getPartitionHandler()).getPartitionStepNames();
}
/* (non-Javadoc)
* @see org.springframework.batch.core.step.StepLocator#getStep(java.lang.String)
*/
@Override
public Step getStep(String stepName) throws NoSuchStepException {
JsrPartitionHandler partitionHandler = (JsrPartitionHandler) getPartitionHandler();
Collection<String> names = partitionHandler.getPartitionStepNames();
if(names.contains(stepName)) {
return partitionHandler.getStep();
} else {
throw new NoSuchStepException(stepName + " was not found");
}
}
}