From ca8087dfc6c4a7914f115d4fcfdbfdfc78355e0d Mon Sep 17 00:00:00 2001
From: Michael Minella <mminella@gopivotal.com>
Date: Wed, 21 May 2014 21:30:05 -0500
Subject: [PATCH] BATCH-2206: Added option to set collection for
 MongoItemReader

---
 .../batch/item/data/MongoItemReader.java      | 29 ++++++++----
 .../batch/item/data/MongoItemReaderTests.java | 46 ++++++++++++++-----
 2 files changed, 55 insertions(+), 20 deletions(-)

diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/data/MongoItemReader.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/data/MongoItemReader.java
index 41a770ce2..23d100c00 100644
--- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/data/MongoItemReader.java
+++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/data/MongoItemReader.java
@@ -16,13 +16,7 @@
 
 package org.springframework.batch.item.data;
 
-import java.util.ArrayList;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
-
+import com.mongodb.util.JSON;
 import org.springframework.batch.item.ExecutionContext;
 import org.springframework.batch.item.ItemReader;
 import org.springframework.beans.factory.InitializingBean;
@@ -36,7 +30,12 @@ import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
 import org.springframework.util.StringUtils;
 
-import com.mongodb.util.JSON;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 /**
  * <p>
@@ -75,6 +74,7 @@ public class MongoItemReader<T> extends AbstractPaginatedDataItemReader<T> imple
 	private Sort sort;
 	private String hint;
 	private String fields;
+	private String collection;
 	private List<Object> parameterValues;
 
 	public MongoItemReader() {
@@ -143,6 +143,13 @@ public class MongoItemReader<T> extends AbstractPaginatedDataItemReader<T> imple
 		this.sort = convertToSort(sorts);
 	}
 
+	/**
+	 * @param collection Mongo collection to be queried.
+	 */
+	public void setCollection(String collection) {
+		this.collection = collection;
+	}
+
 	/**
 	 * JSON String telling MongoDB what index to use.
 	 *
@@ -175,7 +182,11 @@ public class MongoItemReader<T> extends AbstractPaginatedDataItemReader<T> imple
 			mongoQuery.withHint(hint);
 		}
 
-		return (Iterator<T>) template.find(mongoQuery, type).iterator();
+		if(StringUtils.hasText(collection)) {
+			return (Iterator<T>) template.find(mongoQuery, type, collection).iterator();
+		} else {
+			return (Iterator<T>) template.find(mongoQuery, type).iterator();
+		}
 	}
 
 	/**
diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/MongoItemReaderTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/MongoItemReaderTests.java
index 7373b538b..ceb1bc8fb 100644
--- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/MongoItemReaderTests.java
+++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/MongoItemReaderTests.java
@@ -15,17 +15,6 @@
  */
 package org.springframework.batch.item.data;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.fail;
-import static org.mockito.Matchers.eq;
-import static org.mockito.Mockito.when;
-
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.Map;
-
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
@@ -35,6 +24,17 @@ import org.springframework.data.domain.Sort;
 import org.springframework.data.mongodb.core.MongoOperations;
 import org.springframework.data.mongodb.core.query.Query;
 
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.fail;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.when;
+
 public class MongoItemReaderTests {
 
 	private MongoItemReader<String> reader;
@@ -196,4 +196,28 @@ public class MongoItemReaderTests {
 		assertEquals("{ \"name\" : \"foo\"}", query.getQueryObject().toString());
 		assertEquals("{ \"name\" : -1}", query.getSortObject().toString());
 	}
+
+	@SuppressWarnings("serial")
+	@Test
+	public void testQueryWithCollection() {
+		reader.setParameterValues(new ArrayList<Object>(){{
+			add("foo");
+		}});
+
+		reader.setQuery("{ name : ?0 }");
+		reader.setCollection("collection");
+		ArgumentCaptor<Query> queryContainer = ArgumentCaptor.forClass(Query.class);
+		ArgumentCaptor<String> collectionContainer = ArgumentCaptor.forClass(String.class);
+
+		when(template.find(queryContainer.capture(), eq(String.class), collectionContainer.capture())).thenReturn(new ArrayList<String>());
+
+		assertFalse(reader.doPageRead().hasNext());
+
+		Query query = queryContainer.getValue();
+		assertEquals(50, query.getLimit());
+		assertEquals(0, query.getSkip());
+		assertEquals("{ \"name\" : \"foo\"}", query.getQueryObject().toString());
+		assertEquals("{ \"name\" : -1}", query.getSortObject().toString());
+		assertEquals("collection", collectionContainer.getValue());
+	}
 }