Skip to content

Commit 253a78c

Browse files
christophstroblchristophstrobl
committed
Support GeoPage
1 parent fbceb94 commit 253a78c

File tree

13 files changed

+214
-66
lines changed

13 files changed

+214
-66
lines changed

‎spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ExecutableFindOperation.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -225,6 +225,14 @@ interface TerminatingFindNear<T> {
225225
* @return never {@literal null}.
226226
*/
227227
GeoResults<T> all();
228+
229+
/**
230+
* Count matching elements.
231+
*
232+
* @return number of elements matching the query.
233+
* @since 5.0
234+
*/
235+
long count();
228236
}
229237

230238
/**

‎spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ExecutableFindOperationSupport.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,11 @@ public <R> TerminatingFindNear<R> map(QueryResultConverter<? super G, ? extends
243243
public GeoResults<G> all() {
244244
return template.doGeoNear(nearQuery, domainType, getCollectionName(), returnType, resultConverter);
245245
}
246+
247+
@Override
248+
public long count() {
249+
return template.doGeoNearCount(nearQuery, domainType, getCollectionName());
250+
}
246251
}
247252
}
248253

‎spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
import org.springframework.dao.support.PersistenceExceptionTranslator;
4949
import org.springframework.data.convert.EntityReader;
5050
import org.springframework.data.domain.OffsetScrollPosition;
51+
import org.springframework.data.domain.Pageable;
5152
import org.springframework.data.domain.Window;
5253
import org.springframework.data.geo.Distance;
5354
import org.springframework.data.geo.GeoResult;
@@ -1044,6 +1045,31 @@ public <T> GeoResults<T> geoNear(NearQuery near, Class<?> domainType, String col
10441045
return doGeoNear(near, domainType, collectionName, returnType, QueryResultConverter.entity());
10451046
}
10461047

1048+
long doGeoNearCount(NearQuery near, Class<?> domainType, String collectionName) {
1049+
1050+
Builder optionsBuilder = AggregationOptions.builder().collation(near.getCollation());
1051+
1052+
if (near.hasReadPreference()) {
1053+
optionsBuilder.readPreference(near.getReadPreference());
1054+
}
1055+
1056+
if (near.hasReadConcern()) {
1057+
optionsBuilder.readConcern(near.getReadConcern());
1058+
}
1059+
1060+
String distanceField = operations.nearQueryDistanceFieldName(domainType);
1061+
Aggregation $geoNear = TypedAggregation.newAggregation(domainType,
1062+
Aggregation.geoNear(near, distanceField).skip(-1).limit(-1), Aggregation.count().as("_totalCount"))
1063+
.withOptions(optionsBuilder.build());
1064+
1065+
AggregationResults<Document> results = doAggregate($geoNear, collectionName, Document.class,
1066+
queryOperations.createAggregation($geoNear, (AggregationOperationContext) null));
1067+
Iterator<Document> iterator = results.iterator();
1068+
return iterator.hasNext()
1069+
? NumberUtils.convertNumberToTargetClass(iterator.next().get("_totalCount", Integer.class), Long.class)
1070+
: 0L;
1071+
}
1072+
10471073
<T, R> GeoResults<R> doGeoNear(NearQuery near, Class<?> domainType, String collectionName, Class<T> returnType,
10481074
QueryResultConverter<? super T, ? extends R> resultConverter) {
10491075

‎spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/GeoNearOperation.java‎

Lines changed: 44 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,8 @@ public class GeoNearOperation implements AggregationOperation {
4242
private final NearQuery nearQuery;
4343
private final String distanceField;
4444
private final @Nullable String indexKey;
45+
private final @Nullable Long skip;
46+
private final @Nullable Integer limit;
4547

4648
/**
4749
* Creates a new {@link GeoNearOperation} from the given {@link NearQuery} and the given distance field. The
@@ -51,7 +53,7 @@ public class GeoNearOperation implements AggregationOperation {
5153
* @param distanceField must not be {@literal null}.
5254
*/
5355
public GeoNearOperation(NearQuery nearQuery, String distanceField) {
54-
this(nearQuery, distanceField, null);
56+
this(nearQuery, distanceField, null, nearQuery.getSkip(), null);
5557
}
5658

5759
/**
@@ -63,14 +65,17 @@ public GeoNearOperation(NearQuery nearQuery, String distanceField) {
6365
* @param indexKey can be {@literal null};
6466
* @since 2.1
6567
*/
66-
private GeoNearOperation(NearQuery nearQuery, String distanceField, @Nullable String indexKey) {
68+
private GeoNearOperation(NearQuery nearQuery, String distanceField, @Nullable String indexKey, @Nullable Long skip,
69+
@Nullable Integer limit) {
6770

6871
Assert.notNull(nearQuery, "NearQuery must not be null");
6972
Assert.hasLength(distanceField, "Distance field must not be null or empty");
7073

7174
this.nearQuery = nearQuery;
7275
this.distanceField = distanceField;
7376
this.indexKey = indexKey;
77+
this.skip = skip;
78+
this.limit = limit;
7479
}
7580

7681
/**
@@ -83,7 +88,30 @@ private GeoNearOperation(NearQuery nearQuery, String distanceField, @Nullable St
8388
*/
8489
@Contract("_ -> new")
8590
public GeoNearOperation useIndex(String key) {
86-
return new GeoNearOperation(nearQuery, distanceField, key);
91+
return new GeoNearOperation(nearQuery, distanceField, key, skip, limit);
92+
}
93+
94+
/**
95+
* Override potential skip applied via {@link NearQuery#getSkip()}. Adds an additional {@link SkipOperation} if value
96+
* is non negative.
97+
*
98+
* @param skip
99+
* @return new instance of {@link GeoNearOperation}.
100+
* @since 5.0
101+
*/
102+
public GeoNearOperation skip(long skip) {
103+
return new GeoNearOperation(nearQuery, distanceField, indexKey, skip, limit);
104+
}
105+
106+
/**
107+
* Override potential limit value. Adds an additional {@link LimitOperation} if value is non negative.
108+
*
109+
* @param limit
110+
* @return new instance of {@link GeoNearOperation}.
111+
* @since 5.0
112+
*/
113+
public GeoNearOperation limit(Integer limit) {
114+
return new GeoNearOperation(nearQuery, distanceField, indexKey, skip, limit);
87115
}
88116

89117
@Override
@@ -92,7 +120,13 @@ public Document toDocument(AggregationOperationContext context) {
92120
Document command = context.getMappedObject(nearQuery.toDocument());
93121

94122
if (command.containsKey("query")) {
95-
command.replace("query", context.getMappedObject(command.get("query", Document.class)));
123+
Document query = command.get("query", Document.class);
124+
if (query == null || query.isEmpty()) {
125+
command.remove("query");
126+
} else {
127+
command.replace("query", context.getMappedObject(query));
128+
}
129+
96130
}
97131

98132
command.remove("collation");
@@ -115,15 +149,18 @@ public List<Document> toPipelineStages(AggregationOperationContext context) {
115149

116150
Document command = toDocument(context);
117151
Number limit = (Number) command.get("$geoNear", Document.class).remove("num");
152+
if (limit != null && this.limit != null) {
153+
limit = this.limit;
154+
}
118155

119156
List<Document> stages = new ArrayList<>(3);
120157
stages.add(command);
121158

122-
if (nearQuery.getSkip() != null && nearQuery.getSkip() > 0) {
123-
stages.add(new Document("$skip", nearQuery.getSkip()));
159+
if (this.skip != null && this.skip > 0) {
160+
stages.add(new Document("$skip", this.skip));
124161
}
125162

126-
if (limit != null) {
163+
if (limit != null && limit.longValue() > 0) {
127164
stages.add(new Document("$limit", limit.longValue()));
128165
}
129166

‎spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/NearQuery.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -671,7 +671,7 @@ public Document toDocument() {
671671
document.put("distanceMultiplier", getDistanceMultiplier());
672672
}
673673

674-
if (limit != null) {
674+
if (limit != null && limit > 0) {
675675
document.put("num", limit);
676676
}
677677

‎spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/aot/GeoBlocks.java‎

Lines changed: 20 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,14 @@
1515
*/
1616
package org.springframework.data.mongodb.repository.aot;
1717

18-
import java.util.List;
19-
import java.util.stream.Collectors;
20-
2118
import org.springframework.data.geo.Distance;
2219
import org.springframework.data.geo.GeoPage;
2320
import org.springframework.data.geo.GeoResults;
2421
import org.springframework.data.mongodb.core.MongoOperations;
2522
import org.springframework.data.mongodb.core.query.NearQuery;
2623
import org.springframework.data.mongodb.repository.query.MongoQueryMethod;
2724
import org.springframework.data.repository.aot.generate.AotQueryMethodGenerationContext;
25+
import org.springframework.data.support.PageableExecutionUtils;
2826
import org.springframework.javapoet.CodeBlock;
2927
import org.springframework.util.ClassUtils;
3028

@@ -38,14 +36,12 @@ static class GeoNearCodeBlockBuilder {
3836

3937
private final AotQueryMethodGenerationContext context;
4038
private final MongoQueryMethod queryMethod;
41-
private final List<CodeBlock> arguments;
4239

4340
private String variableName;
4441

4542
GeoNearCodeBlockBuilder(AotQueryMethodGenerationContext context, MongoQueryMethod queryMethod) {
4643

4744
this.context = context;
48-
this.arguments = context.getBindableParameterNames().stream().map(CodeBlock::of).collect(Collectors.toList());
4945
this.queryMethod = queryMethod;
5046
}
5147

@@ -119,24 +115,31 @@ CodeBlock build() {
119115
CodeBlock.Builder builder = CodeBlock.builder();
120116
builder.add("\n");
121117

122-
// TODO: move the section below into dedicated executor builder
118+
String executorVar = context.localVariable("nearFinder");
119+
builder.addStatement("var $L = $L.query($T.class).near($L)", executorVar,
120+
context.fieldNameOf(MongoOperations.class), context.getRepositoryInformation().getDomainType(),
121+
queryVariableName);
122+
123123
if (ClassUtils.isAssignable(GeoPage.class, context.getReturnType().getRawClass())) {
124-
builder.addStatement("return new $T<>($L.query($T.class).near($L).all())", GeoPage.class,
125-
context.fieldNameOf(MongoOperations.class), context.getRepositoryInformation().getDomainType(),
126-
queryVariableName);
127-
}
128124

129-
else if (ClassUtils.isAssignable(GeoResults.class, context.getReturnType().getRawClass())) {
125+
String geoResultVar = context.localVariable("geoResult");
126+
builder.addStatement("var $L = $L.all()", geoResultVar, executorVar);
130127

131-
builder.addStatement("return $L.query($T.class).near($L).all()", context.fieldNameOf(MongoOperations.class),
132-
context.getRepositoryInformation().getDomainType(), queryVariableName);
128+
builder.beginControlFlow("if($L.isUnpaged())", context.getPageableParameterName());
129+
builder.addStatement("return new $T<>($L)", GeoPage.class, geoResultVar);
130+
builder.endControlFlow();
131+
132+
String pageVar = context.localVariable("resultPage");
133+
builder.addStatement("var $L = $T.getPage($L.getContent(), $L, () -> $L.count())", pageVar,
134+
PageableExecutionUtils.class, geoResultVar, context.getPageableParameterName(), executorVar);
135+
builder.addStatement("return new $T<>($L, $L, $L.getTotalElements())", GeoPage.class, geoResultVar,
136+
context.getPageableParameterName(), pageVar);
137+
} else if (ClassUtils.isAssignable(GeoResults.class, context.getReturnType().getRawClass())) {
138+
builder.addStatement("return $L.all()", executorVar);
133139
} else {
134-
builder.addStatement("return $L.query($T.class).near($L).all().getContent()",
135-
context.fieldNameOf(MongoOperations.class), context.getRepositoryInformation().getDomainType(),
136-
queryVariableName);
140+
builder.addStatement("return $L.all().getContent()", executorVar);
137141
}
138142
return builder.build();
139143
}
140-
141144
}
142145
}

‎spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/aot/MongoRepositoryContributor.java‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@
3434
import org.apache.commons.logging.LogFactory;
3535
import org.jspecify.annotations.Nullable;
3636
import org.springframework.core.annotation.AnnotatedElementUtils;
37-
import org.springframework.data.geo.GeoPage;
3837
import org.springframework.data.mongodb.core.MongoOperations;
3938
import org.springframework.data.mongodb.core.aggregation.AggregationUpdate;
4039
import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
@@ -52,7 +51,6 @@
5251
import org.springframework.data.repository.query.parser.PartTree;
5352
import org.springframework.javapoet.CodeBlock;
5453
import org.springframework.javapoet.TypeName;
55-
import org.springframework.util.ClassUtils;
5654
import org.springframework.util.ObjectUtils;
5755
import org.springframework.util.StringUtils;
5856

@@ -189,8 +187,7 @@ private static boolean backoff(MongoQueryMethod method) {
189187

190188
// TODO: namedQuery, Regex queries, queries accepting Shapes (e.g. within) or returning arrays.
191189
boolean skip = method.isSearchQuery() || method.getName().toLowerCase(Locale.ROOT).contains("regex")
192-
|| method.getReturnType().getType().isArray()
193-
|| ClassUtils.isAssignable(GeoPage.class, method.getReturnType().getType());
190+
|| method.getReturnType().getType().isArray();
194191

195192
if (skip && logger.isDebugEnabled()) {
196193
logger.debug("Skipping AOT generation for [%s]. Method is either returning an array or a geo-near, regex query"

‎spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/query/MongoQueryExecution.java‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -186,8 +186,11 @@ public Object execute(Query query) {
186186
return isListOfGeoResult(method.getReturnType()) ? results.getContent() : results;
187187
}
188188

189-
@SuppressWarnings({ "unchecked", "NullAway" })
190189
GeoResults<Object> doExecuteQuery(Query query) {
190+
return doExecuteQuery(nearQuery(query));
191+
}
192+
193+
NearQuery nearQuery(Query query) {
191194

192195
Point nearLocation = accessor.getGeoNearLocation();
193196
Assert.notNull(nearLocation, "[query.location] must not be null");
@@ -205,9 +208,12 @@ GeoResults<Object> doExecuteQuery(Query query) {
205208
distances.getUpperBound().getValue().ifPresent(it -> nearQuery.maxDistance(it).in(it.getMetric()));
206209

207210
Pageable pageable = accessor.getPageable();
208-
nearQuery.with(pageable);
211+
return nearQuery.with(pageable);
212+
}
209213

210-
return (GeoResults<Object>) operation.near(nearQuery).all();
214+
@SuppressWarnings({ "unchecked", "NullAway" })
215+
GeoResults<Object> doExecuteQuery(NearQuery query) {
216+
return (GeoResults<Object>) operation.near(query).all();
211217
}
212218

213219
private static boolean isListOfGeoResult(TypeInformation<?> returnType) {
@@ -324,16 +330,11 @@ final class PagingGeoNearExecution extends GeoNearExecution {
324330
@Override
325331
public Object execute(Query query) {
326332

327-
GeoResults<Object> geoResults = doExecuteQuery(query);
333+
NearQuery nearQuery = nearQuery(query);
334+
GeoResults<Object> geoResults = doExecuteQuery(nearQuery);
328335

329336
Page<GeoResult<Object>> page = PageableExecutionUtils.getPage(geoResults.getContent(), accessor.getPageable(),
330-
() -> {
331-
332-
Query countQuery = mongoQuery.createCountQuery(accessor);
333-
countQuery = mongoQuery.applyQueryMetaAttributesWhenPresent(countQuery);
334-
335-
return operation.matching(countQuery).count();
336-
});
337+
() -> operation.near(nearQuery).count());
337338

338339
// transform to GeoPage after applying optimization
339340
return new GeoPage<>(geoResults, accessor.getPageable(), page.getTotalElements());

‎spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ExecutableFindOperationSupportTests.java‎

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,13 @@
1515
*/
1616
package org.springframework.data.mongodb.core;
1717

18-
import static org.assertj.core.api.Assertions.*;
19-
import static org.springframework.data.mongodb.core.query.Criteria.*;
20-
import static org.springframework.data.mongodb.core.query.Query.*;
21-
import static org.springframework.data.mongodb.test.util.DirtiesStateExtension.*;
18+
import static org.assertj.core.api.Assertions.assertThat;
19+
import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
20+
import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
21+
import static org.springframework.data.mongodb.core.query.Criteria.where;
22+
import static org.springframework.data.mongodb.core.query.Query.query;
23+
import static org.springframework.data.mongodb.test.util.DirtiesStateExtension.DirtiesState;
24+
import static org.springframework.data.mongodb.test.util.DirtiesStateExtension.StateFunctions;
2225

2326
import java.util.Date;
2427
import java.util.List;
@@ -47,6 +50,7 @@
4750
import org.springframework.data.mongodb.core.mapping.Field;
4851
import org.springframework.data.mongodb.core.query.BasicQuery;
4952
import org.springframework.data.mongodb.core.query.NearQuery;
53+
import org.springframework.data.mongodb.core.query.Query;
5054
import org.springframework.data.mongodb.test.util.DirtiesStateExtension;
5155
import org.springframework.data.mongodb.test.util.MongoTemplateExtension;
5256
import org.springframework.data.mongodb.test.util.MongoTestTemplate;
@@ -81,7 +85,7 @@ public void clear() {
8185

8286
@Override
8387
public void setupState() {
84-
template.indexOps(Planet.class).ensureIndex(
88+
template.indexOps(Planet.class).createIndex(
8589
new GeospatialIndex("coordinates").typed(GeoSpatialIndexType.GEO_2DSPHERE).named("planet-coordinate-idx"));
8690

8791
initPersons();
@@ -162,7 +166,7 @@ void findAllByWithCollection() {
162166
void findAllAsDocument() {
163167
assertThat(
164168
template.query(Document.class).inCollection(STAR_WARS).matching(query(where("firstname").is("luke"))).all())
165-
.hasSize(1);
169+
.hasSize(1);
166170
}
167171

168172
@Test // DATAMONGO-1563
@@ -324,6 +328,14 @@ void findAllNearBy() {
324328
assertThat(results.getContent().get(0).getDistance()).isNotNull();
325329
}
326330

331+
@Test
332+
void countResultsOfNearQuery() {
333+
334+
Long count = template.query(Planet.class)
335+
.near(NearQuery.near(-73.9667, 40.78).spherical(true).query(new Query(where("name").is("alderan")))).count();
336+
assertThat(count).isEqualTo(1);
337+
}
338+
327339
@Test // DATAMONGO-1563
328340
void findAllNearByWithCollectionAndProjection() {
329341

0 commit comments

Comments
 (0)