Add support for arbitrary where clauses in Single Query Loading.

Closes #1601

Author: schauder

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

@@ -21,15 +21,22 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
+import java.util.function.BiFunction;
 
 import org.springframework.dao.IncorrectResultSizeDataAccessException;
 import org.springframework.data.relational.core.dialect.Dialect;
 import org.springframework.data.relational.core.mapping.AggregatePath;
 import org.springframework.data.relational.core.mapping.RelationalPersistentEntity;
+import org.springframework.data.relational.core.query.CriteriaDefinition;
+import org.springframework.data.relational.core.query.Query;
+import org.springframework.data.relational.core.sql.Condition;
+import org.springframework.data.relational.core.sql.Table;
 import org.springframework.data.relational.core.sqlgeneration.AliasFactory;
 import org.springframework.data.relational.core.sqlgeneration.SingleQuerySqlGenerator;
 import org.springframework.data.relational.core.sqlgeneration.SqlGenerator;
 import org.springframework.data.relational.domain.RowDocument;
+import org.springframework.jdbc.core.namedparam.MapSqlParameterSource;
 import org.springframework.jdbc.core.namedparam.NamedParameterJdbcOperations;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
@@ -89,6 +96,35 @@ public Iterable<T> findAllById(Iterable<?> ids) {
 		return jdbcTemplate.query(sqlGenerator.findAllById(), Map.of("ids", convertedIds), this::extractAll);
 	}
 
+	public Iterable<T> findAllBy(Query query) {
+
+		MapSqlParameterSource parameterSource = new MapSqlParameterSource();
+		BiFunction<Table, RelationalPersistentEntity, Condition> condition = createConditionSource(query, parameterSource);
+		return jdbcTemplate.query(sqlGenerator.findAllByCondition(condition), parameterSource, this::extractAll);
+	}
+
+	public Optional<T> findOneByQuery(Query query) {
+		
+		MapSqlParameterSource parameterSource = new MapSqlParameterSource();
+		BiFunction<Table, RelationalPersistentEntity, Condition> condition = createConditionSource(query, parameterSource);
+
+		return Optional.ofNullable(
+				jdbcTemplate.query(sqlGenerator.findAllByCondition(condition), parameterSource, this::extractZeroOrOne));
+	}
+
+	private BiFunction<Table, RelationalPersistentEntity, Condition> createConditionSource(Query query, MapSqlParameterSource parameterSource) {
+
+		QueryMapper queryMapper = new QueryMapper(converter);
+
+		BiFunction<Table, RelationalPersistentEntity, Condition> condition = (table, aggregate) -> {
+			Optional<CriteriaDefinition> criteria = query.getCriteria();
+			return criteria
+					.map(criteriaDefinition -> queryMapper.getMappedObject(parameterSource, criteriaDefinition, table, aggregate))
+					.orElse(null);
+		};
+		return condition;
+	}
+
 	/**
 	 * Extracts a list of aggregates from the given {@link ResultSet} by utilizing the
 	 * {@link RowDocumentResultSetExtractor} and the {@link JdbcConverter}. When used as a method reference this conforms
@@ -115,7 +151,8 @@ private List<T> extractAll(ResultSet rs) throws SQLException {
 	 * to the {@link org.springframework.jdbc.core.ResultSetExtractor} contract.
 	 *
 	 * @param @param rs the {@link ResultSet} from which to extract the data. Must not be {(}@literal null}.
-	 * @return The single instance when the conversion results in exactly one instance. If the {@literal ResultSet} is empty, null is returned.
+	 * @return The single instance when the conversion results in exactly one instance. If the {@literal ResultSet} is
+	 *         empty, null is returned.
 	 * @throws SQLException
 	 * @throws IncorrectResultSizeDataAccessException when the conversion yields more than one instance.
 	 */
@@ -190,9 +227,15 @@ public String findAllById() {
 			return findAllById;
 		}
 
+		@Override
+		public String findAllByCondition(BiFunction<Table, RelationalPersistentEntity, Condition> conditionSource) {
+			return delegate.findAllByCondition(conditionSource);
+		}
+
 		@Override
 		public AliasFactory getAliasFactory() {
 			return delegate.getAliasFactory();
 		}
+
 	}
 }

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

@@ -77,12 +77,13 @@ public <T> Iterable<T> findAll(Class<T> domainType, Pageable pageable) {
 
 	@Override
 	public <T> Optional<T> findOne(Query query, Class<T> domainType) {
-		return Optional.empty();
+		return getReader(domainType).findOneByQuery(query);
 	}
 
 	@Override
 	public <T> Iterable<T> findAll(Query query, Class<T> domainType) {
-		throw new UnsupportedOperationException();
+
+		return getReader(domainType).findAllBy(query);
 	}
 
 	@Override

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

@@ -16,9 +16,11 @@
 package org.springframework.data.jdbc.core.convert;
 
 import java.util.Collections;
+import java.util.Optional;
 
 import org.springframework.data.mapping.PersistentPropertyPath;
 import org.springframework.data.relational.core.mapping.RelationalPersistentProperty;
+import org.springframework.data.relational.core.query.Query;
 import org.springframework.jdbc.core.namedparam.NamedParameterJdbcOperations;
 import org.springframework.util.Assert;
 
@@ -85,13 +87,37 @@ public <T> Iterable<T> findAllById(Iterable<?> ids, Class<T> domainType) {
 		return super.findAllById(ids, domainType);
 	}
 
+	public <T> Optional<T> findOne(Query query, Class<T> domainType) {
+
+		if (isSingleSelectQuerySupported(domainType) && isSingleSelectQuerySupported(query)) {
+			return singleSelectDelegate.findOne(query, domainType);
+		}
+
+		return super.findOne(query, domainType);
+	}
+
+	@Override
+	public <T> Iterable<T> findAll(Query query, Class<T> domainType) {
+
+		if (isSingleSelectQuerySupported(domainType) && isSingleSelectQuerySupported(query)) {
+			return singleSelectDelegate.findAll(query, domainType);
+		}
+
+		return super.findAll(query, domainType);
+	}
+
+	private static boolean isSingleSelectQuerySupported(Query query) {
+		return !query.isSorted() && !query.isLimited();
+	}
+
 	private boolean isSingleSelectQuerySupported(Class<?> entityType) {
 
-		return sqlGeneratorSource.getDialect().supportsSingleQueryLoading()//
-				&& entityQualifiesForSingleSelectQuery(entityType);
+		return converter.getMappingContext().isSingleQueryLoadingEnabled()
+				&& sqlGeneratorSource.getDialect().supportsSingleQueryLoading()//
+				&& entityQualifiesForSingleQueryLoading(entityType);
 	}
 
-	private boolean entityQualifiesForSingleSelectQuery(Class<?> entityType) {
+	private boolean entityQualifiesForSingleQueryLoading(Class<?> entityType) {
 
 		boolean referenceFound = false;
 		for (PersistentPropertyPath<RelationalPersistentProperty> path : converter.getMappingContext()
@@ -113,9 +139,9 @@ private boolean entityQualifiesForSingleSelectQuery(Class<?> entityType) {
 			}
 
 			// AggregateReferences aren't supported yet
-			if (property.isAssociation()) {
-				return false;
-			}
+			// if (property.isAssociation()) {
+			// return false;
+			// }
 		}
 		return true;
 

spring-data-jdbc/src/test/java/org/springframework/data/jdbc/core/AbstractJdbcAggregateTemplateIntegrationTests.java

@@ -24,15 +24,7 @@
 import static org.springframework.test.context.TestExecutionListeners.MergeMode.*;
 
 import java.time.LocalDateTime;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-import java.util.Set;
+import java.util.*;
 import java.util.function.Function;
 import java.util.stream.IntStream;
 
@@ -67,6 +59,9 @@
 import org.springframework.data.relational.core.mapping.MappedCollection;
 import org.springframework.data.relational.core.mapping.RelationalMappingContext;
 import org.springframework.data.relational.core.mapping.Table;
+import org.springframework.data.relational.core.query.Criteria;
+import org.springframework.data.relational.core.query.CriteriaDefinition;
+import org.springframework.data.relational.core.query.Query;
 import org.springframework.jdbc.core.namedparam.NamedParameterJdbcOperations;
 import org.springframework.test.context.ActiveProfiles;
 import org.springframework.test.context.ContextConfiguration;
@@ -233,6 +228,49 @@ void findAllById() {
 				.containsExactlyInAnyOrder(tuple(entity.id, "entity"), tuple(yetAnother.id, "yetAnother"));
 	}
 
+	@Test // GH-1601
+	void findAllByQuery() {
+
+		ListParent entity = new ListParent();
+		entity = template.save(entity);
+
+		ListParent other = new ListParent();
+		other = template.save(other);
+
+		ListParent yetAnother = new ListParent();
+		yetAnother = template.save(yetAnother);
+
+		CriteriaDefinition criteria  = CriteriaDefinition.from(Criteria.where("id").is(other.id));
+		Query query = Query.query(criteria);
+		Iterable<ListParent> reloadedById = template.findAll(query,
+				ListParent.class);
+
+
+		assertThat(reloadedById).extracting(e -> e.id)
+				.containsExactly(other.id);
+	}
+
+	@Test // GH-1601
+	void findOneByQuery() {
+
+		ListParent entity = new ListParent();
+		entity = template.save(entity);
+
+		ListParent other = new ListParent();
+		other = template.save(other);
+
+		ListParent yetAnother = new ListParent();
+		yetAnother = template.save(yetAnother);
+
+		CriteriaDefinition criteria  = CriteriaDefinition.from(Criteria.where("id").is(other.id));
+		Query query = Query.query(criteria);
+		Optional<ListParent> reloadedById = template.findOne(query,
+				ListParent.class);
+
+
+		assertThat(reloadedById).get().extracting(e -> e.id).isEqualTo(other.id);
+	}
+
 	@Test // DATAJDBC-112
 	@EnabledOnFeature(SUPPORTS_QUOTED_IDS)
 	void saveAndLoadAnEntityWithReferencedEntityById() {

spring-data-relational/src/main/java/org/springframework/data/relational/core/sqlgeneration/SingleQuerySqlGenerator.java

@@ -19,6 +19,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.function.BiFunction;
 
 import org.jetbrains.annotations.NotNull;
 import org.springframework.data.mapping.PersistentProperty;
@@ -81,20 +82,20 @@ public String findAllById() {
 		return createSelect(condition);
 	}
 
+	@Override
+	public String findAllByCondition(BiFunction<Table, RelationalPersistentEntity, Condition> conditionSource) {
+		Condition condition = conditionSource.apply(table, aggregate);
+		return createSelect(condition);
+	}
+
 	/**
 	 * @return The {@link AggregatePath} to the id property of the aggregate root.
 	 */
 	private AggregatePath getRootIdPath() {
 		return context.getAggregatePath(aggregate).append(aggregate.getRequiredIdProperty());
 	}
 
-	/**
-	 * Creates a SQL suitable of loading all the data required for constructing complete aggregates.
-	 *
-	 * @param condition a constraint for limiting the aggregates to be loaded.
-	 * @return a {@literal  String} containing the generated SQL statement
-	 */
-	private String createSelect(Condition condition) {
+	String createSelect(Condition condition) {
 
 		AggregatePath rootPath = context.getAggregatePath(aggregate);
 		QueryMeta queryMeta = createInlineQuery(rootPath, condition);

spring-data-relational/src/main/java/org/springframework/data/relational/core/sqlgeneration/SqlGenerator.java

@@ -15,6 +15,12 @@
  */
 package org.springframework.data.relational.core.sqlgeneration;
 
+import org.springframework.data.relational.core.mapping.RelationalPersistentEntity;
+import org.springframework.data.relational.core.sql.Condition;
+import org.springframework.data.relational.core.sql.Table;
+
+import java.util.function.BiFunction;
+
 /**
  * Generates SQL statements for loading aggregates.
  *
@@ -28,5 +34,7 @@ public interface SqlGenerator {
 
 	String findAllById();
 
+	String findAllByCondition(BiFunction<Table, RelationalPersistentEntity, Condition> conditionSource);
+
 	AliasFactory getAliasFactory();
 }