Simplify ResultStreamWrapper to use native Java stream iterator

schema/src/main/java/com/introproventures/graphql/jpa/query/schema/impl/GraphQLJpaQueryFactory.java

@@ -26,6 +26,7 @@
 import static graphql.introspection.Introspection.SchemaMetaFieldDef;
 import static graphql.introspection.Introspection.TypeMetaFieldDef;
 import static graphql.introspection.Introspection.TypeNameMetaFieldDef;
+import static java.lang.Integer.min;
 import static java.util.stream.Collectors.groupingBy;
 
 import com.introproventures.graphql.jpa.query.annotation.GraphQLDefaultOrderBy;
@@ -248,16 +249,17 @@ public List<Object> queryKeys(
     public List<Object> queryResultList(DataFetchingEnvironment environment, int maxResults, List<Object> keys) {
         // Let's execute query and get result as stream
         Stream<Object> resultStream = queryResultStream(environment, maxResults, keys);
+        var size = keys.isEmpty() ? maxResults : min(keys.size(), maxResults);
         // Let's wrap stream into lazy list to pass it downstream
-        return ResultStreamWrapper.wrap(resultStream, maxResults);
+        return ResultStreamWrapper.wrap(resultStream, size);
     }
 
     protected Stream<Object> queryResultStream(DataFetchingEnvironment environment, int maxResults, List<Object> keys) {
         MergedField queryField = resolveQueryField(environment.getField());
 
         // Override query environment with associated entity object type and
         final DataFetchingEnvironment queryEnvironment = getQueryEnvironment(environment, queryField);
-        final int fetchSize = Integer.min(maxResults, defaultFetchSize);
+        final int fetchSize = min(maxResults, defaultFetchSize);
         final boolean isDistinct = resolveDistinctArgument(queryEnvironment.getField());
 
         final TypedQuery<Object> query = getQuery(
@@ -280,7 +282,7 @@ protected <T> Stream<T> getResultStream(
         boolean isDistinct,
         EntityGraph<?> entityGraph
     ) {
-        // Let' try reduce overhead and disable all caching
+        // Let's try to reduce overhead and disable all caching
         query.setHint(ORG_HIBERNATE_READ_ONLY, true);
         query.setHint(ORG_HIBERNATE_FETCH_SIZE, fetchSize);
         query.setHint(ORG_HIBERNATE_CACHEABLE, false);

schema/src/main/java/com/introproventures/graphql/jpa/query/schema/impl/ResultStreamWrapper.java

@@ -20,7 +20,6 @@
 import java.lang.reflect.Method;
 import java.lang.reflect.Proxy;
 import java.util.Collection;
-import java.util.Iterator;
 import java.util.List;
 import java.util.stream.Stream;
 
@@ -55,7 +54,7 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
             if ("size".equals(method.getName())) {
                 return size;
             } else if ("iterator".equals(method.getName())) {
-                return new ResultIteratorWrapper(stream.iterator(), size);
+                return stream.limit(size).iterator();
             } else if ("equals".equals(method.getName())) {
                 // Only consider equal when proxies are identical.
                 return (proxy == args[0]);
@@ -67,33 +66,5 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
             }
             throw new UnsupportedOperationException(method + " is not supported");
         }
-
-        class ResultIteratorWrapper implements Iterator<T> {
-
-            final Iterator<T> delegate;
-            final int size;
-            int current = 0;
-
-            ResultIteratorWrapper(Iterator<T> delegate, int size) {
-                this.delegate = delegate;
-                this.size = size;
-            }
-
-            @Override
-            public boolean hasNext() {
-                return (current < size) && delegate.hasNext();
-            }
-
-            @Override
-            public T next() {
-                T result = delegate.next();
-
-                try {
-                    return result;
-                } finally {
-                    current++;
-                }
-            }
-        }
     }
 }