GH-2428 - Provide mapping functions via supplier.

michael-simons · michael-simons · commit 2bd28befa9ce · 2021-11-04T19:35:18.000+01:00
This avoids potentially stale objects during caches in mapping functions (like in our own) and fixes #2428.
diff --git a/src/main/java/org/springframework/data/neo4j/core/Neo4jTemplate.java b/src/main/java/org/springframework/data/neo4j/core/Neo4jTemplate.java
@@ -461,7 +461,7 @@ private <T> ExecutableQuery<T> createExecutableQuery(Class<T> domainType, String
 		PreparedQuery<T> preparedQuery = PreparedQuery.queryFor(domainType)
 				.withCypherQuery(cypherStatement)
 				.withParameters(parameters)
-				.usingMappingFunction(neo4jMappingContext.getRequiredMappingFunctionFor(domainType))
+				.usingMappingFunction(() -> neo4jMappingContext.getRequiredMappingFunctionFor(domainType))
 				.build();
 
 		return toExecutableQuery(preparedQuery);
@@ -677,7 +677,7 @@ private <T> ExecutableQuery<T> createExecutableQuery(Class<T> domainType,
 
 		PreparedQuery<T> preparedQuery = PreparedQuery.queryFor(domainType)
 				.withQueryFragmentsAndParameters(queryFragmentsAndParameters)
-				.usingMappingFunction(neo4jMappingContext.getRequiredMappingFunctionFor(domainType))
+				.usingMappingFunction(() -> neo4jMappingContext.getRequiredMappingFunctionFor(domainType))
 				.build();
 		return toExecutableQuery(preparedQuery);
 	}
diff --git a/src/main/java/org/springframework/data/neo4j/core/PreparedQuery.java b/src/main/java/org/springframework/data/neo4j/core/PreparedQuery.java
@@ -40,6 +40,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BiFunction;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 /**
@@ -65,29 +66,33 @@ public static <CT> RequiredBuildStep<CT> queryFor(Class<CT> resultType) {
 
 	private final Class<T> resultType;
 	private final QueryFragmentsAndParameters queryFragmentsAndParameters;
-	private final @Nullable BiFunction<TypeSystem, Record, T> mappingFunction;
+	private final @Nullable Supplier<BiFunction<TypeSystem, MapAccessor, ?>> mappingFunctionSupplier;
+	private volatile Optional<BiFunction<TypeSystem, Record, T>> lastMappingFunction = Optional.empty();
 
 	private PreparedQuery(OptionalBuildSteps<T> optionalBuildSteps) {
 		this.resultType = optionalBuildSteps.resultType;
-		if (optionalBuildSteps.mappingFunction == null) {
-			this.mappingFunction = null;
-		} else {
-			this.mappingFunction = (BiFunction<TypeSystem, Record, T>) new AggregatingMappingFunction(
-					optionalBuildSteps.mappingFunction);
-		}
+		this.mappingFunctionSupplier = optionalBuildSteps.mappingFunctionSupplier;
 		this.queryFragmentsAndParameters = optionalBuildSteps.queryFragmentsAndParameters;
 	}
 
 	public Class<T> getResultType() {
 		return this.resultType;
 	}
 
-	public Optional<BiFunction<TypeSystem, Record, T>> getOptionalMappingFunction() {
-		return Optional.ofNullable(mappingFunction);
+	@SuppressWarnings("unchecked")
+	public synchronized Optional<BiFunction<TypeSystem, Record, T>> getOptionalMappingFunction() {
+		lastMappingFunction = Optional.ofNullable(this.mappingFunctionSupplier)
+				.map(Supplier::get)
+				.map(f -> (BiFunction<TypeSystem, Record, T>) new AggregatingMappingFunction(f));
+		return lastMappingFunction;
 	}
 
-	boolean resultsHaveBeenAggregated() {
-		return this.mappingFunction != null && ((AggregatingMappingFunction) this.mappingFunction).hasAggregated();
+	synchronized boolean resultsHaveBeenAggregated() {
+		return lastMappingFunction
+				.filter(AggregatingMappingFunction.class::isInstance)
+				.map(AggregatingMappingFunction.class::cast)
+				.map(AggregatingMappingFunction::hasAggregated)
+				.orElse(false);
 	}
 
 	public QueryFragmentsAndParameters getQueryFragmentsAndParameters() {
@@ -122,7 +127,7 @@ public static class OptionalBuildSteps<CT> {
 
 		final Class<CT> resultType;
 		final QueryFragmentsAndParameters queryFragmentsAndParameters;
-		@Nullable BiFunction<TypeSystem, MapAccessor, ?> mappingFunction;
+		@Nullable Supplier<BiFunction<TypeSystem, MapAccessor, ?>> mappingFunctionSupplier;
 
 		OptionalBuildSteps(Class<CT> resultType, QueryFragmentsAndParameters queryFragmentsAndParameters) {
 			this.resultType = resultType;
@@ -140,9 +145,21 @@ public OptionalBuildSteps<CT> withParameters(Map<String, Object> newParameters)
 			return this;
 		}
 
+		/**
+		 * @param newMappingFunction A new mapping function
+		 * @return This builder
+		 * @deprecated since 6.1.7, please use {@link #usingMappingFunction(Supplier)}, otherwise your instance will be cached,
+		 * leading to potentially stale results
+		 */
+		@Deprecated
 		public OptionalBuildSteps<CT> usingMappingFunction(
 				@Nullable BiFunction<TypeSystem, MapAccessor, ?> newMappingFunction) {
-			this.mappingFunction = newMappingFunction;
+			this.mappingFunctionSupplier = newMappingFunction == null ? null : () -> newMappingFunction;
+			return this;
+		}
+
+		public OptionalBuildSteps<CT> usingMappingFunction(@Nullable Supplier<BiFunction<TypeSystem, MapAccessor, ?>> newMappingFunction) {
+			this.mappingFunctionSupplier = newMappingFunction;
 			return this;
 		}
 
diff --git a/src/main/java/org/springframework/data/neo4j/core/ReactiveNeo4jTemplate.java b/src/main/java/org/springframework/data/neo4j/core/ReactiveNeo4jTemplate.java
@@ -452,7 +452,7 @@ private <T> Mono<ExecutableQuery<T>> createExecutableQuery(Class<T> domainType,
 		Assert.notNull(neo4jMappingContext.getPersistentEntity(domainType), "Cannot get or create persistent entity.");
 		PreparedQuery<T> preparedQuery = PreparedQuery.queryFor(domainType).withCypherQuery(cypherQuery)
 				.withParameters(parameters)
-				.usingMappingFunction(this.neo4jMappingContext.getRequiredMappingFunctionFor(domainType)).build();
+				.usingMappingFunction(() -> this.neo4jMappingContext.getRequiredMappingFunctionFor(domainType)).build();
 		return this.toExecutableQuery(preparedQuery);
 	}
 
diff --git a/src/main/java/org/springframework/data/neo4j/repository/query/AbstractNeo4jQuery.java b/src/main/java/org/springframework/data/neo4j/repository/query/AbstractNeo4jQuery.java
@@ -141,7 +141,7 @@ private Slice<?> createSlice(boolean incrementLimit, Neo4jParameterAccessor para
 	protected abstract <T extends Object> PreparedQuery<T> prepareQuery(Class<T> returnedType,
 			List<String> includedProperties, Neo4jParameterAccessor parameterAccessor,
 			@Nullable Neo4jQueryType queryType,
-			@Nullable BiFunction<TypeSystem, MapAccessor, ?> mappingFunction,
+			@Nullable Supplier<BiFunction<TypeSystem, MapAccessor, ?>> mappingFunction,
 			@Nullable UnaryOperator<Integer> limitModifier);
 
 	protected Optional<PreparedQuery<Long>> getCountQuery(Neo4jParameterAccessor parameterAccessor) {
diff --git a/src/main/java/org/springframework/data/neo4j/repository/query/AbstractReactiveNeo4jQuery.java b/src/main/java/org/springframework/data/neo4j/repository/query/AbstractReactiveNeo4jQuery.java
@@ -17,6 +17,7 @@
 
 import java.util.List;
 import java.util.function.BiFunction;
+import java.util.function.Supplier;
 
 import org.neo4j.driver.types.MapAccessor;
 import org.neo4j.driver.types.TypeSystem;
@@ -83,5 +84,5 @@ public final Object execute(Object[] parameters) {
 
 	protected abstract <T extends Object> PreparedQuery<T> prepareQuery(Class<T> returnedType,
 			List<String> includedProperties, Neo4jParameterAccessor parameterAccessor, @Nullable Neo4jQueryType queryType,
-			@Nullable BiFunction<TypeSystem, MapAccessor, ?> mappingFunction);
+			@Nullable Supplier<BiFunction<TypeSystem, MapAccessor, ?>> mappingFunction);
 }
diff --git a/src/main/java/org/springframework/data/neo4j/repository/query/Neo4jQuerySupport.java b/src/main/java/org/springframework/data/neo4j/repository/query/Neo4jQuerySupport.java
@@ -95,25 +95,28 @@ static Class<?> getDomainType(QueryMethod queryMethod) {
 		this.queryType = queryType;
 	}
 
-	protected final BiFunction<TypeSystem, MapAccessor, ?> getMappingFunction(final ResultProcessor resultProcessor) {
-
-		final ReturnedType returnedTypeMetadata = resultProcessor.getReturnedType();
-		final Class<?> returnedType = returnedTypeMetadata.getReturnedType();
-		final Class<?> domainType = returnedTypeMetadata.getDomainType();
-
-		final BiFunction<TypeSystem, MapAccessor, ?> mappingFunction;
-
-		if (mappingContext.getConversionService().isSimpleType(returnedType)) {
-			// Clients automatically selects a single value mapping function.
-			// It will throw an error if the query contains more than one column.
-			mappingFunction = null;
-		} else if (returnedTypeMetadata.isProjecting()) {
-			BiFunction<TypeSystem, MapAccessor, ?> target = this.mappingContext.getRequiredMappingFunctionFor(domainType);
-			mappingFunction = (t, r) -> new EntityInstanceWithSource(target.apply(t, r), t, r);
-		} else {
-			mappingFunction = this.mappingContext.getRequiredMappingFunctionFor(domainType);
-		}
-		return mappingFunction;
+	protected final Supplier<BiFunction<TypeSystem, MapAccessor, ?>> getMappingFunction(final ResultProcessor resultProcessor) {
+
+		return () -> {
+			final ReturnedType returnedTypeMetadata = resultProcessor.getReturnedType();
+			final Class<?> returnedType = returnedTypeMetadata.getReturnedType();
+			final Class<?> domainType = returnedTypeMetadata.getDomainType();
+
+			final BiFunction<TypeSystem, MapAccessor, ?> mappingFunction;
+
+			if (mappingContext.getConversionService().isSimpleType(returnedType)) {
+				// Clients automatically selects a single value mapping function.
+				// It will throw an error if the query contains more than one column.
+				mappingFunction = null;
+			} else if (returnedTypeMetadata.isProjecting()) {
+				BiFunction<TypeSystem, MapAccessor, ?> target = this.mappingContext.getRequiredMappingFunctionFor(
+						domainType);
+				mappingFunction = (t, r) -> new EntityInstanceWithSource(target.apply(t, r), t, r);
+			} else {
+				mappingFunction = this.mappingContext.getRequiredMappingFunctionFor(domainType);
+			}
+			return mappingFunction;
+		};
 	}
 
 	protected final List<String> getInputProperties(final ResultProcessor resultProcessor) {
diff --git a/src/main/java/org/springframework/data/neo4j/repository/query/PartTreeNeo4jQuery.java b/src/main/java/org/springframework/data/neo4j/repository/query/PartTreeNeo4jQuery.java
@@ -18,6 +18,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.function.BiFunction;
+import java.util.function.Supplier;
 import java.util.function.UnaryOperator;
 
 import org.neo4j.driver.types.MapAccessor;
@@ -60,7 +61,7 @@ private PartTreeNeo4jQuery(Neo4jOperations neo4jOperations, Neo4jMappingContext
 	@Override
 	protected <T extends Object> PreparedQuery<T> prepareQuery(Class<T> returnedType, List<String> includedProperties,
 			Neo4jParameterAccessor parameterAccessor, @Nullable Neo4jQueryType queryType,
-			@Nullable BiFunction<TypeSystem, MapAccessor, ?> mappingFunction, UnaryOperator<Integer> limitModifier) {
+			@Nullable Supplier<BiFunction<TypeSystem, MapAccessor, ?>> mappingFunction, UnaryOperator<Integer> limitModifier) {
 
 		CypherQueryCreator queryCreator = new CypherQueryCreator(mappingContext, getDomainType(queryMethod),
 				Optional.ofNullable(queryType).orElseGet(() -> Neo4jQueryType.fromPartTree(tree)), tree, parameterAccessor,
diff --git a/src/main/java/org/springframework/data/neo4j/repository/query/ReactivePartTreeNeo4jQuery.java b/src/main/java/org/springframework/data/neo4j/repository/query/ReactivePartTreeNeo4jQuery.java
@@ -18,6 +18,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.function.BiFunction;
+import java.util.function.Supplier;
 import java.util.function.UnaryOperator;
 
 import org.neo4j.driver.types.MapAccessor;
@@ -60,7 +61,7 @@ private ReactivePartTreeNeo4jQuery(ReactiveNeo4jOperations neo4jOperations, Neo4
 	@Override
 	protected <T extends Object> PreparedQuery<T> prepareQuery(Class<T> returnedType, List<String> includedProperties,
 			Neo4jParameterAccessor parameterAccessor, @Nullable Neo4jQueryType queryType,
-			@Nullable BiFunction<TypeSystem, MapAccessor, ?> mappingFunction) {
+			@Nullable Supplier<BiFunction<TypeSystem, MapAccessor, ?>> mappingFunction) {
 
 		CypherQueryCreator queryCreator = new CypherQueryCreator(mappingContext, getDomainType(queryMethod),
 				Optional.ofNullable(queryType).orElseGet(() -> Neo4jQueryType.fromPartTree(tree)), tree, parameterAccessor,
diff --git a/src/main/java/org/springframework/data/neo4j/repository/query/ReactiveStringBasedNeo4jQuery.java b/src/main/java/org/springframework/data/neo4j/repository/query/ReactiveStringBasedNeo4jQuery.java
@@ -20,6 +20,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.function.BiFunction;
+import java.util.function.Supplier;
 
 import org.neo4j.driver.types.MapAccessor;
 import org.neo4j.driver.types.TypeSystem;
@@ -124,7 +125,7 @@ private ReactiveStringBasedNeo4jQuery(ReactiveNeo4jOperations neo4jOperations, N
 	@Override
 	protected <T extends Object> PreparedQuery<T> prepareQuery(Class<T> returnedType, List<String> includedProperties,
 			Neo4jParameterAccessor parameterAccessor, @Nullable Neo4jQueryType queryType,
-			@Nullable BiFunction<TypeSystem, MapAccessor, ?> mappingFunction) {
+			@Nullable Supplier<BiFunction<TypeSystem, MapAccessor, ?>> mappingFunction) {
 
 		Map<String, Object> boundParameters = bindParameters(parameterAccessor);
 		QueryContext queryContext = new QueryContext(
diff --git a/src/main/java/org/springframework/data/neo4j/repository/query/StringBasedNeo4jQuery.java b/src/main/java/org/springframework/data/neo4j/repository/query/StringBasedNeo4jQuery.java
@@ -21,6 +21,7 @@
 import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.function.BiFunction;
+import java.util.function.Supplier;
 import java.util.function.UnaryOperator;
 import java.util.regex.Pattern;
 
@@ -171,7 +172,7 @@ private StringBasedNeo4jQuery(Neo4jOperations neo4jOperations, Neo4jMappingConte
 	@Override
 	protected <T extends Object> PreparedQuery<T> prepareQuery(Class<T> returnedType, List<String> includedProperties,
 			Neo4jParameterAccessor parameterAccessor, @Nullable Neo4jQueryType queryType,
-			@Nullable BiFunction<TypeSystem, MapAccessor, ?> mappingFunction,
+			@Nullable Supplier<BiFunction<TypeSystem, MapAccessor, ?>> mappingFunction,
 			UnaryOperator<Integer> limitModifier
 	) {
 
diff --git a/src/test/java/org/springframework/data/neo4j/repository/query/ReactiveRepositoryQueryTest.java b/src/test/java/org/springframework/data/neo4j/repository/query/ReactiveRepositoryQueryTest.java
@@ -142,7 +142,7 @@ void shouldWarnWhenUsingSortedAndCustomQuery(LogbackCapture logbackCapture) {
 					Collections.emptyList(),
 					parameterAccessor,
 					Neo4jQueryType.DEFAULT,
-					(typeSystem, mapAccessor) -> new TestEntity()
+					() -> (typeSystem, mapAccessor) -> new TestEntity()
 			);
 			assertThat(logbackCapture.getFormattedMessages())
 					.anyMatch(s -> s.matches(
@@ -176,7 +176,7 @@ void orderBySpelShouldWork(LogbackCapture logbackCapture) {
 					Collections.emptyList(),
 					parameterAccessor,
 					Neo4jQueryType.DEFAULT,
-					(typeSystem, mapAccessor) -> new TestEntity()
+					() -> (typeSystem, mapAccessor) -> new TestEntity()
 			);
 			assertThat(pq.getQueryFragmentsAndParameters().getCypherQuery())
 					.isEqualTo("MATCH (n:Test) RETURN n ORDER BY name ASC SKIP $skip LIMIT $limit");
@@ -212,7 +212,7 @@ void literalReplacementsShouldWork() {
 						Collections.emptyList(),
 						parameterAccessor,
 						Neo4jQueryType.DEFAULT,
-						(typeSystem, mapAccessor) -> new TestEntity()
+						() -> (typeSystem, mapAccessor) -> new TestEntity()
 				);
 				return pq.getQueryFragmentsAndParameters().getCypherQuery();
 			}).block();
diff --git a/src/test/java/org/springframework/data/neo4j/repository/query/RepositoryQueryTest.java b/src/test/java/org/springframework/data/neo4j/repository/query/RepositoryQueryTest.java
@@ -333,7 +333,7 @@ void shouldWarnWhenUsingSortedAndCustomQuery(LogbackCapture logbackCapture) {
 					Collections.emptyList(),
 					parameterAccessor,
 					Neo4jQueryType.DEFAULT,
-					(typeSystem, mapAccessor) -> new TestEntity(),
+					() -> (typeSystem, mapAccessor) -> new TestEntity(),
 					UnaryOperator.identity()
 			);
 			assertThat(logbackCapture.getFormattedMessages())
@@ -359,7 +359,7 @@ void shouldWarnWhenUsingSortedPageable(LogbackCapture logbackCapture) {
 					Collections.emptyList(),
 					parameterAccessor,
 					Neo4jQueryType.DEFAULT,
-					(typeSystem, mapAccessor) -> new TestEntity(),
+					() -> (typeSystem, mapAccessor) -> new TestEntity(),
 					UnaryOperator.identity()
 			);
 			assertThat(logbackCapture.getFormattedMessages())
@@ -391,7 +391,7 @@ void orderBySpelShouldWork(LogbackCapture logbackCapture) {
 					Collections.emptyList(),
 					parameterAccessor,
 					Neo4jQueryType.DEFAULT,
-					(typeSystem, mapAccessor) -> new TestEntity(),
+					() -> (typeSystem, mapAccessor) -> new TestEntity(),
 					UnaryOperator.identity()
 			);
 			assertThat(pq.getQueryFragmentsAndParameters().getCypherQuery())
@@ -424,7 +424,7 @@ void literalReplacementsShouldWork() {
 					Collections.emptyList(),
 					parameterAccessor,
 					Neo4jQueryType.DEFAULT,
-					(typeSystem, mapAccessor) -> new TestEntity(),
+					() -> (typeSystem, mapAccessor) -> new TestEntity(),
 					UnaryOperator.identity()
 			);
 			assertThat(pq.getQueryFragmentsAndParameters().getCypherQuery())

Original file line number	Diff line number	Diff line change
`@@ -452,7 +452,7 @@ private <T> Mono<ExecutableQuery<T>> createExecutableQuery(Class<T> domainType,`
`452`	`452`	`Assert.notNull(neo4jMappingContext.getPersistentEntity(domainType), "Cannot get or create persistent entity.");`
`453`	`453`	`PreparedQuery<T> preparedQuery = PreparedQuery.queryFor(domainType).withCypherQuery(cypherQuery)`
`454`	`454`	`.withParameters(parameters)`
`455`		`- .usingMappingFunction(this.neo4jMappingContext.getRequiredMappingFunctionFor(domainType)).build();`
	`455`	`+ .usingMappingFunction(() -> this.neo4jMappingContext.getRequiredMappingFunctionFor(domainType)).build();`
`456`	`456`	`return this.toExecutableQuery(preparedQuery);`
`457`	`457`	`}`
`458`	`458`