Consider supporting interface-based, class-based and dynamic projecti…

…ons for result mapping of query methods Closes gh-26
evgeniycheban · Jan 1, 2025 · 2bbd24d · 2bbd24d
1 parent c22b4cf
commit 2bbd24d
Show file tree

Hide file tree

Showing 7 changed files with 472 additions and 7 deletions.
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
@@ -44,6 +44,10 @@ public final class ReindexerQueryMethod extends QueryMethod {
 
 	private final Lazy<Query> queryAnnotationExtractor;
 
+	private final Class<?> returnType;
+
+	private final ProjectionFactory factory;
+
 	/**
 	 * Creates a new {@link QueryMethod} from the given parameters. Looks up the correct query to use for following
 	 * invocations of the method given.
@@ -54,11 +58,13 @@ 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.isIteratorQuery = Lazy.of(() -> Iterator.class.isAssignableFrom(method.getReturnType()));
 		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));
+		this.returnType = method.getReturnType();
+		this.factory = factory;
 	}
 
 	/**
@@ -129,4 +135,30 @@ public boolean isUpdateQuery() {
 		return query.update();
 	}
 
+	/**
+	 * Returns method's return type
+	 *
+	 * @return the method's return type
+	 */
+	Class<?> getReturnType() {
+		return this.returnType;
+	}
+
+	/**
+	 * Returns a {@link ProjectionFactory} to be used.
+	 *
+	 * @return the {@link ProjectionFactory} to use
+	 */
+	ProjectionFactory getFactory() {
+		return this.factory;
+	}
+
+	/**
+	 * {@inheritDoc}
+	 */
+	@Override
+	public Class<?> getDomainClass() {
+		return super.getDomainClass();
+	}
+
 }
diff --git a/...in/java/org/springframework/data/reindexer/repository/query/ReindexerRepositoryQuery.java b/...in/java/org/springframework/data/reindexer/repository/query/ReindexerRepositoryQuery.java
@@ -15,22 +15,33 @@
  */
 package org.springframework.data.reindexer.repository.query;
 
+import java.beans.PropertyDescriptor;
 import java.lang.reflect.Array;
+import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
+import java.util.Spliterator;
+import java.util.Spliterators;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
+import ru.rt.restream.reindexer.AggregationResult;
 import org.springframework.data.reindexer.repository.support.TransactionalNamespace;
 import ru.rt.restream.reindexer.FieldType;
 import ru.rt.restream.reindexer.Namespace;
 import ru.rt.restream.reindexer.Query;
 import ru.rt.restream.reindexer.Reindexer;
 import ru.rt.restream.reindexer.ReindexerIndex;
 import ru.rt.restream.reindexer.ReindexerNamespace;
+import ru.rt.restream.reindexer.ResultIterator;
+import ru.rt.restream.reindexer.util.BeanPropertyUtils;
 
+import org.springframework.core.CollectionFactory;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.domain.Sort.Order;
 import org.springframework.data.repository.query.ParametersParameterAccessor;
@@ -55,6 +66,8 @@ public class ReindexerRepositoryQuery implements RepositoryQuery {
 
 	private final Map<String, ReindexerIndex> indexes;
 
+	private final String[] selectFields;
+
 	/**
 	 * Creates an instance.
 	 *
@@ -72,22 +85,33 @@ public ReindexerRepositoryQuery(ReindexerQueryMethod queryMethod, ReindexerEntit
 		for (ReindexerIndex index : namespace.getIndexes()) {
 			this.indexes.put(index.getName(), index);
 		}
+		this.selectFields = getSelectFields();
+	}
+
+	private String[] getSelectFields() {
+		if (this.queryMethod.getDomainClass() == this.queryMethod.getReturnedObjectType()
+				|| this.queryMethod.getParameters().hasDynamicProjection()) {
+			return new String[0];
+		}
+		return getSelectFields(this.queryMethod.getReturnedObjectType());
 	}
 
 	@Override
 	public Object execute(Object[] parameters) {
 		Query<?> query = createQuery(parameters);
 		if (this.queryMethod.isCollectionQuery()) {
-			return query.toList();
+			return toCollection(query, parameters);
 		}
 		if (this.queryMethod.isStreamQuery()) {
-			return query.stream();
+			return toStream(query, parameters);
 		}
 		if (this.queryMethod.isIteratorQuery()) {
-			return query.execute();
+			return new ProjectingResultIterator(query, parameters);
 		}
 		if (this.queryMethod.isQueryForEntity()) {
-			return query.getOne();
+			Object entity = toEntity(query, parameters);
+			Assert.state(entity != null, "Exactly one item expected, but there is zero");
+			return entity;
 		}
 		if (this.tree.isExistsProjection()) {
 			return query.exists();
@@ -99,13 +123,13 @@ public Object execute(Object[] parameters) {
 			query.delete();
 			return null;
 		}
-		return query.findOne();
+		return Optional.ofNullable(toEntity(query, parameters));
 	}
 
 	private Query<?> createQuery(Object[] parameters) {
 		ParametersParameterAccessor accessor =
 				new ParametersParameterAccessor(this.queryMethod.getParameters(), parameters);
-		Query<?> base = this.namespace.query();
+		Query<?> base = this.namespace.query().select(getSelectFields(parameters));
 		Iterator<Object> iterator = accessor.iterator();
 		for (OrPart node : this.tree) {
 			Iterator<Part> parts = node.iterator();
@@ -125,6 +149,34 @@ private Query<?> createQuery(Object[] parameters) {
 		return base;
 	}
 
+	private String[] getSelectFields(Object[] parameters) {
+		if (this.queryMethod.getParameters().hasDynamicProjection()) {
+			Class<?> type = (Class<?>) parameters[this.queryMethod.getParameters().getDynamicProjectionIndex()];
+			return getSelectFields(type);
+		}
+		return this.selectFields;
+	}
+
+	private String[] getSelectFields(Class<?> type) {
+		if (type.isInterface()) {
+			List<PropertyDescriptor> inputProperties = this.queryMethod.getFactory()
+					.getProjectionInformation(type).getInputProperties();
+			String[] result = new String[inputProperties.size()];
+			for (int i = 0; i < result.length; i++) {
+				result[i] = inputProperties.get(i).getName();
+			}
+			return result;
+		}
+		else {
+			List<Field> inheritedFields = BeanPropertyUtils.getInheritedFields(type);
+			String[] result = new String[inheritedFields.size()];
+			for (int i = 0; i < result.length; i++) {
+				result[i] = inheritedFields.get(i).getName();
+			}
+			return result;
+		}
+	}
+
 	private Query<?> where(Part part, Query<?> criteria, Iterator<Object> parameters) {
 		String indexName = part.getProperty().toDotPath();
 		switch (part.getType()) {
@@ -188,9 +240,110 @@ private Object getParameterValue(String indexName, Object value) {
 		return value;
 	}
 
+	private Collection<?> toCollection(Query<?> query, Object[] parameters) {
+		try (ResultIterator<?> iterator = new ProjectingResultIterator(query, parameters)) {
+			Collection<Object> result = CollectionFactory.createCollection(this.queryMethod.getReturnType(),
+					this.queryMethod.getReturnedObjectType(), (int) iterator.size());
+			while (iterator.hasNext()) {
+				result.add(iterator.next());
+			}
+			return result;
+		}
+	}
+
+	private Stream<?> toStream(Query<?> query, Object[] parameters) {
+		ResultIterator<?> iterator = new ProjectingResultIterator(query, parameters);
+		Spliterator<?> spliterator = Spliterators.spliterator(iterator, iterator.size(), Spliterator.NONNULL);
+		return StreamSupport.stream(spliterator, false).onClose(iterator::close);
+	}
+
+	private Object toEntity(Query<?> query, Object[] parameters) {
+		try (ResultIterator<?> iterator = new ProjectingResultIterator(query, parameters)) {
+			Object item = null;
+			if (iterator.hasNext()) {
+				item = iterator.next();
+			}
+			if (iterator.hasNext()) {
+				throw new IllegalStateException("Exactly one item expected, but there are more");
+			}
+			return item;
+		} catch (Exception e) {
+			throw new RuntimeException(e.getMessage(), e);
+		}
+	}
+
 	@Override
 	public ReindexerQueryMethod getQueryMethod() {
 		return this.queryMethod;
 	}
 
+	private final class ProjectingResultIterator implements ResultIterator<Object> {
+
+		private final Object[] parameters;
+
+		private final ResultIterator<?> delegate;
+
+		private ProjectingResultIterator(Query<?> query, Object[] parameters) {
+			this.parameters = parameters;
+			this.delegate = query.execute(determineReturnType());
+		}
+
+		private Class<?> determineReturnType() {
+			if (ReindexerRepositoryQuery.this.queryMethod.getParameters().hasDynamicProjection()) {
+				Class<?> type = (Class<?>) this.parameters[ReindexerRepositoryQuery.this.queryMethod.getParameters().getDynamicProjectionIndex()];
+				if (type.isInterface()) {
+					return ReindexerRepositoryQuery.this.queryMethod.getDomainClass();
+				}
+				return type;
+			}
+			if (ReindexerRepositoryQuery.this.queryMethod.getDomainClass() != ReindexerRepositoryQuery.this.queryMethod.getReturnedObjectType()
+					&& ReindexerRepositoryQuery.this.queryMethod.getReturnedObjectType().isInterface()) {
+				return ReindexerRepositoryQuery.this.queryMethod.getDomainClass();
+			}
+			return ReindexerRepositoryQuery.this.queryMethod.getReturnedObjectType();
+		}
+
+		@Override
+		public long getTotalCount() {
+			return this.delegate.getTotalCount();
+		}
+
+		@Override
+		public long size() {
+			return this.delegate.size();
+		}
+
+		@Override
+		public List<AggregationResult> aggResults() {
+			return this.delegate.aggResults();
+		}
+
+		@Override
+		public void close() {
+			this.delegate.close();
+		}
+
+		@Override
+		public boolean hasNext() {
+			return this.delegate.hasNext();
+		}
+
+		@Override
+		public Object next() {
+			Object item = this.delegate.next();
+			if (ReindexerRepositoryQuery.this.queryMethod.getParameters().hasDynamicProjection()) {
+				Class<?> type = (Class<?>) this.parameters[ReindexerRepositoryQuery.this.queryMethod.getParameters().getDynamicProjectionIndex()];
+				if (type.isInterface()) {
+					return ReindexerRepositoryQuery.this.queryMethod.getFactory().createProjection(type, item);
+				}
+			}
+			else if (ReindexerRepositoryQuery.this.queryMethod.getReturnedObjectType().isInterface()) {
+				return ReindexerRepositoryQuery.this.queryMethod.getFactory()
+						.createProjection(ReindexerRepositoryQuery.this.queryMethod.getReturnedObjectType(), item);
+			}
+			return item;
+		}
+
+	}
+
 }
diff --git a/...ingframework/data/reindexer/repository/support/ReindexerAnnotationRepositoryMetadata.java b/...ingframework/data/reindexer/repository/support/ReindexerAnnotationRepositoryMetadata.java
@@ -0,0 +1,50 @@
+/*
+ * Copyright 2022 evgeniycheban
+ *
+ * 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.reindexer.repository.support;
+
+import java.lang.reflect.Method;
+
+import org.springframework.data.reindexer.repository.util.ReindexerQueryExecutionConverters;
+import org.springframework.data.repository.core.support.AnnotationRepositoryMetadata;
+import org.springframework.data.repository.util.ReactiveWrapperConverters;
+import org.springframework.data.util.TypeInformation;
+
+/**
+ * {@link org.springframework.data.repository.core.RepositoryMetadata} implementation inspecting the given
+ * repository interface for a {@link org.springframework.data.repository.RepositoryDefinition} annotation.
+ *
+ * @author Evgeniy Cheban
+ */
+public class ReindexerAnnotationRepositoryMetadata extends AnnotationRepositoryMetadata {
+
+	/**
+	 * Creates a new {@link ReindexerAnnotationRepositoryMetadata} instance looking up repository types from a
+	 * {@link org.springframework.data.repository.RepositoryDefinition} annotation.
+	 *
+	 * @param repositoryInterface must not be {@literal null}.
+	 */
+	public ReindexerAnnotationRepositoryMetadata(Class<?> repositoryInterface) {
+		super(repositoryInterface);
+	}
+
+	@Override
+	public Class<?> getReturnedDomainClass(Method method) {
+		TypeInformation<?> returnType = getReturnType(method);
+		returnType = ReactiveWrapperConverters.unwrapWrapperTypes(returnType);
+		return ReindexerQueryExecutionConverters.unwrapWrapperTypes(returnType, getDomainTypeInformation()).getType();
+	}
+
+}