BATCH-2568: MongoItemReader should support Query object by default

JIRA: https://jira.spring.io/browse/BATCH-2568

Author: ninja-inc

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012 the original author or authors.
+ * Copyright 2012-2017 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.
@@ -25,6 +25,8 @@
 
 import com.mongodb.util.JSON;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.batch.item.ExecutionContext;
 import org.springframework.batch.item.ItemReader;
 import org.springframework.beans.factory.InitializingBean;
@@ -45,14 +47,23 @@
  * </p>
  *
  * <p>
- * It executes the JSON {@link #setQuery(String)} to retrieve the requested
- * documents.  The query is executed using paged requests specified in the
+ * If you set JSON String query {@link #setQuery(String)} then
+ * it executes the JSON to retrieve the requested documents.
+ * </p>
+ * 
+ * <p>
+ * If you set Query object {@link #setQuery(Query)} then
+ * it executes the Query to retrieve the requested documents.
+ * </p>
+ * 
+ * <p>
+ * The query is executed using paged requests specified in the
  * {@link #setPageSize(int)}.  Additional pages are requested as needed to
  * provide data when the {@link #read()} method is called.
  * </p>
  *
  * <p>
- * The JSON query provided supports parameter substitution via ?&lt;index&gt;
+ * The JSON String query provided supports parameter substitution via ?&lt;index&gt;
  * placeholders where the &lt;index&gt; indicates the index of the
  * parameterValue to substitute.
  * </p>
@@ -65,12 +76,16 @@
  *
  *
  * @author Michael Minella
+ * @author Takaaki Iida
  */
 public class MongoItemReader<T> extends AbstractPaginatedDataItemReader<T> implements InitializingBean {
-
+	
+	private static final Logger log = LoggerFactory.getLogger(MongoItemReader.class);
+	
 	private static final Pattern PLACEHOLDER = Pattern.compile("\\?(\\d+)");
 	private MongoOperations template;
-	private String query;
+	private Query query;
+	private String queryString;
 	private Class<? extends T> type;
 	private Sort sort;
 	private String hint;
@@ -82,6 +97,15 @@ public MongoItemReader() {
 		super();
 		setName(ClassUtils.getShortName(MongoItemReader.class));
 	}
+	
+	/**
+	 * A Mongo Query to be used.
+	 *
+	 * @param query Mongo Query to be used.
+	 */
+	public void setQuery(Query query) {
+		this.query = query;
+	}
 
 	/**
 	 * Used to perform operations against the MongoDB instance.  Also
@@ -99,10 +123,10 @@ public void setTemplate(MongoOperations template) {
 	 * via ?&lt;index&gt; placeholders where the &lt;index&gt; indicates the index of the
 	 * parameterValue to substitute.
 	 *
-	 * @param query JSON formatted Mongo query
+	 * @param queryString JSON formatted Mongo query
 	 */
-	public void setQuery(String query) {
-		this.query = query;
+	public void setQuery(String queryString) {
+		this.queryString = queryString;
 	}
 
 	/**
@@ -163,30 +187,41 @@ public void setHint(String hint) {
 	@Override
 	@SuppressWarnings("unchecked")
 	protected Iterator<T> doPageRead() {
-
-		Pageable pageRequest = PageRequest.of(page, pageSize, sort);
-
-		String populatedQuery = replacePlaceholders(query, parameterValues);
-
-		Query mongoQuery;
-
-		if(StringUtils.hasText(fields)) {
-			mongoQuery = new BasicQuery(populatedQuery, fields);
-		}
-		else {
-			mongoQuery = new BasicQuery(populatedQuery);
-		}
-
-		mongoQuery.with(pageRequest);
-
-		if(StringUtils.hasText(hint)) {
-			mongoQuery.withHint(hint);
-		}
-
-		if(StringUtils.hasText(collection)) {
-			return (Iterator<T>) template.find(mongoQuery, type, collection).iterator();
+		if (queryString != null) {
+			Pageable pageRequest = new PageRequest(page, pageSize, sort);
+	
+			String populatedQuery = replacePlaceholders(queryString, parameterValues);
+	
+			Query mongoQuery = null;
+	
+			if(StringUtils.hasText(fields)) {
+				mongoQuery = new BasicQuery(populatedQuery, fields);
+			}
+			else {
+				mongoQuery = new BasicQuery(populatedQuery);
+			}
+	
+			mongoQuery.with(pageRequest);
+	
+			if(StringUtils.hasText(hint)) {
+				mongoQuery.withHint(hint);
+			}
+	
+			if(StringUtils.hasText(collection)) {
+				return (Iterator<T>) template.find(mongoQuery, type, collection).iterator();
+			} else {
+				return (Iterator<T>) template.find(mongoQuery, type).iterator();
+			}
+			
 		} else {
-			return (Iterator<T>) template.find(mongoQuery, type).iterator();
+			Pageable pageRequest = new PageRequest(page, pageSize);
+			query.with(pageRequest);
+			
+			if(StringUtils.hasText(collection)) {
+				return (Iterator<T>) template.find(query, type, collection).iterator();
+			} else {
+				return (Iterator<T>) template.find(query, type).iterator();
+			}
 		}
 	}
 
@@ -199,8 +234,18 @@ protected Iterator<T> doPageRead() {
 	public void afterPropertiesSet() throws Exception {
 		Assert.state(template != null, "An implementation of MongoOperations is required.");
 		Assert.state(type != null, "A type to convert the input into is required.");
-		Assert.state(query != null, "A query is required.");
-		Assert.state(sort != null, "A sort is required.");
+		Assert.state(queryString != null || query != null, "A query is required.");
+		
+		if (queryString != null) {
+			Assert.state(sort != null, "A sort is required.");
+		}
+		if (query != null) {
+			Assert.state(query.getSortObject() != null, "A Sort in Query object is required.");
+		}
+		
+		if (query != null && query.getLimit() != 0) {
+			log.warn("PageSize in Query object was ignored. Please set it by MongoItemReader.setPageSize().");
+		}
 	}
 
 	// Copied from StringBasedMongoQuery...is there a place where this type of logic is already exposed?

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

@@ -25,15 +25,17 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import org.springframework.data.domain.PageRequest;
 import org.springframework.data.domain.Sort;
+import org.springframework.data.domain.Sort.Order;
 import org.springframework.data.mongodb.core.MongoOperations;
 import org.springframework.data.mongodb.core.query.Query;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
-import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.when;
 
 public class MongoItemReaderTests {
@@ -60,8 +62,8 @@ public void setUp() throws Exception {
 	}
 
 	@Test
-	public void testAfterPropertiesSet() throws Exception{
-		reader = new MongoItemReader<>();
+	public void testAfterPropertiesSetForQueryString() throws Exception{
+		reader = new MongoItemReader<String>();
 
 		try {
 			reader.afterPropertiesSet();
@@ -109,6 +111,31 @@ public void testAfterPropertiesSet() throws Exception{
 
 		reader.afterPropertiesSet();
 	}
+	
+	@Test
+	public void testAfterPropertiesSetForQueryObject() throws Exception{
+		reader = new MongoItemReader<String>();
+		
+		reader.setTemplate(template);
+		reader.setTargetType(String.class);
+		
+		Query query1 = new Query();
+		reader.setQuery(query1);
+
+		try {
+			reader.afterPropertiesSet();
+			fail("Sort was not set but exception was not thrown.");
+		} catch (IllegalStateException iae) {
+			assertEquals("A Sort in Query object is required.", iae.getMessage());
+		} catch (Throwable t) {
+			fail("Wrong exception was thrown.");
+		}
+
+		Query query2 = new Query().with(new Sort(new Order(Sort.Direction.ASC, "_id")));
+		reader.setQuery(query2);
+
+		reader.afterPropertiesSet();
+	}
 
 	@Test
 	public void testBasicQueryFirstPage() {
@@ -223,4 +250,102 @@ public void testQueryWithCollection() {
 		assertEquals("{ \"name\" : -1 }", query.getSortObject().toJson());
 		assertEquals("collection", collectionContainer.getValue());
 	}
+	
+	@Test
+	public void testQueryObject() throws Exception {
+		reader = new MongoItemReader<String>();
+		reader.setTemplate(template);
+		
+		Query query = new Query()
+				.with(new Sort(new Order(Sort.Direction.ASC, "_id")));
+		reader.setQuery(query);
+		reader.setTargetType(String.class);
+		
+		reader.afterPropertiesSet();
+		
+		ArgumentCaptor<Query> queryContainer = ArgumentCaptor.forClass(Query.class);
+		when(template.find(queryContainer.capture(), eq(String.class))).thenReturn(new ArrayList<String>());
+		
+		assertFalse(reader.doPageRead().hasNext());
+		
+		Query actualQuery = queryContainer.getValue();
+		assertFalse(reader.doPageRead().hasNext());
+		assertEquals(10, actualQuery.getLimit());
+		assertEquals(0, actualQuery.getSkip());
+	}
+	
+	@Test
+	public void testQueryObjectWithIgnoredPageSize() throws Exception {
+		reader = new MongoItemReader<String>();
+		reader.setTemplate(template);
+		
+		Query query = new Query()
+				.with(new Sort(new Order(Sort.Direction.ASC, "_id")))
+				.with(new PageRequest(0, 50));
+		reader.setQuery(query);
+		reader.setTargetType(String.class);
+		
+		reader.afterPropertiesSet();
+		
+		ArgumentCaptor<Query> queryContainer = ArgumentCaptor.forClass(Query.class);
+		when(template.find(queryContainer.capture(), eq(String.class))).thenReturn(new ArrayList<String>());
+		
+		assertFalse(reader.doPageRead().hasNext());
+		
+		Query actualQuery = queryContainer.getValue();
+		assertFalse(reader.doPageRead().hasNext());
+		assertEquals(10, actualQuery.getLimit());
+		assertEquals(0, actualQuery.getSkip());
+	}
+	
+	@Test
+	public void testQueryObjectWithPageSize() throws Exception {
+		reader = new MongoItemReader<String>();
+		reader.setTemplate(template);
+		
+		Query query = new Query()
+				.with(new Sort(new Order(Sort.Direction.ASC, "_id")))
+				.with(new PageRequest(30, 50));
+		reader.setQuery(query);
+		reader.setTargetType(String.class);
+		reader.setPageSize(100);
+		
+		reader.afterPropertiesSet();
+		
+		ArgumentCaptor<Query> queryContainer = ArgumentCaptor.forClass(Query.class);
+		when(template.find(queryContainer.capture(), eq(String.class))).thenReturn(new ArrayList<String>());
+		
+		assertFalse(reader.doPageRead().hasNext());
+		
+		Query actualQuery = queryContainer.getValue();
+		assertFalse(reader.doPageRead().hasNext());
+		assertEquals(100, actualQuery.getLimit());
+		assertEquals(0, actualQuery.getSkip());
+	}
+	
+	@Test
+	public void testQueryObjectWithCollection() throws Exception {
+		reader = new MongoItemReader<String>();
+		reader.setTemplate(template);
+		
+		Query query = new Query()
+				.with(new Sort(new Order(Sort.Direction.ASC, "_id")));
+		reader.setQuery(query);
+		reader.setTargetType(String.class);
+		reader.setCollection("collection");
+		
+		reader.afterPropertiesSet();
+		
+		ArgumentCaptor<Query> queryContainer = ArgumentCaptor.forClass(Query.class);
+		ArgumentCaptor<String> stringContainer = ArgumentCaptor.forClass(String.class);
+		when(template.find(queryContainer.capture(), eq(String.class), stringContainer.capture())).thenReturn(new ArrayList<String>());
+		
+		assertFalse(reader.doPageRead().hasNext());
+		
+		Query actualQuery = queryContainer.getValue();
+		assertFalse(reader.doPageRead().hasNext());
+		assertEquals(10, actualQuery.getLimit());
+		assertEquals(0, actualQuery.getSkip());
+		assertEquals("collection", stringContainer.getValue());
+	}
 }