Add option to provide a custom JobKeyGenerator in JdbcJobInstanceDao

Resolves #3926
2023-04-03 14:42:30 -04:00
parent 2c97974366
commit e36a44788d
12 changed files with 349 additions and 3 deletions
--- a/spring-batch-core/src/test/java/org/springframework/batch/core/configuration/annotation/BatchRegistrarTests.java
+++ b/spring-batch-core/src/test/java/org/springframework/batch/core/configuration/annotation/BatchRegistrarTests.java
@@ -24,6 +24,8 @@ import org.mockito.Mockito;

 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 @@ class BatchRegistrarTests {
 		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 @@ class BatchRegistrarTests {

 	}

+	@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
--- a/spring-batch-core/src/test/java/org/springframework/batch/core/explore/support/JobExplorerFactoryBeanTests.java
+++ b/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.mockito.Mockito;

 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 @@ class JobExplorerFactoryBeanTests {
 		}
 	}

+	@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";
+		}
+
+	}
+
 }
--- a/spring-batch-core/src/test/java/org/springframework/batch/core/repository/dao/JdbcJobInstanceDaoCustomTests.java
+++ b/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
+	private ApplicationContext applicationContext;
+
+	@Autowired
+	private 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";
+	}
+
+}
--- a/spring-batch-core/src/test/java/org/springframework/batch/core/repository/dao/JdbcJobInstanceDaoTests.java
+++ b/spring-batch-core/src/test/java/org/springframework/batch/core/repository/dao/JdbcJobInstanceDaoTests.java
@@ -27,13 +27,17 @@ import javax.sql.DataSource;

 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
+
+import org.springframework.batch.core.DefaultJobKeyGenerator;
 import org.springframework.batch.core.JobExecution;
 import org.springframework.batch.core.JobInstance;
+import org.springframework.batch.core.JobKeyGenerator;
 import org.springframework.batch.core.JobParameters;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.test.context.junit.jupiter.SpringJUnitConfig;
 import org.springframework.test.jdbc.JdbcTestUtils;
+import org.springframework.test.util.ReflectionTestUtils;
 import org.springframework.transaction.annotation.Transactional;

@SpringJUnitConfig(locations = "sql-dao-test.xml")
@@ -115,4 +119,11 @@ public class JdbcJobInstanceDaoTests extends AbstractJobInstanceDaoTests {
 		Assertions.assertNull(dao.getJobInstance(jobInstance.getId()));
 	}

+	@Test
+	void testDefaultJobKeyGeneratorIsUsed() {
+		JobKeyGenerator jobKeyGenerator = (JobKeyGenerator) ReflectionTestUtils.getField(jobInstanceDao,
+				"jobKeyGenerator");
+		Assertions.assertEquals(DefaultJobKeyGenerator.class, jobKeyGenerator.getClass());
+	}
+
 }
--- a/spring-batch-core/src/test/java/org/springframework/batch/core/repository/support/JobRepositoryFactoryBeanTests.java
+++ b/spring-batch-core/src/test/java/org/springframework/batch/core/repository/support/JobRepositoryFactoryBeanTests.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.
@@ -30,6 +30,8 @@ import org.mockito.Mockito;

 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.JobParameters;
 import org.springframework.batch.core.repository.ExecutionContextSerializer;
 import org.springframework.batch.core.repository.JobRepository;
@@ -373,6 +375,30 @@ class JobRepositoryFactoryBeanTests {
 		assertNotNull(repository);
 	}

+	@Test
+	public void testDefaultJobKeyGenerator() throws Exception {
+		testCreateRepository();
+		JobKeyGenerator jobKeyGenerator = (JobKeyGenerator) ReflectionTestUtils.getField(factory, "jobKeyGenerator");
+		Assertions.assertEquals(DefaultJobKeyGenerator.class, jobKeyGenerator.getClass());
+	}
+
+	@Test
+	public void testCustomJobKeyGenerator() throws Exception {
+		factory.setJobKeyGenerator(new CustomJobKeyGenerator());
+		testCreateRepository();
+		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";
+		}
+
+	}
+
 	private static class StubIncrementer implements DataFieldMaxValueIncrementer {

 		@Override
--- a/spring-batch-core/src/test/resources/org/springframework/batch/core/repository/dao/sql-dao-custom-key-generator-test.xml
+++ b/spring-batch-core/src/test/resources/org/springframework/batch/core/repository/dao/sql-dao-custom-key-generator-test.xml
@@ -0,0 +1,75 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<beans xmlns="http://www.springframework.org/schema/beans" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+	   xmlns:jdbc="http://www.springframework.org/schema/jdbc"
+	   xsi:schemaLocation="http://www.springframework.org/schema/beans https://www.springframework.org/schema/beans/spring-beans.xsd http://www.springframework.org/schema/jdbc http://www.springframework.org/schema/jdbc/spring-jdbc.xsd">
+
+	<bean id="dataSource" class="org.apache.commons.dbcp2.BasicDataSource">
+		<property name="driverClassName" value="org.hsqldb.jdbcDriver" />
+		<property name="url" value="jdbc:hsqldb:mem:test;sql.enforce_strict_size=true;hsqldb.tx=mvcc" />
+		<property name="username" value="sa" />
+		<property name="password" value="" />
+	</bean>
+
+	<jdbc:initialize-database>
+		<jdbc:script location="org/springframework/batch/core/schema-drop-hsqldb.sql"/>
+		<jdbc:script location="org/springframework/batch/core/schema-hsqldb.sql"/>
+	</jdbc:initialize-database>
+
+	<bean id="transactionManager" class="org.springframework.jdbc.support.JdbcTransactionManager">
+		<property name="dataSource" ref="dataSource" />
+	</bean>
+
+	<bean id="jobRepository" class="org.springframework.batch.core.repository.support.SimpleJobRepository">
+		<constructor-arg ref="jobInstanceDao" />
+		<constructor-arg ref="jobExecutionDao" />
+		<constructor-arg ref="stepExecutionDao" />
+		<constructor-arg ref="executionContextDao" />
+	</bean>
+
+	<bean id="executionContextDao" class="org.springframework.batch.core.repository.dao.JdbcExecutionContextDao">
+		<property name="jdbcTemplate" ref="jdbcTemplate" />
+		<property name="serializer" ref="serializer"/>
+	</bean>
+
+	<bean id="jobInstanceDao" class="org.springframework.batch.core.repository.dao.JdbcJobInstanceDao">
+		<property name="jdbcTemplate" ref="jdbcTemplate" />
+		<property name="jobInstanceIncrementer" ref="jobInstanceIncrementer" />
+		<property name="jobKeyGenerator" ref="jobKeyGenerator" />
+	</bean>
+
+	<bean id="jobKeyGenerator" class="org.springframework.batch.core.repository.dao.CustomJobKeyGenerator" />
+
+	<bean id="jobExecutionDao" class="org.springframework.batch.core.repository.dao.JdbcJobExecutionDao">
+		<property name="jdbcTemplate" ref="jdbcTemplate" />
+		<property name="jobExecutionIncrementer" ref="jobExecutionIncrementer" />
+	</bean>
+
+	<bean id="stepExecutionDao" class="org.springframework.batch.core.repository.dao.JdbcStepExecutionDao">
+		<property name="jdbcTemplate" ref="jdbcTemplate" />
+		<property name="stepExecutionIncrementer" ref="stepExecutionIncrementer" />
+	</bean>
+
+	<bean id="jdbcTemplate" class="org.springframework.jdbc.core.JdbcTemplate">
+		<constructor-arg ref="dataSource" />
+	</bean>
+
+	<bean id="serializer" class="org.springframework.batch.core.repository.dao.DefaultExecutionContextSerializer"/>
+
+	<bean id="jobInstanceIncrementer" class="org.springframework.jdbc.support.incrementer.HsqlMaxValueIncrementer">
+		<constructor-arg name="dataSource" ref="dataSource"/>
+		<constructor-arg name="incrementerName" value="BATCH_JOB_SEQ"/>
+		<constructor-arg name="columnName" value="ID"/>
+	</bean>
+
+	<bean id="jobExecutionIncrementer" class="org.springframework.jdbc.support.incrementer.HsqlMaxValueIncrementer">
+		<constructor-arg name="dataSource" ref="dataSource"/>
+		<constructor-arg name="incrementerName" value="BATCH_JOB_EXECUTION_SEQ"/>
+		<constructor-arg name="columnName" value="ID"/>
+	</bean>
+
+	<bean id="stepExecutionIncrementer" class="org.springframework.jdbc.support.incrementer.HsqlMaxValueIncrementer">
+		<constructor-arg name="dataSource" ref="dataSource"/>
+		<constructor-arg name="incrementerName" value="BATCH_STEP_EXECUTION_SEQ"/>
+		<constructor-arg name="columnName" value="ID"/>
+	</bean>
+</beans>