From 4b6915976ce7e3d688711c49179a85942716dbab Mon Sep 17 00:00:00 2001
From: Igor Motov <igor@motovs.org>
Date: Wed, 23 May 2018 15:15:19 -0400
Subject: [PATCH] Add support for indexed shape routing in geo_shape query
 (#30760)

Adds ability to specify the routing value for the indexed shape in the
geo_shape query.

Closes #7663
---
 .../query-dsl/geo-shape-query.asciidoc        |  1 +
 .../index/query/GeoShapeQueryBuilder.java     | 47 ++++++++++++++++++-
 .../query/GeoShapeQueryBuilderTests.java      |  7 +++
 .../search/geo/GeoShapeIntegrationIT.java     | 38 +++++++++++++++
 4 files changed, 92 insertions(+), 1 deletion(-)

diff --git a/docs/reference/query-dsl/geo-shape-query.asciidoc b/docs/reference/query-dsl/geo-shape-query.asciidoc
index 08b504951e1..4e00a2f49b4 100644
--- a/docs/reference/query-dsl/geo-shape-query.asciidoc
+++ b/docs/reference/query-dsl/geo-shape-query.asciidoc
@@ -93,6 +93,7 @@ to 'shapes'.
 * `type` - Index type where the pre-indexed shape is.
 * `path` - The field specified as path containing the pre-indexed shape.
 Defaults to 'shape'.
+* `routing` - The routing of the shape document if required.
 
 The following is an example of using the Filter with a pre-indexed
 shape:
diff --git a/server/src/main/java/org/elasticsearch/index/query/GeoShapeQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/GeoShapeQueryBuilder.java
index d7d9797dad8..03fb2098180 100644
--- a/server/src/main/java/org/elasticsearch/index/query/GeoShapeQueryBuilder.java
+++ b/server/src/main/java/org/elasticsearch/index/query/GeoShapeQueryBuilder.java
@@ -29,6 +29,7 @@ import org.apache.lucene.spatial.prefix.RecursivePrefixTreeStrategy;
 import org.apache.lucene.spatial.query.SpatialArgs;
 import org.apache.lucene.spatial.query.SpatialOperation;
 import org.apache.lucene.util.SetOnce;
+import org.elasticsearch.Version;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.get.GetRequest;
 import org.elasticsearch.action.get.GetResponse;
@@ -77,6 +78,7 @@ public class GeoShapeQueryBuilder extends AbstractQueryBuilder<GeoShapeQueryBuil
     private static final ParseField SHAPE_TYPE_FIELD = new ParseField("type");
     private static final ParseField SHAPE_INDEX_FIELD = new ParseField("index");
     private static final ParseField SHAPE_PATH_FIELD = new ParseField("path");
+    private static final ParseField SHAPE_ROUTING_FIELD = new ParseField("routing");
     private static final ParseField IGNORE_UNMAPPED_FIELD = new ParseField("ignore_unmapped");
 
     private final String fieldName;
@@ -89,8 +91,10 @@ public class GeoShapeQueryBuilder extends AbstractQueryBuilder<GeoShapeQueryBuil
     private final String indexedShapeId;
     private final String indexedShapeType;
 
+
     private String indexedShapeIndex = DEFAULT_SHAPE_INDEX_NAME;
     private String indexedShapePath = DEFAULT_SHAPE_FIELD_NAME;
+    private String indexedShapeRouting;
 
     private ShapeRelation relation = DEFAULT_SHAPE_RELATION;
 
@@ -166,6 +170,11 @@ public class GeoShapeQueryBuilder extends AbstractQueryBuilder<GeoShapeQueryBuil
             indexedShapeType = in.readOptionalString();
             indexedShapeIndex = in.readOptionalString();
             indexedShapePath = in.readOptionalString();
+            if (in.getVersion().onOrAfter(Version.V_7_0_0_alpha1)) {
+                indexedShapeRouting = in.readOptionalString();
+            } else {
+                indexedShapeRouting = null;
+            }
         }
         relation = ShapeRelation.readFromStream(in);
         strategy = in.readOptionalWriteable(SpatialStrategy::readFromStream);
@@ -188,6 +197,11 @@ public class GeoShapeQueryBuilder extends AbstractQueryBuilder<GeoShapeQueryBuil
             out.writeOptionalString(indexedShapeType);
             out.writeOptionalString(indexedShapeIndex);
             out.writeOptionalString(indexedShapePath);
+            if (out.getVersion().onOrAfter(Version.V_7_0_0_alpha1)) {
+                out.writeOptionalString(indexedShapeRouting);
+            } else if (indexedShapeRouting != null) {
+                throw new IllegalStateException("indexed shape routing cannot be serialized to older nodes");
+            }
         }
         relation.writeTo(out);
         out.writeOptionalWriteable(strategy);
@@ -285,6 +299,26 @@ public class GeoShapeQueryBuilder extends AbstractQueryBuilder<GeoShapeQueryBuil
         return indexedShapePath;
     }
 
+    /**
+     * Sets the optional routing to the indexed Shape that will be used in the query
+     *
+     * @param indexedShapeRouting indexed shape routing
+     * @return this
+     */
+    public GeoShapeQueryBuilder indexedShapeRouting(String indexedShapeRouting) {
+        this.indexedShapeRouting = indexedShapeRouting;
+        return this;
+    }
+
+
+    /**
+     * @return the optional routing to the indexed Shape that will be used in the
+     *         Query
+     */
+    public String indexedShapeRouting() {
+        return indexedShapeRouting;
+    }
+
     /**
      * Sets the relation of query shape and indexed shape.
      *
@@ -473,6 +507,9 @@ public class GeoShapeQueryBuilder extends AbstractQueryBuilder<GeoShapeQueryBuil
             if (indexedShapePath != null) {
                 builder.field(SHAPE_PATH_FIELD.getPreferredName(), indexedShapePath);
             }
+            if (indexedShapeRouting != null) {
+                builder.field(SHAPE_ROUTING_FIELD.getPreferredName(), indexedShapeRouting);
+            }
             builder.endObject();
         }
 
@@ -498,6 +535,7 @@ public class GeoShapeQueryBuilder extends AbstractQueryBuilder<GeoShapeQueryBuil
         String type = null;
         String index = null;
         String shapePath = null;
+        String shapeRouting = null;
 
         XContentParser.Token token;
         String currentFieldName = null;
@@ -544,6 +582,8 @@ public class GeoShapeQueryBuilder extends AbstractQueryBuilder<GeoShapeQueryBuil
                                         index = parser.text();
                                     } else if (SHAPE_PATH_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
                                         shapePath = parser.text();
+                                    } else if (SHAPE_ROUTING_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
+                                        shapeRouting = parser.text();
                                     }
                                 } else {
                                     throw new ParsingException(parser.getTokenLocation(), "[" + GeoShapeQueryBuilder.NAME +
@@ -581,6 +621,9 @@ public class GeoShapeQueryBuilder extends AbstractQueryBuilder<GeoShapeQueryBuil
         if (shapePath != null) {
             builder.indexedShapePath(shapePath);
         }
+        if (shapeRouting != null) {
+            builder.indexedShapeRouting(shapeRouting);
+        }
         if (shapeRelation != null) {
             builder.relation(shapeRelation);
         }
@@ -602,6 +645,7 @@ public class GeoShapeQueryBuilder extends AbstractQueryBuilder<GeoShapeQueryBuil
                 && Objects.equals(indexedShapeIndex, other.indexedShapeIndex)
                 && Objects.equals(indexedShapePath, other.indexedShapePath)
                 && Objects.equals(indexedShapeType, other.indexedShapeType)
+                && Objects.equals(indexedShapeRouting, other.indexedShapeRouting)
                 && Objects.equals(relation, other.relation)
                 && Objects.equals(shape, other.shape)
                 && Objects.equals(supplier, other.supplier)
@@ -612,7 +656,7 @@ public class GeoShapeQueryBuilder extends AbstractQueryBuilder<GeoShapeQueryBuil
     @Override
     protected int doHashCode() {
         return Objects.hash(fieldName, indexedShapeId, indexedShapeIndex,
-                indexedShapePath, indexedShapeType, relation, shape, strategy, ignoreUnmapped, supplier);
+                indexedShapePath, indexedShapeType, indexedShapeRouting, relation, shape, strategy, ignoreUnmapped, supplier);
     }
 
     @Override
@@ -629,6 +673,7 @@ public class GeoShapeQueryBuilder extends AbstractQueryBuilder<GeoShapeQueryBuil
             SetOnce<ShapeBuilder> supplier = new SetOnce<>();
             queryRewriteContext.registerAsyncAction((client, listener) -> {
                 GetRequest getRequest = new GetRequest(indexedShapeIndex, indexedShapeType, indexedShapeId);
+                getRequest.routing(indexedShapeRouting);
                 fetch(client, getRequest, indexedShapePath, ActionListener.wrap(builder-> {
                     supplier.set(builder);
                     listener.onResponse(null);
diff --git a/server/src/test/java/org/elasticsearch/index/query/GeoShapeQueryBuilderTests.java b/server/src/test/java/org/elasticsearch/index/query/GeoShapeQueryBuilderTests.java
index 3282077ba6a..6356b2122ed 100644
--- a/server/src/test/java/org/elasticsearch/index/query/GeoShapeQueryBuilderTests.java
+++ b/server/src/test/java/org/elasticsearch/index/query/GeoShapeQueryBuilderTests.java
@@ -59,6 +59,7 @@ public class GeoShapeQueryBuilderTests extends AbstractQueryTestCase<GeoShapeQue
     private static String indexedShapeType;
     private static String indexedShapePath;
     private static String indexedShapeIndex;
+    private static String indexedShapeRouting;
     private static ShapeBuilder indexedShapeToReturn;
 
     @Override
@@ -85,6 +86,10 @@ public class GeoShapeQueryBuilderTests extends AbstractQueryTestCase<GeoShapeQue
                 indexedShapePath = randomAlphaOfLengthBetween(3, 20);
                 builder.indexedShapePath(indexedShapePath);
             }
+            if (randomBoolean()) {
+                indexedShapeRouting = randomAlphaOfLengthBetween(3, 20);
+                builder.indexedShapeRouting(indexedShapeRouting);
+            }
         }
         if (randomBoolean()) {
             SpatialStrategy strategy = randomFrom(SpatialStrategy.values());
@@ -112,6 +117,7 @@ public class GeoShapeQueryBuilderTests extends AbstractQueryTestCase<GeoShapeQue
         assertThat(indexedShapeType, notNullValue());
         assertThat(getRequest.id(), equalTo(indexedShapeId));
         assertThat(getRequest.type(), equalTo(indexedShapeType));
+        assertThat(getRequest.routing(), equalTo(indexedShapeRouting));
         String expectedShapeIndex = indexedShapeIndex == null ? GeoShapeQueryBuilder.DEFAULT_SHAPE_INDEX_NAME : indexedShapeIndex;
         assertThat(getRequest.index(), equalTo(expectedShapeIndex));
         String expectedShapePath = indexedShapePath == null ? GeoShapeQueryBuilder.DEFAULT_SHAPE_FIELD_NAME : indexedShapePath;
@@ -136,6 +142,7 @@ public class GeoShapeQueryBuilderTests extends AbstractQueryTestCase<GeoShapeQue
         indexedShapeType = null;
         indexedShapePath = null;
         indexedShapeIndex = null;
+        indexedShapeRouting = null;
     }
 
     @Override
diff --git a/server/src/test/java/org/elasticsearch/search/geo/GeoShapeIntegrationIT.java b/server/src/test/java/org/elasticsearch/search/geo/GeoShapeIntegrationIT.java
index a8f559ce35e..83b79562118 100644
--- a/server/src/test/java/org/elasticsearch/search/geo/GeoShapeIntegrationIT.java
+++ b/server/src/test/java/org/elasticsearch/search/geo/GeoShapeIntegrationIT.java
@@ -31,6 +31,7 @@ import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.indices.IndicesService;
 import org.elasticsearch.test.ESIntegTestCase;
 
+import static org.elasticsearch.index.query.QueryBuilders.geoShapeQuery;
 import static org.elasticsearch.index.query.QueryBuilders.matchAllQuery;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.hamcrest.Matchers.equalTo;
@@ -121,6 +122,43 @@ public class GeoShapeIntegrationIT extends ESIntegTestCase {
         assertThat(searchResponse.getHits().getTotalHits(), equalTo(1L));
     }
 
+    /**
+     * Test that the indexed shape routing can be provided if it is required
+     */
+    public void testIndexShapeRouting() throws Exception {
+        String mapping = "{\n" +
+            "    \"_routing\": {\n" +
+            "      \"required\": true\n" +
+            "    },\n" +
+            "    \"properties\": {\n" +
+            "      \"shape\": {\n" +
+            "        \"type\": \"geo_shape\"\n" +
+            "      }\n" +
+            "    }\n" +
+            "  }";
+
+
+        // create index
+        assertAcked(client().admin().indices().prepareCreate("test").addMapping("doc", mapping, XContentType.JSON).get());
+        ensureGreen();
+
+        String source = "{\n" +
+            "    \"shape\" : {\n" +
+            "        \"type\" : \"circle\",\n" +
+            "        \"coordinates\" : [-45.0, 45.0],\n" +
+            "        \"radius\" : \"100m\"\n" +
+            "    }\n" +
+            "}";
+
+        indexRandom(true, client().prepareIndex("test", "doc", "0").setSource(source, XContentType.JSON).setRouting("ABC"));
+
+        SearchResponse searchResponse = client().prepareSearch("test").setQuery(
+            geoShapeQuery("shape", "0", "doc").indexedShapeIndex("test").indexedShapeRouting("ABC")
+        ).get();
+
+        assertThat(searchResponse.getHits().getTotalHits(), equalTo(1L));
+    }
+
     private String findNodeName(String index) {
         ClusterState state = client().admin().cluster().prepareState().get().getState();
         IndexShardRoutingTable shard = state.getRoutingTable().index(index).shard(0);