From 0c9fefe5bce4e5151d50cb4e72149f1fd36cc8cb Mon Sep 17 00:00:00 2001
From: Julie Tibshirani <julie.tibshirani@elastic.co>
Date: Thu, 27 Sep 2018 14:16:57 -0700
Subject: [PATCH] Support prefix and regexp queries on _index fields.

---
 .../index/mapper/IndexFieldMapper.java        | 34 ++++++++++++-
 .../index/mapper/IndexFieldTypeTests.java     | 48 ++++++++++++++++++-
 2 files changed, 78 insertions(+), 4 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/index/mapper/IndexFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/IndexFieldMapper.java
index 456805e64160e..7e8ac563cacc1 100644
--- a/server/src/main/java/org/elasticsearch/index/mapper/IndexFieldMapper.java
+++ b/server/src/main/java/org/elasticsearch/index/mapper/IndexFieldMapper.java
@@ -38,6 +38,7 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
+import java.util.regex.Pattern;
 
 
 public class IndexFieldMapper extends MetadataFieldMapper {
@@ -151,14 +152,43 @@ public Query termsQuery(List values, QueryShardContext context) {
                 + " vs. " + values);
         }
 
+        @Override
+        public Query prefixQuery(String value,
+                                 @Nullable MultiTermQuery.RewriteMethod method,
+                                 QueryShardContext context) {
+            String indexName = context.getFullyQualifiedIndex().getName();
+            if (indexName.startsWith(value)) {
+                return Queries.newMatchAllQuery();
+            } else {
+                return Queries.newMatchNoDocsQuery("The index [" + indexName +
+                    "] doesn't match the provided prefix [" + value + "].");
+            }
+        }
+
+        @Override
+        public Query regexpQuery(String value, int flags, int maxDeterminizedStates,
+                                 MultiTermQuery.RewriteMethod method, QueryShardContext context) {
+            String indexName = context.getFullyQualifiedIndex().getName();
+            Pattern pattern = Regex.compile(value, Regex.flagsToString(flags));
+
+            if (pattern.matcher(indexName).matches()) {
+                return Queries.newMatchAllQuery();
+            } else {
+                return Queries.newMatchNoDocsQuery("The index [" + indexName +
+                    "] doesn't match the provided pattern [" + value + "].");
+            }
+        }
+
         @Override
         public Query wildcardQuery(String value,
                                    @Nullable MultiTermQuery.RewriteMethod method,
                                    QueryShardContext context) {
-            if (isSameIndex(value, context.getFullyQualifiedIndex().getName())) {
+            String indexName = context.getFullyQualifiedIndex().getName();
+            if (isSameIndex(value, indexName)) {
                 return Queries.newMatchAllQuery();
             } else {
-                return Queries.newMatchNoDocsQuery("Index didn't match. Index queried: " + context.index().getName() + " vs. " + value);
+                return Queries.newMatchNoDocsQuery("The index [" + indexName +
+                    "] doesn't match the provided pattern [" + value + "].");
             }
         }
 
diff --git a/server/src/test/java/org/elasticsearch/index/mapper/IndexFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/IndexFieldTypeTests.java
index ecea620f11cf1..82f0edf24f4da 100644
--- a/server/src/test/java/org/elasticsearch/index/mapper/IndexFieldTypeTests.java
+++ b/server/src/test/java/org/elasticsearch/index/mapper/IndexFieldTypeTests.java
@@ -18,12 +18,56 @@
  */
 package org.elasticsearch.index.mapper;
 
-import org.elasticsearch.index.mapper.IndexFieldMapper;
-import org.elasticsearch.index.mapper.MappedFieldType;
+import org.apache.lucene.index.IndexOptions;
+import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.MatchNoDocsQuery;
+import org.elasticsearch.index.Index;
+import org.elasticsearch.index.query.QueryShardContext;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 public class IndexFieldTypeTests extends FieldTypeTestCase {
+
     @Override
     protected MappedFieldType createDefaultFieldType() {
         return new IndexFieldMapper.IndexFieldType();
     }
+
+    public void testPrefixQuery() {
+        MappedFieldType ft = createDefaultFieldType();
+        ft.setName("field");
+        ft.setIndexOptions(IndexOptions.DOCS);
+
+        assertEquals(new MatchAllDocsQuery(), ft.prefixQuery("ind", null, createContext()));
+        assertEquals(new MatchNoDocsQuery(), ft.prefixQuery("other_ind", null, createContext()));
+    }
+
+    public void testRegexpQuery() {
+        MappedFieldType ft = createDefaultFieldType();
+        ft.setName("field");
+        ft.setIndexOptions(IndexOptions.DOCS);
+
+        assertEquals(new MatchAllDocsQuery(), ft.regexpQuery("ind.x", 0, 10, null, createContext()));
+        assertEquals(new MatchNoDocsQuery(), ft.regexpQuery("ind?x", 0, 10, null, createContext()));
+    }
+
+    public void testWildcardQuery() {
+        MappedFieldType ft = createDefaultFieldType();
+        ft.setName("field");
+        ft.setIndexOptions(IndexOptions.DOCS);
+
+        assertEquals(new MatchAllDocsQuery(), ft.wildcardQuery("ind*x", null, createContext()));
+        assertEquals(new MatchNoDocsQuery(), ft.wildcardQuery("other_ind*x", null, createContext()));
+    }
+
+    private QueryShardContext createContext() {
+        QueryShardContext context = mock(QueryShardContext.class);
+
+        Index index = new Index("index", "123");
+        when(context.getFullyQualifiedIndex()).thenReturn(index);
+        when(context.index()).thenReturn(index);
+
+        return context;
+    }
 }