DATAJPA-677 - Add support for Java 8 Stream in repository finder methods.

Added support for Streaming query results in Hibernate, EclipseLink and OpenJPA by falling back to store specific APIs to execute the query.

Original pull request: #136.

src/main/java/org/springframework/data/jpa/provider/PersistenceProvider.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2008-2014 the original author or authors.
+ * Copyright 2008-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.
@@ -21,17 +21,31 @@ import static org.springframework.data.jpa.provider.PersistenceProvider.Constant
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
 
 import javax.persistence.EntityManager;
 import javax.persistence.Query;
 import javax.persistence.metamodel.Metamodel;
 
 import org.apache.openjpa.enhance.PersistenceCapable;
+import org.apache.openjpa.persistence.OpenJPAPersistence;
 import org.apache.openjpa.persistence.OpenJPAQuery;
+import org.apache.openjpa.persistence.jdbc.FetchDirection;
+import org.apache.openjpa.persistence.jdbc.JDBCFetchPlan;
+import org.apache.openjpa.persistence.jdbc.LRSSizeAlgorithm;
+import org.apache.openjpa.persistence.jdbc.ResultSetType;
 import org.eclipse.persistence.jpa.JpaQuery;
+import org.eclipse.persistence.queries.ScrollableCursor;
+import org.hibernate.ScrollMode;
+import org.hibernate.ScrollableResults;
 import org.hibernate.ejb.HibernateQuery;
 import org.hibernate.proxy.HibernateProxy;
+import org.springframework.beans.DirectFieldAccessor;
+import org.springframework.data.util.CloseableIterator;
+import org.springframework.transaction.support.TransactionSynchronizationManager;
 import org.springframework.util.Assert;
+import org.springframework.util.ClassUtils;
 
 /**
  * Enumeration representing persistence providers to be used.
@@ -94,6 +108,11 @@ public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor {
 		public <T> Collection<T> potentiallyConvertEmptyCollection(Collection<T> collection) {
 			return collection == null || collection.isEmpty() ? null : collection;
 		}
+
+		@Override
+		public CloseableIterator<Object> executeQueryWithResultStream(Query jpaQuery) {
+			return new HibernateScrollableResultsIterator<Object>(jpaQuery);
+		}
 	},
 
 	/**
@@ -131,6 +150,11 @@ public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor {
 		public <T> Collection<T> potentiallyConvertEmptyCollection(Collection<T> collection) {
 			return collection == null || collection.isEmpty() ? null : collection;
 		}
+
+		@Override
+		public CloseableIterator<Object> executeQueryWithResultStream(Query jpaQuery) {
+			return new EclipseLinkScrollableResultsIterator<Object>(jpaQuery);
+		}
 	},
 
 	/**
@@ -164,6 +188,11 @@ public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor {
 		public Object getIdentifierFrom(Object entity) {
 			return ((PersistenceCapable) entity).pcFetchObjectId();
 		}
+
+		@Override
+		public CloseableIterator<Object> executeQueryWithResultStream(Query jpaQuery) {
+			return new OpenJpaResultStreamingIterator<Object>(jpaQuery);
+		}
 	},
 
 	/**
@@ -310,4 +339,195 @@ public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor {
 	public <T> Collection<T> potentiallyConvertEmptyCollection(Collection<T> collection) {
 		return collection;
 	}
+
+	public CloseableIterator<Object> executeQueryWithResultStream(Query jpaQuery) {
+		throw new UnsupportedOperationException("Streaming results is not implement for this PersistenceProvider: "
+				+ name());
+	}
+
+	/**
+	 * @author Thomas Darimont
+	 * @param <T>
+	 * @since 1.8
+	 */
+	@SuppressWarnings("unchecked")
+	static class HibernateScrollableResultsIterator<T> implements CloseableIterator<T> {
+
+		private ScrollableResults scrollableResults;
+
+		private static final boolean IS_HIBERNATE3 = ClassUtils.isPresent("org.hibernate.ejb.QueryImpl",
+				HibernateScrollableResultsIterator.class.getClassLoader());
+
+		public HibernateScrollableResultsIterator(Query jpaQuery) {
+
+			// see http://java.dzone.com/articles/bulk-fetching-hibernate
+			// we could also use a Hibernate stateless session here for constructing a query
+
+			org.hibernate.Query qry = IS_HIBERNATE3 ? extractHibernate3QueryFrom(jpaQuery) : extractHibernate4Query(jpaQuery);
+
+			ScrollableResults scrollableResults = qry.setReadOnly(
+					TransactionSynchronizationManager.isCurrentTransactionReadOnly()).scroll(ScrollMode.FORWARD_ONLY);
+
+			this.scrollableResults = scrollableResults;
+		}
+
+		private org.hibernate.Query extractHibernate4Query(Query jpaQuery) {
+
+			Object queryImpl = jpaQuery;
+			if (jpaQuery.getClass().getName().equals("org.hibernate.jpa.criteria.compile.CriteriaQueryTypeQueryAdapter")) {
+				queryImpl = new DirectFieldAccessor(jpaQuery).getPropertyValue("jpqlQuery");
+			}
+
+			return extractHibernateQueryFromQueryImpl(queryImpl);
+		}
+
+		private org.hibernate.Query extractHibernate3QueryFrom(Query jpaQuery) {
+
+			Object queryImpl = jpaQuery;
+			if (jpaQuery.getClass().isAnonymousClass()
+					&& jpaQuery.getClass().getEnclosingClass().getName()
+							.equals("org.hibernate.ejb.criteria.CriteriaQueryCompiler")) {
+				queryImpl = new DirectFieldAccessor(jpaQuery).getPropertyValue("val$jpaqlQuery");
+			}
+
+			return extractHibernateQueryFromQueryImpl(queryImpl);
+		}
+
+		private org.hibernate.Query extractHibernateQueryFromQueryImpl(Object queryImpl) {
+			return (org.hibernate.Query) new DirectFieldAccessor(queryImpl).getPropertyValue("query");
+		}
+
+		@Override
+		public T next() {
+
+			Object item = scrollableResults.get()[0];
+
+			return (T) item;
+		}
+
+		@Override
+		public boolean hasNext() {
+
+			if (scrollableResults == null) {
+				return false;
+			}
+
+			return scrollableResults.next();
+		}
+
+		@Override
+		public void close() {
+
+			if (scrollableResults == null) {
+				return;
+			}
+
+			try {
+				scrollableResults.close();
+			} finally {
+				scrollableResults = null;
+			}
+		}
+	}
+
+	/**
+	 * @author Thomas Darimont
+	 * @param <T>
+	 * @since 1.8
+	 */
+	@SuppressWarnings("unchecked")
+	static class EclipseLinkScrollableResultsIterator<T> implements CloseableIterator<T> {
+
+		private ScrollableCursor scrollableCursor;
+
+		public EclipseLinkScrollableResultsIterator(Query jpaQuery) {
+
+			jpaQuery.setHint("eclipselink.cursor.scrollable", true);
+			this.scrollableCursor = (ScrollableCursor) jpaQuery.getSingleResult();
+		}
+
+		@Override
+		public boolean hasNext() {
+
+			if (scrollableCursor == null) {
+				return false;
+			}
+
+			return scrollableCursor.hasNext();
+		}
+
+		@Override
+		public T next() {
+
+			Object item = scrollableCursor.next();
+			return (T) item;
+		}
+
+		@Override
+		public void close() {
+
+			if (scrollableCursor == null) {
+				return;
+			}
+
+			try {
+				scrollableCursor.close();
+			} finally {
+				scrollableCursor = null;
+			}
+		}
+	}
+
+	/**
+	 * @author Thomas Darimont
+	 * @param <T>
+	 * @since 1.8
+	 */
+	static class OpenJpaResultStreamingIterator<T> implements CloseableIterator<T> {
+
+		private Iterator<T> iterator;
+
+		@SuppressWarnings({ "rawtypes", "unchecked" })
+		public OpenJpaResultStreamingIterator(Query jpaQuery) {
+
+			OpenJPAQuery kq = OpenJPAPersistence.cast(jpaQuery);
+			JDBCFetchPlan fetch = (JDBCFetchPlan) kq.getFetchPlan();
+			fetch.setFetchBatchSize(20);
+			fetch.setResultSetType(ResultSetType.SCROLL_SENSITIVE);
+			fetch.setFetchDirection(FetchDirection.FORWARD);
+			fetch.setLRSSizeAlgorithm(LRSSizeAlgorithm.LAST);
+
+			List<T> resultList = kq.getResultList();
+			iterator = resultList.iterator();
+		}
+
+		@Override
+		public boolean hasNext() {
+
+			if (iterator == null) {
+				return false;
+			}
+
+			return iterator.hasNext();
+		}
+
+		@Override
+		public T next() {
+			return iterator.next();
+		}
+
+		@Override
+		public void close() {
+			if (iterator == null) {
+				return;
+			}
+
+			try {
+				OpenJPAPersistence.close(iterator);
+			} finally {
+				iterator = null;
+			}
+		}
+
+	}
 }

src/main/java/org/springframework/data/jpa/repository/query/AbstractJpaQuery.java

@@ -28,6 +28,7 @@ import org.springframework.data.jpa.repository.query.JpaQueryExecution.PagedExec
 import org.springframework.data.jpa.repository.query.JpaQueryExecution.ProcedureExecution;
 import org.springframework.data.jpa.repository.query.JpaQueryExecution.SingleEntityExecution;
 import org.springframework.data.jpa.repository.query.JpaQueryExecution.SlicedExecution;
+import org.springframework.data.jpa.repository.query.JpaQueryExecution.StreamExecution;
 import org.springframework.data.repository.query.RepositoryQuery;
 import org.springframework.util.Assert;
 
@@ -99,7 +100,9 @@ public abstract class AbstractJpaQuery implements RepositoryQuery {
 
 	protected JpaQueryExecution getExecution() {
 
-		if (method.isProcedureQuery()) {
+		if (method.isStreamQuery()) {
+			return new StreamExecution();
+		} else if (method.isProcedureQuery()) {
 			return new ProcedureExecution();
 		} else if (method.isCollectionQuery()) {
 			return new CollectionExecution();

src/main/java/org/springframework/data/jpa/repository/query/JpaQueryExecution.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2008-2014 the original author or authors.
+ * Copyright 2008-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,9 +31,12 @@ import org.springframework.data.domain.PageImpl;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Slice;
 import org.springframework.data.domain.SliceImpl;
+import org.springframework.data.jpa.provider.PersistenceProvider;
 import org.springframework.data.repository.query.ParameterAccessor;
 import org.springframework.data.repository.query.Parameters;
 import org.springframework.data.repository.query.ParametersParameterAccessor;
+import org.springframework.data.util.CloseableIterator;
+import org.springframework.data.util.Java8StreamUtils;
 import org.springframework.util.Assert;
 
 /**
@@ -298,4 +301,23 @@ public abstract class JpaQueryExecution {
 			return storedProcedureJpaQuery.extractOutputValue(storedProcedure);
 		}
 	}
+
+	/**
+	 * {@link Execution} executing a Java 8 Stream.
+	 * 
+	 * @author Thomas Darimont
+	 * @since 1.8
+	 */
+	static class StreamExecution extends JpaQueryExecution {
+
+		@Override
+		protected Object doExecute(final AbstractJpaQuery query, Object[] values) {
+
+			Query jpaQuery = query.createQuery(values);
+			PersistenceProvider persistenceProvider = PersistenceProvider.fromEntityManager(query.getEntityManager());
+			CloseableIterator<Object> iter = persistenceProvider.executeQueryWithResultStream(jpaQuery);
+
+			return Java8StreamUtils.createStreamFromIterator(iter);
+		}
+	}
 }

src/test/java/org/springframework/data/jpa/repository/UserRepositoryTests.java

@@ -28,6 +28,8 @@ import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.function.Consumer;
+import java.util.stream.Stream;
 
 import javax.persistence.EntityManager;
 import javax.persistence.PersistenceContext;
@@ -1751,8 +1753,8 @@ public class UserRepositoryTests {
 		assertThat(users.getContent(), hasSize(2));
 		assertThat(users.getContent().get(0), is(thirdUser));
 		assertThat(users.getContent().get(1), is(fourthUser));
-	}
-
+	}
+
 	/**
 	 * @see DATAJPA-629
 	 */
@@ -1816,6 +1818,66 @@ public class UserRepositoryTests {
 		assertThat(users.get(0), is(secondUser));
 	}
 
+	/**
+	 * @see DATAJPA-677
+	 */
+	@Test
+	public void shouldSupportJava8StreamsForRepositoryFinderMethods() {
+
+		flushTestUsers();
+
+		Stream<User> stream = repository.findAllByCustomQueryAndStream();
+
+		final List<User> users = new ArrayList<User>();
+
+		try {
+			stream.forEach(new Consumer<User>() {
+
+				@Override
+				public void accept(User user) {
+
+					// System.out.printf("%s%n", user);
+					users.add(user);
+				}
+
+			});
+		} finally {
+			stream.close();
+		}
+
+		assertThat(users, hasSize(4));
+	}
+
+	/**
+	 * @see DATAJPA-677
+	 */
+	@Test
+	public void shouldSupportJava8StreamsForRepositoryDerivedFinderMethods() {
+
+		flushTestUsers();
+
+		Stream<User> stream = repository.readAllByFirstnameNotNull();
+
+		final List<User> users = new ArrayList<User>();
+
+		try {
+			stream.forEach(new Consumer<User>() {
+
+				@Override
+				public void accept(User user) {
+
+					// System.out.printf("%s%n", user);
+					users.add(user);
+				}
+
+			});
+		} finally {
+			stream.close();
+		}
+
+		assertThat(users, hasSize(4));
+	}
+
 	private Page<User> executeSpecWithSort(Sort sort) {
 
 		flushTestUsers();

src/test/java/org/springframework/data/jpa/repository/sample/UserRepository.java

@@ -19,6 +19,7 @@ import java.util.Collection;
 import java.util.Date;
 import java.util.List;
 import java.util.Set;
+import java.util.stream.Stream;
 
 import javax.persistence.EntityManager;
 import javax.persistence.QueryHint;
@@ -557,4 +558,15 @@ public interface UserRepository extends JpaRepository<User, Integer>, JpaSpecifi
 	 * DATAJPA-606
 	 */
 	List<User> queryByAgeInOrFirstname(Integer[] ages, String firstname);
+
+	/**
+	 * DATAJPA-677
+	 */
+	@Query("select u from User u")
+	Stream<User> findAllByCustomQueryAndStream();
+
+	/**
+	 * DATAJPA-677
+	 */
+	Stream<User> readAllByFirstnameNotNull();
 }