Add support for fluent Query by Example query definition.

mp911de · mp911de · commit a53f4ac37598 · 2021-10-08T13:56:34.000+02:00
Closes: #663
diff --git a/src/main/java/org/springframework/data/r2dbc/repository/support/ReactiveFluentQuerySupport.java b/src/main/java/org/springframework/data/r2dbc/repository/support/ReactiveFluentQuerySupport.java
@@ -0,0 +1,100 @@
+/*
+ * Copyright 2021 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.r2dbc.repository.support;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+
+import org.springframework.data.domain.Sort;
+import org.springframework.data.repository.query.FluentQuery;
+import org.springframework.util.Assert;
+
+/**
+ * Support class for {@link org.springframework.data.repository.query.FluentQuery.ReactiveFluentQuery} implementations.
+ *
+ * @author Mark Paluch
+ * @since 1.4
+ */
+abstract class ReactiveFluentQuerySupport<P, T> implements FluentQuery.ReactiveFluentQuery<T> {
+
+	private final P predicate;
+	private final Sort sort;
+	private final Class<T> resultType;
+	private final List<String> fieldsToInclude;
+
+	ReactiveFluentQuerySupport(P predicate, Sort sort, Class<T> resultType, List<String> fieldsToInclude) {
+		this.predicate = predicate;
+		this.sort = sort;
+		this.resultType = resultType;
+		this.fieldsToInclude = fieldsToInclude;
+	}
+
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.repository.query.FluentQuery.ReactiveFluentQuery#sortBy(org.springframework.data.domain.Sort)
+	 */
+	@Override
+	public ReactiveFluentQuery<T> sortBy(Sort sort) {
+
+		Assert.notNull(sort, "Sort must not be null!");
+
+		return create(predicate, sort, resultType, fieldsToInclude);
+	}
+
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.repository.query.FluentQuery.ReactiveFluentQuery#as(java.lang.Class)
+	 */
+	@Override
+	public <R> ReactiveFluentQuery<R> as(Class<R> projection) {
+
+		Assert.notNull(projection, "Projection target type must not be null!");
+
+		return create(predicate, sort, projection, fieldsToInclude);
+	}
+
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.repository.query.FluentQuery.ReactiveFluentQuery#project(java.util.Collection)
+	 */
+	@Override
+	public ReactiveFluentQuery<T> project(Collection<String> properties) {
+
+		Assert.notNull(properties, "Projection properties must not be null!");
+
+		return create(predicate, sort, resultType, new ArrayList<>(properties));
+	}
+
+	protected abstract <R> ReactiveFluentQuerySupport<P, R> create(P predicate, Sort sort, Class<R> resultType,
+			List<String> fieldsToInclude);
+
+	P getPredicate() {
+		return predicate;
+	}
+
+	Sort getSort() {
+		return sort;
+	}
+
+	Class<T> getResultType() {
+		return resultType;
+	}
+
+	List<String> getFieldsToInclude() {
+		return fieldsToInclude;
+	}
+}
diff --git a/src/main/java/org/springframework/data/r2dbc/repository/support/ReactivePageableExecutionUtils.java b/src/main/java/org/springframework/data/r2dbc/repository/support/ReactivePageableExecutionUtils.java
@@ -0,0 +1,69 @@
+/*
+ * Copyright 2021 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.r2dbc.repository.support;
+
+import reactor.core.publisher.Mono;
+
+import java.util.List;
+
+import org.springframework.data.domain.Page;
+import org.springframework.data.domain.PageImpl;
+import org.springframework.data.domain.Pageable;
+import org.springframework.util.Assert;
+
+/**
+ * Support for query execution using {@link Pageable}. Using {@link ReactivePageableExecutionUtils} assumes that data
+ * queries are cheaper than {@code COUNT} queries and so some cases can take advantage of optimizations.
+ *
+ * @author Mark Paluch
+ * @since 1.4
+ */
+abstract class ReactivePageableExecutionUtils {
+
+	private ReactivePageableExecutionUtils() {}
+
+	/**
+	 * Constructs a {@link Page} based on the given {@code content}, {@link Pageable} and {@link Mono} applying
+	 * optimizations. The construction of {@link Page} omits a count query if the total can be determined based on the
+	 * result size and {@link Pageable}.
+	 *
+	 * @param content must not be {@literal null}.
+	 * @param pageable must not be {@literal null}.
+	 * @param totalSupplier must not be {@literal null}.
+	 * @return the {@link Page}.
+	 */
+	public static <T> Mono<Page<T>> getPage(List<T> content, Pageable pageable, Mono<Long> totalSupplier) {
+
+		Assert.notNull(content, "Content must not be null!");
+		Assert.notNull(pageable, "Pageable must not be null!");
+		Assert.notNull(totalSupplier, "TotalSupplier must not be null!");
+
+		if (pageable.isUnpaged() || pageable.getOffset() == 0) {
+
+			if (pageable.isUnpaged() || pageable.getPageSize() > content.size()) {
+				return Mono.just(new PageImpl<>(content, pageable, content.size()));
+			}
+
+			return totalSupplier.map(total -> new PageImpl<>(content, pageable, total));
+		}
+
+		if (content.size() != 0 && pageable.getPageSize() > content.size()) {
+			return Mono.just(new PageImpl<>(content, pageable, pageable.getOffset() + content.size()));
+		}
+
+		return totalSupplier.map(total -> new PageImpl<>(content, pageable, total));
+	}
+}
diff --git a/src/main/java/org/springframework/data/r2dbc/repository/support/SimpleR2dbcRepository.java b/src/main/java/org/springframework/data/r2dbc/repository/support/SimpleR2dbcRepository.java
@@ -18,21 +18,28 @@
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
+import java.util.Collections;
 import java.util.List;
+import java.util.function.Function;
+import java.util.function.UnaryOperator;
 
 import org.reactivestreams.Publisher;
 import org.springframework.data.domain.Example;
+import org.springframework.data.domain.Page;
+import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.r2dbc.convert.R2dbcConverter;
 import org.springframework.data.r2dbc.core.R2dbcEntityOperations;
 import org.springframework.data.r2dbc.core.R2dbcEntityTemplate;
 import org.springframework.data.r2dbc.core.ReactiveDataAccessStrategy;
+import org.springframework.data.r2dbc.core.ReactiveSelectOperation;
 import org.springframework.data.r2dbc.repository.R2dbcRepository;
 import org.springframework.data.relational.core.mapping.RelationalPersistentProperty;
 import org.springframework.data.relational.core.query.Criteria;
 import org.springframework.data.relational.core.query.Query;
 import org.springframework.data.relational.repository.query.RelationalEntityInformation;
 import org.springframework.data.relational.repository.query.RelationalExampleMapper;
+import org.springframework.data.repository.query.FluentQuery;
 import org.springframework.data.repository.reactive.ReactiveSortingRepository;
 import org.springframework.data.util.Lazy;
 import org.springframework.data.util.Streamable;
@@ -432,11 +439,130 @@ public <S extends T> Mono<Boolean> exists(Example<S> example) {
 		return this.entityOperations.exists(query, example.getProbeType());
 	}
 
+	@Override
+	public <S extends T, R, P extends Publisher<R>> P findBy(Example<S> example,
+			Function<FluentQuery.ReactiveFluentQuery<S>, P> queryFunction) {
+
+		Assert.notNull(example, "Sample must not be null!");
+		Assert.notNull(queryFunction, "Query function must not be null!");
+
+		return queryFunction.apply(new ReactiveFluentQueryByExample<>(example, example.getProbeType()));
+	}
+
 	private RelationalPersistentProperty getIdProperty() {
 		return this.idProperty.get();
 	}
 
 	private Query getIdQuery(Object id) {
 		return Query.query(Criteria.where(getIdProperty().getName()).is(id));
 	}
+
+	/**
+	 * {@link org.springframework.data.repository.query.FluentQuery.ReactiveFluentQuery} using {@link Example}.
+	 *
+	 * @author Mark Paluch
+	 * @since 1.4
+	 */
+	class ReactiveFluentQueryByExample<S, T> extends ReactiveFluentQuerySupport<Example<S>, T> {
+
+		ReactiveFluentQueryByExample(Example<S> example, Class<T> resultType) {
+			this(example, Sort.unsorted(), resultType, Collections.emptyList());
+		}
+
+		ReactiveFluentQueryByExample(Example<S> example, Sort sort, Class<T> resultType, List<String> fieldsToInclude) {
+			super(example, sort, resultType, fieldsToInclude);
+		}
+
+		@Override
+		protected <R> ReactiveFluentQueryByExample<S, R> create(Example<S> predicate, Sort sort, Class<R> resultType,
+				List<String> fieldsToInclude) {
+			return new ReactiveFluentQueryByExample<>(predicate, sort, resultType, fieldsToInclude);
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see org.springframework.data.repository.query.FluentQuery.ReactiveFluentQuery#one()
+		 */
+		@Override
+		public Mono<T> one() {
+			return createQuery().one();
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see org.springframework.data.repository.query.FluentQuery.ReactiveFluentQuery#first()
+		 */
+		@Override
+		public Mono<T> first() {
+			return createQuery().first();
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see org.springframework.data.repository.query.FluentQuery.ReactiveFluentQuery#all()
+		 */
+		@Override
+		public Flux<T> all() {
+			return createQuery().all();
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see org.springframework.data.repository.query.FluentQuery.ReactiveFluentQuery#page(org.springframework.data.domain.Pageable)
+		 */
+		@Override
+		public Mono<Page<T>> page(Pageable pageable) {
+
+			Assert.notNull(pageable, "Pageable must not be null!");
+
+			Mono<List<T>> items = createQuery(q -> q.with(pageable)).all().collectList();
+
+			return items.flatMap(content -> ReactivePageableExecutionUtils.getPage(content, pageable, this.count()));
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see org.springframework.data.repository.query.FluentQuery.ReactiveFluentQuery#count()
+		 */
+		@Override
+		public Mono<Long> count() {
+			return createQuery().count();
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see org.springframework.data.repository.query.FluentQuery.ReactiveFluentQuery#exists()
+		 */
+		@Override
+		public Mono<Boolean> exists() {
+			return createQuery().exists();
+		}
+
+		private ReactiveSelectOperation.TerminatingSelect<T> createQuery() {
+			return createQuery(UnaryOperator.identity());
+		}
+
+		@SuppressWarnings("unchecked")
+		private ReactiveSelectOperation.TerminatingSelect<T> createQuery(UnaryOperator<Query> queryCustomizer) {
+
+			Query query = exampleMapper.getMappedExample(getPredicate());
+
+			if (getSort().isSorted()) {
+				query = query.sort(getSort());
+			}
+
+			if (!getFieldsToInclude().isEmpty()) {
+				query = query.columns(getFieldsToInclude().toArray(new String[0]));
+			}
+
+			query = queryCustomizer.apply(query);
+
+			ReactiveSelectOperation.ReactiveSelect<S> select = entityOperations.select(getPredicate().getProbeType());
+
+			if (getResultType() != getPredicate().getProbeType()) {
+				return select.as(getResultType()).matching(query);
+			}
+			return (ReactiveSelectOperation.TerminatingSelect<T>) select.matching(query);
+		}
+	}
 }
diff --git a/src/test/java/org/springframework/data/r2dbc/repository/support/AbstractSimpleR2dbcRepositoryIntegrationTests.java b/src/test/java/org/springframework/data/r2dbc/repository/support/AbstractSimpleR2dbcRepositoryIntegrationTests.java
