Use BulkOperations API for Remove Operations in MongoItemWriter

Issue #3737
This commit is contained in:
Parikshit Dutta
2020-07-15 19:16:16 +05:30
committed by Mahmoud Ben Hassine
parent e7ea16c9bc
commit 6c5cb2b13d
3 changed files with 157 additions and 81 deletions

View File

@@ -143,27 +143,22 @@ public class MongoItemWriter<T> implements ItemWriter<T>, InitializingBean {
}
private void delete(List<? extends T> 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);
BulkOperations bulkOperations = initBulkOperations(BulkMode.ORDERED, items.get(0));
MongoConverter mongoConverter = this.template.getConverter();
for (Object item : items) {
Document document = new Document();
mongoConverter.write(item, document);
Object objectId = document.get(ID_KEY);
if (objectId != null) {
Query query = new Query().addCriteria(Criteria.where(ID_KEY).is(objectId));
bulkOperations.remove(query);
}
}
bulkOperations.execute();
}
private void saveOrUpdate(List<? extends T> 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)));
}
BulkOperations bulkOperations = initBulkOperations(BulkMode.ORDERED, items.get(0));
MongoConverter mongoConverter = this.template.getConverter();
FindAndReplaceOptions upsert = new FindAndReplaceOptions().upsert();
for (Object item : items) {
@@ -176,6 +171,17 @@ public class MongoItemWriter<T> implements ItemWriter<T>, InitializingBean {
bulkOperations.execute();
}
private BulkOperations initBulkOperations(BulkMode bulkMode, Object item) {
BulkOperations bulkOperations;
if (StringUtils.hasText(this.collection)) {
bulkOperations = this.template.bulkOps(bulkMode, this.collection);
}
else {
bulkOperations = this.template.bulkOps(bulkMode, ClassUtils.getUserClass(item));
}
return bulkOperations;
}
private boolean transactionActive() {
return TransactionSynchronizationManager.isActualTransactionActive();
}

View File

@@ -23,12 +23,24 @@ import org.bson.Document;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mock;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.when;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyZeroInteractions;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.never;
import org.mockito.MockitoAnnotations;
import org.springframework.batch.support.transaction.ResourcelessTransactionManager;
import org.springframework.data.mapping.context.MappingContext;
import org.springframework.data.mongodb.core.BulkOperations;
import org.springframework.data.mongodb.core.MongoOperations;
import org.springframework.data.mongodb.core.convert.DbRefResolver;
import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
import org.springframework.data.mongodb.core.convert.MongoConverter;
import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.support.TransactionCallback;
@@ -39,12 +51,6 @@ 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;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verifyZeroInteractions;
/**
* @author Michael Minella
@@ -60,15 +66,19 @@ public class MongoItemWriterTests {
@Mock
private BulkOperations bulkOperations;
@Mock
private MongoConverter mongoConverter;
DbRefResolver dbRefResolver;
private PlatformTransactionManager transactionManager = new ResourcelessTransactionManager();
@Before
public void setUp() throws Exception {
MockitoAnnotations.initMocks(this);
when(template.bulkOps(any(), anyString())).thenReturn(bulkOperations);
when(template.bulkOps(any(), any(Class.class))).thenReturn(bulkOperations);
when(template.getConverter()).thenReturn(mongoConverter);
when(this.template.bulkOps(any(), anyString())).thenReturn(this.bulkOperations);
when(this.template.bulkOps(any(), any(Class.class))).thenReturn(this.bulkOperations);
MappingContext mappingContext = new MongoMappingContext();
MappingMongoConverter mongoConverter = spy(new MappingMongoConverter(this.dbRefResolver, mappingContext));
when(this.template.getConverter()).thenReturn(mongoConverter);
writer = new MongoItemWriter<>();
writer.setTemplate(template);
@@ -91,9 +101,9 @@ public class MongoItemWriterTests {
@Test
public void testWriteNoTransactionNoCollection() throws Exception {
List<Object> items = new ArrayList<Object>() {{
add(new Object());
add(new Object());
List<Item> items = new ArrayList<Item>() {{
add(new Item("Foo"));
add(new Item("Bar"));
}};
writer.write(items);
@@ -105,8 +115,8 @@ public class MongoItemWriterTests {
@Test
public void testWriteNoTransactionWithCollection() throws Exception {
List<Object> items = new ArrayList<Object>() {{
add(new Object());
add(new Object());
add(new Item("Foo"));
add(new Item("Bar"));
}};
writer.setCollection("collection");
@@ -128,8 +138,8 @@ public class MongoItemWriterTests {
@Test
public void testWriteTransactionNoCollection() throws Exception {
final List<Object> items = new ArrayList<Object>() {{
add(new Object());
add(new Object());
add(new Item("Foo"));
add(new Item("Bar"));
}};
new TransactionTemplate(transactionManager).execute((TransactionCallback<Void>) status -> {
@@ -149,8 +159,8 @@ public class MongoItemWriterTests {
@Test
public void testWriteTransactionWithCollection() throws Exception {
final List<Object> items = new ArrayList<Object>() {{
add(new Object());
add(new Object());
add(new Item("Foo"));
add(new Item("Bar"));
}};
writer.setCollection("collection");
@@ -172,8 +182,8 @@ public class MongoItemWriterTests {
@Test
public void testWriteTransactionFails() throws Exception {
final List<Object> items = new ArrayList<Object>() {{
add(new Object());
add(new Object());
add(new Item("Foo"));
add(new Item("Bar"));
}};
writer.setCollection("collection");
@@ -204,8 +214,8 @@ public class MongoItemWriterTests {
@Test
public void testWriteTransactionReadOnly() throws Exception {
final List<Object> items = new ArrayList<Object>() {{
add(new Object());
add(new Object());
add(new Item("Foo"));
add(new Item("Bar"));
}};
writer.setCollection("collection");
@@ -230,33 +240,62 @@ public class MongoItemWriterTests {
}
@Test
public void testRemoveNoTransactionNoCollection() throws Exception {
public void testRemoveNoObjectIdNoCollection() throws Exception {
writer.setDelete(true);
List<Object> items = new ArrayList<Object>() {{
add(new Object());
add(new Object());
add(new Item("Foo"));
add(new Item("Bar"));
}};
writer.write(items);
verify(template).remove(items.get(0));
verify(template).remove(items.get(1));
verify(template).bulkOps(any(), any(Class.class));
verify(bulkOperations, never()).remove(any(Query.class));
}
@Test
public void testRemoveNoObjectIdWithCollection() throws Exception {
writer.setDelete(true);
List<Object> items = new ArrayList<Object>() {{
add(new Item("Foo"));
add(new Item("Bar"));
}};
writer.setCollection("collection");
writer.write(items);
verify(template).bulkOps(any(), eq("collection"));
verify(bulkOperations, never()).remove(any(Query.class));
}
@Test
public void testRemoveNoTransactionNoCollection() throws Exception {
writer.setDelete(true);
List<Object> items = new ArrayList<Object>() {{
add(new Item(1));
add(new Item(2));
}};
writer.write(items);
verify(template).bulkOps(any(), any(Class.class));
verify(bulkOperations, times(2)).remove(any(Query.class));
}
@Test
public void testRemoveNoTransactionWithCollection() throws Exception {
writer.setDelete(true);
List<Object> items = new ArrayList<Object>() {{
add(new Object());
add(new Object());
add(new Item(1));
add(new Item(2));
}};
writer.setCollection("collection");
writer.write(items);
verify(template).remove(items.get(0), "collection");
verify(template).remove(items.get(1), "collection");
verify(template).bulkOps(any(), eq("collection"));
verify(bulkOperations, times(2)).remove(any(Query.class));
}
// BATCH-2018, test code updated to pass BATCH-3713
@@ -312,3 +351,14 @@ public class MongoItemWriterTests {
}
}
}
class Item {
Integer id;
String name;
public Item(Integer id){
this.id = id;
}
public Item(String name) {
this.name = name;
}
}

View File

@@ -23,12 +23,21 @@ import org.bson.Document;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.when;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import org.mockito.MockitoAnnotations;
import org.springframework.batch.item.data.MongoItemWriter;
import org.springframework.data.mapping.context.MappingContext;
import org.springframework.data.mongodb.core.BulkOperations;
import org.springframework.data.mongodb.core.MongoOperations;
import org.springframework.data.mongodb.core.convert.DbRefResolver;
import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
import org.springframework.data.mongodb.core.convert.MongoConverter;
import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
import org.springframework.data.mongodb.core.query.Query;
import static org.junit.Assert.assertEquals;
@@ -36,14 +45,11 @@ 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
* @author Parikshit Dutta
*/
public class MongoItemWriterBuilderTests {
@Mock
@@ -51,60 +57,63 @@ public class MongoItemWriterBuilderTests {
@Mock
private BulkOperations bulkOperations;
@Mock
DbRefResolver dbRefResolver;
private MongoConverter mongoConverter;
private List<String> items;
private List<Item> saveItems;
private List<Item> removeItems;
@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");
MappingContext mappingContext = new MongoMappingContext();
mongoConverter = spy(new MappingMongoConverter(this.dbRefResolver, mappingContext));
when(this.template.getConverter()).thenReturn(mongoConverter);
this.saveItems = Arrays.asList(new Item("Foo"), new Item("Bar"));
this.removeItems = Arrays.asList(new Item(1), new Item(2));
}
@Test
public void testBasicWrite() throws Exception {
MongoItemWriter<String> writer = new MongoItemWriterBuilder<String>().template(this.template).build();
writer.write(this.items);
MongoItemWriter<Item> writer = new MongoItemWriterBuilder<Item>().template(this.template).build();
writer.write(this.saveItems);
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.mongoConverter).write(eq(this.saveItems.get(0)), any(Document.class));
verify(this.mongoConverter).write(eq(this.saveItems.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));
}
@Test
public void testDelete() throws Exception {
MongoItemWriter<String> writer = new MongoItemWriterBuilder<String>().template(this.template)
.delete(true)
.build();
writer.write(this.items);
verify(this.template).remove(this.items.get(0));
verify(this.template).remove(this.items.get(1));
verify(this.template, never()).bulkOps(any(), any(Class.class));
verify(this.mongoConverter, never()).write(any(), any());
verify(this.bulkOperations, never()).remove(any(Query.class));
}
@Test
public void testWriteToCollection() throws Exception {
MongoItemWriter<String> writer = new MongoItemWriterBuilder<String>().collection("collection")
MongoItemWriter<Item> writer = new MongoItemWriterBuilder<Item>().collection("collection")
.template(this.template)
.build();
writer.write(this.items);
writer.write(this.saveItems);
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.mongoConverter).write(eq(this.saveItems.get(0)), any(Document.class));
verify(this.mongoConverter).write(eq(this.saveItems.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");
verify(this.bulkOperations, never()).remove(any(Query.class));
}
@Test
public void testDelete() throws Exception {
MongoItemWriter<Item> writer = new MongoItemWriterBuilder<Item>().template(this.template)
.delete(true)
.build();
writer.write(this.removeItems);
verify(this.template).bulkOps(any(), any(Class.class));
verify(this.bulkOperations, times(2)).remove(any(Query.class));
}
@Test
@@ -119,3 +128,14 @@ public class MongoItemWriterBuilderTests {
}
}
}
class Item {
Integer id;
String name;
public Item(Integer id){
this.id = id;
}
public Item(String name) {
this.name = name;
}
}