IncrementalColumnRangePartitioner.java

/*
 * Copyright 2015 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.integration.x;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

import javax.sql.DataSource;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import org.springframework.batch.core.ExitStatus;
import org.springframework.batch.core.JobExecution;
import org.springframework.batch.core.JobInstance;
import org.springframework.batch.core.StepExecution;
import org.springframework.batch.core.StepExecutionListener;
import org.springframework.batch.core.explore.JobExplorer;
import org.springframework.batch.core.partition.support.Partitioner;
import org.springframework.batch.item.ExecutionContext;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.jdbc.core.JdbcOperations;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.util.StringUtils;

/**
 * Provides the ability to partition a job and append additional conditions to support
 * incremental imports.  Incremental imports are supported via the checkColumn attribute.
 * The partitionMax value processed in the current run will be set as the minimum value for the
 * next run.
 *
 * @author Michael Minella
 * @since 1.2
 */
public class IncrementalColumnRangePartitioner implements Partitioner, StepExecutionListener, InitializingBean {

	private static final Log log = LogFactory.getLog(IncrementalColumnRangePartitioner.class);

	public static final String BATCH_INCREMENTAL_MAX_ID = "batch.incremental.maxId";

	private JdbcOperations jdbcTemplate;

	private String table;

	private String column;

	private int partitions;

	private long partitionMax = Long.MAX_VALUE;

	private long partitionMin = Long.MIN_VALUE;

	private long incrementalMin = Long.MIN_VALUE;

	private JobExplorer jobExplorer;

	private String checkColumn;

	private Long overrideValue;

	/**
	 * The data source for connecting to the database.
	 *
	 * @param dataSource a {@link DataSource}
	 */
	public void setDataSource(DataSource dataSource) {
		jdbcTemplate = new JdbcTemplate(dataSource);
	}

	/**
	 * The name of the SQL table the data are in.
	 *
	 * @param table the name of the table
	 */
	public void setTable(String table) {
		this.table = table;
	}

	/**
	 * The name of the column to partition.
	 *
	 * @param column the column name.
	 */
	public void setColumn(String column) {
		this.column = column;
	}

	/**
	 * The number of partitions to create.
	 *
	 * @param partitions the number of partitions.
	 */
	public void setPartitions(int partitions) {
		this.partitions = partitions;
	}

	public void setJobExplorer(JobExplorer jobExplorer) {
		this.jobExplorer = jobExplorer;
	}

	public void setCheckColumn(String checkColumn) {
		this.checkColumn = checkColumn;
	}

	public void setOverrideValue(Long overrideValue) {
		this.overrideValue = overrideValue;
	}

	/**
	 * Partition a database table assuming that the data in the column specified
	 * are uniformly distributed. The execution context values will have keys
	 * <code>minValue</code> and <code>maxValue</code> specifying the range of
	 * values to consider in each partition.
	 *
	 * @see Partitioner#partition(int)
	 */
	@Override
	public Map<String, ExecutionContext> partition(int gridSize) {
		StringBuilder incrementalClause = new StringBuilder();
		Map<String, ExecutionContext> result = new HashMap<>();

		if(!StringUtils.hasText(checkColumn) && !StringUtils.hasText(column)) {
			ExecutionContext value = new ExecutionContext();
			value.put("partClause", "");
			result.put("partition0", value);
			value.put("partSuffix", "");
		}
		else {
			if(StringUtils.hasText(checkColumn)) {
				incrementalClause.append(checkColumn).append(" > ").append(this.incrementalMin);
			}

			long targetSize = (this.partitionMax - this.partitionMin) / partitions + 1;

			int number = 0;
			long start = this.partitionMin;
			long end = start + targetSize - 1;

			while (start >= 0 && start <= this.partitionMax) {
				ExecutionContext value = new ExecutionContext();
				result.put("partition" + number, value);

				if (end >= this.partitionMax) {
					end = this.partitionMax;
				}

				if(StringUtils.hasText(checkColumn)) {
					value.putString("partClause", String.format("WHERE (%s BETWEEN %s AND %s) AND %s", column, start, end, incrementalClause.toString()));
				}
				else {
					value.putString("partClause", String.format("WHERE (%s BETWEEN %s AND %s)", column, start, end));
				}

				value.putString("partSuffix", "-p"+number);
				start += targetSize;
				end += targetSize;
				number++;

				log.debug("Current ExecutionContext = " + value);
			}
		}

		return result;
	}

	@Override
	public void beforeStep(StepExecution stepExecution) {
		if(StringUtils.hasText(checkColumn)) {

			if(overrideValue != null && overrideValue >= 0) {
				this.incrementalMin = overrideValue;
			}
			else {
				String jobName = stepExecution.getJobExecution().getJobInstance().getJobName();

				// Get the last jobInstance...not the current one
				List<JobInstance> jobInstances = jobExplorer.getJobInstances(jobName, 1, 1);

				if(jobInstances.size() > 0) {
					JobInstance lastInstance = jobInstances.get(jobInstances.size() - 1);

					List<JobExecution> executions = jobExplorer.getJobExecutions(lastInstance);

					JobExecution lastExecution = executions.get(0);

					for (JobExecution execution : executions) {
						if(lastExecution.getEndTime().getTime() < execution.getEndTime().getTime()) {
							lastExecution = execution;
						}
					}

					if(lastExecution.getExecutionContext().containsKey(BATCH_INCREMENTAL_MAX_ID)) {
						this.incrementalMin = lastExecution.getExecutionContext().getLong(BATCH_INCREMENTAL_MAX_ID);
					}
					else {
						this.incrementalMin = Long.MIN_VALUE;
					}
				}
				else {
					this.incrementalMin = Long.MIN_VALUE;
				}
			}

			long newMin = jdbcTemplate.queryForObject(String.format("select max(%s) from %s", checkColumn, table), Integer.class);

			stepExecution.getExecutionContext().put(BATCH_INCREMENTAL_MAX_ID, newMin);
		}

		if(StringUtils.hasText(column) && StringUtils.hasText(table)) {
			if(StringUtils.hasText(checkColumn)) {
				Long minResult = jdbcTemplate.queryForObject("SELECT MIN(" + column + ") from " + table + " where " + checkColumn + " > " + this.incrementalMin, Long.class);
				Long maxResult = jdbcTemplate.queryForObject("SELECT MAX(" + column + ") from " + table + " where " + checkColumn + " > " + this.incrementalMin, Long.class);
				this.partitionMin = minResult != null ? minResult : Long.MIN_VALUE;
				this.partitionMax = maxResult != null ? maxResult : Long.MAX_VALUE;
			}
			else {
				Long minResult = jdbcTemplate.queryForObject("SELECT MIN(" + column + ") from " + table, Long.class);
				Long maxResult = jdbcTemplate.queryForObject("SELECT MAX(" + column + ") from " + table, Long.class);
				this.partitionMin = minResult != null ? minResult : Long.MIN_VALUE;
				this.partitionMax = maxResult != null ? maxResult : Long.MAX_VALUE;
			}
		}
	}

	@Override
	public ExitStatus afterStep(StepExecution stepExecution) {
		return stepExecution.getExitStatus();
	}

	@Override
	public void afterPropertiesSet() throws Exception {
		if(!StringUtils.hasText(this.column)) {
			this.column = this.checkColumn;
		}
	}
}