Skip to content

Commit 34eda1f

Browse files
baezzysfmbenhassine
authored andcommitted
Add option to clear the persistence context in JpaItemWriter
Resolves #804
1 parent 58a4134 commit 34eda1f

File tree

3 files changed

+53
-2
lines changed

3 files changed

+53
-2
lines changed

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
*
4444
* @author Thomas Risberg
4545
* @author Mahmoud Ben Hassine
46+
* @author Jinwoo Bae
4647
*
4748
*/
4849
public class JpaItemWriter<T> implements ItemWriter<T>, InitializingBean {
@@ -53,6 +54,8 @@ public class JpaItemWriter<T> implements ItemWriter<T>, InitializingBean {
5354

5455
private boolean usePersist = false;
5556

57+
private boolean clearEntityManager = true;
58+
5659
/**
5760
* Set the EntityManager to be used internally.
5861
* @param entityManagerFactory the entityManagerFactory to set
@@ -69,6 +72,15 @@ public void setUsePersist(boolean usePersist) {
6972
this.usePersist = usePersist;
7073
}
7174

75+
/**
76+
* Flag to indicate that the EntityManager should be cleared and flushed at the end of
77+
* the write (default true).
78+
* @param clearEntityManager the flag value to set
79+
*/
80+
public void setClearEntityManager(boolean clearEntityManager) {
81+
this.clearEntityManager = clearEntityManager;
82+
}
83+
7284
/**
7385
* Check mandatory properties - there must be an entityManagerFactory.
7486
*/
@@ -91,6 +103,9 @@ public void write(Chunk<? extends T> items) {
91103
}
92104
doWrite(entityManager, items);
93105
entityManager.flush();
106+
if (clearEntityManager) {
107+
entityManager.clear();
108+
}
94109
}
95110

96111
/**

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

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2018-2021 the original author or authors.
2+
* Copyright 2018-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,7 @@
2424
* A builder for the {@link JpaItemWriter}.
2525
*
2626
* @author Mahmoud Ben Hassine
27+
* @author Jinwoo Bae
2728
* @since 4.1
2829
* @see JpaItemWriter
2930
*/
@@ -33,6 +34,8 @@ public class JpaItemWriterBuilder<T> {
3334

3435
private boolean usePersist = false;
3536

37+
private boolean clearEntityManager = true;
38+
3639
/**
3740
* The JPA {@link EntityManagerFactory} to obtain an entity manager from. Required.
3841
* @param entityManagerFactory the {@link EntityManagerFactory}
@@ -57,6 +60,19 @@ public JpaItemWriterBuilder<T> usePersist(boolean usePersist) {
5760
return this;
5861
}
5962

63+
/**
64+
* If set to false, the {@link jakarta.persistence.EntityManager} will not be cleared
65+
* at the end of the chunk.
66+
* @param clearEntityManager defaults to true
67+
* @return this instance for method chaining
68+
* @see org.springframework.batch.item.database.JpaItemWriter#setClearEntityManager(boolean)
69+
*/
70+
public JpaItemWriterBuilder<T> clearEntityManager(boolean clearEntityManager) {
71+
this.clearEntityManager = clearEntityManager;
72+
73+
return this;
74+
}
75+
6076
/**
6177
* Returns a fully built {@link JpaItemWriter}.
6278
* @return the writer
@@ -67,6 +83,7 @@ public JpaItemWriter<T> build() {
6783
JpaItemWriter<T> writer = new JpaItemWriter<>();
6884
writer.setEntityManagerFactory(this.entityManagerFactory);
6985
writer.setUsePersist(this.usePersist);
86+
writer.setClearEntityManager(this.clearEntityManager);
7087

7188
return writer;
7289
}

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

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2018-2022 the original author or authors.
2+
* Copyright 2018-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.
@@ -32,10 +32,12 @@
3232

3333
import static org.junit.jupiter.api.Assertions.assertEquals;
3434
import static org.junit.jupiter.api.Assertions.assertThrows;
35+
import static org.mockito.Mockito.never;
3536
import static org.mockito.Mockito.verify;
3637

3738
/**
3839
* @author Mahmoud Ben Hassine
40+
* @author Jinwoo Bae
3941
*/
4042
@ExtendWith(MockitoExtension.class)
4143
class JpaItemWriterBuilderTests {
@@ -71,6 +73,23 @@ void testConfiguration() throws Exception {
7173

7274
verify(this.entityManager).merge(chunk.getItems().get(0));
7375
verify(this.entityManager).merge(chunk.getItems().get(1));
76+
verify(this.entityManager).clear();
77+
}
78+
79+
@Test
80+
void testConfigurationClearEntityManager() throws Exception {
81+
JpaItemWriter<String> itemWriter = new JpaItemWriterBuilder<String>().clearEntityManager(false)
82+
.entityManagerFactory(this.entityManagerFactory).build();
83+
84+
itemWriter.afterPropertiesSet();
85+
86+
Chunk<String> chunk = Chunk.of("foo", "bar");
87+
88+
itemWriter.write(chunk);
89+
90+
verify(this.entityManager).merge(chunk.getItems().get(0));
91+
verify(this.entityManager).merge(chunk.getItems().get(1));
92+
verify(this.entityManager, never()).clear();
7493
}
7594

7695
@Test

0 commit comments

Comments
 (0)