Skip to content

Commit b5bc432

Browse files
mp911demp911de
committed
DATAMONGO-2059 - Polishing.
Move query rewriting into CountQuery. Consider existing $and items during query rewrite. Original pull request: #604.
1 parent 909c51d commit b5bc432

File tree

6 files changed

+405
-187
lines changed

6 files changed

+405
-187
lines changed

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

+227
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,227 @@
1+
/*
2+
* Copyright 2019 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.core;
17+
18+
import java.util.ArrayList;
19+
import java.util.Arrays;
20+
import java.util.Collection;
21+
import java.util.Collections;
22+
import java.util.List;
23+
import java.util.Map;
24+
25+
import org.bson.Document;
26+
27+
import org.springframework.data.geo.Point;
28+
import org.springframework.lang.Nullable;
29+
import org.springframework.util.ObjectUtils;
30+
31+
/**
32+
* Value object representing a count query. Count queries using {@code $near} or {@code $nearSphere} require a rewrite
33+
* to {@code $geoWithin}.
34+
*
35+
* @author Christoph Strobl
36+
* @author Mark Paluch
37+
* @since 2.3
38+
*/
39+
class CountQuery {
40+
41+
private Document source;
42+
43+
private CountQuery(Document source) {
44+
this.source = source;
45+
}
46+
47+
public static CountQuery of(Document source) {
48+
return new CountQuery(source);
49+
}
50+
51+
/**
52+
* Returns the query {@link Document} that can be used with {@code countDocuments()}. Potentially rewrites the query
53+
* to be usable with {@code countDocuments()}.
54+
*
55+
* @return the query {@link Document} that can be used with {@code countDocuments()}.
56+
*/
57+
public Document toQueryDocument() {
58+
59+
if (!requiresRewrite(source)) {
60+
return source;
61+
}
62+
63+
Document target = new Document();
64+
65+
for (Map.Entry<String, Object> entry : source.entrySet()) {
66+
67+
if (entry.getValue() instanceof Document && requiresRewrite(entry.getValue())) {
68+
69+
Document theValue = (Document) entry.getValue();
70+
target.putAll(createGeoWithin(entry.getKey(), theValue, source.get("$and")));
71+
continue;
72+
}
73+
74+
if (entry.getValue() instanceof Collection && requiresRewrite(entry.getValue())) {
75+
76+
Collection<?> source = (Collection<?>) entry.getValue();
77+
78+
target.put(entry.getKey(), rewriteCollection(source));
79+
continue;
80+
}
81+
82+
if ("$and".equals(entry.getKey()) && target.containsKey("$and")) {
83+
// Expect $and to be processed with Document and createGeoWithin.
84+
continue;
85+
}
86+
87+
target.put(entry.getKey(), entry.getValue());
88+
}
89+
90+
return target;
91+
}
92+
93+
/**
94+
* @param valueToInspect
95+
* @return {@code true} if the enclosing element needs to be rewritten.
96+
*/
97+
private boolean requiresRewrite(Object valueToInspect) {
98+
99+
if (valueToInspect instanceof Document) {
100+
return requiresRewrite((Document) valueToInspect);
101+
}
102+
103+
if (valueToInspect instanceof Collection) {
104+
return requiresRewrite((Collection) valueToInspect);
105+
}
106+
107+
return false;
108+
}
109+
110+
private boolean requiresRewrite(Collection<?> collection) {
111+
112+
for (Object o : collection) {
113+
if (o instanceof Document && requiresRewrite((Document) o)) {
114+
return true;
115+
}
116+
}
117+
118+
return false;
119+
}
120+
121+
private boolean requiresRewrite(Document document) {
122+
123+
if (containsNear(document)) {
124+
return true;
125+
}
126+
127+
for (Object entry : document.values()) {
128+
129+
if (requiresRewrite(entry)) {
130+
return true;
131+
}
132+
}
133+
134+
return false;
135+
}
136+
137+
private Collection<Object> rewriteCollection(Collection<?> source) {
138+
139+
Collection<Object> rewrittenCollection = new ArrayList<>(source.size());
140+
141+
for (Object item : source) {
142+
if (item instanceof Document && requiresRewrite(item)) {
143+
rewrittenCollection.add(CountQuery.of((Document) item).toQueryDocument());
144+
} else {
145+
rewrittenCollection.add(item);
146+
}
147+
}
148+
149+
return rewrittenCollection;
150+
}
151+
152+
/**
153+
* Rewrite the near query for field {@code key} to {@code $geoWithin}.
154+
*
155+
* @param key the queried field.
156+
* @param source source {@link Document}.
157+
* @param $and potentially existing {@code $and} condition.
158+
* @return the rewritten query {@link Document}.
159+
*/
160+
private static Document createGeoWithin(String key, Document source, @Nullable Object $and) {
161+
162+
boolean spheric = source.containsKey("$nearSphere");
163+
Object $near = spheric ? source.get("$nearSphere") : source.get("$near");
164+
165+
Number maxDistance = source.containsKey("$maxDistance") ? (Number) source.get("$maxDistance") : Double.MAX_VALUE;
166+
List<Object> $centerMax = Arrays.asList(toCenterCoordinates($near), maxDistance);
167+
Document $geoWithinMax = new Document("$geoWithin",
168+
new Document(spheric ? "$centerSphere" : "$center", $centerMax));
169+
170+
if (!containsNearWithMinDistance(source)) {
171+
return new Document(key, $geoWithinMax);
172+
}
173+
174+
Number minDistance = (Number) source.get("$minDistance");
175+
List<Object> $centerMin = Arrays.asList(toCenterCoordinates($near), minDistance);
176+
Document $geoWithinMin = new Document("$geoWithin",
177+
new Document(spheric ? "$centerSphere" : "$center", $centerMin));
178+
179+
List<Document> criteria = new ArrayList<>();
180+
181+
if ($and != null) {
182+
if ($and instanceof Collection) {
183+
criteria.addAll((Collection) $and);
184+
} else {
185+
throw new IllegalArgumentException(
186+
"Cannot rewrite query as it contains an '$and' element that is not a Collection!: Offending element: "
187+
+ $and);
188+
}
189+
}
190+
191+
criteria.add(new Document("$nor", Collections.singletonList(new Document(key, $geoWithinMin))));
192+
criteria.add(new Document(key, $geoWithinMax));
193+
return new Document("$and", criteria);
194+
}
195+
196+
private static boolean containsNear(Document source) {
197+
return source.containsKey("$near") || source.containsKey("$nearSphere");
198+
}
199+
200+
private static boolean containsNearWithMinDistance(Document source) {
201+
202+
if (!containsNear(source)) {
203+
return false;
204+
}
205+
206+
return source.containsKey("$minDistance");
207+
}
208+
209+
private static Object toCenterCoordinates(Object value) {
210+
211+
if (ObjectUtils.isArray(value)) {
212+
return value;
213+
}
214+
215+
if (value instanceof Point) {
216+
return Arrays.asList(((Point) value).getX(), ((Point) value).getY());
217+
}
218+
219+
if (value instanceof Document && ((Document) value).containsKey("x")) {
220+
221+
Document point = (Document) value;
222+
return Arrays.asList(point.get("x"), point.get("y"));
223+
}
224+
225+
return value;
226+
}
227+
}

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

+2-1
Original file line numberDiff line numberDiff line change
@@ -1192,7 +1192,8 @@ protected long doCount(String collectionName, Document filter, CountOptions opti
11921192
LOGGER.debug("Executing count: {} in collection: {}", serializeToJsonSafely(filter), collectionName);
11931193
}
11941194

1195-
return execute(collectionName, collection -> collection.countDocuments(QueryMapper.processCountFilter(filter), options));
1195+
return execute(collectionName,
1196+
collection -> collection.countDocuments(CountQuery.of(filter).toQueryDocument(), options));
11961197
}
11971198

11981199
/*

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -1301,7 +1301,7 @@ public Mono<Long> count(Query query, @Nullable Class<?> entityClass, String coll
13011301
protected Mono<Long> doCount(String collectionName, Document filter, CountOptions options) {
13021302

13031303
return createMono(collectionName,
1304-
collection -> collection.countDocuments(QueryMapper.processCountFilter(filter), options));
1304+
collection -> collection.countDocuments(CountQuery.of(filter).toQueryDocument(), options));
13051305
}
13061306

13071307
/*

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

+11-95
Original file line numberDiff line numberDiff line change
@@ -15,19 +15,28 @@
1515
*/
1616
package org.springframework.data.mongodb.core.convert;
1717

18-
import java.util.*;
18+
import java.util.ArrayList;
19+
import java.util.Arrays;
20+
import java.util.Collections;
21+
import java.util.HashSet;
22+
import java.util.Iterator;
23+
import java.util.LinkedHashMap;
24+
import java.util.List;
25+
import java.util.Map;
1926
import java.util.Map.Entry;
27+
import java.util.Optional;
28+
import java.util.Set;
2029
import java.util.regex.Matcher;
2130
import java.util.regex.Pattern;
2231

2332
import org.bson.BsonValue;
2433
import org.bson.Document;
2534
import org.bson.conversions.Bson;
2635
import org.bson.types.ObjectId;
36+
2737
import org.springframework.core.convert.ConversionService;
2838
import org.springframework.core.convert.converter.Converter;
2939
import org.springframework.data.domain.Example;
30-
import org.springframework.data.geo.Point;
3140
import org.springframework.data.mapping.Association;
3241
import org.springframework.data.mapping.MappingException;
3342
import org.springframework.data.mapping.PersistentEntity;
@@ -1282,97 +1291,4 @@ public String convert(MongoPersistentProperty source) {
12821291
public MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> getMappingContext() {
12831292
return mappingContext;
12841293
}
1285-
1286-
public static Document processCountFilter(Document source) {
1287-
1288-
Document target = new Document();
1289-
for (Entry<String, Object> entry : source.entrySet()) {
1290-
1291-
if (entry.getValue() instanceof Document) {
1292-
1293-
Document theValue = (Document) entry.getValue();
1294-
if (containsNear(theValue)) {
1295-
target.putAll(createGeoWithin(entry.getKey(), theValue));
1296-
} else {
1297-
target.put(entry.getKey(), entry.getValue());
1298-
}
1299-
} else if (entry.getValue() instanceof Collection) {
1300-
1301-
Collection<Object> tmp = new ArrayList<>();
1302-
for (Object val : (Collection) entry.getValue()) {
1303-
if (val instanceof Document) {
1304-
tmp.add(processCountFilter((Document) val));
1305-
} else {
1306-
tmp.add(val);
1307-
}
1308-
}
1309-
target.put(entry.getKey(), tmp);
1310-
} else {
1311-
target.put(entry.getKey(), entry.getValue());
1312-
}
1313-
}
1314-
return target;
1315-
}
1316-
1317-
private static Document createGeoWithin(String key, Document source) {
1318-
1319-
boolean spheric = source.containsKey("$nearSphere");
1320-
Object $near = spheric ? source.get("$nearSphere") : source.get("$near");
1321-
1322-
Number maxDistance = source.containsKey("$maxDistance") ? (Number) source.get("$maxDistance") : Double.MAX_VALUE;
1323-
List<Object> $centerMax = Arrays.asList(toCenterCoordinates($near), maxDistance);
1324-
Document $geoWithinMax = new Document("$geoWithin",
1325-
new Document(spheric ? "$centerSphere" : "$center", $centerMax));
1326-
1327-
if (!containsNearWithMinDistance(source)) {
1328-
return new Document(key, $geoWithinMax);
1329-
}
1330-
1331-
Number minDistance = (Number) source.get("$minDistance");
1332-
List<Object> $centerMin = Arrays.asList(toCenterCoordinates($near), minDistance);
1333-
Document $geoWithinMin = new Document("$geoWithin",
1334-
new Document(spheric ? "$centerSphere" : "$center", $centerMin));
1335-
1336-
List<Document> criteria = new ArrayList<>();
1337-
criteria.add(new Document("$nor", Arrays.asList(new Document(key, $geoWithinMin))));
1338-
criteria.add(new Document(key, $geoWithinMax));
1339-
return new Document("$and", criteria);
1340-
}
1341-
1342-
private static boolean containsNear(Document source) {
1343-
1344-
if (source.containsKey("$near") || source.containsKey("$nearSphere")) {
1345-
return true;
1346-
}
1347-
1348-
return false;
1349-
}
1350-
1351-
private static boolean containsNearWithMinDistance(Document source) {
1352-
1353-
if (!containsNear(source)) {
1354-
return false;
1355-
}
1356-
1357-
return source.containsKey("$minDistance");
1358-
}
1359-
1360-
private static Object toCenterCoordinates(Object value) {
1361-
1362-
if (ObjectUtils.isArray(value)) {
1363-
return value;
1364-
}
1365-
1366-
if (value instanceof Point) {
1367-
return Arrays.asList(((Point) value).getX(), ((Point) value).getY());
1368-
}
1369-
1370-
if (value instanceof Document && ((Document) value).containsKey("x")) {
1371-
1372-
Document point = (Document) value;
1373-
return Arrays.asList(point.get("x"), point.get("y"));
1374-
}
1375-
1376-
return value;
1377-
}
13781294
}

0 commit comments

Comments
 (0)