001/*
002 * Copyright 2006-2018 the original author or authors.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *      https://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016package org.springframework.batch.test;
017
018import java.sql.ResultSet;
019import java.sql.SQLException;
020import java.util.ArrayList;
021import java.util.Collection;
022import java.util.Collections;
023import java.util.List;
024
025import javax.sql.DataSource;
026
027import org.springframework.batch.core.JobExecution;
028import org.springframework.batch.core.JobInstance;
029import org.springframework.batch.core.JobParameter;
030import org.springframework.batch.core.JobParameters;
031import org.springframework.batch.core.JobParametersIncrementer;
032import org.springframework.batch.core.StepExecution;
033import org.springframework.batch.core.repository.JobExecutionAlreadyRunningException;
034import org.springframework.batch.core.repository.JobInstanceAlreadyCompleteException;
035import org.springframework.batch.core.repository.JobRepository;
036import org.springframework.batch.core.repository.JobRestartException;
037import org.springframework.batch.core.repository.dao.AbstractJdbcBatchMetadataDao;
038import org.springframework.beans.factory.InitializingBean;
039import org.springframework.beans.factory.annotation.Autowired;
040import org.springframework.dao.DataAccessException;
041import org.springframework.jdbc.core.JdbcOperations;
042import org.springframework.jdbc.core.JdbcTemplate;
043import org.springframework.jdbc.core.RowMapper;
044import org.springframework.lang.Nullable;
045import org.springframework.util.Assert;
046
047/**
048 * Convenience class for creating and removing {@link JobExecution} instances
049 * from a database. Typical usage in test case would be to create instances
050 * before a transaction, save the result, and then use it to remove them after
051 * the transaction.
052 *
053 * @author Dave Syer
054 * @author Mahmoud Ben Hassine
055 */
056public class JobRepositoryTestUtils extends AbstractJdbcBatchMetadataDao implements InitializingBean {
057
058        private JobRepository jobRepository;
059
060        private JobParametersIncrementer jobParametersIncrementer = new JobParametersIncrementer() {
061
062                Long count = 0L;
063
064                @Override
065                public JobParameters getNext(@Nullable JobParameters parameters) {
066                        return new JobParameters(Collections.singletonMap("count", new JobParameter(count++)));
067                }
068
069        };
070
071        private JdbcOperations jdbcTemplate;
072
073        /**
074         * @see InitializingBean#afterPropertiesSet()
075         */
076        @Override
077        public void afterPropertiesSet() throws Exception {
078                Assert.notNull(jobRepository, "JobRepository must be set");
079                Assert.notNull(jdbcTemplate, "DataSource must be set");
080        }
081
082        /**
083         * Default constructor.
084         */
085        public JobRepositoryTestUtils() {
086        }
087
088        /**
089         * Create a {@link JobRepositoryTestUtils} with all its mandatory
090         * properties.
091         *
092         * @param jobRepository a {@link JobRepository} backed by a database
093         * @param dataSource a {@link DataSource}
094         */
095        public JobRepositoryTestUtils(JobRepository jobRepository, DataSource dataSource) {
096                super();
097                this.jobRepository = jobRepository;
098                setDataSource(dataSource);
099        }
100
101        @Autowired
102        public final void setDataSource(DataSource dataSource) {
103                jdbcTemplate = new JdbcTemplate(dataSource);
104        }
105
106        /**
107         * @param jobParametersIncrementer the jobParametersIncrementer to set
108         */
109        public void setJobParametersIncrementer(JobParametersIncrementer jobParametersIncrementer) {
110                this.jobParametersIncrementer = jobParametersIncrementer;
111        }
112
113        /**
114         * @param jobRepository the jobRepository to set
115         */
116        @Autowired
117        public void setJobRepository(JobRepository jobRepository) {
118                this.jobRepository = jobRepository;
119        }
120
121        /**
122         * Use the {@link JobRepository} to create some {@link JobExecution}
123         * instances each with the given job name and each having step executions
124         * with the given step names.
125         *
126         * @param jobName the name of the job
127         * @param stepNames the names of the step executions
128         * @param count the required number of instances of {@link JobExecution} to
129         * create
130         * @return a collection of {@link JobExecution}
131         *
132         * @throws JobExecutionAlreadyRunningException thrown if Job is already running.
133         * @throws JobRestartException thrown if Job is not restartable.
134         * @throws JobInstanceAlreadyCompleteException thrown if Job Instance is already complete.
135         */
136        public List<JobExecution> createJobExecutions(String jobName, String[] stepNames, int count)
137                        throws JobExecutionAlreadyRunningException, JobRestartException, JobInstanceAlreadyCompleteException {
138                List<JobExecution> list = new ArrayList<JobExecution>();
139                JobParameters jobParameters = new JobParameters();
140                for (int i = 0; i < count; i++) {
141                        JobExecution jobExecution = jobRepository.createJobExecution(jobName, jobParametersIncrementer
142                                        .getNext(jobParameters));
143                        list.add(jobExecution);
144                        for (String stepName : stepNames) {
145                                jobRepository.add(jobExecution.createStepExecution(stepName));
146                        }
147                }
148                return list;
149        }
150
151        /**
152         * Use the {@link JobRepository} to create some {@link JobExecution}
153         * instances each with a single step execution.
154         *
155         * @param count the required number of instances of {@link JobExecution} to
156         * create
157         * @return a collection of {@link JobExecution}
158         * @throws JobExecutionAlreadyRunningException thrown if Job is already running.
159         * @throws JobRestartException thrown if Job is not restartable.
160         * @throws JobInstanceAlreadyCompleteException thrown if Job Instance is already complete.
161         */
162        public List<JobExecution> createJobExecutions(int count) throws JobExecutionAlreadyRunningException,
163        JobRestartException, JobInstanceAlreadyCompleteException {
164                return createJobExecutions("job", new String[] { "step" }, count);
165        }
166
167        /**
168         * Remove the {@link JobExecution} instances, and all associated
169         * {@link JobInstance} and {@link StepExecution} instances from the standard
170         * RDBMS locations used by Spring Batch.
171         *
172         * @param list a list of {@link JobExecution}
173         * @throws DataAccessException if there is a problem
174         */
175        public void removeJobExecutions(Collection<JobExecution> list) throws DataAccessException {
176                for (JobExecution jobExecution : list) {
177                        List<Long> stepExecutionIds = jdbcTemplate.query(
178                                        getQuery("select STEP_EXECUTION_ID from %PREFIX%STEP_EXECUTION where JOB_EXECUTION_ID=?"),
179                                        new RowMapper<Long>() {
180                                                @Override
181                                                public Long mapRow(ResultSet rs, int rowNum) throws SQLException {
182                                                        return rs.getLong(1);
183                                                }
184                                        }, jobExecution.getId());
185                        for (Long stepExecutionId : stepExecutionIds) {
186                                jdbcTemplate.update(getQuery("delete from %PREFIX%STEP_EXECUTION_CONTEXT where STEP_EXECUTION_ID=?"),
187                                                stepExecutionId);
188                                jdbcTemplate.update(getQuery("delete from %PREFIX%STEP_EXECUTION where STEP_EXECUTION_ID=?"),
189                                                stepExecutionId);
190                        }
191                        jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_EXECUTION_CONTEXT where JOB_EXECUTION_ID=?"),
192                                        jobExecution.getId());
193                        jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_EXECUTION_PARAMS where JOB_EXECUTION_ID=?"), jobExecution
194                                        .getId());
195                        jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_EXECUTION where JOB_EXECUTION_ID=?"), jobExecution
196                                        .getId());
197                }
198                for (JobExecution jobExecution : list) {
199                        jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_INSTANCE where JOB_INSTANCE_ID=?"), jobExecution
200                                        .getJobId());
201                }
202        }
203
204        /**
205         * Remove all the {@link JobExecution} instances, and all associated
206         * {@link JobInstance} and {@link StepExecution} instances from the standard
207         * RDBMS locations used by Spring Batch.
208         *
209         * @throws DataAccessException if there is a problem
210         */
211        public void removeJobExecutions() throws DataAccessException {
212                jdbcTemplate.update(getQuery("delete from %PREFIX%STEP_EXECUTION_CONTEXT"));
213                jdbcTemplate.update(getQuery("delete from %PREFIX%STEP_EXECUTION"));
214                jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_EXECUTION_CONTEXT"));
215                jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_EXECUTION_PARAMS"));
216                jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_EXECUTION"));
217                jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_INSTANCE"));
218
219        }
220
221}