Add option to provide a custom JobKeyGenerator in JdbcJobInstanceDao

Resolves #3926

spring-batch-core/src/main/java/org/springframework/batch/core/configuration/annotation/BatchRegistrar.java

@@ -113,6 +113,11 @@ private void registerJobRepository(BeanDefinitionRegistry registry, EnableBatchP
 			beanDefinitionBuilder.addPropertyReference("incrementerFactory", incrementerFactoryRef);
 		}
 
+		String jobKeyGeneratorRef = batchAnnotation.jobKeyGeneratorRef();
+		if (registry.containsBeanDefinition(jobKeyGeneratorRef)) {
+			beanDefinitionBuilder.addPropertyReference("jobKeyGenerator", jobKeyGeneratorRef);
+		}
+
 		String charset = batchAnnotation.charset();
 		if (charset != null) {
 			beanDefinitionBuilder.addPropertyValue("charset", Charset.forName(charset));
@@ -165,6 +170,11 @@ private void registerJobExplorer(BeanDefinitionRegistry registry, EnableBatchPro
 			beanDefinitionBuilder.addPropertyReference("conversionService", conversionServiceRef);
 		}
 
+		String jobKeyGeneratorRef = batchAnnotation.jobKeyGeneratorRef();
+		if (registry.containsBeanDefinition(jobKeyGeneratorRef)) {
+			beanDefinitionBuilder.addPropertyReference("jobKeyGenerator", jobKeyGeneratorRef);
+		}
+
 		String charset = batchAnnotation.charset();
 		if (charset != null) {
 			beanDefinitionBuilder.addPropertyValue("charset", Charset.forName(charset));

spring-batch-core/src/main/java/org/springframework/batch/core/configuration/annotation/EnableBatchProcessing.java

@@ -214,6 +214,15 @@
 	 */
 	String incrementerFactoryRef() default "incrementerFactory";
 
+	/**
+	 * The generator that determines a unique key for identifying job instance objects
+	 * @return the bean name of the job key generator to use. Defaults to
+	 * {@literal jobKeyGenerator}.
+	 *
+	 * @since 5.1
+	 */
+	String jobKeyGeneratorRef() default "jobKeyGenerator";
+
 	/**
 	 * The large object handler to use in job repository and job explorer.
 	 * @return the bean name of the lob handler to use. Defaults to {@literal lobHandler}.

spring-batch-core/src/main/java/org/springframework/batch/core/configuration/support/DefaultBatchConfiguration.java

@@ -21,6 +21,9 @@
 
 import javax.sql.DataSource;
 
+import org.springframework.batch.core.DefaultJobKeyGenerator;
+import org.springframework.batch.core.JobInstance;
+import org.springframework.batch.core.JobKeyGenerator;
 import org.springframework.batch.core.configuration.BatchConfigurationException;
 import org.springframework.batch.core.configuration.JobRegistry;
 import org.springframework.batch.core.converter.DateToStringConverter;
@@ -129,6 +132,7 @@ public JobRepository jobRepository() throws BatchConfigurationException {
 			jobRepositoryFactoryBean.setTransactionManager(getTransactionManager());
 			jobRepositoryFactoryBean.setDatabaseType(getDatabaseType());
 			jobRepositoryFactoryBean.setIncrementerFactory(getIncrementerFactory());
+			jobRepositoryFactoryBean.setJobKeyGenerator(getJobKeyGenerator());
 			jobRepositoryFactoryBean.setClobType(getClobType());
 			jobRepositoryFactoryBean.setTablePrefix(getTablePrefix());
 			jobRepositoryFactoryBean.setSerializer(getExecutionContextSerializer());
@@ -167,6 +171,7 @@ public JobExplorer jobExplorer() throws BatchConfigurationException {
 		jobExplorerFactoryBean.setDataSource(getDataSource());
 		jobExplorerFactoryBean.setTransactionManager(getTransactionManager());
 		jobExplorerFactoryBean.setJdbcOperations(getJdbcOperations());
+		jobExplorerFactoryBean.setJobKeyGenerator(getJobKeyGenerator());
 		jobExplorerFactoryBean.setCharset(getCharset());
 		jobExplorerFactoryBean.setTablePrefix(getTablePrefix());
 		jobExplorerFactoryBean.setLobHandler(getLobHandler());
@@ -348,6 +353,16 @@ protected DataFieldMaxValueIncrementerFactory getIncrementerFactory() {
 		return new DefaultDataFieldMaxValueIncrementerFactory(getDataSource());
 	}
 
+	/**
+	 * A custom implementation of the {@link JobKeyGenerator}. The default, if not
+	 * injected, is the {@link DefaultJobKeyGenerator}.
+	 * @return the generator that creates the key used in identifying {@link JobInstance}
+	 * objects
+	 */
+	protected JobKeyGenerator getJobKeyGenerator() {
+		return new DefaultJobKeyGenerator();
+	}
+
 	/**
 	 * Return the database type. The default will be introspected from the JDBC meta-data
 	 * of the data source.

spring-batch-core/src/main/java/org/springframework/batch/core/explore/support/JobExplorerFactoryBean.java

@@ -21,6 +21,8 @@
 
 import javax.sql.DataSource;
 
+import org.springframework.batch.core.DefaultJobKeyGenerator;
+import org.springframework.batch.core.JobKeyGenerator;
 import org.springframework.batch.core.converter.DateToStringConverter;
 import org.springframework.batch.core.converter.LocalDateTimeToStringConverter;
 import org.springframework.batch.core.converter.LocalDateToStringConverter;
@@ -77,6 +79,8 @@ protected long getNextKey() {
 		}
 	};
 
+	private JobKeyGenerator jobKeyGenerator;
+
 	private LobHandler lobHandler;
 
 	private ExecutionContextSerializer serializer;
@@ -124,6 +128,16 @@ public void setTablePrefix(String tablePrefix) {
 		this.tablePrefix = tablePrefix;
 	}
 
+	/**
+	 * * Sets the generator for creating the key used in identifying unique {link
+	 * JobInstance} objects
+	 * @param jobKeyGenerator a {@link JobKeyGenerator}
+	 * @since 5.1
+	 */
+	public void setJobKeyGenerator(JobKeyGenerator jobKeyGenerator) {
+		this.jobKeyGenerator = jobKeyGenerator;
+	}
+
 	/**
 	 * The lob handler to use when saving {@link ExecutionContext} instances. Defaults to
 	 * {@code null}, which works for most databases.
@@ -166,6 +180,10 @@ public void afterPropertiesSet() throws Exception {
 			jdbcOperations = new JdbcTemplate(dataSource);
 		}
 
+		if (jobKeyGenerator == null) {
+			jobKeyGenerator = new DefaultJobKeyGenerator();
+		}
+
 		if (serializer == null) {
 			serializer = new DefaultExecutionContextSerializer();
 		}
@@ -203,6 +221,7 @@ protected JobInstanceDao createJobInstanceDao() throws Exception {
 		JdbcJobInstanceDao dao = new JdbcJobInstanceDao();
 		dao.setJdbcTemplate(jdbcOperations);
 		dao.setJobInstanceIncrementer(incrementer);
+		dao.setJobKeyGenerator(jobKeyGenerator);
 		dao.setTablePrefix(tablePrefix);
 		dao.afterPropertiesSet();
 		return dao;

spring-batch-core/src/main/java/org/springframework/batch/core/repository/dao/JdbcJobInstanceDao.java

@@ -129,7 +129,7 @@ SELECT COUNT(*)
 
 	private DataFieldMaxValueIncrementer jobInstanceIncrementer;
 
-	private final JobKeyGenerator<JobParameters> jobKeyGenerator = new DefaultJobKeyGenerator();
+	private JobKeyGenerator<JobParameters> jobKeyGenerator = new DefaultJobKeyGenerator();
 
 	/**
 	 * In this JDBC implementation a job instance id is obtained by asking the
@@ -341,6 +341,18 @@ public void setJobInstanceIncrementer(DataFieldMaxValueIncrementer jobInstanceIn
 		this.jobInstanceIncrementer = jobInstanceIncrementer;
 	}
 
+	/**
+	 * Setter for {@link JobKeyGenerator} to be used when generating unique identifiers
+	 * for {@link JobInstance} objects.
+	 * @param jobKeyGenerator the {@link JobKeyGenerator}
+	 *
+	 * @since 5.1
+	 */
+	public void setJobKeyGenerator(JobKeyGenerator jobKeyGenerator) {
+		Assert.notNull(jobKeyGenerator, "jobKeyGenerator must not be null.");
+		this.jobKeyGenerator = jobKeyGenerator;
+	}
+
 	@Override
 	public void afterPropertiesSet() throws Exception {
 		super.afterPropertiesSet();

spring-batch-core/src/main/java/org/springframework/batch/core/repository/support/JobRepositoryFactoryBean.java

@@ -26,6 +26,8 @@
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
+import org.springframework.batch.core.DefaultJobKeyGenerator;
+import org.springframework.batch.core.JobKeyGenerator;
 import org.springframework.batch.core.converter.DateToStringConverter;
 import org.springframework.batch.core.converter.LocalDateTimeToStringConverter;
 import org.springframework.batch.core.converter.LocalDateToStringConverter;
@@ -87,6 +89,8 @@ public class JobRepositoryFactoryBean extends AbstractJobRepositoryFactoryBean i
 
 	private DataFieldMaxValueIncrementerFactory incrementerFactory;
 
+	private JobKeyGenerator jobKeyGenerator;
+
 	private int maxVarCharLength = AbstractJdbcBatchMetadataDao.DEFAULT_EXIT_MESSAGE_LENGTH;
 
 	private LobHandler lobHandler;
@@ -182,6 +186,16 @@ public void setIncrementerFactory(DataFieldMaxValueIncrementerFactory incremente
 		this.incrementerFactory = incrementerFactory;
 	}
 
+	/**
+	 * * Sets the generator for creating the key used in identifying unique {link
+	 * JobInstance} objects
+	 * @param jobKeyGenerator a {@link JobKeyGenerator}
+	 * @since 5.1
+	 */
+	public void setJobKeyGenerator(JobKeyGenerator jobKeyGenerator) {
+		this.jobKeyGenerator = jobKeyGenerator;
+	}
+
 	/**
 	 * Set the {@link Charset} to use when serializing/deserializing the execution
 	 * context. Defaults to "UTF-8". Must not be {@code null}.
@@ -218,6 +232,10 @@ public void afterPropertiesSet() throws Exception {
 			incrementerFactory = new DefaultDataFieldMaxValueIncrementerFactory(dataSource);
 		}
 
+		if (jobKeyGenerator == null) {
+			jobKeyGenerator = new DefaultJobKeyGenerator();
+		}
+
 		if (databaseType == null) {
 			databaseType = DatabaseType.fromMetaData(dataSource).name();
 			if (logger.isInfoEnabled()) {
@@ -262,6 +280,7 @@ protected JobInstanceDao createJobInstanceDao() throws Exception {
 		JdbcJobInstanceDao dao = new JdbcJobInstanceDao();
 		dao.setJdbcTemplate(jdbcOperations);
 		dao.setJobInstanceIncrementer(incrementerFactory.getIncrementer(databaseType, tablePrefix + "JOB_SEQ"));
+		dao.setJobKeyGenerator(jobKeyGenerator);
 		dao.setTablePrefix(tablePrefix);
 		dao.afterPropertiesSet();
 		return dao;

spring-batch-core/src/test/java/org/springframework/batch/core/configuration/annotation/BatchRegistrarTests.java

@@ -24,6 +24,8 @@
 
 import org.springframework.aop.Advisor;
 import org.springframework.aop.framework.Advised;
+import org.springframework.batch.core.DefaultJobKeyGenerator;
+import org.springframework.batch.core.JobKeyGenerator;
 import org.springframework.batch.core.configuration.JobRegistry;
 import org.springframework.batch.core.explore.JobExplorer;
 import org.springframework.batch.core.launch.JobLauncher;
@@ -167,6 +169,35 @@ void testDefaultInfrastructureBeansRegistration() {
 		Assertions.assertNotNull(jobOperator);
 	}
 
+	@Test
+	@DisplayName("When no JobKeyGenerator is provided the default implementation should be used")
+	public void testDefaultJobKeyGeneratorConfiguration() {
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(JobConfiguration.class);
+
+		JobRepository jobRepository = context.getBean(JobRepository.class);
+		JdbcJobInstanceDao jobInstanceDao = (JdbcJobInstanceDao) ReflectionTestUtils.getField(jobRepository,
+				"jobInstanceDao");
+		JobKeyGenerator jobKeyGenerator = (JobKeyGenerator) ReflectionTestUtils.getField(jobInstanceDao,
+				"jobKeyGenerator");
+
+		Assertions.assertEquals(DefaultJobKeyGenerator.class, jobKeyGenerator.getClass());
+	}
+
+	@Test
+	@DisplayName("When a custom JobKeyGenerator implementation is found that should be used")
+	public void testCustomJobKeyGeneratorConfiguration() {
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(
+				CustomJobKeyGeneratorConfiguration.class);
+
+		JobRepository jobRepository = context.getBean(JobRepository.class);
+		JdbcJobInstanceDao jobInstanceDao = (JdbcJobInstanceDao) ReflectionTestUtils.getField(jobRepository,
+				"jobInstanceDao");
+		JobKeyGenerator jobKeyGenerator = (JobKeyGenerator) ReflectionTestUtils.getField(jobInstanceDao,
+				"jobKeyGenerator");
+		Assertions.assertEquals(CustomJobKeyGeneratorConfiguration.TestCustomJobKeyGenerator.class,
+				jobKeyGenerator.getClass());
+	}
+
 	@Configuration
 	@EnableBatchProcessing
 	public static class JobConfigurationWithoutDataSource {
@@ -253,6 +284,39 @@ public JdbcTransactionManager batchTransactionManager(DataSource dataSource) {
 
 	}
 
+	@Configuration
+	@EnableBatchProcessing
+	public static class CustomJobKeyGeneratorConfiguration {
+
+		@Bean
+		public DataSource dataSource() {
+			return new EmbeddedDatabaseBuilder().setType(EmbeddedDatabaseType.HSQL)
+				.addScript("/org/springframework/batch/core/schema-hsqldb.sql")
+				.generateUniqueName(true)
+				.build();
+		}
+
+		@Bean
+		public JdbcTransactionManager transactionManager(DataSource dataSource) {
+			return new JdbcTransactionManager(dataSource);
+		}
+
+		@Bean
+		public JobKeyGenerator jobKeyGenerator() {
+			return new TestCustomJobKeyGenerator();
+		}
+
+		private class TestCustomJobKeyGenerator implements JobKeyGenerator {
+
+			@Override
+			public String generateKey(Object source) {
+				return "1";
+			}
+
+		}
+
+	}
+
 	private PlatformTransactionManager getTransactionManagerSetOnJobRepository(JobRepository jobRepository) {
 		Advised target = (Advised) jobRepository; // proxy created by
 													// AbstractJobRepositoryFactoryBean

spring-batch-core/src/test/java/org/springframework/batch/core/explore/support/JobExplorerFactoryBeanTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2006-2022 the original author or authors.
+ * Copyright 2006-2023 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.
@@ -24,6 +24,8 @@
 
 import org.springframework.aop.Advisor;
 import org.springframework.aop.framework.Advised;
+import org.springframework.batch.core.DefaultJobKeyGenerator;
+import org.springframework.batch.core.JobKeyGenerator;
 import org.springframework.batch.core.explore.JobExplorer;
 import org.springframework.jdbc.core.JdbcOperations;
 import org.springframework.jdbc.core.JdbcTemplate;
@@ -129,4 +131,28 @@ public void testCustomTransactionAttributesSource() throws Exception {
 		}
 	}
 
+	@Test
+	public void testDefaultJobKeyGenerator() throws Exception {
+		this.factory.afterPropertiesSet();
+		JobKeyGenerator jobKeyGenerator = (JobKeyGenerator) ReflectionTestUtils.getField(factory, "jobKeyGenerator");
+		Assertions.assertEquals(DefaultJobKeyGenerator.class, jobKeyGenerator.getClass());
+	}
+
+	@Test
+	public void testCustomJobKeyGenerator() throws Exception {
+		factory.setJobKeyGenerator(new CustomJobKeyGenerator());
+		this.factory.afterPropertiesSet();
+		JobKeyGenerator jobKeyGenerator = (JobKeyGenerator) ReflectionTestUtils.getField(factory, "jobKeyGenerator");
+		Assertions.assertEquals(CustomJobKeyGenerator.class, jobKeyGenerator.getClass());
+	}
+
+	class CustomJobKeyGenerator implements JobKeyGenerator<String> {
+
+		@Override
+		public String generateKey(String source) {
+			return "1";
+		}
+
+	}
+
 }