Skip to content

Commit 435ceab

Browse files
christophstroblchristophstrobl
committed
ReadPreference for geo queries
1 parent 364fa66 commit 435ceab

File tree

6 files changed

+182
-15
lines changed

6 files changed

+182
-15
lines changed

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@ CodeBlock build() {
8585
builder.addStatement("$L.with($L)", variableName, context.getPageableParameterName());
8686
}
8787

88+
MongoCodeBlocks.appendReadPreference(context, builder, variableName);
89+
8890
return builder.build();
8991
}
9092

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121

2222
import org.bson.Document;
2323
import org.jspecify.annotations.Nullable;
24+
import org.springframework.core.annotation.MergedAnnotation;
25+
import org.springframework.data.mongodb.repository.ReadPreference;
2426
import org.springframework.data.mongodb.repository.aot.AggregationBlocks.AggregationCodeBlockBuilder;
2527
import org.springframework.data.mongodb.repository.aot.AggregationBlocks.AggregationExecutionCodeBlockBuilder;
2628
import org.springframework.data.mongodb.repository.aot.DeleteBlocks.DeleteExecutionCodeBlockBuilder;
@@ -189,4 +191,15 @@ static CodeBlock renderExpressionToDocument(@Nullable String source, String vari
189191
static boolean containsPlaceholder(String source) {
190192
return PARAMETER_BINDING_PATTERN.matcher(source).find();
191193
}
194+
195+
static void appendReadPreference(AotQueryMethodGenerationContext context, Builder builder, String queryVariableName) {
196+
197+
MergedAnnotation<ReadPreference> readPreferenceAnnotation = context.getAnnotation(ReadPreference.class);
198+
String readPreference = readPreferenceAnnotation.isPresent() ? readPreferenceAnnotation.getString("value") : null;
199+
200+
if (StringUtils.hasText(readPreference)) {
201+
builder.addStatement("$L.withReadPreference($T.valueOf($S))", queryVariableName,
202+
com.mongodb.ReadPreference.class, readPreference);
203+
}
204+
}
192205
}

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

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -131,8 +131,6 @@ protected void customizeConstructor(AotRepositoryConstructorBuilder constructorB
131131
}
132132
}
133133

134-
135-
136134
if (query.isDelete()) {
137135
return deleteMethodContributor(queryMethod, query);
138136
}
@@ -216,7 +214,7 @@ private static MethodContributor<MongoQueryMethod> nearQueryMethodContributor(Mo
216214
});
217215
}
218216

219-
private static MethodContributor<MongoQueryMethod> aggregationMethodContributor(MongoQueryMethod queryMethod,
217+
static MethodContributor<MongoQueryMethod> aggregationMethodContributor(MongoQueryMethod queryMethod,
220218
AggregationInteraction aggregation) {
221219

222220
return MethodContributor.forQueryMethod(queryMethod).withMetadata(aggregation).contribute(context -> {
@@ -232,7 +230,7 @@ private static MethodContributor<MongoQueryMethod> aggregationMethodContributor(
232230
});
233231
}
234232

235-
private static MethodContributor<MongoQueryMethod> updateMethodContributor(MongoQueryMethod queryMethod,
233+
static MethodContributor<MongoQueryMethod> updateMethodContributor(MongoQueryMethod queryMethod,
236234
UpdateInteraction update) {
237235

238236
return MethodContributor.forQueryMethod(queryMethod).withMetadata(update).contribute(context -> {
@@ -261,7 +259,7 @@ private static MethodContributor<MongoQueryMethod> updateMethodContributor(Mongo
261259
});
262260
}
263261

264-
private static MethodContributor<MongoQueryMethod> aggregationUpdateMethodContributor(MongoQueryMethod queryMethod,
262+
static MethodContributor<MongoQueryMethod> aggregationUpdateMethodContributor(MongoQueryMethod queryMethod,
265263
AggregationUpdateInteraction update) {
266264

267265
return MethodContributor.forQueryMethod(queryMethod).withMetadata(update).contribute(context -> {
@@ -287,7 +285,7 @@ private static MethodContributor<MongoQueryMethod> aggregationUpdateMethodContri
287285
});
288286
}
289287

290-
private static MethodContributor<MongoQueryMethod> deleteMethodContributor(MongoQueryMethod queryMethod,
288+
static MethodContributor<MongoQueryMethod> deleteMethodContributor(MongoQueryMethod queryMethod,
291289
QueryInteraction query) {
292290

293291
return MethodContributor.forQueryMethod(queryMethod).withMetadata(query).contribute(context -> {
@@ -302,7 +300,7 @@ private static MethodContributor<MongoQueryMethod> deleteMethodContributor(Mongo
302300
});
303301
}
304302

305-
private static MethodContributor<MongoQueryMethod> queryMethodContributor(MongoQueryMethod queryMethod,
303+
static MethodContributor<MongoQueryMethod> queryMethodContributor(MongoQueryMethod queryMethod,
306304
QueryInteraction query) {
307305

308306
return MethodContributor.forQueryMethod(queryMethod).withMetadata(query).contribute(context -> {

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

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@
3434
import org.springframework.data.mongodb.core.query.BasicQuery;
3535
import org.springframework.data.mongodb.repository.Hint;
3636
import org.springframework.data.mongodb.repository.Meta;
37-
import org.springframework.data.mongodb.repository.ReadPreference;
3837
import org.springframework.data.mongodb.repository.query.MongoParameters.MongoParameter;
3938
import org.springframework.data.mongodb.repository.query.MongoQueryExecution.PagedExecution;
4039
import org.springframework.data.mongodb.repository.query.MongoQueryExecution.SlicedExecution;
@@ -249,13 +248,7 @@ CodeBlock build() {
249248
builder.addStatement("$L.withHint($S)", queryVariableName, hint);
250249
}
251250

252-
MergedAnnotation<ReadPreference> readPreferenceAnnotation = context.getAnnotation(ReadPreference.class);
253-
String readPreference = readPreferenceAnnotation.isPresent() ? readPreferenceAnnotation.getString("value") : null;
254-
255-
if (StringUtils.hasText(readPreference)) {
256-
builder.addStatement("$L.withReadPreference($T.valueOf($S))", queryVariableName,
257-
com.mongodb.ReadPreference.class, readPreference);
258-
}
251+
MongoCodeBlocks.appendReadPreference(context, builder, queryVariableName);
259252

260253
MergedAnnotation<Meta> metaAnnotation = context.getAnnotation(Meta.class);
261254
if (metaAnnotation.isPresent()) {
@@ -297,6 +290,8 @@ private CodeBlock renderExpressionToQuery(@Nullable String source, String variab
297290
builder.add(iterator.next());
298291
if (iterator.hasNext()) {
299292
builder.add(", ");
293+
} else {
294+
builder.add(" ");
300295
}
301296
}
302297
builder.add("});\n");

‎spring-data-mongodb/src/test/java/example/aot/UserRepository.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
import org.springframework.data.geo.Polygon;
4444
import org.springframework.data.mongodb.core.aggregation.AggregationResults;
4545
import org.springframework.data.mongodb.core.geo.GeoJsonPolygon;
46+
import org.springframework.data.mongodb.core.geo.Sphere;
4647
import org.springframework.data.mongodb.repository.Aggregation;
4748
import org.springframework.data.mongodb.repository.Hint;
4849
import org.springframework.data.mongodb.repository.Person;
@@ -118,6 +119,8 @@ public interface UserRepository extends CrudRepository<User, String> {
118119

119120
List<User> findByLocationCoordinatesWithin(Circle circle);
120121

122+
List<User> findByLocationCoordinatesWithin(Sphere circle);
123+
121124
List<User> findByLocationCoordinatesWithin(Box box);
122125

123126
List<User> findByLocationCoordinatesWithin(Polygon polygon);

‎spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/aot/QueryMethodContributionUnitTests.java

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
/*
2+
* Copyright 2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.data.mongodb.repository.aot;
17+
18+
import static org.assertj.core.api.Assertions.assertThat;
19+
20+
import example.aot.User;
21+
import example.aot.UserRepository;
22+
23+
import java.lang.reflect.Method;
24+
25+
import javax.lang.model.element.Modifier;
26+
27+
import org.junit.jupiter.api.Test;
28+
import org.springframework.data.geo.Box;
29+
import org.springframework.data.geo.Circle;
30+
import org.springframework.data.geo.Distance;
31+
import org.springframework.data.geo.GeoResults;
32+
import org.springframework.data.geo.Point;
33+
import org.springframework.data.geo.Polygon;
34+
import org.springframework.data.mongodb.core.MongoOperations;
35+
import org.springframework.data.mongodb.core.geo.Sphere;
36+
import org.springframework.data.mongodb.repository.ReadPreference;
37+
import org.springframework.data.repository.Repository;
38+
import org.springframework.data.repository.aot.generate.AotQueryMethodGenerationContext;
39+
import org.springframework.data.repository.aot.generate.AotRepositoryFragmentMetadata;
40+
import org.springframework.data.repository.aot.generate.MethodContributor;
41+
import org.springframework.data.repository.core.RepositoryInformation;
42+
import org.springframework.data.repository.query.QueryMethod;
43+
import org.springframework.javapoet.ClassName;
44+
import org.springframework.javapoet.FieldSpec;
45+
import org.springframework.javapoet.MethodSpec;
46+
47+
/**
48+
* @author Christoph Strobl
49+
*/
50+
public class QueryMethodContributionUnitTests {
51+
52+
@Test
53+
void rendersQueryForNearUsingPoint() throws NoSuchMethodException {
54+
55+
MethodSpec methodSpec = codeOf(UserRepository.class, "findByLocationCoordinatesNear", Point.class);
56+
57+
assertThat(methodSpec.toString()) //
58+
.contains("{'location.coordinates':{'$near':?0}}") //
59+
.contains("Object[]{ location }") //
60+
.contains("return finder.matching(filterQuery).all()");
61+
}
62+
63+
@Test
64+
void rendersQueryForWithinUsingCircle() throws NoSuchMethodException {
65+
66+
MethodSpec methodSpec = codeOf(UserRepository.class, "findByLocationCoordinatesWithin", Circle.class);
67+
68+
assertThat(methodSpec.toString()) //
69+
.contains("{'location.coordinates':{'$geoWithin':{'$center':?0}}") //
70+
.contains(
71+
"List.of(circle.getCenter().getX(), circle.getCenter().getY()), circle.getRadius().getNormalizedValue())") //
72+
.contains("return finder.matching(filterQuery).all()");
73+
}
74+
75+
@Test
76+
void rendersQueryForWithinUsingSphere() throws NoSuchMethodException {
77+
78+
MethodSpec methodSpec = codeOf(UserRepository.class, "findByLocationCoordinatesWithin", Sphere.class);
79+
80+
assertThat(methodSpec.toString()) //
81+
.contains("{'location.coordinates':{'$geoWithin':{'$centerSphere':?0}}") //
82+
.contains(
83+
"List.of(circle.getCenter().getX(), circle.getCenter().getY()), circle.getRadius().getNormalizedValue())") //
84+
.contains("return finder.matching(filterQuery).all()");
85+
}
86+
87+
@Test
88+
void rendersQueryForWithinUsingBox() throws NoSuchMethodException {
89+
90+
MethodSpec methodSpec = codeOf(UserRepository.class, "findByLocationCoordinatesWithin", Box.class);
91+
92+
assertThat(methodSpec.toString()) //
93+
.contains("{'location.coordinates':{'$geoWithin':{'$box':?0}}") //
94+
.contains("List.of(box.getFirst().getX(), box.getFirst().getY())") //
95+
.contains("List.of(box.getSecond().getX(), box.getSecond().getY())") //
96+
.contains("return finder.matching(filterQuery).all()");
97+
}
98+
99+
@Test
100+
void rendersQueryForWithinUsingPolygon() throws NoSuchMethodException {
101+
102+
MethodSpec methodSpec = codeOf(UserRepository.class, "findByLocationCoordinatesWithin", Polygon.class);
103+
104+
assertThat(methodSpec.toString()) //
105+
.contains("{'location.coordinates':{'$geoWithin':{'$polygon':?0}}") //
106+
.contains("polygon.getPoints().stream().map(_p ->") //
107+
.contains("List.of(_p.getX(), _p.getY())") //
108+
.contains("return finder.matching(filterQuery).all()");
109+
}
110+
111+
@Test
112+
void rendersNearQueryForGeoResults() throws NoSuchMethodException {
113+
114+
MethodSpec methodSpec = codeOf(UserRepoWithMeta.class, "findByLocationCoordinatesNear", Point.class,
115+
Distance.class);
116+
117+
assertThat(methodSpec.toString()) //
118+
.contains("NearQuery.near(point)") //
119+
.contains("nearQuery.maxDistance(maxDistance).in(maxDistance.getMetric())") //
120+
.contains(".withReadPreference(com.mongodb.ReadPreference.valueOf(\"NEAREST\")") //
121+
.contains(".near(nearQuery).all()");
122+
}
123+
124+
private static MethodSpec codeOf(Class<?> repository, String methodName, Class<?>... args)
125+
throws NoSuchMethodException {
126+
127+
Method method = repository.getMethod(methodName, args);
128+
129+
TestMongoAotRepositoryContext repoContext = new TestMongoAotRepositoryContext(repository, null);
130+
MongoRepositoryContributor contributor = new MongoRepositoryContributor(repoContext);
131+
MethodContributor<? extends QueryMethod> methodContributor = contributor.contributeQueryMethod(method);
132+
133+
AotRepositoryFragmentMetadata metadata = new AotRepositoryFragmentMetadata(ClassName.get(UserRepository.class));
134+
metadata.addField(
135+
FieldSpec.builder(MongoOperations.class, "mongoOperations", Modifier.PRIVATE, Modifier.FINAL).build());
136+
137+
TestQueryMethodGenerationContext methodContext = new TestQueryMethodGenerationContext(
138+
repoContext.getRepositoryInformation(), method, methodContributor.getQueryMethod(), metadata);
139+
return methodContributor.contribute(methodContext);
140+
}
141+
142+
static class TestQueryMethodGenerationContext extends AotQueryMethodGenerationContext {
143+
144+
protected TestQueryMethodGenerationContext(RepositoryInformation repositoryInformation, Method method,
145+
QueryMethod queryMethod, AotRepositoryFragmentMetadata targetTypeMetadata) {
146+
super(repositoryInformation, method, queryMethod, targetTypeMetadata);
147+
}
148+
}
149+
150+
interface UserRepoWithMeta extends Repository<User, String> {
151+
152+
@ReadPreference("NEAREST")
153+
GeoResults<User> findByLocationCoordinatesNear(Point point, Distance maxDistance);
154+
155+
}
156+
}

0 commit comments

Comments
 (0)