From d5e2c250485621cead0001dcfe9a0b0d506db6d3 Mon Sep 17 00:00:00 2001
From: Evgeniy Cheban <mister.cheban@gmail.com>
Date: Wed, 3 Aug 2022 21:55:43 +0200
Subject: [PATCH] Consider adding more return types for a native query (#5)

Consider adding more return types for a native query

Closes gh-1
---
 .../query/ReindexerQueryMethod.java           | 42 ++++++++++++
 .../StringBasedReindexerRepositoryQuery.java  | 67 +++++++++++++++++++
 2 files changed, 109 insertions(+)

diff --git a/src/main/java/org/springframework/data/reindexer/repository/query/ReindexerQueryMethod.java b/src/main/java/org/springframework/data/reindexer/repository/query/ReindexerQueryMethod.java
index 900f2c4..444a33d 100644
--- a/src/main/java/org/springframework/data/reindexer/repository/query/ReindexerQueryMethod.java
+++ b/src/main/java/org/springframework/data/reindexer/repository/query/ReindexerQueryMethod.java
@@ -17,6 +17,9 @@
 
 import java.lang.reflect.Method;
 import java.util.Iterator;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
 
 import org.springframework.data.projection.ProjectionFactory;
 import org.springframework.data.reindexer.core.mapping.Query;
@@ -33,6 +36,12 @@ public final class ReindexerQueryMethod extends QueryMethod {
 
 	private final Lazy<Boolean> isIteratorQuery;
 
+	private final Lazy<Boolean> isOptionalQuery;
+
+	private final Lazy<Boolean> isListQuery;
+
+	private final Lazy<Boolean> isSetQuery;
+
 	private final Lazy<Query> queryAnnotationExtractor;
 
 	/**
@@ -46,6 +55,9 @@ public final class ReindexerQueryMethod extends QueryMethod {
 	public ReindexerQueryMethod(Method method, RepositoryMetadata metadata, ProjectionFactory factory) {
 		super(method, metadata, factory);
 		this.isIteratorQuery = Lazy.of(() -> Iterator.class.isAssignableFrom(getReturnedObjectType()));
+		this.isOptionalQuery = Lazy.of(() -> Optional.class.isAssignableFrom(method.getReturnType()));
+		this.isListQuery = Lazy.of(() -> List.class.isAssignableFrom(method.getReturnType()));
+		this.isSetQuery = Lazy.of(() -> Set.class.isAssignableFrom(method.getReturnType()));
 		this.queryAnnotationExtractor = Lazy.of(() -> method.getAnnotation(Query.class));
 	}
 
@@ -58,6 +70,36 @@ public boolean isIteratorQuery() {
 		return this.isIteratorQuery.get();
 	}
 
+	/**
+	 * Returns true if the method returns {@link Optional}.
+	 *
+	 * @return true if the method returns {@link Optional}
+	 * @since 1.1
+	 */
+	public boolean isOptionalQuery() {
+		return this.isOptionalQuery.get();
+	}
+
+	/**
+	 * Returns true if the method returns {@link List}.
+	 *
+	 * @return true if the method returns {@link List}
+	 * @since 1.1
+	 */
+	public boolean isListQuery() {
+		return this.isListQuery.get();
+	}
+
+	/**
+	 * Returns true if the method returns {@link Set}.
+	 *
+	 * @return true if the method returns {@link Set}
+	 * @since 1.1
+	 */
+	public boolean isSetQuery() {
+		return this.isSetQuery.get();
+	}
+
 	/**
 	 * Returns true if the method has {@link Query} annotation.
 	 *
diff --git a/src/main/java/org/springframework/data/reindexer/repository/query/StringBasedReindexerRepositoryQuery.java b/src/main/java/org/springframework/data/reindexer/repository/query/StringBasedReindexerRepositoryQuery.java
index ecee8b1..c596126 100644
--- a/src/main/java/org/springframework/data/reindexer/repository/query/StringBasedReindexerRepositoryQuery.java
+++ b/src/main/java/org/springframework/data/reindexer/repository/query/StringBasedReindexerRepositoryQuery.java
@@ -15,6 +15,17 @@
  */
 package org.springframework.data.reindexer.repository.query;
 
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Optional;
+import java.util.Spliterator;
+import java.util.Spliterators;
+import java.util.function.Supplier;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
+
+import ru.rt.restream.reindexer.CloseableIterator;
 import ru.rt.restream.reindexer.Namespace;
 import ru.rt.restream.reindexer.Reindexer;
 
@@ -55,9 +66,65 @@ public Object execute(Object[] parameters) {
 		if (this.queryMethod.isIteratorQuery()) {
 			return this.namespace.execSql(query);
 		}
+		if (this.queryMethod.isStreamQuery()) {
+			return toStream(this.namespace.execSql(query));
+		}
+		if (this.queryMethod.isListQuery()) {
+			return toCollection(this.namespace.execSql(query), ArrayList::new);
+		}
+		if (this.queryMethod.isSetQuery()) {
+			return toCollection(this.namespace.execSql(query), HashSet::new);
+		}
+		if (this.queryMethod.isOptionalQuery()) {
+			return toOptionalEntity(this.namespace.execSql(query));
+		}
+		if (this.queryMethod.isQueryForEntity()) {
+			return toEntity(this.namespace.execSql(query));
+		}
 		throw new IllegalStateException("Unsupported method return type " + this.queryMethod.getReturnedObjectType());
 	}
 
+	private <T> Stream<T> toStream(CloseableIterator<T> iterator) {
+		Spliterator<T> spliterator = Spliterators.spliterator(iterator, iterator.size(), Spliterator.NONNULL);
+		return StreamSupport.stream(spliterator, false);
+	}
+
+	private <T> Collection<T> toCollection(CloseableIterator<T> iterator, Supplier<Collection<T>> collectionSupplier) {
+		Collection<T> result = collectionSupplier.get();
+		try (CloseableIterator<T> it = iterator) {
+			while (it.hasNext()) {
+				result.add(it.next());
+			}
+		}
+		return result;
+	}
+
+	public <T> Optional<T> toOptionalEntity(CloseableIterator<T> iterator) {
+		T item = getOneInternal(iterator);
+		return Optional.ofNullable(item);
+	}
+
+	public <T> T toEntity(CloseableIterator<T> iterator) {
+		T item = getOneInternal(iterator);
+		if (item == null) {
+			throw new IllegalStateException("Exactly one item expected, but there is zero");
+		}
+		return item;
+	}
+
+	private <T> T getOneInternal(CloseableIterator<T> iterator) {
+		try (CloseableIterator<T> it = iterator) {
+			T item = null;
+			if (it.hasNext()) {
+				item = it.next();
+			}
+			if (it.hasNext()) {
+				throw new IllegalStateException("Exactly one item expected, but there are more");
+			}
+			return item;
+		}
+	}
+
 	@Override
 	public QueryMethod getQueryMethod() {
 		return this.queryMethod;