Allow configuration of QueryEngine for Map repositories

pom.xml

@@ -5,7 +5,7 @@
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-keyvalue</artifactId>
-	<version>3.4.0-SNAPSHOT</version>
+	<version>3.4.0-GH-576-SNAPSHOT</version>
 
 	<name>Spring Data KeyValue</name>
 

src/main/java/org/springframework/data/keyvalue/core/PredicateQueryEngine.java

@@ -32,7 +32,7 @@
  * @author Christoph Strobl
  * @since 3.3
  */
-class PredicateQueryEngine extends QueryEngine<KeyValueAdapter, Predicate<?>, Comparator<?>> {
+public class PredicateQueryEngine extends QueryEngine<KeyValueAdapter, Predicate<?>, Comparator<?>> {
 
 	/**
 	 * Creates a new {@link PredicateQueryEngine}.

src/main/java/org/springframework/data/keyvalue/core/QueryEngineFactory.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2024 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
+ *
+ *      https://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.keyvalue.core;
+
+/**
+ * Interface for {@code QueryEngineFactory} implementations that provide a {@link QueryEngine} object as part of the
+ * configuration.
+ * <p>
+ * The factory is used during configuration to supply the query engine to be used. When configured, a
+ * {@code QueryEngineFactory} can be instantiated by accepting a {@link SortAccessor} in its constructor. Otherwise,
+ * implementations are expected to declare a no-args constructor.
+ *
+ * @author Mark Paluch
+ * @since 3.3.1
+ */
+public interface QueryEngineFactory {
+
+	/**
+	 * Factory method for creating a {@link QueryEngine}.
+	 *
+	 * @return the query engine.
+	 */
+	QueryEngine<?, ?, ?> create();
+}

src/main/java/org/springframework/data/keyvalue/core/SpelQueryEngine.java

@@ -34,9 +34,8 @@
  * @author Christoph Strobl
  * @author Oliver Gierke
  * @author Mark Paluch
- * @param <T>
  */
-class SpelQueryEngine extends QueryEngine<KeyValueAdapter, SpelCriteria, Comparator<?>> {
+public class SpelQueryEngine extends QueryEngine<KeyValueAdapter, SpelCriteria, Comparator<?>> {
 
 	private static final SpelExpressionParser PARSER = new SpelExpressionParser();
 

src/main/java/org/springframework/data/keyvalue/repository/config/KeyValueRepositoryConfigurationExtension.java

@@ -25,6 +25,7 @@
 import org.springframework.beans.factory.support.BeanDefinitionRegistry;
 import org.springframework.beans.factory.support.RootBeanDefinition;
 import org.springframework.core.annotation.AnnotationAttributes;
+import org.springframework.core.annotation.MergedAnnotation;
 import org.springframework.core.type.AnnotationMetadata;
 import org.springframework.data.keyvalue.core.mapping.context.KeyValueMappingContext;
 import org.springframework.data.keyvalue.repository.KeyValueRepository;
@@ -36,7 +37,6 @@
 import org.springframework.data.repository.config.RepositoryConfigurationExtensionSupport;
 import org.springframework.data.repository.config.RepositoryConfigurationSource;
 import org.springframework.lang.Nullable;
-import org.springframework.util.CollectionUtils;
 
 /**
  * {@link RepositoryConfigurationExtension} for {@link KeyValueRepository}.
@@ -90,17 +90,15 @@ public void postProcess(BeanDefinitionBuilder builder, AnnotationRepositoryConfi
 	 */
 	private static Class<?> getQueryCreatorType(AnnotationRepositoryConfigurationSource config) {
 
-		AnnotationMetadata metadata = config.getEnableAnnotationMetadata();
+		AnnotationMetadata amd = (AnnotationMetadata) config.getSource();
+		MergedAnnotation<QueryCreatorType> queryCreator = amd.getAnnotations().get(QueryCreatorType.class);
+		Class<?> queryCreatorType = queryCreator.isPresent() ? queryCreator.getClass("value") : Class.class;
 
-		Map<String, Object> queryCreatorAnnotationAttributes = metadata
-				.getAnnotationAttributes(QueryCreatorType.class.getName());
-
-		if (CollectionUtils.isEmpty(queryCreatorAnnotationAttributes)) {
+		if (queryCreatorType == Class.class) {
 			return SpelQueryCreator.class;
 		}
 
-		AnnotationAttributes queryCreatorAttributes = new AnnotationAttributes(queryCreatorAnnotationAttributes);
-		return queryCreatorAttributes.getClass("value");
+		return queryCreatorType;
 	}
 
 	/**
@@ -132,10 +130,10 @@ public void registerBeansForRoot(BeanDefinitionRegistry registry, RepositoryConf
 
 		registerIfNotAlreadyRegistered(() -> {
 
-			RootBeanDefinition definitionefinition = new RootBeanDefinition(KeyValueMappingContext.class);
-			definitionefinition.setSource(configurationSource.getSource());
+			RootBeanDefinition mappingContext = new RootBeanDefinition(KeyValueMappingContext.class);
+			mappingContext.setSource(configurationSource.getSource());
 
-			return definitionefinition;
+			return mappingContext;
 
 		}, registry, getMappingContextBeanRef(), configurationSource);
 

src/main/java/org/springframework/data/keyvalue/repository/query/KeyValuePartTreeQuery.java

@@ -89,7 +89,7 @@ public KeyValuePartTreeQuery(QueryMethod queryMethod, QueryMethodEvaluationConte
 			QueryCreatorFactory<AbstractQueryCreator<KeyValueQuery<?>, ?>> queryCreatorFactory) {
 
 		Assert.notNull(queryMethod, "Query method must not be null");
-		Assert.notNull(evaluationContextProvider, "EvaluationContextprovider must not be null");
+		Assert.notNull(evaluationContextProvider, "EvaluationContextProvider must not be null");
 		Assert.notNull(keyValueOperations, "KeyValueOperations must not be null");
 		Assert.notNull(queryCreatorFactory, "QueryCreatorFactory type must not be null");
 

src/main/java/org/springframework/data/keyvalue/repository/query/PredicateQueryCreator.java

@@ -38,6 +38,8 @@
 import org.springframework.util.comparator.NullSafeComparator;
 
 /**
+ * {@link AbstractQueryCreator} to create {@link Predicate}-based {@link KeyValueQuery}s.
+ *
  * @author Christoph Strobl
  * @since 3.3
  */

src/main/java/org/springframework/data/keyvalue/repository/query/SpelQueryCreator.java

@@ -32,7 +32,7 @@
 import org.springframework.util.StringUtils;
 
 /**
- * {@link AbstractQueryCreator} to create {@link SpelExpression} based {@link KeyValueQuery}s.
+ * {@link AbstractQueryCreator} to create {@link SpelExpression}-based {@link KeyValueQuery}s.
  *
  * @author Christoph Strobl
  * @author Oliver Gierke
@@ -180,8 +180,8 @@ protected SpelExpression toPredicateExpression(PartTree tree) {
 					case NEGATING_SIMPLE_PROPERTY:
 					case EXISTS:
 					default:
-						throw new InvalidDataAccessApiUsageException(String.format("Found invalid part '%s' in query",
-								part.getType()));
+						throw new InvalidDataAccessApiUsageException(
+								String.format("Found invalid part '%s' in query", part.getType()));
 				}
 
 				if (partIter.hasNext()) {

src/main/java/org/springframework/data/map/repository/config/EnableMapRepositories.java

@@ -28,30 +28,33 @@
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.ComponentScan.Filter;
 import org.springframework.context.annotation.Import;
+import org.springframework.core.annotation.AliasFor;
 import org.springframework.data.keyvalue.core.KeyValueOperations;
 import org.springframework.data.keyvalue.core.KeyValueTemplate;
+import org.springframework.data.keyvalue.core.QueryEngineFactory;
 import org.springframework.data.keyvalue.core.SortAccessor;
 import org.springframework.data.keyvalue.repository.config.QueryCreatorType;
-import org.springframework.data.keyvalue.repository.query.KeyValuePartTreeQuery;
 import org.springframework.data.keyvalue.repository.query.PredicateQueryCreator;
 import org.springframework.data.keyvalue.repository.support.KeyValueRepositoryFactoryBean;
 import org.springframework.data.repository.config.DefaultRepositoryBaseClass;
 import org.springframework.data.repository.query.QueryLookupStrategy;
 import org.springframework.data.repository.query.QueryLookupStrategy.Key;
+import org.springframework.data.repository.query.parser.AbstractQueryCreator;
 
 /**
  * Annotation to activate Map repositories. If no base package is configured through either {@link #value()},
  * {@link #basePackages()} or {@link #basePackageClasses()} it will trigger scanning of the package of annotated class.
  *
  * @author Christoph Strobl
  * @author Oliver Gierke
+ * @author Mark Paluch
  */
 @Target(ElementType.TYPE)
 @Retention(RetentionPolicy.RUNTIME)
 @Documented
 @Inherited
 @Import(MapRepositoriesRegistrar.class)
-@QueryCreatorType(value = PredicateQueryCreator.class, repositoryQueryType = KeyValuePartTreeQuery.class)
+@QueryCreatorType(PredicateQueryCreator.class)
 public @interface EnableMapRepositories {
 
 	/**
@@ -145,10 +148,29 @@
 	@SuppressWarnings("rawtypes")
 	Class<? extends Map> mapType() default ConcurrentHashMap.class;
 
+	/**
+	 * Configures the {@link QueryEngineFactory} to create the QueryEngine. When both, the query engine and sort accessors
+	 * are configured, the query engine is instantiated using the configured sort accessor.
+	 *
+	 * @return {@link QueryEngineFactory} to configure the QueryEngine.
+	 * @since 3.3.1
+	 */
+	Class<? extends QueryEngineFactory> queryEngineFactory() default QueryEngineFactory.class;
+
+	/**
+	 * Configures the {@code QueryCreator} to create Part-Tree queries. The QueryCreator must create queries supported by
+	 * the underlying {@code QueryEngine}.
+	 *
+	 * @return {@link AbstractQueryCreator}
+	 * @since 3.3.1
+	 */
+	@AliasFor(annotation = QueryCreatorType.class, value = "value")
+	Class<? extends AbstractQueryCreator<?, ?>> queryCreator() default PredicateQueryCreator.class;
+
 	/**
 	 * Configures the {@link SortAccessor accessor} for sorting results.
 	 *
-	 * @return {@link SortAccessor} to indicate usage of default implementation.
+	 * @return the configured {@link SortAccessor}.
 	 * @since 3.1.10
 	 */
 	Class<? extends SortAccessor> sortAccessor() default SortAccessor.class;

src/main/java/org/springframework/data/map/repository/config/MapRepositoryConfigurationExtension.java

@@ -15,6 +15,7 @@
  */
 package org.springframework.data.map.repository.config;
 
+import java.lang.reflect.Constructor;
 import java.util.Map;
 
 import org.springframework.beans.BeanUtils;
@@ -24,15 +25,23 @@
 import org.springframework.core.type.AnnotationMetadata;
 import org.springframework.data.config.ParsingUtils;
 import org.springframework.data.keyvalue.core.KeyValueTemplate;
+import org.springframework.data.keyvalue.core.QueryEngine;
+import org.springframework.data.keyvalue.core.QueryEngineFactory;
 import org.springframework.data.keyvalue.core.SortAccessor;
 import org.springframework.data.keyvalue.repository.config.KeyValueRepositoryConfigurationExtension;
 import org.springframework.data.map.MapKeyValueAdapter;
+import org.springframework.data.repository.config.RepositoryConfigurationExtension;
 import org.springframework.data.repository.config.RepositoryConfigurationSource;
 import org.springframework.lang.Nullable;
+import org.springframework.util.ClassUtils;
 
 /**
+ * {@link RepositoryConfigurationExtension} for Map-based repositories.
+ *
  * @author Christoph Strobl
+ * @author Mark Paluch
  */
+@SuppressWarnings("unchecked")
 public class MapRepositoryConfigurationExtension extends KeyValueRepositoryConfigurationExtension {
 
 	@Override
@@ -58,7 +67,11 @@ protected AbstractBeanDefinition getDefaultKeyValueTemplateBeanDefinition(
 		adapterBuilder.addConstructorArgValue(getMapTypeToUse(configurationSource));
 
 		SortAccessor<?> sortAccessor = getSortAccessor(configurationSource);
-		if(sortAccessor != null) {
+		QueryEngine<?, ?, ?> queryEngine = getQueryEngine(sortAccessor, configurationSource);
+
+		if (queryEngine != null) {
+			adapterBuilder.addConstructorArgValue(queryEngine);
+		} else if (sortAccessor != null) {
 			adapterBuilder.addConstructorArgValue(sortAccessor);
 		}
 
@@ -73,20 +86,60 @@ protected AbstractBeanDefinition getDefaultKeyValueTemplateBeanDefinition(
 	@SuppressWarnings({ "unchecked", "rawtypes" })
 	private static Class<? extends Map> getMapTypeToUse(RepositoryConfigurationSource source) {
 
-		return (Class<? extends Map>) ((AnnotationMetadata) source.getSource()).getAnnotationAttributes(
-				EnableMapRepositories.class.getName()).get("mapType");
+		return (Class<? extends Map>) getAnnotationAttributes(source).get("mapType");
 	}
 
 	@Nullable
 	private static SortAccessor<?> getSortAccessor(RepositoryConfigurationSource source) {
 
-		Class<? extends SortAccessor<?>> sortAccessorType = (Class<? extends SortAccessor<?>>) ((AnnotationMetadata) source.getSource()).getAnnotationAttributes(
-				EnableMapRepositories.class.getName()).get("sortAccessor");
+		Class<? extends SortAccessor<?>> sortAccessorType = (Class<? extends SortAccessor<?>>) getAnnotationAttributes(
+				source).get("sortAccessor");
 
-		if(sortAccessorType != null && !sortAccessorType.isInterface()) {
+		if (sortAccessorType != null && !sortAccessorType.isInterface()) {
 			return BeanUtils.instantiateClass(sortAccessorType);
 		}
 
 		return null;
 	}
+
+	@Nullable
+	private static QueryEngine<?, ?, ?> getQueryEngine(@Nullable SortAccessor<?> sortAccessor,
+			RepositoryConfigurationSource source) {
+
+		Class<? extends QueryEngineFactory> queryEngineFactoryType = (Class<? extends QueryEngineFactory>) getAnnotationAttributes(
+				source).get("queryEngineFactory");
+
+		if (queryEngineFactoryType != null && !queryEngineFactoryType.isInterface()) {
+
+			if (sortAccessor != null) {
+				Constructor<? extends QueryEngineFactory> constructor = ClassUtils
+						.getConstructorIfAvailable(queryEngineFactoryType, SortAccessor.class);
+				if (constructor != null) {
+					return BeanUtils.instantiateClass(constructor, sortAccessor).create();
+				}
+			}
+
+			return BeanUtils.instantiateClass(queryEngineFactoryType).create();
+		}
+
+		return null;
+	}
+
+	private static Map<String, Object> getAnnotationAttributes(RepositoryConfigurationSource source) {
+
+		AnnotationMetadata annotationSource = (AnnotationMetadata) source.getSource();
+
+		if (annotationSource == null) {
+			throw new IllegalArgumentException("AnnotationSource not available");
+		}
+
+		Map<String, Object> annotationAttributes = annotationSource
+				.getAnnotationAttributes(EnableMapRepositories.class.getName());
+
+		if (annotationAttributes == null) {
+			throw new IllegalStateException("No annotation attributes for @EnableMapRepositories");
+		}
+
+		return annotationAttributes;
+	}
 }