simplified forQueryByIds

Author: schauder

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/convert/SqlParametersFactory.java

@@ -19,6 +19,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Predicate;
 
@@ -148,31 +149,27 @@ <T> SqlIdentifierParameterSource forQueryByIds(Iterable<?> ids, Class<T> domainT
 
 		SqlIdentifierParameterSource parameterSource = new SqlIdentifierParameterSource();
 
-		RelationalPersistentEntity<?> entity = context.getPersistentEntity(domainType);
+		RelationalPersistentEntity<?> entity = context.getRequiredPersistentEntity(domainType);
 		RelationalPersistentProperty singleIdProperty = entity.getRequiredIdProperty();
+		RelationalPersistentEntity<?> complexId = context.getPersistentEntity(singleIdProperty);
+		AggregatePath.ColumnInfos idColumnInfos = context.getAggregatePath(entity).getTableInfo().idColumnInfos();
 
-		if (singleIdProperty.isEntity()) {
-
-			RelationalPersistentEntity<?> complexId = context.getPersistentEntity(singleIdProperty);
-
-			AggregatePath.ColumnInfos idColumnInfos = context.getAggregatePath(entity).getTableInfo().idColumnInfos();
+		BiFunction<Object, AggregatePath, Object> valueExtractor = complexId == null
+				? (id, ap) -> id
+				: (id, ap) -> complexId.getPropertyPathAccessor(id).getProperty(ap.getRequiredPersistentPropertyPath());
 
-			List<Object[]> parameterValues = new ArrayList<>();
-			for (Object id : ids) {
+		List<Object[]> parameterValues = new ArrayList<>();
+		for (Object id : ids) {
 
-				PersistentPropertyPathAccessor<Object> accessor = complexId.getPropertyPathAccessor(id);
+			List<Object> tupleList = new ArrayList<>();
+			idColumnInfos.forEach((ap, ci) -> {
+				tupleList.add(valueExtractor.apply(id, ap));
+			});
+			parameterValues.add(tupleList.toArray(new Object[0]));
+		}
 
-				List<Object> tupleList = new ArrayList<>();
-				idColumnInfos.forEach((ap, ci) -> {
-					tupleList.add(accessor.getProperty(ap.getRequiredPersistentPropertyPath()));
-				});
-				parameterValues.add(tupleList.toArray(new Object[0]));
-			}
+		parameterSource.addValue(SqlGenerator.IDS_SQL_PARAMETER, parameterValues);
 
-			parameterSource.addValue(SqlGenerator.IDS_SQL_PARAMETER, parameterValues);
-		} else {
-			addConvertedPropertyValuesAsList(parameterSource, getRequiredPersistentEntity(domainType).getRequiredIdProperty(), ids);
-		}
 		return parameterSource;
 	}