Add option to clear the persistence context in JpaItemWriter

Resolves #804

Author: baezzys

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/JpaItemWriter.java

@@ -43,6 +43,7 @@
  *
  * @author Thomas Risberg
  * @author Mahmoud Ben Hassine
+ * @author Jinwoo Bae
  *
  */
 public class JpaItemWriter<T> implements ItemWriter<T>, InitializingBean {
@@ -53,6 +54,8 @@ public class JpaItemWriter<T> implements ItemWriter<T>, InitializingBean {
 
 	private boolean usePersist = false;
 
+	private boolean clearEntityManager = true;
+
 	/**
 	 * Set the EntityManager to be used internally.
 	 * @param entityManagerFactory the entityManagerFactory to set
@@ -69,6 +72,15 @@ public void setUsePersist(boolean usePersist) {
 		this.usePersist = usePersist;
 	}
 
+	/**
+	 * Flag to indicate that the EntityManager should be cleared and flushed at the end of
+	 * the write (default true).
+	 * @param clearEntityManager the flag value to set
+	 */
+	public void setClearEntityManager(boolean clearEntityManager) {
+		this.clearEntityManager = clearEntityManager;
+	}
+
 	/**
 	 * Check mandatory properties - there must be an entityManagerFactory.
 	 */
@@ -91,6 +103,9 @@ public void write(Chunk<? extends T> items) {
 		}
 		doWrite(entityManager, items);
 		entityManager.flush();
+		if (clearEntityManager) {
+			entityManager.clear();
+		}
 	}
 
 	/**

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/builder/JpaItemWriterBuilder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2021 the original author or authors.
+ * Copyright 2018-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,7 @@
  * A builder for the {@link JpaItemWriter}.
  *
  * @author Mahmoud Ben Hassine
+ * @author Jinwoo Bae
  * @since 4.1
  * @see JpaItemWriter
  */
@@ -33,6 +34,8 @@ public class JpaItemWriterBuilder<T> {
 
 	private boolean usePersist = false;
 
+	private boolean clearEntityManager = true;
+
 	/**
 	 * The JPA {@link EntityManagerFactory} to obtain an entity manager from. Required.
 	 * @param entityManagerFactory the {@link EntityManagerFactory}
@@ -57,6 +60,19 @@ public JpaItemWriterBuilder<T> usePersist(boolean usePersist) {
 		return this;
 	}
 
+	/**
+	 * If set to false, the {@link jakarta.persistence.EntityManager} will not be cleared
+	 * at the end of the chunk.
+	 * @param clearEntityManager defaults to true
+	 * @return this instance for method chaining
+	 * @see org.springframework.batch.item.database.JpaItemWriter#setClearEntityManager(boolean)
+	 */
+	public JpaItemWriterBuilder<T> clearEntityManager(boolean clearEntityManager) {
+		this.clearEntityManager = clearEntityManager;
+
+		return this;
+	}
+
 	/**
 	 * Returns a fully built {@link JpaItemWriter}.
 	 * @return the writer
@@ -67,6 +83,7 @@ public JpaItemWriter<T> build() {
 		JpaItemWriter<T> writer = new JpaItemWriter<>();
 		writer.setEntityManagerFactory(this.entityManagerFactory);
 		writer.setUsePersist(this.usePersist);
+		writer.setClearEntityManager(this.clearEntityManager);
 
 		return writer;
 	}

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/builder/JpaItemWriterBuilderTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2022 the original author or authors.
+ * Copyright 2018-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.
@@ -32,10 +32,12 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 
 /**
  * @author Mahmoud Ben Hassine
+ * @author Jinwoo Bae
  */
 @ExtendWith(MockitoExtension.class)
 class JpaItemWriterBuilderTests {
@@ -71,6 +73,23 @@ void testConfiguration() throws Exception {
 
 		verify(this.entityManager).merge(chunk.getItems().get(0));
 		verify(this.entityManager).merge(chunk.getItems().get(1));
+		verify(this.entityManager).clear();
+	}
+
+	@Test
+	void testConfigurationClearEntityManager() throws Exception {
+		JpaItemWriter<String> itemWriter = new JpaItemWriterBuilder<String>().clearEntityManager(false)
+				.entityManagerFactory(this.entityManagerFactory).build();
+
+		itemWriter.afterPropertiesSet();
+
+		Chunk<String> chunk = Chunk.of("foo", "bar");
+
+		itemWriter.write(chunk);
+
+		verify(this.entityManager).merge(chunk.getItems().get(0));
+		verify(this.entityManager).merge(chunk.getItems().get(1));
+		verify(this.entityManager, never()).clear();
 	}
 
 	@Test