Add option to provide custom JobKeyGenerator in JdbcJobInstanceDao

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

@@ -114,6 +114,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));

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

@@ -214,6 +214,13 @@
 	 */
 	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}.
+	 */
+	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

@@ -24,6 +24,9 @@
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
+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;
@@ -136,6 +139,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());
@@ -355,6 +359,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 JobKeyGenerator<JobParameters> jobKeyGenerator = new DefaultJobKeyGenerator();
+	private JobKeyGenerator<JobParameters> jobKeyGenerator;
 
 	/**
 	 * In this JDBC implementation a job instance id is obtained by asking the
@@ -348,10 +348,24 @@ 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) {
+		this.jobKeyGenerator = jobKeyGenerator;
+	}
+
 	@Override
 	public void afterPropertiesSet() throws Exception {
 		super.afterPropertiesSet();
 		Assert.state(jobInstanceIncrementer != null, "jobInstanceIncrementer is required");
+		if (jobKeyGenerator == null) {
+			jobKeyGenerator = new DefaultJobKeyGenerator();
+		}
 	}
 
 	/**

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()) {
@@ -264,6 +282,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

@@ -25,6 +25,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;
@@ -182,6 +184,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 {
@@ -274,6 +305,37 @@ 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

@@ -30,6 +30,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.batch.core.repository.JobRepository;
 import org.springframework.jdbc.core.JdbcOperations;
@@ -130,4 +132,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";
+		}
+
+	}
+
 }

spring-batch-core/src/test/java/org/springframework/batch/core/repository/dao/JdbcJobInstanceDaoCustomTests.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2008-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.
+ * You may obtain a copy of the License at
+ *
+ *      https://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.core.repository.dao;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.batch.core.JobKeyGenerator;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.ApplicationContext;
+import org.springframework.test.context.junit.jupiter.SpringJUnitConfig;
+import org.springframework.test.util.ReflectionTestUtils;
+
+@SpringJUnitConfig(locations = "sql-dao-custom-key-generator-test.xml")
+public class JdbcJobInstanceDaoCustomTests {
+
+	@Autowired
+	ApplicationContext applicationContext;
+
+	@Autowired
+	JobInstanceDao jobInstanceDao;
+
+	@Test
+	public void testCustomJobKeyGeneratorIsWired() {
+		Object jobKeyGenerator = applicationContext.getBean("jobKeyGenerator");
+
+		Assertions.assertTrue(jobKeyGenerator != null);
+		Assertions.assertEquals(CustomJobKeyGenerator.class, jobKeyGenerator.getClass());
+	}
+
+	@Test
+	public void testCustomJobKeyGeneratorIsUsed() {
+		JobKeyGenerator jobKeyGenerator = (JobKeyGenerator) ReflectionTestUtils.getField(jobInstanceDao,
+				"jobKeyGenerator");
+		Assertions.assertEquals(CustomJobKeyGenerator.class, jobKeyGenerator.getClass());
+	}
+
+}
+
+class CustomJobKeyGenerator implements JobKeyGenerator<String> {
+
+	@Override
+	public String generateKey(String source) {
+		return "1";
+	}
+
+}