DATAMONGO-1165 - Add support for Streaming large result lists.

Initial PoC for a MongoDB Cursor backed Iterator that allows fetching of large result sets.

Author: Thomas Darimont

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoOperations.java

@@ -33,6 +33,7 @@
 import org.springframework.data.mongodb.core.query.NearQuery;
 import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.core.query.Update;
+import org.springframework.data.mongodb.util.CloseableIterator;
 
 import com.mongodb.CommandResult;
 import com.mongodb.DBCollection;
@@ -959,4 +960,6 @@ <T> T findAndModify(Query query, Update update, FindAndModifyOptions options, Cl
 	 * @return
 	 */
 	MongoConverter getConverter();
+
+	CloseableIterator<Object> getStreamingMappingCursor(Query query, Class<?> entityType);
 }

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2014 the original author or authors.
+ * Copyright 2010-2015 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.
@@ -95,6 +95,7 @@
 import org.springframework.data.mongodb.core.query.NearQuery;
 import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.core.query.Update;
+import org.springframework.data.mongodb.util.CloseableIterator;
 import org.springframework.jca.cci.core.ConnectionCallback;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
@@ -103,6 +104,7 @@
 
 import com.mongodb.BasicDBObject;
 import com.mongodb.CommandResult;
+import com.mongodb.Cursor;
 import com.mongodb.DB;
 import com.mongodb.DBCollection;
 import com.mongodb.DBCursor;
@@ -2262,4 +2264,81 @@ public GeoResult<T> doWith(DBObject object) {
 		}
 	}
 
+	@Override
+	public CloseableIterator<Object> getStreamingMappingCursor(Query query, Class<?> entityClass) {
+
+		MongoPersistentEntity<?> entity = mappingContext.getPersistentEntity(entityClass);
+
+		DBCollection collection = getCollection(getCollectionName(entityClass));
+
+		DBObject mappedFields = queryMapper.getMappedFields(query.getFieldsObject(), entity);
+		DBObject mappedQuery = queryMapper.getMappedObject(query.getQueryObject(), entity);
+
+		return new CloseableIterableCusorAdapter(collection.find(mappedQuery, mappedFields), mongoConverter, entityClass);
+	}
+
+	static class CloseableIterableCusorAdapter implements CloseableIterator<Object> {
+
+		private volatile Cursor cursor;
+		private MongoConverter converter;
+		private Class<?> entityClass;
+
+		public CloseableIterableCusorAdapter(Cursor cursor, MongoConverter converter, Class<?> entityClass) {
+
+			this.cursor = cursor;
+			this.converter = converter;
+			this.entityClass = entityClass;
+		}
+
+		@Override
+		public boolean hasNext() {
+
+			if (cursor == null) {
+				return false;
+			}
+
+			try {
+				return cursor.hasNext();
+			} catch (Exception ex) {
+				throw new RuntimeException("error on hasNext");
+			}
+		}
+
+		@Override
+		public Object next() {
+
+			if (cursor == null) {
+				return null;
+			}
+
+			try {
+
+				DBObject item = cursor.next();
+				Object converted = converter.read(entityClass, item);
+
+				return converted;
+			} catch (Exception ex) {
+				throw new RuntimeException("error on next");
+			}
+		}
+
+		@Override
+		public void close() throws IOException {
+
+			Cursor c = cursor;
+			try {
+				c.close();
+			} finally {
+
+				cursor = null;
+				converter = null;
+				entityClass = null;
+			}
+		}
+
+		@Override
+		public Iterator<Object> iterator() {
+			return this;
+		}
+	}
 }

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/query/AbstractMongoQuery.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2014 the original author or authors.
+ * Copyright 2010-2015 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.
@@ -18,8 +18,6 @@
 import java.util.Collections;
 import java.util.List;
 
-import org.springframework.core.convert.ConversionService;
-import org.springframework.core.convert.support.DefaultConversionService;
 import org.springframework.data.domain.PageImpl;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Slice;
@@ -32,10 +30,12 @@
 import org.springframework.data.mongodb.core.MongoOperations;
 import org.springframework.data.mongodb.core.query.NearQuery;
 import org.springframework.data.mongodb.core.query.Query;
+import org.springframework.data.mongodb.util.CloseableIterator;
 import org.springframework.data.repository.query.ParameterAccessor;
 import org.springframework.data.repository.query.RepositoryQuery;
 import org.springframework.data.util.TypeInformation;
 import org.springframework.util.Assert;
+import org.springframework.util.ClassUtils;
 
 import com.mongodb.WriteResult;
 
@@ -48,8 +48,6 @@
  */
 public abstract class AbstractMongoQuery implements RepositoryQuery {
 
-	private static final ConversionService CONVERSION_SERVICE = new DefaultConversionService();
-
 	private final MongoQueryMethod method;
 	private final MongoOperations operations;
 
@@ -87,6 +85,10 @@ public Object execute(Object[] parameters) {
 
 		applyQueryMetaAttributesWhenPresent(query);
 
+		if (method.isStreamQuery()) {
+			return new StreamExecution().execute(query);
+		}
+
 		if (isDeleteQuery()) {
 			return new DeleteExecution().execute(query);
 		} else if (method.isGeoNearQuery() && method.isPageQuery()) {
@@ -415,4 +417,22 @@ private Object deleteAndConvertResult(Query query, MongoEntityMetadata<?> metada
 			return writeResult != null ? writeResult.getN() : 0L;
 		}
 	}
+
+	final class StreamExecution extends Execution {
+
+		private final boolean IS_JAVA_7 = ClassUtils.isPresent("java.lang.AutoCloseable",
+				StreamExecution.class.getClassLoader());
+
+		@Override
+		Object execute(Query query) {
+
+			Class<?> entityType = getQueryMethod().getEntityInformation().getJavaType();
+			CloseableIterator<Object> result = operations.getStreamingMappingCursor(query, entityType);
+			return IS_JAVA_7 ? wrapInAutoCloseableIterator(result) : result;
+		}
+
+		private CloseableIterator<Object> wrapInAutoCloseableIterator(CloseableIterator<Object> iter) {
+			return new ForwardingAutoCloseableIterator<Object>(iter);
+		}
+	}
 }

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/query/ForwardingAutoCloseableIterator.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2015 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.repository.query;
+
+import java.io.IOException;
+import java.util.Iterator;
+
+import org.springframework.data.mongodb.util.AutoCloseableIterator;
+import org.springframework.data.mongodb.util.CloseableIterator;
+
+/**
+ * An {@link CloseableIterator} that can be used with TRW in Java 7.
+ * 
+ * @author Thomas Darimont
+ * @param <T>
+ */
+class ForwardingAutoCloseableIterator<T> implements AutoCloseableIterator<T> {
+
+	private final CloseableIterator<T> target;
+
+	public ForwardingAutoCloseableIterator(CloseableIterator<T> target) {
+		this.target = target;
+	}
+
+	@Override
+	public boolean hasNext() {
+		return target.hasNext();
+	}
+
+	@Override
+	public T next() {
+		return target.next();
+	}
+
+	@Override
+	public void close() throws IOException {
+		this.target.close();
+	}
+
+	@Override
+	public Iterator<T> iterator() {
+		return this;
+	}
+}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/query/MongoQueryMethod.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2011-2014 the original author or authors.
+ * Copyright 2011-2015 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.
@@ -29,6 +29,7 @@
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
 import org.springframework.data.mongodb.repository.Meta;
 import org.springframework.data.mongodb.repository.Query;
+import org.springframework.data.mongodb.util.CloseableIterator;
 import org.springframework.data.repository.core.RepositoryMetadata;
 import org.springframework.data.repository.query.QueryMethod;
 import org.springframework.data.util.ClassTypeInformation;
@@ -233,4 +234,10 @@ public org.springframework.data.mongodb.core.query.Meta getQueryMetaAttributes()
 
 		return metaAttributes;
 	}
+
+	public boolean isStreamQuery() {
+
+		Class<?> returnType = method.getReturnType();
+		return org.springframework.util.ClassUtils.isAssignable(CloseableIterator.class, returnType);
+	}
 }

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/util/AutoCloseableIterator.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2015 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.util;
+
+/**
+ * @author Thomas Darimont
+ * @param <T>
+ */
+public interface AutoCloseableIterator<T> extends CloseableIterator<T>, AutoCloseable {}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/util/CloseableIterator.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2015 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.util;
+
+import java.io.Closeable;
+import java.util.Iterator;
+
+/**
+ * @author Thomas Darimont
+ * @param <T>
+ */
+public interface CloseableIterator<T> extends Iterator<T>, Closeable, Iterable<T> {
+
+}

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/PersonRepository.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2014 the original author or authors.
+ * Copyright 2010-2015 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.
@@ -31,6 +31,7 @@
 import org.springframework.data.geo.Point;
 import org.springframework.data.geo.Polygon;
 import org.springframework.data.mongodb.repository.Person.Sex;
+import org.springframework.data.mongodb.util.AutoCloseableIterator;
 import org.springframework.data.querydsl.QueryDslPredicateExecutor;
 
 /**
@@ -320,4 +321,7 @@ public interface PersonRepository extends MongoRepository<Person, String>, Query
 
 	@Query("{ ?0 : ?1 }")
 	List<Person> findByKeyValue(String key, String value);
+
+	@Query("{firstname:{$in:?0}}")
+	AutoCloseableIterator<Person> findByCustomQueryWithCursorByFirstnames(List<String> firstnames);
 }

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/PersonRepositoryIntegrationTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2013 the original author or authors.
+ * Copyright 2010-2015 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.
@@ -15,6 +15,11 @@
  */
 package org.springframework.data.mongodb.repository;
 
+import java.util.Arrays;
+import java.util.List;
+
+import org.junit.Test;
+import org.springframework.data.mongodb.util.AutoCloseableIterator;
 import org.springframework.test.context.ContextConfiguration;
 
 /**
@@ -24,4 +29,25 @@
  * @author Thomas Darimont
  */
 @ContextConfiguration
-public class PersonRepositoryIntegrationTests extends AbstractPersonRepositoryIntegrationTests {}
+public class PersonRepositoryIntegrationTests extends AbstractPersonRepositoryIntegrationTests {
+
+	@Test
+	public void findsAllMusiciansWithCursor() throws Exception {
+
+		repository.save(new Person("foo", "bar"));
+		repository.save(new Person("bar", "bar"));
+		repository.save(new Person("fuu", "bar"));
+		repository.save(new Person("notfound", "bar"));
+
+		List<String> firstNames = Arrays.asList("bar", "foo", "fuu");
+		AutoCloseableIterator<Person> result = repository.findByCustomQueryWithCursorByFirstnames(firstNames);
+
+		try {
+			for (Person person : result) {
+				System.out.printf("%s%n", person);
+			}
+		} finally {
+			result.close();
+		}
+	}
+}