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 @@
@@ -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 @@
@@ -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