Allow IndexCoordinates in repositories

It is common to have multiple indexes with the same document structure. E.g. rolling indices over time and having a wildcard alias. Sometimes one wants to select different IndexCoordinates than specified in `@Document(indexName = "...", ...)`. Also you might want to use the same Repository to do cross cluster searches and for this need to specify the IndexCoordinates. It would therefore be helpful to support IndexCoordinates as an argument to repository methods the same way ScrollPosition, Sort and Pageable are handled. This PR is an attempt to introduce this change. I keep it as draft to get some comments I can add tests to it if we can agree on the change being valid.
spring-projects · Mar 27, 2023 · f8eaa18 · f8eaa18
1 parent 675b779
commit f8eaa18
Show file tree

Hide file tree

Showing 9 changed files with 94 additions and 37 deletions.
diff --git a/...ngframework/data/elasticsearch/repository/query/AbstractElasticsearchRepositoryQuery.java b/...ngframework/data/elasticsearch/repository/query/AbstractElasticsearchRepositoryQuery.java
@@ -15,12 +15,10 @@
  */
 package org.springframework.data.elasticsearch.repository.query;
 
+import java.util.Collections;
+
 import org.springframework.data.domain.PageRequest;
-import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
-import org.springframework.data.elasticsearch.core.SearchHitSupport;
-import org.springframework.data.elasticsearch.core.SearchHits;
-import org.springframework.data.elasticsearch.core.SearchHitsImpl;
-import org.springframework.data.elasticsearch.core.TotalHitsRelation;
+import org.springframework.data.elasticsearch.core.*;
 import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverter;
 import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
 import org.springframework.data.elasticsearch.core.query.Query;
@@ -32,8 +30,6 @@
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
 
-import java.util.Collections;
-
 /**
  * AbstractElasticsearchRepositoryQuery
  *
@@ -74,15 +70,17 @@ public QueryMethod getQueryMethod() {
 	@Override
 	public Object execute(Object[] parameters) {
 
-		ParametersParameterAccessor parameterAccessor = getParameterAccessor(parameters);
+		ElasticsearchParametersParameterAccessor parameterAccessor = getParameterAccessor(parameters);
 		Class<?> clazz = getResultClass();
 		Query query = createQuery(parameters);
 
-		IndexCoordinates index = elasticsearchOperations.getIndexCoordinatesFor(clazz);
+		IndexCoordinates index = parameterAccessor
+				.getIndexCoordinatesOrDefaults(elasticsearchOperations.getIndexCoordinatesFor(clazz));
 
 		Object result = null;
 
 		if (isDeleteQuery()) {
+			index = elasticsearchOperations.getIndexCoordinatesFor(clazz);
 			result = countOrGetDocumentsForDelete(query, parameterAccessor);
 			elasticsearchOperations.delete(query, clazz, index);
 			elasticsearchOperations.indexOps(index).refresh();
@@ -158,8 +156,8 @@ private Class<?> getResultClass() {
 		return queryMethod.getResultProcessor().getReturnedType().getDomainType();
 	}
 
-	private ParametersParameterAccessor getParameterAccessor(Object[] parameters) {
-		return new ParametersParameterAccessor(queryMethod.getParameters(), parameters);
+	private ElasticsearchParametersParameterAccessor getParameterAccessor(Object[] parameters) {
+		return new ElasticsearchParametersParameterAccessor(queryMethod, parameters);
 	}
 
 	@Nullable

diff --git a/...ork/data/elasticsearch/repository/query/AbstractReactiveElasticsearchRepositoryQuery.java b/...ork/data/elasticsearch/repository/query/AbstractReactiveElasticsearchRepositoryQuery.java
@@ -95,7 +95,7 @@ private Object execute(ElasticsearchParameterAccessor parameterAccessor) {
 		}
 
 		String indexName = queryMethod.getEntityInformation().getIndexName();
-		IndexCoordinates index = IndexCoordinates.of(indexName);
+		IndexCoordinates index = parameterAccessor.getIndexCoordinatesOrDefaults(IndexCoordinates.of(indexName));
 
 		ReactiveElasticsearchQueryExecution execution = getExecution(parameterAccessor,
 				new ResultProcessingConverter(processor));

diff --git a/...g/springframework/data/elasticsearch/repository/query/ElasticsearchParameterAccessor.java b/...g/springframework/data/elasticsearch/repository/query/ElasticsearchParameterAccessor.java
@@ -15,7 +15,9 @@
  */
 package org.springframework.data.elasticsearch.repository.query;
 
+import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
 import org.springframework.data.repository.query.ParameterAccessor;
+import org.springframework.lang.NonNull;
 
 /**
  * @author Christoph Strobl
@@ -26,7 +28,9 @@ public interface ElasticsearchParameterAccessor extends ParameterAccessor {
 	/**
 	 * Returns the raw parameter values of the underlying query method.
 	 *
-	 * @return
+	 * @return values.
 	 */
 	Object[] getValues();
+
+	IndexCoordinates getIndexCoordinatesOrDefaults(@NonNull IndexCoordinates defaults);
 }
diff --git a/...java/org/springframework/data/elasticsearch/repository/query/ElasticsearchParameters.java b/...java/org/springframework/data/elasticsearch/repository/query/ElasticsearchParameters.java
@@ -16,54 +16,89 @@
 package org.springframework.data.elasticsearch.repository.query;
 
 import java.lang.reflect.Method;
+import java.util.ArrayList;
 import java.util.List;
 
 import org.springframework.core.MethodParameter;
+import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
 import org.springframework.data.elasticsearch.repository.query.ElasticsearchParameters.ElasticsearchParameter;
 import org.springframework.data.geo.Distance;
 import org.springframework.data.repository.query.Parameter;
 import org.springframework.data.repository.query.Parameters;
+import org.springframework.data.util.TypeInformation;
 
 /**
  * @author Christoph Strobl
  * @author Peter-Josef Meisch
  * @since 3.2
  */
 public class ElasticsearchParameters extends Parameters<ElasticsearchParameters, ElasticsearchParameter> {
-
-	public ElasticsearchParameters(Method method) {
-		super(method);
-	}
+	private final int indexCoordinatesIndex;
 
 	private ElasticsearchParameters(List<ElasticsearchParameter> parameters) {
 		super(parameters);
+		this.indexCoordinatesIndex = initIndexCoordinatesIndex();
 	}
 
-	@Override
-	protected ElasticsearchParameter createParameter(MethodParameter parameter) {
-		return new ElasticsearchParameter(parameter);
+	public ElasticsearchParameters(Method method, TypeInformation<?> aggregateType) {
+		super(method, (param) -> new ElasticsearchParameter(param, aggregateType));
+		this.indexCoordinatesIndex = initIndexCoordinatesIndex();
+	}
+
+	private int initIndexCoordinatesIndex() {
+		int index = 0;
+		List<Integer> foundIndices = new ArrayList<>();
+		for (ElasticsearchParameter parameter : this) {
+			if (parameter.isIndexCoordinatesParameter()) {
+				foundIndices.add(index);
+			}
+			index++;
+		}
+		if (foundIndices.size() > 1) {
+			throw new IllegalArgumentException(this + " can only contain at most one IndexCoordinates parameter.");
+		}
+		return foundIndices.isEmpty() ? -1 : foundIndices.get(0);
 	}
 
 	@Override
 	protected ElasticsearchParameters createFrom(List<ElasticsearchParameter> parameters) {
 		return new ElasticsearchParameters(parameters);
 	}
 
+	public boolean hasIndexCoordinatesParameter() {
+		return this.indexCoordinatesIndex != -1;
+	}
+
+	public int getIndexCoordinatesIndex() {
+		return indexCoordinatesIndex;
+	}
+
 	/**
 	 * Custom {@link Parameter} implementation adding parameters of type {@link Distance} to the special ones.
 	 *
 	 * @author Christoph Strobl
 	 */
-	class ElasticsearchParameter extends Parameter {
+	static class ElasticsearchParameter extends Parameter {
+
+		private final boolean indexCoordinatesParameter;
 
 		/**
 		 * Creates a new {@link ElasticsearchParameter}.
 		 *
 		 * @param parameter must not be {@literal null}.
 		 */
-		ElasticsearchParameter(MethodParameter parameter) {
-			super(parameter);
+		ElasticsearchParameter(MethodParameter parameter, TypeInformation<?> aggregateType) {
+			super(parameter, aggregateType);
+			this.indexCoordinatesParameter = parameter.getParameter().getType().isAssignableFrom(IndexCoordinates.class);
 		}
 
+		@Override
+		public boolean isSpecialParameter() {
+			return isIndexCoordinatesParameter() || super.isSpecialParameter();
+		}
+
+		public boolean isIndexCoordinatesParameter() {
+			return this.indexCoordinatesParameter;
+		}
 	}
 }
diff --git a/...amework/data/elasticsearch/repository/query/ElasticsearchParametersParameterAccessor.java b/...amework/data/elasticsearch/repository/query/ElasticsearchParametersParameterAccessor.java
@@ -18,7 +18,9 @@
 import java.util.Arrays;
 import java.util.List;
 
+import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
 import org.springframework.data.repository.query.ParametersParameterAccessor;
+import org.springframework.lang.NonNull;
 
 /**
  * @author Christoph Strobl
@@ -36,13 +38,25 @@ class ElasticsearchParametersParameterAccessor extends ParametersParameterAccess
 	 * @param values must not be {@literal null}.
 	 */
 	ElasticsearchParametersParameterAccessor(ElasticsearchQueryMethod method, Object... values) {
-
 		super(method.getParameters(), values);
 		this.values = Arrays.asList(values);
 	}
 
+	@Override
+	public ElasticsearchParameters getParameters() {
+		return (ElasticsearchParameters) super.getParameters();
+	}
+
 	@Override
 	public Object[] getValues() {
 		return values.toArray();
 	}
+
+	@Override
+	public IndexCoordinates getIndexCoordinatesOrDefaults(@NonNull IndexCoordinates defaults) {
+		if (!getParameters().hasIndexCoordinatesParameter()) {
+			return defaults;
+		}
+		return (IndexCoordinates) getValues()[getParameters().getIndexCoordinatesIndex()];
+	}
 }
diff --git a/...ava/org/springframework/data/elasticsearch/repository/query/ElasticsearchQueryMethod.java b/...ava/org/springframework/data/elasticsearch/repository/query/ElasticsearchQueryMethod.java
@@ -322,4 +322,14 @@ private String[] mapParameters(String[] source, ParameterAccessor parameterAcces
 
 		return fieldNames.toArray(new String[0]);
 	}
+
+	@Override
+	protected ElasticsearchParameters createParameters(Method method, TypeInformation<?> domainType) {
+		return new ElasticsearchParameters(method, domainType);
+	}
+
+	@Override
+	public ElasticsearchParameters getParameters() {
+		return (ElasticsearchParameters) super.getParameters();
+	}
 }
diff --git a/...data/elasticsearch/repository/query/ReactiveElasticsearchParametersParameterAccessor.java b/...data/elasticsearch/repository/query/ReactiveElasticsearchParametersParameterAccessor.java
@@ -81,7 +81,7 @@ protected <T> T getValue(int index) {
 	@Override
 	public Object[] getValues() {
 
-		Object[] result = new Object[getValues().length];
+		Object[] result = new Object[super.getValues().length];
 		for (int i = 0; i < result.length; i++) {
 			result[i] = getValue(i);
 		}

diff --git a/...springframework/data/elasticsearch/repository/query/ReactiveElasticsearchQueryMethod.java b/...springframework/data/elasticsearch/repository/query/ReactiveElasticsearchQueryMethod.java
@@ -15,7 +15,7 @@
  */
 package org.springframework.data.elasticsearch.repository.query;
 
-import static org.springframework.data.repository.util.ClassUtils.*;
+import static org.springframework.data.repository.util.ClassUtils.hasParameterOfType;
 
 import reactor.core.publisher.Mono;
 
@@ -102,11 +102,6 @@ protected void verifyCountQueryTypes() {
 		}
 	}
 
-	@Override
-	protected ElasticsearchParameters createParameters(Method method) {
-		return new ElasticsearchParameters(method);
-	}
-
 	/**
 	 * Check if the given {@link org.springframework.data.repository.query.QueryMethod} receives a reactive parameter
 	 * wrapper as one of its parameters.
@@ -143,11 +138,6 @@ public boolean isStreamQuery() {
 		return true;
 	}
 
-	@Override
-	public ElasticsearchParameters getParameters() {
-		return (ElasticsearchParameters) super.getParameters();
-	}
-
 	@Override
 	protected boolean isAllowedGenericType(ParameterizedType methodGenericReturnType) {
 		return super.isAllowedGenericType(methodGenericReturnType)

diff --git a/...t/java/org/springframework/data/elasticsearch/repository/query/StubParameterAccessor.java b/...t/java/org/springframework/data/elasticsearch/repository/query/StubParameterAccessor.java
@@ -17,12 +17,13 @@
 
 import java.util.Arrays;
 import java.util.Iterator;
-import java.util.Optional;
 
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.ScrollPosition;
 import org.springframework.data.domain.Sort;
+import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
 import org.springframework.data.repository.query.ParameterAccessor;
+import org.springframework.lang.NonNull;
 
 /**
  * Simple {@link ParameterAccessor} that returns the given parameters unfiltered.
@@ -97,6 +98,11 @@ public Object[] getValues() {
 		return this.values;
 	}
 
+	@Override
+	public IndexCoordinates getIndexCoordinatesOrDefaults(@NonNull IndexCoordinates defaults) {
+		return defaults;
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.repository.query.ParameterAccessor#findDynamicProjection()