1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16 package org.springframework.batch.test;
17
18 import java.sql.ResultSet;
19 import java.sql.SQLException;
20 import java.util.ArrayList;
21 import java.util.Collection;
22 import java.util.Collections;
23 import java.util.List;
24
25 import javax.sql.DataSource;
26
27 import org.springframework.batch.core.JobExecution;
28 import org.springframework.batch.core.JobInstance;
29 import org.springframework.batch.core.JobParameter;
30 import org.springframework.batch.core.JobParameters;
31 import org.springframework.batch.core.JobParametersIncrementer;
32 import org.springframework.batch.core.StepExecution;
33 import org.springframework.batch.core.repository.JobExecutionAlreadyRunningException;
34 import org.springframework.batch.core.repository.JobInstanceAlreadyCompleteException;
35 import org.springframework.batch.core.repository.JobRepository;
36 import org.springframework.batch.core.repository.JobRestartException;
37 import org.springframework.batch.core.repository.dao.AbstractJdbcBatchMetadataDao;
38 import org.springframework.beans.factory.InitializingBean;
39 import org.springframework.dao.DataAccessException;
40 import org.springframework.jdbc.core.simple.ParameterizedRowMapper;
41 import org.springframework.jdbc.core.JdbcOperations;
42 import org.springframework.jdbc.core.JdbcTemplate;
43 import org.springframework.util.Assert;
44
45
46
47
48
49
50
51
52
53 public class JobRepositoryTestUtils extends AbstractJdbcBatchMetadataDao implements InitializingBean {
54
55 private JobRepository jobRepository;
56
57 private JobParametersIncrementer jobParametersIncrementer = new JobParametersIncrementer() {
58
59 Long count = 0L;
60
61 @Override
62 public JobParameters getNext(JobParameters parameters) {
63 return new JobParameters(Collections.singletonMap("count", new JobParameter(count++)));
64 }
65
66 };
67
68 private JdbcOperations jdbcTemplate;
69
70
71
72
73 @Override
74 public void afterPropertiesSet() throws Exception {
75 Assert.notNull(jobRepository, "JobRepository must be set");
76 Assert.notNull(jdbcTemplate, "DataSource must be set");
77 }
78
79
80
81
82 public JobRepositoryTestUtils() {
83 }
84
85
86
87
88
89
90
91
92 public JobRepositoryTestUtils(JobRepository jobRepository, DataSource dataSource) {
93 super();
94 this.jobRepository = jobRepository;
95 setDataSource(dataSource);
96 }
97
98 public final void setDataSource(DataSource dataSource) {
99 jdbcTemplate = new JdbcTemplate(dataSource);
100 }
101
102
103
104
105 public void setJobParametersIncrementer(JobParametersIncrementer jobParametersIncrementer) {
106 this.jobParametersIncrementer = jobParametersIncrementer;
107 }
108
109
110
111
112 public void setJobRepository(JobRepository jobRepository) {
113 this.jobRepository = jobRepository;
114 }
115
116
117
118
119
120
121
122
123
124
125
126
127
128 public List<JobExecution> createJobExecutions(String jobName, String[] stepNames, int count)
129 throws JobExecutionAlreadyRunningException, JobRestartException, JobInstanceAlreadyCompleteException {
130 List<JobExecution> list = new ArrayList<JobExecution>();
131 JobParameters jobParameters = new JobParameters();
132 for (int i = 0; i < count; i++) {
133 JobExecution jobExecution = jobRepository.createJobExecution(jobName, jobParametersIncrementer
134 .getNext(jobParameters));
135 list.add(jobExecution);
136 for (String stepName : stepNames) {
137 jobRepository.add(jobExecution.createStepExecution(stepName));
138 }
139 }
140 return list;
141 }
142
143
144
145
146
147
148
149
150
151
152 public List<JobExecution> createJobExecutions(int count) throws JobExecutionAlreadyRunningException,
153 JobRestartException, JobInstanceAlreadyCompleteException {
154 return createJobExecutions("job", new String[] { "step" }, count);
155 }
156
157
158
159
160
161
162
163
164
165 public void removeJobExecutions(Collection<JobExecution> list) throws DataAccessException {
166 for (JobExecution jobExecution : list) {
167 List<Long> stepExecutionIds = jdbcTemplate.query(
168 getQuery("select STEP_EXECUTION_ID from %PREFIX%STEP_EXECUTION where JOB_EXECUTION_ID=?"),
169 new ParameterizedRowMapper<Long>() {
170 @Override
171 public Long mapRow(ResultSet rs, int rowNum) throws SQLException {
172 return rs.getLong(1);
173 }
174 }, jobExecution.getId());
175 for (Long stepExecutionId : stepExecutionIds) {
176 jdbcTemplate.update(getQuery("delete from %PREFIX%STEP_EXECUTION_CONTEXT where STEP_EXECUTION_ID=?"),
177 stepExecutionId);
178 jdbcTemplate.update(getQuery("delete from %PREFIX%STEP_EXECUTION where STEP_EXECUTION_ID=?"),
179 stepExecutionId);
180 }
181 jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_EXECUTION_CONTEXT where JOB_EXECUTION_ID=?"),
182 jobExecution.getId());
183 jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_EXECUTION where JOB_EXECUTION_ID=?"), jobExecution
184 .getId());
185 }
186 for (JobExecution jobExecution : list) {
187 jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_PARAMS where JOB_INSTANCE_ID=?"), jobExecution
188 .getJobId());
189 jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_INSTANCE where JOB_INSTANCE_ID=?"), jobExecution
190 .getJobId());
191 }
192 }
193
194
195
196
197
198
199
200
201 public void removeJobExecutions() throws DataAccessException {
202 jdbcTemplate.update(getQuery("delete from %PREFIX%STEP_EXECUTION_CONTEXT"));
203 jdbcTemplate.update(getQuery("delete from %PREFIX%STEP_EXECUTION"));
204 jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_EXECUTION_CONTEXT"));
205 jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_EXECUTION"));
206 jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_PARAMS"));
207 jdbcTemplate.update(getQuery("delete from %PREFIX%JOB_INSTANCE"));
208
209 }
210
211 }