Browse Source
Move query rewriting into CountQuery. Consider existing $and items during query rewrite. Original pull request: #604.pull/802/head
Mark Paluch
6 years ago
6 changed files with 405 additions and 187 deletions
@ -0,0 +1,227 @@ |
|||||||
|
/* |
||||||
|
* Copyright 2019 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.mongodb.core; |
||||||
|
|
||||||
|
import java.util.ArrayList; |
||||||
|
import java.util.Arrays; |
||||||
|
import java.util.Collection; |
||||||
|
import java.util.Collections; |
||||||
|
import java.util.List; |
||||||
|
import java.util.Map; |
||||||
|
|
||||||
|
import org.bson.Document; |
||||||
|
|
||||||
|
import org.springframework.data.geo.Point; |
||||||
|
import org.springframework.lang.Nullable; |
||||||
|
import org.springframework.util.ObjectUtils; |
||||||
|
|
||||||
|
/** |
||||||
|
* Value object representing a count query. Count queries using {@code $near} or {@code $nearSphere} require a rewrite |
||||||
|
* to {@code $geoWithin}. |
||||||
|
* |
||||||
|
* @author Christoph Strobl |
||||||
|
* @author Mark Paluch |
||||||
|
* @since 2.3 |
||||||
|
*/ |
||||||
|
class CountQuery { |
||||||
|
|
||||||
|
private Document source; |
||||||
|
|
||||||
|
private CountQuery(Document source) { |
||||||
|
this.source = source; |
||||||
|
} |
||||||
|
|
||||||
|
public static CountQuery of(Document source) { |
||||||
|
return new CountQuery(source); |
||||||
|
} |
||||||
|
|
||||||
|
/** |
||||||
|
* Returns the query {@link Document} that can be used with {@code countDocuments()}. Potentially rewrites the query |
||||||
|
* to be usable with {@code countDocuments()}. |
||||||
|
* |
||||||
|
* @return the query {@link Document} that can be used with {@code countDocuments()}. |
||||||
|
*/ |
||||||
|
public Document toQueryDocument() { |
||||||
|
|
||||||
|
if (!requiresRewrite(source)) { |
||||||
|
return source; |
||||||
|
} |
||||||
|
|
||||||
|
Document target = new Document(); |
||||||
|
|
||||||
|
for (Map.Entry<String, Object> entry : source.entrySet()) { |
||||||
|
|
||||||
|
if (entry.getValue() instanceof Document && requiresRewrite(entry.getValue())) { |
||||||
|
|
||||||
|
Document theValue = (Document) entry.getValue(); |
||||||
|
target.putAll(createGeoWithin(entry.getKey(), theValue, source.get("$and"))); |
||||||
|
continue; |
||||||
|
} |
||||||
|
|
||||||
|
if (entry.getValue() instanceof Collection && requiresRewrite(entry.getValue())) { |
||||||
|
|
||||||
|
Collection<?> source = (Collection<?>) entry.getValue(); |
||||||
|
|
||||||
|
target.put(entry.getKey(), rewriteCollection(source)); |
||||||
|
continue; |
||||||
|
} |
||||||
|
|
||||||
|
if ("$and".equals(entry.getKey()) && target.containsKey("$and")) { |
||||||
|
// Expect $and to be processed with Document and createGeoWithin.
|
||||||
|
continue; |
||||||
|
} |
||||||
|
|
||||||
|
target.put(entry.getKey(), entry.getValue()); |
||||||
|
} |
||||||
|
|
||||||
|
return target; |
||||||
|
} |
||||||
|
|
||||||
|
/** |
||||||
|
* @param valueToInspect |
||||||
|
* @return {@code true} if the enclosing element needs to be rewritten. |
||||||
|
*/ |
||||||
|
private boolean requiresRewrite(Object valueToInspect) { |
||||||
|
|
||||||
|
if (valueToInspect instanceof Document) { |
||||||
|
return requiresRewrite((Document) valueToInspect); |
||||||
|
} |
||||||
|
|
||||||
|
if (valueToInspect instanceof Collection) { |
||||||
|
return requiresRewrite((Collection) valueToInspect); |
||||||
|
} |
||||||
|
|
||||||
|
return false; |
||||||
|
} |
||||||
|
|
||||||
|
private boolean requiresRewrite(Collection<?> collection) { |
||||||
|
|
||||||
|
for (Object o : collection) { |
||||||
|
if (o instanceof Document && requiresRewrite((Document) o)) { |
||||||
|
return true; |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
return false; |
||||||
|
} |
||||||
|
|
||||||
|
private boolean requiresRewrite(Document document) { |
||||||
|
|
||||||
|
if (containsNear(document)) { |
||||||
|
return true; |
||||||
|
} |
||||||
|
|
||||||
|
for (Object entry : document.values()) { |
||||||
|
|
||||||
|
if (requiresRewrite(entry)) { |
||||||
|
return true; |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
return false; |
||||||
|
} |
||||||
|
|
||||||
|
private Collection<Object> rewriteCollection(Collection<?> source) { |
||||||
|
|
||||||
|
Collection<Object> rewrittenCollection = new ArrayList<>(source.size()); |
||||||
|
|
||||||
|
for (Object item : source) { |
||||||
|
if (item instanceof Document && requiresRewrite(item)) { |
||||||
|
rewrittenCollection.add(CountQuery.of((Document) item).toQueryDocument()); |
||||||
|
} else { |
||||||
|
rewrittenCollection.add(item); |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
return rewrittenCollection; |
||||||
|
} |
||||||
|
|
||||||
|
/** |
||||||
|
* Rewrite the near query for field {@code key} to {@code $geoWithin}. |
||||||
|
* |
||||||
|
* @param key the queried field. |
||||||
|
* @param source source {@link Document}. |
||||||
|
* @param $and potentially existing {@code $and} condition. |
||||||
|
* @return the rewritten query {@link Document}. |
||||||
|
*/ |
||||||
|
private static Document createGeoWithin(String key, Document source, @Nullable Object $and) { |
||||||
|
|
||||||
|
boolean spheric = source.containsKey("$nearSphere"); |
||||||
|
Object $near = spheric ? source.get("$nearSphere") : source.get("$near"); |
||||||
|
|
||||||
|
Number maxDistance = source.containsKey("$maxDistance") ? (Number) source.get("$maxDistance") : Double.MAX_VALUE; |
||||||
|
List<Object> $centerMax = Arrays.asList(toCenterCoordinates($near), maxDistance); |
||||||
|
Document $geoWithinMax = new Document("$geoWithin", |
||||||
|
new Document(spheric ? "$centerSphere" : "$center", $centerMax)); |
||||||
|
|
||||||
|
if (!containsNearWithMinDistance(source)) { |
||||||
|
return new Document(key, $geoWithinMax); |
||||||
|
} |
||||||
|
|
||||||
|
Number minDistance = (Number) source.get("$minDistance"); |
||||||
|
List<Object> $centerMin = Arrays.asList(toCenterCoordinates($near), minDistance); |
||||||
|
Document $geoWithinMin = new Document("$geoWithin", |
||||||
|
new Document(spheric ? "$centerSphere" : "$center", $centerMin)); |
||||||
|
|
||||||
|
List<Document> criteria = new ArrayList<>(); |
||||||
|
|
||||||
|
if ($and != null) { |
||||||
|
if ($and instanceof Collection) { |
||||||
|
criteria.addAll((Collection) $and); |
||||||
|
} else { |
||||||
|
throw new IllegalArgumentException( |
||||||
|
"Cannot rewrite query as it contains an '$and' element that is not a Collection!: Offending element: " |
||||||
|
+ $and); |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
criteria.add(new Document("$nor", Collections.singletonList(new Document(key, $geoWithinMin)))); |
||||||
|
criteria.add(new Document(key, $geoWithinMax)); |
||||||
|
return new Document("$and", criteria); |
||||||
|
} |
||||||
|
|
||||||
|
private static boolean containsNear(Document source) { |
||||||
|
return source.containsKey("$near") || source.containsKey("$nearSphere"); |
||||||
|
} |
||||||
|
|
||||||
|
private static boolean containsNearWithMinDistance(Document source) { |
||||||
|
|
||||||
|
if (!containsNear(source)) { |
||||||
|
return false; |
||||||
|
} |
||||||
|
|
||||||
|
return source.containsKey("$minDistance"); |
||||||
|
} |
||||||
|
|
||||||
|
private static Object toCenterCoordinates(Object value) { |
||||||
|
|
||||||
|
if (ObjectUtils.isArray(value)) { |
||||||
|
return value; |
||||||
|
} |
||||||
|
|
||||||
|
if (value instanceof Point) { |
||||||
|
return Arrays.asList(((Point) value).getX(), ((Point) value).getY()); |
||||||
|
} |
||||||
|
|
||||||
|
if (value instanceof Document && ((Document) value).containsKey("x")) { |
||||||
|
|
||||||
|
Document point = (Document) value; |
||||||
|
return Arrays.asList(point.get("x"), point.get("y")); |
||||||
|
} |
||||||
|
|
||||||
|
return value; |
||||||
|
} |
||||||
|
} |
||||||
@ -0,0 +1,164 @@ |
|||||||
|
/* |
||||||
|
* Copyright 2019 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.mongodb.core; |
||||||
|
|
||||||
|
import static org.mockito.Mockito.*; |
||||||
|
import static org.springframework.data.mongodb.core.query.Criteria.*; |
||||||
|
import static org.springframework.data.mongodb.core.query.Query.*; |
||||||
|
import static org.springframework.data.mongodb.test.util.Assertions.*; |
||||||
|
|
||||||
|
import org.junit.Before; |
||||||
|
import org.junit.Test; |
||||||
|
|
||||||
|
import org.springframework.data.geo.Point; |
||||||
|
import org.springframework.data.mongodb.MongoDbFactory; |
||||||
|
import org.springframework.data.mongodb.core.convert.DefaultDbRefResolver; |
||||||
|
import org.springframework.data.mongodb.core.convert.MappingMongoConverter; |
||||||
|
import org.springframework.data.mongodb.core.convert.QueryMapper; |
||||||
|
import org.springframework.data.mongodb.core.mapping.MongoMappingContext; |
||||||
|
import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity; |
||||||
|
import org.springframework.data.mongodb.core.query.Criteria; |
||||||
|
import org.springframework.data.mongodb.core.query.Query; |
||||||
|
|
||||||
|
/** |
||||||
|
* Unit tests for {@link CountQuery}. |
||||||
|
* |
||||||
|
* @author Mark Paluch |
||||||
|
* @author Christoph Strobl |
||||||
|
*/ |
||||||
|
public class CountQueryUnitTests { |
||||||
|
|
||||||
|
QueryMapper mapper; |
||||||
|
MongoMappingContext context; |
||||||
|
MappingMongoConverter converter; |
||||||
|
|
||||||
|
MongoDbFactory factory = mock(MongoDbFactory.class); |
||||||
|
|
||||||
|
@Before |
||||||
|
public void setUp() { |
||||||
|
|
||||||
|
this.context = new MongoMappingContext(); |
||||||
|
|
||||||
|
this.converter = new MappingMongoConverter(new DefaultDbRefResolver(factory), context); |
||||||
|
this.converter.afterPropertiesSet(); |
||||||
|
|
||||||
|
this.mapper = new QueryMapper(converter); |
||||||
|
} |
||||||
|
|
||||||
|
@Test // DATAMONGO-2059
|
||||||
|
public void nearToGeoWithinWithoutDistance() { |
||||||
|
|
||||||
|
Query source = query(where("location").near(new Point(-73.99171, 40.738868))); |
||||||
|
org.bson.Document target = postProcessQueryForCount(source); |
||||||
|
|
||||||
|
assertThat(target).isEqualTo(org.bson.Document |
||||||
|
.parse("{\"location\": {\"$geoWithin\": {\"$center\": [[-73.99171, 40.738868], 1.7976931348623157E308]}}}")); |
||||||
|
} |
||||||
|
|
||||||
|
@Test // DATAMONGO-2059
|
||||||
|
public void nearAndExisting$and() { |
||||||
|
|
||||||
|
Query source = query(where("location").near(new Point(-73.99171, 40.738868)).minDistance(0.01)) |
||||||
|
.addCriteria(new Criteria().andOperator(where("foo").is("bar"))); |
||||||
|
org.bson.Document target = postProcessQueryForCount(source); |
||||||
|
|
||||||
|
assertThat(target).isEqualTo(org.bson.Document.parse("{\"$and\":[" //
|
||||||
|
+ "{\"foo\":\"bar\"}" //
|
||||||
|
+ "{\"$nor\":[{\"location\":{\"$geoWithin\":{\"$center\":[ [ -73.99171, 40.738868 ], 0.01]}}}]},"//
|
||||||
|
+ " {\"location\":{\"$geoWithin\":{\"$center\":[ [ -73.99171, 40.738868 ], 1.7976931348623157E308]}}},"//
|
||||||
|
+ "]}")); |
||||||
|
} |
||||||
|
|
||||||
|
@Test // DATAMONGO-2059
|
||||||
|
public void nearSphereToGeoWithinWithoutDistance() { |
||||||
|
|
||||||
|
Query source = query(where("location").nearSphere(new Point(-73.99171, 40.738868))); |
||||||
|
org.bson.Document target = postProcessQueryForCount(source); |
||||||
|
|
||||||
|
assertThat(target).isEqualTo(org.bson.Document.parse( |
||||||
|
"{\"location\": {\"$geoWithin\": {\"$centerSphere\": [[-73.99171, 40.738868], 1.7976931348623157E308]}}}")); |
||||||
|
} |
||||||
|
|
||||||
|
@Test // DATAMONGO-2059
|
||||||
|
public void nearToGeoWithinWithMaxDistance() { |
||||||
|
|
||||||
|
Query source = query(where("location").near(new Point(-73.99171, 40.738868)).maxDistance(10)); |
||||||
|
org.bson.Document target = postProcessQueryForCount(source); |
||||||
|
|
||||||
|
assertThat(target).isEqualTo( |
||||||
|
org.bson.Document.parse("{\"location\": {\"$geoWithin\": {\"$center\": [[-73.99171, 40.738868], 10.0]}}}")); |
||||||
|
} |
||||||
|
|
||||||
|
@Test // DATAMONGO-2059
|
||||||
|
public void nearSphereToGeoWithinWithMaxDistance() { |
||||||
|
|
||||||
|
Query source = query(where("location").nearSphere(new Point(-73.99171, 40.738868)).maxDistance(10)); |
||||||
|
org.bson.Document target = postProcessQueryForCount(source); |
||||||
|
|
||||||
|
assertThat(target).isEqualTo(org.bson.Document |
||||||
|
.parse("{\"location\": {\"$geoWithin\": {\"$centerSphere\": [[-73.99171, 40.738868], 10.0]}}}")); |
||||||
|
} |
||||||
|
|
||||||
|
@Test // DATAMONGO-2059
|
||||||
|
public void nearToGeoWithinWithMinDistance() { |
||||||
|
|
||||||
|
Query source = query(where("location").near(new Point(-73.99171, 40.738868)).minDistance(0.01)); |
||||||
|
org.bson.Document target = postProcessQueryForCount(source); |
||||||
|
|
||||||
|
assertThat(target).isEqualTo(org.bson.Document.parse( |
||||||
|
"{\"$and\":[{\"$nor\":[{\"location\":{\"$geoWithin\":{\"$center\":[ [ -73.99171, 40.738868 ], 0.01]}}}]}," |
||||||
|
+ " {\"location\":{\"$geoWithin\":{\"$center\":[ [ -73.99171, 40.738868 ], 1.7976931348623157E308]}}}]}")); |
||||||
|
} |
||||||
|
|
||||||
|
@Test // DATAMONGO-2059
|
||||||
|
public void nearToGeoWithinWithMaxDistanceAndCombinedWithOtherCriteria() { |
||||||
|
|
||||||
|
Query source = query( |
||||||
|
where("name").is("food").and("location").near(new Point(-73.99171, 40.738868)).maxDistance(10)); |
||||||
|
org.bson.Document target = postProcessQueryForCount(source); |
||||||
|
|
||||||
|
assertThat(target).isEqualTo(org.bson.Document |
||||||
|
.parse("{\"name\": \"food\", \"location\": {\"$geoWithin\": {\"$center\": [[-73.99171, 40.738868], 10.0]}}}")); |
||||||
|
} |
||||||
|
|
||||||
|
@Test // DATAMONGO-2059
|
||||||
|
public void nearToGeoWithinWithMinDistanceOrCombinedWithOtherCriteria() { |
||||||
|
|
||||||
|
Query source = query(new Criteria().orOperator(where("name").is("food"), |
||||||
|
where("location").near(new Point(-73.99171, 40.738868)).minDistance(0.01))); |
||||||
|
org.bson.Document target = postProcessQueryForCount(source); |
||||||
|
|
||||||
|
assertThat(target).isEqualTo(org.bson.Document.parse( |
||||||
|
"{\"$or\" : [ { \"name\": \"food\" }, {\"$and\":[{\"$nor\":[{\"location\":{\"$geoWithin\":{\"$center\":[ [ -73.99171, 40.738868 ], 0.01]}}}]},{\"location\":{\"$geoWithin\":{\"$center\":[ [ -73.99171, 40.738868 ], 1.7976931348623157E308]}}}]} ]}")); |
||||||
|
} |
||||||
|
|
||||||
|
@Test // DATAMONGO-2059
|
||||||
|
public void nearToGeoWithinWithMaxDistanceOrCombinedWithOtherCriteria() { |
||||||
|
|
||||||
|
Query source = query(new Criteria().orOperator(where("name").is("food"), |
||||||
|
where("location").near(new Point(-73.99171, 40.738868)).maxDistance(10))); |
||||||
|
org.bson.Document target = postProcessQueryForCount(source); |
||||||
|
|
||||||
|
assertThat(target).isEqualTo(org.bson.Document.parse( |
||||||
|
"{\"$or\" : [ { \"name\": \"food\" }, {\"location\": {\"$geoWithin\": {\"$center\": [[-73.99171, 40.738868], 10.0]}}} ]}")); |
||||||
|
} |
||||||
|
|
||||||
|
private org.bson.Document postProcessQueryForCount(Query source) { |
||||||
|
|
||||||
|
org.bson.Document intermediate = mapper.getMappedObject(source.getQueryObject(), (MongoPersistentEntity<?>) null); |
||||||
|
return CountQuery.of(intermediate).toQueryDocument(); |
||||||
|
} |
||||||
|
} |
||||||
Loading…
Reference in new issue