Skip to content

Commit e36a447

Browse files
Robert McNeesfmbenhassine
Robert McNees
authored and
fmbenhassine
committed
Add option to provide a custom JobKeyGenerator in JdbcJobInstanceDao
Resolves #3926
1 parent 2c97974 commit e36a447

File tree

12 files changed

+349
-3
lines changed

12 files changed

+349
-3
lines changed

Diff for: ‎spring-batch-core/src/main/java/org/springframework/batch/core/configuration/annotation/BatchRegistrar.java

+10
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,11 @@ private void registerJobRepository(BeanDefinitionRegistry registry, EnableBatchP
113113
beanDefinitionBuilder.addPropertyReference("incrementerFactory", incrementerFactoryRef);
114114
}
115115

116+
String jobKeyGeneratorRef = batchAnnotation.jobKeyGeneratorRef();
117+
if (registry.containsBeanDefinition(jobKeyGeneratorRef)) {
118+
beanDefinitionBuilder.addPropertyReference("jobKeyGenerator", jobKeyGeneratorRef);
119+
}
120+
116121
String charset = batchAnnotation.charset();
117122
if (charset != null) {
118123
beanDefinitionBuilder.addPropertyValue("charset", Charset.forName(charset));
@@ -165,6 +170,11 @@ private void registerJobExplorer(BeanDefinitionRegistry registry, EnableBatchPro
165170
beanDefinitionBuilder.addPropertyReference("conversionService", conversionServiceRef);
166171
}
167172

173+
String jobKeyGeneratorRef = batchAnnotation.jobKeyGeneratorRef();
174+
if (registry.containsBeanDefinition(jobKeyGeneratorRef)) {
175+
beanDefinitionBuilder.addPropertyReference("jobKeyGenerator", jobKeyGeneratorRef);
176+
}
177+
168178
String charset = batchAnnotation.charset();
169179
if (charset != null) {
170180
beanDefinitionBuilder.addPropertyValue("charset", Charset.forName(charset));

Diff for: ‎spring-batch-core/src/main/java/org/springframework/batch/core/configuration/annotation/EnableBatchProcessing.java

+9
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,15 @@
214214
*/
215215
String incrementerFactoryRef() default "incrementerFactory";
216216

217+
/**
218+
* The generator that determines a unique key for identifying job instance objects
219+
* @return the bean name of the job key generator to use. Defaults to
220+
* {@literal jobKeyGenerator}.
221+
*
222+
* @since 5.1
223+
*/
224+
String jobKeyGeneratorRef() default "jobKeyGenerator";
225+
217226
/**
218227
* The large object handler to use in job repository and job explorer.
219228
* @return the bean name of the lob handler to use. Defaults to {@literal lobHandler}.

Diff for: ‎spring-batch-core/src/main/java/org/springframework/batch/core/configuration/support/DefaultBatchConfiguration.java

+15
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@
2121

2222
import javax.sql.DataSource;
2323

24+
import org.springframework.batch.core.DefaultJobKeyGenerator;
25+
import org.springframework.batch.core.JobInstance;
26+
import org.springframework.batch.core.JobKeyGenerator;
2427
import org.springframework.batch.core.configuration.BatchConfigurationException;
2528
import org.springframework.batch.core.configuration.JobRegistry;
2629
import org.springframework.batch.core.converter.DateToStringConverter;
@@ -129,6 +132,7 @@ public JobRepository jobRepository() throws BatchConfigurationException {
129132
jobRepositoryFactoryBean.setTransactionManager(getTransactionManager());
130133
jobRepositoryFactoryBean.setDatabaseType(getDatabaseType());
131134
jobRepositoryFactoryBean.setIncrementerFactory(getIncrementerFactory());
135+
jobRepositoryFactoryBean.setJobKeyGenerator(getJobKeyGenerator());
132136
jobRepositoryFactoryBean.setClobType(getClobType());
133137
jobRepositoryFactoryBean.setTablePrefix(getTablePrefix());
134138
jobRepositoryFactoryBean.setSerializer(getExecutionContextSerializer());
@@ -167,6 +171,7 @@ public JobExplorer jobExplorer() throws BatchConfigurationException {
167171
jobExplorerFactoryBean.setDataSource(getDataSource());
168172
jobExplorerFactoryBean.setTransactionManager(getTransactionManager());
169173
jobExplorerFactoryBean.setJdbcOperations(getJdbcOperations());
174+
jobExplorerFactoryBean.setJobKeyGenerator(getJobKeyGenerator());
170175
jobExplorerFactoryBean.setCharset(getCharset());
171176
jobExplorerFactoryBean.setTablePrefix(getTablePrefix());
172177
jobExplorerFactoryBean.setLobHandler(getLobHandler());
@@ -348,6 +353,16 @@ protected DataFieldMaxValueIncrementerFactory getIncrementerFactory() {
348353
return new DefaultDataFieldMaxValueIncrementerFactory(getDataSource());
349354
}
350355

356+
/**
357+
* A custom implementation of the {@link JobKeyGenerator}. The default, if not
358+
* injected, is the {@link DefaultJobKeyGenerator}.
359+
* @return the generator that creates the key used in identifying {@link JobInstance}
360+
* objects
361+
*/
362+
protected JobKeyGenerator getJobKeyGenerator() {
363+
return new DefaultJobKeyGenerator();
364+
}
365+
351366
/**
352367
* Return the database type. The default will be introspected from the JDBC meta-data
353368
* of the data source.

Diff for: ‎spring-batch-core/src/main/java/org/springframework/batch/core/explore/support/JobExplorerFactoryBean.java

+19
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121

2222
import javax.sql.DataSource;
2323

24+
import org.springframework.batch.core.DefaultJobKeyGenerator;
25+
import org.springframework.batch.core.JobKeyGenerator;
2426
import org.springframework.batch.core.converter.DateToStringConverter;
2527
import org.springframework.batch.core.converter.LocalDateTimeToStringConverter;
2628
import org.springframework.batch.core.converter.LocalDateToStringConverter;
@@ -77,6 +79,8 @@ protected long getNextKey() {
7779
}
7880
};
7981

82+
private JobKeyGenerator jobKeyGenerator;
83+
8084
private LobHandler lobHandler;
8185

8286
private ExecutionContextSerializer serializer;
@@ -124,6 +128,16 @@ public void setTablePrefix(String tablePrefix) {
124128
this.tablePrefix = tablePrefix;
125129
}
126130

131+
/**
132+
* * Sets the generator for creating the key used in identifying unique {link
133+
* JobInstance} objects
134+
* @param jobKeyGenerator a {@link JobKeyGenerator}
135+
* @since 5.1
136+
*/
137+
public void setJobKeyGenerator(JobKeyGenerator jobKeyGenerator) {
138+
this.jobKeyGenerator = jobKeyGenerator;
139+
}
140+
127141
/**
128142
* The lob handler to use when saving {@link ExecutionContext} instances. Defaults to
129143
* {@code null}, which works for most databases.
@@ -166,6 +180,10 @@ public void afterPropertiesSet() throws Exception {
166180
jdbcOperations = new JdbcTemplate(dataSource);
167181
}
168182

183+
if (jobKeyGenerator == null) {
184+
jobKeyGenerator = new DefaultJobKeyGenerator();
185+
}
186+
169187
if (serializer == null) {
170188
serializer = new DefaultExecutionContextSerializer();
171189
}
@@ -203,6 +221,7 @@ protected JobInstanceDao createJobInstanceDao() throws Exception {
203221
JdbcJobInstanceDao dao = new JdbcJobInstanceDao();
204222
dao.setJdbcTemplate(jdbcOperations);
205223
dao.setJobInstanceIncrementer(incrementer);
224+
dao.setJobKeyGenerator(jobKeyGenerator);
206225
dao.setTablePrefix(tablePrefix);
207226
dao.afterPropertiesSet();
208227
return dao;

Diff for: ‎spring-batch-core/src/main/java/org/springframework/batch/core/repository/dao/JdbcJobInstanceDao.java

+13-1
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@ SELECT COUNT(*)
129129

130130
private DataFieldMaxValueIncrementer jobInstanceIncrementer;
131131

132-
private final JobKeyGenerator<JobParameters> jobKeyGenerator = new DefaultJobKeyGenerator();
132+
private JobKeyGenerator<JobParameters> jobKeyGenerator = new DefaultJobKeyGenerator();
133133

134134
/**
135135
* In this JDBC implementation a job instance id is obtained by asking the
@@ -341,6 +341,18 @@ public void setJobInstanceIncrementer(DataFieldMaxValueIncrementer jobInstanceIn
341341
this.jobInstanceIncrementer = jobInstanceIncrementer;
342342
}
343343

344+
/**
345+
* Setter for {@link JobKeyGenerator} to be used when generating unique identifiers
346+
* for {@link JobInstance} objects.
347+
* @param jobKeyGenerator the {@link JobKeyGenerator}
348+
*
349+
* @since 5.1
350+
*/
351+
public void setJobKeyGenerator(JobKeyGenerator jobKeyGenerator) {
352+
Assert.notNull(jobKeyGenerator, "jobKeyGenerator must not be null.");
353+
this.jobKeyGenerator = jobKeyGenerator;
354+
}
355+
344356
@Override
345357
public void afterPropertiesSet() throws Exception {
346358
super.afterPropertiesSet();

Diff for: ‎spring-batch-core/src/main/java/org/springframework/batch/core/repository/support/JobRepositoryFactoryBean.java

+19
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
import org.apache.commons.logging.Log;
2727
import org.apache.commons.logging.LogFactory;
2828

29+
import org.springframework.batch.core.DefaultJobKeyGenerator;
30+
import org.springframework.batch.core.JobKeyGenerator;
2931
import org.springframework.batch.core.converter.DateToStringConverter;
3032
import org.springframework.batch.core.converter.LocalDateTimeToStringConverter;
3133
import org.springframework.batch.core.converter.LocalDateToStringConverter;
@@ -87,6 +89,8 @@ public class JobRepositoryFactoryBean extends AbstractJobRepositoryFactoryBean i
8789

8890
private DataFieldMaxValueIncrementerFactory incrementerFactory;
8991

92+
private JobKeyGenerator jobKeyGenerator;
93+
9094
private int maxVarCharLength = AbstractJdbcBatchMetadataDao.DEFAULT_EXIT_MESSAGE_LENGTH;
9195

9296
private LobHandler lobHandler;
@@ -182,6 +186,16 @@ public void setIncrementerFactory(DataFieldMaxValueIncrementerFactory incremente
182186
this.incrementerFactory = incrementerFactory;
183187
}
184188

189+
/**
190+
* * Sets the generator for creating the key used in identifying unique {link
191+
* JobInstance} objects
192+
* @param jobKeyGenerator a {@link JobKeyGenerator}
193+
* @since 5.1
194+
*/
195+
public void setJobKeyGenerator(JobKeyGenerator jobKeyGenerator) {
196+
this.jobKeyGenerator = jobKeyGenerator;
197+
}
198+
185199
/**
186200
* Set the {@link Charset} to use when serializing/deserializing the execution
187201
* context. Defaults to "UTF-8". Must not be {@code null}.
@@ -218,6 +232,10 @@ public void afterPropertiesSet() throws Exception {
218232
incrementerFactory = new DefaultDataFieldMaxValueIncrementerFactory(dataSource);
219233
}
220234

235+
if (jobKeyGenerator == null) {
236+
jobKeyGenerator = new DefaultJobKeyGenerator();
237+
}
238+
221239
if (databaseType == null) {
222240
databaseType = DatabaseType.fromMetaData(dataSource).name();
223241
if (logger.isInfoEnabled()) {
@@ -262,6 +280,7 @@ protected JobInstanceDao createJobInstanceDao() throws Exception {
262280
JdbcJobInstanceDao dao = new JdbcJobInstanceDao();
263281
dao.setJdbcTemplate(jdbcOperations);
264282
dao.setJobInstanceIncrementer(incrementerFactory.getIncrementer(databaseType, tablePrefix + "JOB_SEQ"));
283+
dao.setJobKeyGenerator(jobKeyGenerator);
265284
dao.setTablePrefix(tablePrefix);
266285
dao.afterPropertiesSet();
267286
return dao;

Diff for: ‎spring-batch-core/src/test/java/org/springframework/batch/core/configuration/annotation/BatchRegistrarTests.java

+64
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424

2525
import org.springframework.aop.Advisor;
2626
import org.springframework.aop.framework.Advised;
27+
import org.springframework.batch.core.DefaultJobKeyGenerator;
28+
import org.springframework.batch.core.JobKeyGenerator;
2729
import org.springframework.batch.core.configuration.JobRegistry;
2830
import org.springframework.batch.core.explore.JobExplorer;
2931
import org.springframework.batch.core.launch.JobLauncher;
@@ -167,6 +169,35 @@ void testDefaultInfrastructureBeansRegistration() {
167169
Assertions.assertNotNull(jobOperator);
168170
}
169171

172+
@Test
173+
@DisplayName("When no JobKeyGenerator is provided the default implementation should be used")
174+
public void testDefaultJobKeyGeneratorConfiguration() {
175+
AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(JobConfiguration.class);
176+
177+
JobRepository jobRepository = context.getBean(JobRepository.class);
178+
JdbcJobInstanceDao jobInstanceDao = (JdbcJobInstanceDao) ReflectionTestUtils.getField(jobRepository,
179+
"jobInstanceDao");
180+
JobKeyGenerator jobKeyGenerator = (JobKeyGenerator) ReflectionTestUtils.getField(jobInstanceDao,
181+
"jobKeyGenerator");
182+
183+
Assertions.assertEquals(DefaultJobKeyGenerator.class, jobKeyGenerator.getClass());
184+
}
185+
186+
@Test
187+
@DisplayName("When a custom JobKeyGenerator implementation is found that should be used")
188+
public void testCustomJobKeyGeneratorConfiguration() {
189+
AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(
190+
CustomJobKeyGeneratorConfiguration.class);
191+
192+
JobRepository jobRepository = context.getBean(JobRepository.class);
193+
JdbcJobInstanceDao jobInstanceDao = (JdbcJobInstanceDao) ReflectionTestUtils.getField(jobRepository,
194+
"jobInstanceDao");
195+
JobKeyGenerator jobKeyGenerator = (JobKeyGenerator) ReflectionTestUtils.getField(jobInstanceDao,
196+
"jobKeyGenerator");
197+
Assertions.assertEquals(CustomJobKeyGeneratorConfiguration.TestCustomJobKeyGenerator.class,
198+
jobKeyGenerator.getClass());
199+
}
200+
170201
@Configuration
171202
@EnableBatchProcessing
172203
public static class JobConfigurationWithoutDataSource {
@@ -253,6 +284,39 @@ public JdbcTransactionManager batchTransactionManager(DataSource dataSource) {
253284

254285
}
255286

287+
@Configuration
288+
@EnableBatchProcessing
289+
public static class CustomJobKeyGeneratorConfiguration {
290+
291+
@Bean
292+
public DataSource dataSource() {
293+
return new EmbeddedDatabaseBuilder().setType(EmbeddedDatabaseType.HSQL)
294+
.addScript("/org/springframework/batch/core/schema-hsqldb.sql")
295+
.generateUniqueName(true)
296+
.build();
297+
}
298+
299+
@Bean
300+
public JdbcTransactionManager transactionManager(DataSource dataSource) {
301+
return new JdbcTransactionManager(dataSource);
302+
}
303+
304+
@Bean
305+
public JobKeyGenerator jobKeyGenerator() {
306+
return new TestCustomJobKeyGenerator();
307+
}
308+
309+
private class TestCustomJobKeyGenerator implements JobKeyGenerator {
310+
311+
@Override
312+
public String generateKey(Object source) {
313+
return "1";
314+
}
315+
316+
}
317+
318+
}
319+
256320
private PlatformTransactionManager getTransactionManagerSetOnJobRepository(JobRepository jobRepository) {
257321
Advised target = (Advised) jobRepository; // proxy created by
258322
// AbstractJobRepositoryFactoryBean

Diff for: ‎spring-batch-core/src/test/java/org/springframework/batch/core/explore/support/JobExplorerFactoryBeanTests.java

+27-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2006-2022 the original author or authors.
2+
* Copyright 2006-2023 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -24,6 +24,8 @@
2424

2525
import org.springframework.aop.Advisor;
2626
import org.springframework.aop.framework.Advised;
27+
import org.springframework.batch.core.DefaultJobKeyGenerator;
28+
import org.springframework.batch.core.JobKeyGenerator;
2729
import org.springframework.batch.core.explore.JobExplorer;
2830
import org.springframework.jdbc.core.JdbcOperations;
2931
import org.springframework.jdbc.core.JdbcTemplate;
@@ -129,4 +131,28 @@ public void testCustomTransactionAttributesSource() throws Exception {
129131
}
130132
}
131133

134+
@Test
135+
public void testDefaultJobKeyGenerator() throws Exception {
136+
this.factory.afterPropertiesSet();
137+
JobKeyGenerator jobKeyGenerator = (JobKeyGenerator) ReflectionTestUtils.getField(factory, "jobKeyGenerator");
138+
Assertions.assertEquals(DefaultJobKeyGenerator.class, jobKeyGenerator.getClass());
139+
}
140+
141+
@Test
142+
public void testCustomJobKeyGenerator() throws Exception {
143+
factory.setJobKeyGenerator(new CustomJobKeyGenerator());
144+
this.factory.afterPropertiesSet();
145+
JobKeyGenerator jobKeyGenerator = (JobKeyGenerator) ReflectionTestUtils.getField(factory, "jobKeyGenerator");
146+
Assertions.assertEquals(CustomJobKeyGenerator.class, jobKeyGenerator.getClass());
147+
}
148+
149+
class CustomJobKeyGenerator implements JobKeyGenerator<String> {
150+
151+
@Override
152+
public String generateKey(String source) {
153+
return "1";
154+
}
155+
156+
}
157+
132158
}

0 commit comments

Comments
 (0)