Skip to content

Commit b5fc922

Browse files
christophstroblmp911de
authored andcommitted
Support generating geospatial queries during AOT.
Closes #5004 Original pull request: #5005
1 parent b1ac55c commit b5fc922

34 files changed

+2185
-819
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/CriteriaDefinition.java

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -46,33 +46,37 @@ public interface CriteriaDefinition {
4646
* @since 5.0
4747
* @author Christoph Strobl
4848
*/
49-
class Placeholder {
50-
51-
private final Object expression;
49+
interface Placeholder {
5250

5351
/**
5452
* Create a new placeholder for index bindable parameter.
5553
*
5654
* @param position the index of the parameter to bind.
5755
* @return new instance of {@link Placeholder}.
5856
*/
59-
public static Placeholder indexed(int position) {
60-
return new Placeholder("?%s".formatted(position));
57+
static Placeholder indexed(int position) {
58+
return new PlaceholderImpl("?%s".formatted(position));
6159
}
6260

63-
public static Placeholder placeholder(String expression) {
64-
return new Placeholder(expression);
61+
static Placeholder placeholder(String expression) {
62+
return new PlaceholderImpl(expression);
6563
}
6664

67-
Placeholder(Object value) {
68-
this.expression = value;
65+
Object getValue();
66+
}
67+
68+
static class PlaceholderImpl implements Placeholder {
69+
private final Object expression;
70+
71+
public PlaceholderImpl(Object expression) {
72+
this.expression = expression;
6973
}
7074

75+
@Override
7176
public Object getValue() {
7277
return expression;
7378
}
7479

75-
@Override
7680
public String toString() {
7781
return getValue().toString();
7882
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import org.springframework.data.geo.Circle;
2323
import org.springframework.data.geo.Polygon;
2424
import org.springframework.data.geo.Shape;
25+
import org.springframework.data.mongodb.core.geo.GeoJson;
2526
import org.springframework.data.mongodb.core.geo.Sphere;
2627
import org.springframework.util.Assert;
2728

@@ -75,6 +76,9 @@ private String getCommand(Shape shape) {
7576

7677
Assert.notNull(shape, "Shape must not be null");
7778

79+
if(shape instanceof GeoJson<?>) {
80+
return "$geometry";
81+
}
7882
if (shape instanceof Box) {
7983
return "$box";
8084
} else if (shape instanceof Circle) {

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

0 commit comments

Comments
 (0)