Skip to content

Commit 6c5cb2b

Browse files
parikshitduttafmbenhassine
authored andcommitted
Use BulkOperations API for Remove Operations in MongoItemWriter
Issue #3737
1 parent e7ea16c commit 6c5cb2b

File tree

3 files changed

+154
-78
lines changed

3 files changed

+154
-78
lines changed

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/data/MongoItemWriter.java

Lines changed: 22 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -143,27 +143,22 @@ protected void doWrite(List<? extends T> items) {
143143
}
144144

145145
private void delete(List<? extends T> items) {
146-
if (StringUtils.hasText(this.collection)) {
147-
for (Object item : items) {
148-
this.template.remove(item, this.collection);
149-
}
150-
}
151-
else {
152-
for (Object item : items) {
153-
this.template.remove(item);
146+
BulkOperations bulkOperations = initBulkOperations(BulkMode.ORDERED, items.get(0));
147+
MongoConverter mongoConverter = this.template.getConverter();
148+
for (Object item : items) {
149+
Document document = new Document();
150+
mongoConverter.write(item, document);
151+
Object objectId = document.get(ID_KEY);
152+
if (objectId != null) {
153+
Query query = new Query().addCriteria(Criteria.where(ID_KEY).is(objectId));
154+
bulkOperations.remove(query);
154155
}
155156
}
157+
bulkOperations.execute();
156158
}
157159

158160
private void saveOrUpdate(List<? extends T> items) {
159-
BulkOperations bulkOperations;
160-
BulkMode bulkMode = BulkMode.ORDERED;
161-
if (StringUtils.hasText(this.collection)) {
162-
bulkOperations = this.template.bulkOps(bulkMode, this.collection);
163-
}
164-
else {
165-
bulkOperations = this.template.bulkOps(bulkMode, ClassUtils.getUserClass(items.get(0)));
166-
}
161+
BulkOperations bulkOperations = initBulkOperations(BulkMode.ORDERED, items.get(0));
167162
MongoConverter mongoConverter = this.template.getConverter();
168163
FindAndReplaceOptions upsert = new FindAndReplaceOptions().upsert();
169164
for (Object item : items) {
@@ -176,6 +171,17 @@ private void saveOrUpdate(List<? extends T> items) {
176171
bulkOperations.execute();
177172
}
178173

174+
private BulkOperations initBulkOperations(BulkMode bulkMode, Object item) {
175+
BulkOperations bulkOperations;
176+
if (StringUtils.hasText(this.collection)) {
177+
bulkOperations = this.template.bulkOps(bulkMode, this.collection);
178+
}
179+
else {
180+
bulkOperations = this.template.bulkOps(bulkMode, ClassUtils.getUserClass(item));
181+
}
182+
return bulkOperations;
183+
}
184+
179185
private boolean transactionActive() {
180186
return TransactionSynchronizationManager.isActualTransactionActive();
181187
}

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/MongoItemWriterTests.java

Lines changed: 81 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -23,12 +23,24 @@
2323
import org.junit.Before;
2424
import org.junit.Test;
2525
import org.mockito.Mock;
26+
import static org.mockito.Mockito.mock;
27+
import static org.mockito.Mockito.spy;
28+
import static org.mockito.Mockito.when;
29+
import static org.mockito.Mockito.doAnswer;
30+
import static org.mockito.Mockito.verify;
31+
import static org.mockito.Mockito.verifyZeroInteractions;
32+
import static org.mockito.Mockito.times;
33+
import static org.mockito.Mockito.never;
2634
import org.mockito.MockitoAnnotations;
2735

2836
import org.springframework.batch.support.transaction.ResourcelessTransactionManager;
37+
import org.springframework.data.mapping.context.MappingContext;
2938
import org.springframework.data.mongodb.core.BulkOperations;
3039
import org.springframework.data.mongodb.core.MongoOperations;
40+
import org.springframework.data.mongodb.core.convert.DbRefResolver;
41+
import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
3142
import org.springframework.data.mongodb.core.convert.MongoConverter;
43+
import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
3244
import org.springframework.data.mongodb.core.query.Query;
3345
import org.springframework.transaction.PlatformTransactionManager;
3446
import org.springframework.transaction.support.TransactionCallback;
@@ -39,12 +51,6 @@
3951
import static org.mockito.ArgumentMatchers.any;
4052
import static org.mockito.ArgumentMatchers.anyString;
4153
import static org.mockito.ArgumentMatchers.eq;
42-
import static org.mockito.Mockito.when;
43-
import static org.mockito.Mockito.doAnswer;
44-
import static org.mockito.Mockito.mock;
45-
import static org.mockito.Mockito.verify;
46-
import static org.mockito.Mockito.times;
47-
import static org.mockito.Mockito.verifyZeroInteractions;
4854

4955
/**
5056
* @author Michael Minella
@@ -60,15 +66,19 @@ public class MongoItemWriterTests {
6066
@Mock
6167
private BulkOperations bulkOperations;
6268
@Mock
63-
private MongoConverter mongoConverter;
69+
DbRefResolver dbRefResolver;
70+
6471
private PlatformTransactionManager transactionManager = new ResourcelessTransactionManager();
6572

6673
@Before
6774
public void setUp() throws Exception {
6875
MockitoAnnotations.initMocks(this);
69-
when(template.bulkOps(any(), anyString())).thenReturn(bulkOperations);
70-
when(template.bulkOps(any(), any(Class.class))).thenReturn(bulkOperations);
71-
when(template.getConverter()).thenReturn(mongoConverter);
76+
when(this.template.bulkOps(any(), anyString())).thenReturn(this.bulkOperations);
77+
when(this.template.bulkOps(any(), any(Class.class))).thenReturn(this.bulkOperations);
78+
79+
MappingContext mappingContext = new MongoMappingContext();
80+
MappingMongoConverter mongoConverter = spy(new MappingMongoConverter(this.dbRefResolver, mappingContext));
81+
when(this.template.getConverter()).thenReturn(mongoConverter);
7282

7383
writer = new MongoItemWriter<>();
7484
writer.setTemplate(template);
@@ -91,9 +101,9 @@ public void testAfterPropertiesSet() throws Exception {
91101

92102
@Test
93103
public void testWriteNoTransactionNoCollection() throws Exception {
94-
List<Object> items = new ArrayList<Object>() {{
95-
add(new Object());
96-
add(new Object());
104+
List<Item> items = new ArrayList<Item>() {{
105+
add(new Item("Foo"));
106+
add(new Item("Bar"));
97107
}};
98108

99109
writer.write(items);
@@ -105,8 +115,8 @@ public void testWriteNoTransactionNoCollection() throws Exception {
105115
@Test
106116
public void testWriteNoTransactionWithCollection() throws Exception {
107117
List<Object> items = new ArrayList<Object>() {{
108-
add(new Object());
109-
add(new Object());
118+
add(new Item("Foo"));
119+
add(new Item("Bar"));
110120
}};
111121

112122
writer.setCollection("collection");
@@ -128,8 +138,8 @@ public void testWriteNoTransactionNoItems() throws Exception {
128138
@Test
129139
public void testWriteTransactionNoCollection() throws Exception {
130140
final List<Object> items = new ArrayList<Object>() {{
131-
add(new Object());
132-
add(new Object());
141+
add(new Item("Foo"));
142+
add(new Item("Bar"));
133143
}};
134144

135145
new TransactionTemplate(transactionManager).execute((TransactionCallback<Void>) status -> {
@@ -149,8 +159,8 @@ public void testWriteTransactionNoCollection() throws Exception {
149159
@Test
150160
public void testWriteTransactionWithCollection() throws Exception {
151161
final List<Object> items = new ArrayList<Object>() {{
152-
add(new Object());
153-
add(new Object());
162+
add(new Item("Foo"));
163+
add(new Item("Bar"));
154164
}};
155165

156166
writer.setCollection("collection");
@@ -172,8 +182,8 @@ public void testWriteTransactionWithCollection() throws Exception {
172182
@Test
173183
public void testWriteTransactionFails() throws Exception {
174184
final List<Object> items = new ArrayList<Object>() {{
175-
add(new Object());
176-
add(new Object());
185+
add(new Item("Foo"));
186+
add(new Item("Bar"));
177187
}};
178188

179189
writer.setCollection("collection");
@@ -204,8 +214,8 @@ public void testWriteTransactionFails() throws Exception {
204214
@Test
205215
public void testWriteTransactionReadOnly() throws Exception {
206216
final List<Object> items = new ArrayList<Object>() {{
207-
add(new Object());
208-
add(new Object());
217+
add(new Item("Foo"));
218+
add(new Item("Bar"));
209219
}};
210220

211221
writer.setCollection("collection");
@@ -229,34 +239,63 @@ public void testWriteTransactionReadOnly() throws Exception {
229239
verifyZeroInteractions(bulkOperations);
230240
}
231241

242+
@Test
243+
public void testRemoveNoObjectIdNoCollection() throws Exception {
244+
writer.setDelete(true);
245+
List<Object> items = new ArrayList<Object>() {{
246+
add(new Item("Foo"));
247+
add(new Item("Bar"));
248+
}};
249+
250+
writer.write(items);
251+
252+
verify(template).bulkOps(any(), any(Class.class));
253+
verify(bulkOperations, never()).remove(any(Query.class));
254+
}
255+
256+
@Test
257+
public void testRemoveNoObjectIdWithCollection() throws Exception {
258+
writer.setDelete(true);
259+
List<Object> items = new ArrayList<Object>() {{
260+
add(new Item("Foo"));
261+
add(new Item("Bar"));
262+
}};
263+
264+
writer.setCollection("collection");
265+
writer.write(items);
266+
267+
verify(template).bulkOps(any(), eq("collection"));
268+
verify(bulkOperations, never()).remove(any(Query.class));
269+
}
270+
232271
@Test
233272
public void testRemoveNoTransactionNoCollection() throws Exception {
234273
writer.setDelete(true);
235274
List<Object> items = new ArrayList<Object>() {{
236-
add(new Object());
237-
add(new Object());
275+
add(new Item(1));
276+
add(new Item(2));
238277
}};
239278

240279
writer.write(items);
241280

242-
verify(template).remove(items.get(0));
243-
verify(template).remove(items.get(1));
281+
verify(template).bulkOps(any(), any(Class.class));
282+
verify(bulkOperations, times(2)).remove(any(Query.class));
244283
}
245284

246285
@Test
247286
public void testRemoveNoTransactionWithCollection() throws Exception {
248287
writer.setDelete(true);
249288
List<Object> items = new ArrayList<Object>() {{
250-
add(new Object());
251-
add(new Object());
289+
add(new Item(1));
290+
add(new Item(2));
252291
}};
253292

254293
writer.setCollection("collection");
255294

256295
writer.write(items);
257296

258-
verify(template).remove(items.get(0), "collection");
259-
verify(template).remove(items.get(1), "collection");
297+
verify(template).bulkOps(any(), eq("collection"));
298+
verify(bulkOperations, times(2)).remove(any(Query.class));
260299
}
261300

262301
// BATCH-2018, test code updated to pass BATCH-3713
@@ -312,3 +351,14 @@ public void testResourceKeyCollision() throws Exception {
312351
}
313352
}
314353
}
354+
355+
class Item {
356+
Integer id;
357+
String name;
358+
public Item(Integer id){
359+
this.id = id;
360+
}
361+
public Item(String name) {
362+
this.name = name;
363+
}
364+
}

0 commit comments

Comments
 (0)