diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/data/MongoItemWriter.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/data/MongoItemWriter.java index ed11a152d..819e7027e 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/data/MongoItemWriter.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/data/MongoItemWriter.java @@ -53,10 +53,12 @@ import org.springframework.util.StringUtils; * * @author Michael Minella * @author Parikshit Dutta + * @author Mahmoud Ben Hassine * */ public class MongoItemWriter implements ItemWriter, InitializingBean { + private static final String ID_KEY = "_id"; private MongoOperations template; private final Object bufferKey; private String collection; @@ -130,46 +132,50 @@ public class MongoItemWriter implements ItemWriter, InitializingBean { * @param items the list of items to be persisted. */ protected void doWrite(List items) { - if(! CollectionUtils.isEmpty(items)) { - if(delete) { - if(StringUtils.hasText(collection)) { - for (Object object : items) { - template.remove(object, collection); - } - } - else { - for (Object object : items) { - template.remove(object); - } - } + if (!CollectionUtils.isEmpty(items)) { + if (this.delete) { + delete(items); } else { - BulkOperations bulkOperations = null; - - if(StringUtils.hasText(collection)) { - bulkOperations = template.bulkOps(BulkMode.ORDERED, collection); - } - else { - bulkOperations = template.bulkOps(BulkMode.ORDERED, ClassUtils.getUserClass(items.get(0))); - } - - for (Object object : items) { - Document document = new Document(); - - MongoConverter mongoConverter = template.getConverter(); - mongoConverter.write(object, document); - - Query query = new Query(); - query.addCriteria(Criteria.where("_id").is((document.get("_id") != null) - ? document.get("_id") : new ObjectId())); - - bulkOperations.replaceOne(query, document, new FindAndReplaceOptions().upsert()); - } - bulkOperations.execute(); + saveOrUpdate(items); } } } + private void delete(List items) { + if (StringUtils.hasText(this.collection)) { + for (Object item : items) { + this.template.remove(item, this.collection); + } + } + else { + for (Object item : items) { + this.template.remove(item); + } + } + } + + private void saveOrUpdate(List items) { + BulkOperations bulkOperations; + BulkMode bulkMode = BulkMode.ORDERED; + if (StringUtils.hasText(this.collection)) { + bulkOperations = this.template.bulkOps(bulkMode, this.collection); + } + else { + bulkOperations = this.template.bulkOps(bulkMode, ClassUtils.getUserClass(items.get(0))); + } + MongoConverter mongoConverter = this.template.getConverter(); + FindAndReplaceOptions upsert = new FindAndReplaceOptions().upsert(); + for (Object item : items) { + Document document = new Document(); + mongoConverter.write(item, document); + Object objectId = document.get(ID_KEY) != null ? document.get(ID_KEY) : new ObjectId(); + Query query = new Query().addCriteria(Criteria.where(ID_KEY).is(objectId)); + bulkOperations.replaceOne(query, document, upsert); + } + bulkOperations.execute(); + } + private boolean transactionActive() { return TransactionSynchronizationManager.isActualTransactionActive(); } diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/MongoItemWriterTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/MongoItemWriterTests.java index 3f6e95fb5..b80fb0b59 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/MongoItemWriterTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/MongoItemWriterTests.java @@ -38,6 +38,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.fail; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.when; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; @@ -48,6 +49,7 @@ import static org.mockito.Mockito.verifyZeroInteractions; /** * @author Michael Minella * @author Parikshit Dutta + * @author Mahmoud Ben Hassine */ @SuppressWarnings("serial") public class MongoItemWriterTests { @@ -111,7 +113,7 @@ public class MongoItemWriterTests { writer.write(items); - verify(template).bulkOps(any(), anyString()); + verify(template).bulkOps(any(), eq("collection")); verify(bulkOperations, times(2)).replaceOne(any(Query.class), any(Object.class), any()); } @@ -163,7 +165,7 @@ public class MongoItemWriterTests { return null; }); - verify(template).bulkOps(any(), anyString()); + verify(template).bulkOps(any(), eq("collection")); verify(bulkOperations, times(2)).replaceOne(any(Query.class), any(Object.class), any()); } diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/builder/MongoItemWriterBuilderTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/builder/MongoItemWriterBuilderTests.java index 7ec9fe84e..39838eb08 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/builder/MongoItemWriterBuilderTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/builder/MongoItemWriterBuilderTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2017 the original author or authors. + * Copyright 2017-2020 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. @@ -19,31 +19,48 @@ package org.springframework.batch.item.data.builder; import java.util.Arrays; import java.util.List; +import org.bson.Document; import org.junit.Before; import org.junit.Test; import org.mockito.Mock; import org.mockito.MockitoAnnotations; import org.springframework.batch.item.data.MongoItemWriter; +import org.springframework.data.mongodb.core.BulkOperations; import org.springframework.data.mongodb.core.MongoOperations; +import org.springframework.data.mongodb.core.convert.MongoConverter; +import org.springframework.data.mongodb.core.query.Query; import static org.junit.Assert.assertEquals; import static org.junit.Assert.fail; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; /** * @author Glenn Renfro + * @author Mahmoud Ben Hassine */ public class MongoItemWriterBuilderTests { @Mock private MongoOperations template; + @Mock + private BulkOperations bulkOperations; + @Mock + private MongoConverter mongoConverter; private List items; @Before public void setUp() throws Exception { MockitoAnnotations.initMocks(this); + when(this.template.bulkOps(any(), anyString())).thenReturn(this.bulkOperations); + when(this.template.bulkOps(any(), any(Class.class))).thenReturn(this.bulkOperations); + when(this.template.getConverter()).thenReturn(this.mongoConverter); this.items = Arrays.asList("foo", "bar"); } @@ -52,8 +69,10 @@ public class MongoItemWriterBuilderTests { MongoItemWriter writer = new MongoItemWriterBuilder().template(this.template).build(); writer.write(this.items); - verify(this.template).save(this.items.get(0)); - verify(this.template).save(this.items.get(1)); + verify(this.template).bulkOps(any(), any(Class.class)); + verify(this.mongoConverter).write(eq(this.items.get(0)), any(Document.class)); + verify(this.mongoConverter).write(eq(this.items.get(1)), any(Document.class)); + verify(this.bulkOperations, times(2)).replaceOne(any(Query.class), any(Object.class), any()); verify(this.template, never()).remove(this.items.get(0)); verify(this.template, never()).remove(this.items.get(1)); } @@ -68,8 +87,8 @@ public class MongoItemWriterBuilderTests { verify(this.template).remove(this.items.get(0)); verify(this.template).remove(this.items.get(1)); - verify(this.template, never()).save(this.items.get(0)); - verify(this.template, never()).save(this.items.get(1)); + verify(this.template, never()).bulkOps(any(), any(Class.class)); + verify(this.mongoConverter, never()).write(any(), any()); } @Test @@ -80,8 +99,10 @@ public class MongoItemWriterBuilderTests { writer.write(this.items); - verify(this.template).save(this.items.get(0), "collection"); - verify(this.template).save(this.items.get(1), "collection"); + verify(this.template).bulkOps(any(), eq("collection")); + verify(this.mongoConverter).write(eq(this.items.get(0)), any(Document.class)); + verify(this.mongoConverter).write(eq(this.items.get(1)), any(Document.class)); + verify(this.bulkOperations, times(2)).replaceOne(any(Query.class), any(Object.class), any()); verify(this.template, never()).remove(this.items.get(0), "collection"); verify(this.template, never()).remove(this.items.get(1), "collection"); }