From 73e503338d51754a5930e382ed415c011caad429 Mon Sep 17 00:00:00 2001
From: Mark Paluch <mark.paluch@broadcom.com>
Date: Fri, 25 Oct 2024 14:54:33 +0200
Subject: [PATCH] Polishing.

Eliminate potential NoSuchElementException from unchecked Optional.get usage. Simplify stream. Return Staged value, fix Nullability annotations.

See #1907
Original pull request: #1920
---
 .../JdbcAggregateChangeExecutionContext.java  | 35 ++++++++++---------
 .../relational/core/conversion/DbAction.java  |  9 +++--
 2 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/JdbcAggregateChangeExecutionContext.java b/spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/JdbcAggregateChangeExecutionContext.java
index df87d3813d..c53e3770b1 100644
--- a/spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/JdbcAggregateChangeExecutionContext.java
+++ b/spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/JdbcAggregateChangeExecutionContext.java
@@ -49,6 +49,7 @@
  * @author Umut Erturk
  * @author Myeonghyeon Lee
  * @author Chirag Tailor
+ * @author Mark Paluch
  */
 @SuppressWarnings("rawtypes")
 class JdbcAggregateChangeExecutionContext {
@@ -268,12 +269,10 @@ <T> List<T> populateIdsIfNecessary() {
 
 					cascadingValues.stage(insert.getDependingOn(), insert.getPropertyPath(),
 							qualifierValue, newEntity);
-
 				} else if (insert.getPropertyPath().getLeafProperty().isCollectionLike()) {
 
 					cascadingValues.gather(insert.getDependingOn(), insert.getPropertyPath(),
 							qualifierValue, newEntity);
-
 				}
 			}
 		}
@@ -289,6 +288,7 @@ <T> List<T> populateIdsIfNecessary() {
 		return roots;
 	}
 
+	@SuppressWarnings("unchecked")
 	private <S> Object setIdAndCascadingProperties(DbAction.WithEntity<S> action, @Nullable Object generatedId,
 			StagedValues cascadingValues) {
 
@@ -328,6 +328,7 @@ private PersistentPropertyPath<?> getRelativePath(DbAction<?> action, Persistent
 		throw new IllegalArgumentException(String.format("DbAction of type %s is not supported", action.getClass()));
 	}
 
+	@SuppressWarnings("unchecked")
 	private <T> RelationalPersistentEntity<T> getRequiredPersistentEntity(Class<T> type) {
 		return (RelationalPersistentEntity<T>) context.getRequiredPersistentEntity(type);
 	}
@@ -358,7 +359,7 @@ private <T> void updateWithVersion(DbAction.UpdateRoot<T> update) {
 	 */
 	private static class StagedValues {
 
-		static final List<MultiValueAggregator> aggregators = Arrays.asList(SetAggregator.INSTANCE, MapAggregator.INSTANCE,
+		static final List<MultiValueAggregator<?>> aggregators = Arrays.asList(SetAggregator.INSTANCE, MapAggregator.INSTANCE,
 				ListAggregator.INSTANCE, SingleElementAggregator.INSTANCE);
 
 		Map<DbAction, Map<PersistentPropertyPath, StagedValue>> values = new HashMap<>();
@@ -374,13 +375,14 @@ private static class StagedValues {
 		 *          be {@literal null}.
 		 * @param value The value to be set. Must not be {@literal null}.
 		 */
-		@SuppressWarnings("unchecked")
-		<T> void stage(DbAction<?> action, PersistentPropertyPath path, @Nullable Object qualifier, Object value) {
-			gather(action, path, qualifier, value);
-			values.get(action).get(path).isStaged = true;
+		void stage(DbAction<?> action, PersistentPropertyPath path, @Nullable Object qualifier, Object value) {
+
+			StagedValue gather = gather(action, path, qualifier, value);
+			gather.isStaged = true;
 		}
 
-		<T> void gather(DbAction<?> action, PersistentPropertyPath path, @Nullable Object qualifier, Object value) {
+		@SuppressWarnings("unchecked")
+		<T> StagedValue gather(DbAction<?> action, PersistentPropertyPath path, @Nullable Object qualifier, Object value) {
 
 			MultiValueAggregator<T> aggregator = getAggregatorFor(path);
 
@@ -391,19 +393,20 @@ <T> void gather(DbAction<?> action, PersistentPropertyPath path, @Nullable Objec
 					persistentPropertyPath -> new StagedValue(aggregator.createEmptyInstance()));
 			T currentValue = (T) stagedValue.value;
 
-			Object newValue = aggregator.add(currentValue, qualifier, value);
-
-			stagedValue.value = newValue;
+			stagedValue.value = aggregator.add(currentValue, qualifier, value);
 
 			valuesForPath.put(path, stagedValue);
+
+			return stagedValue;
 		}
 
-		private MultiValueAggregator getAggregatorFor(PersistentPropertyPath path) {
+		@SuppressWarnings("unchecked")
+		private <T> MultiValueAggregator<T> getAggregatorFor(PersistentPropertyPath path) {
 
 			PersistentProperty property = path.getLeafProperty();
-			for (MultiValueAggregator aggregator : aggregators) {
+			for (MultiValueAggregator<?> aggregator : aggregators) {
 				if (aggregator.handles(property)) {
-					return aggregator;
+					return (MultiValueAggregator<T>) aggregator;
 				}
 			}
 
@@ -427,10 +430,10 @@ void forEachPath(DbAction<?> dbAction, BiConsumer<PersistentPropertyPath, Object
 	}
 
 	private static class StagedValue {
-		Object value;
+		@Nullable Object value;
 		boolean isStaged;
 
-		public StagedValue(Object value) {
+		public StagedValue(@Nullable Object value) {
 			this.value = value;
 		}
 	}
diff --git a/spring-data-relational/src/main/java/org/springframework/data/relational/core/conversion/DbAction.java b/spring-data-relational/src/main/java/org/springframework/data/relational/core/conversion/DbAction.java
index 1ccca7355e..013bf65b91 100644
--- a/spring-data-relational/src/main/java/org/springframework/data/relational/core/conversion/DbAction.java
+++ b/spring-data-relational/src/main/java/org/springframework/data/relational/core/conversion/DbAction.java
@@ -482,16 +482,21 @@ interface WithDependingOn<T> extends WithPropertyPath<T>, WithEntity<T> {
 		default Pair<PersistentPropertyPath<RelationalPersistentProperty>, Object> getQualifier() {
 
 			Map<PersistentPropertyPath<RelationalPersistentProperty>, Object> qualifiers = getQualifiers();
+
 			if (qualifiers.isEmpty()) {
 				return null;
 			}
 
 			Set<Map.Entry<PersistentPropertyPath<RelationalPersistentProperty>, Object>> entries = qualifiers.entrySet();
-			Map.Entry<PersistentPropertyPath<RelationalPersistentProperty>, Object> entry = entries.stream().sorted(Comparator.comparing(e -> -e.getKey().getLength())).findFirst().get();
+			Optional<Map.Entry<PersistentPropertyPath<RelationalPersistentProperty>, Object>> optionalEntry = entries.stream()
+					.filter(e -> e.getValue() != null).min(Comparator.comparing(e -> -e.getKey().getLength()));
+
+			Map.Entry<PersistentPropertyPath<RelationalPersistentProperty>, Object> entry = optionalEntry.orElse(null);
 
-			if (entry.getValue() == null) {
+			if (entry == null) {
 				return null;
 			}
+
 			return Pair.of(entry.getKey(), entry.getValue());
 		}