From 01decb1814fbe8852f26854e1ee0f6b620a3556f Mon Sep 17 00:00:00 2001
From: Rishabh Maurya <rishabhmaurya05@gmail.com>
Date: Wed, 20 Mar 2024 18:28:17 -0700
Subject: [PATCH] Add tests for all supported query types for derived fields

Signed-off-by: Rishabh Maurya <rishabhmaurya05@gmail.com>
---
 .../index/mapper/DerivedFieldMapper.java      |  21 +-
 .../opensearch/index/mapper/ObjectMapper.java |   6 +-
 .../index/mapper/ParametrizedFieldMapper.java |   4 +-
 .../mapper/DerivedFieldMapperQueryTests.java  | 279 ++++++++++++++++++
 .../index/mapper/DerivedFieldMapperTests.java |  57 ++--
 .../index/mapper/ObjectMapperTests.java       |   3 +-
 6 files changed, 326 insertions(+), 44 deletions(-)
 create mode 100644 server/src/test/java/org/opensearch/index/mapper/DerivedFieldMapperQueryTests.java

diff --git a/server/src/main/java/org/opensearch/index/mapper/DerivedFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/DerivedFieldMapper.java
index c02b618bfccd5..5d6d5ad2e73e8 100644
--- a/server/src/main/java/org/opensearch/index/mapper/DerivedFieldMapper.java
+++ b/server/src/main/java/org/opensearch/index/mapper/DerivedFieldMapper.java
@@ -56,12 +56,7 @@ private static DerivedFieldMapper toType(FieldMapper in) {
      */
     public static class Builder extends ParametrizedFieldMapper.Builder {
         // TODO: The type of parameter may change here if the actual underlying FieldType object is needed
-        private final Parameter<String> type = Parameter.stringParam(
-            "type",
-            false,
-            m -> toType(m).type,
-            "text"
-        );
+        private final Parameter<String> type = Parameter.stringParam("type", false, m -> toType(m).type, "text");
 
         private final Parameter<Script> script = new Parameter<>(
             "script",
@@ -83,9 +78,17 @@ protected List<Parameter<?>> getParameters() {
         @Override
         public DerivedFieldMapper build(BuilderContext context) {
             FieldMapper fieldMapper = DerivedFieldSupportedTypes.getFieldMapperFromType(type.getValue(), name, context);
-            Function<Object, IndexableField> fieldFunction =
-                DerivedFieldSupportedTypes.getIndexableFieldGeneratorType(type.getValue(), name);
-            DerivedFieldType ft = new DerivedFieldType(buildFullName(context), type.getValue(), script.getValue(), fieldMapper, fieldFunction);
+            Function<Object, IndexableField> fieldFunction = DerivedFieldSupportedTypes.getIndexableFieldGeneratorType(
+                type.getValue(),
+                name
+            );
+            DerivedFieldType ft = new DerivedFieldType(
+                buildFullName(context),
+                type.getValue(),
+                script.getValue(),
+                fieldMapper,
+                fieldFunction
+            );
             return new DerivedFieldMapper(name, ft, multiFieldsBuilder.build(this, context), copyTo.build(), this);
         }
     }
diff --git a/server/src/main/java/org/opensearch/index/mapper/ObjectMapper.java b/server/src/main/java/org/opensearch/index/mapper/ObjectMapper.java
index 33c4b3393f7d2..2c73a53f70526 100644
--- a/server/src/main/java/org/opensearch/index/mapper/ObjectMapper.java
+++ b/server/src/main/java/org/opensearch/index/mapper/ObjectMapper.java
@@ -721,7 +721,8 @@ public void toXContent(XContentBuilder builder, Params params, ToXContent custom
         doXContent(builder, params);
 
         // sort the mappers so we get consistent serialization format
-        Mapper[] derivedSortedMappers = mappers.values().stream()
+        Mapper[] derivedSortedMappers = mappers.values()
+            .stream()
             .filter(m -> m instanceof DerivedFieldMapper)
             .toArray(size -> new Mapper[size]);
         Arrays.sort(derivedSortedMappers, new Comparator<Mapper>() {
@@ -731,7 +732,8 @@ public int compare(Mapper o1, Mapper o2) {
             }
         });
 
-        Mapper[] sortedMappers = mappers.values().stream()
+        Mapper[] sortedMappers = mappers.values()
+            .stream()
             .filter(m -> !(m instanceof DerivedFieldMapper))
             .toArray(size -> new Mapper[size]);
         Arrays.sort(sortedMappers, new Comparator<Mapper>() {
diff --git a/server/src/main/java/org/opensearch/index/mapper/ParametrizedFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/ParametrizedFieldMapper.java
index cde351f923767..4f4ec45001a54 100644
--- a/server/src/main/java/org/opensearch/index/mapper/ParametrizedFieldMapper.java
+++ b/server/src/main/java/org/opensearch/index/mapper/ParametrizedFieldMapper.java
@@ -671,7 +671,9 @@ public final void parse(String name, ParserContext parserContext, Map<String, Ob
                 }
             }
             String type = (String) fieldNode.get("type");
-            if (paramsMap.get("type") == null) { fieldNode.remove("type"); }
+            if (paramsMap.get("type") == null) {
+                fieldNode.remove("type");
+            }
 
             for (Iterator<Map.Entry<String, Object>> iterator = fieldNode.entrySet().iterator(); iterator.hasNext();) {
                 Map.Entry<String, Object> entry = iterator.next();
diff --git a/server/src/test/java/org/opensearch/index/mapper/DerivedFieldMapperQueryTests.java b/server/src/test/java/org/opensearch/index/mapper/DerivedFieldMapperQueryTests.java
new file mode 100644
index 0000000000000..64f4f1f3e083e
--- /dev/null
+++ b/server/src/test/java/org/opensearch/index/mapper/DerivedFieldMapperQueryTests.java
@@ -0,0 +1,279 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index.mapper;
+
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.Field;
+import org.apache.lucene.document.TextField;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.store.Directory;
+import org.opensearch.common.lucene.Lucene;
+import org.opensearch.core.index.Index;
+import org.opensearch.index.query.QueryBuilders;
+import org.opensearch.index.query.QueryShardContext;
+import org.opensearch.script.DerivedFieldScript;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+
+import org.mockito.Mockito;
+
+import static org.mockito.Mockito.when;
+
+public class DerivedFieldMapperQueryTests extends MapperServiceTestCase {
+
+    // First element is the document ingested, other elements corresponds to value returned against a given derived field script
+    // Raw Message, Request Succeeded (boolean), Timestamp (long), Client IP, Method, Request Size (double), Duration (long)
+    private static final Object[][] raw_requests = new Object[][] {
+        {
+            "40.135.0.0 GET /images/hm_bg.jpg?size=1.5KB HTTP/1.0 200 2024-03-20T08:30:45 1500",
+            true,
+            1710923445000L,
+            "40.135.0.0",
+            "GET",
+            1.5,
+            1500L },
+        {
+            "232.0.0.0 GET /images/hm_bg.jpg?size=2.3KB HTTP/1.0 400 2024-03-20T09:15:20 2300",
+            false,
+            1710926120000L,
+            "232.0.0.0",
+            "GET",
+            2.3,
+            2300L },
+        {
+            "26.1.0.0 DELETE /images/hm_bg.jpg?size=3.7KB HTTP/1.0 200 2024-03-20T10:05:55 3700",
+            true,
+            1710929155000L,
+            "26.1.0.0",
+            "DELETE",
+            3.7,
+            3700L },
+        {
+            "247.37.0.0 GET /french/splash_inet.html?size=4.1KB HTTP/1.0 400 2024-03-20T11:20:10 4100",
+            false,
+            1710933610000L,
+            "247.37.0.0",
+            "GET",
+            4.1,
+            4100L },
+        {
+            "247.37.0.0 DELETE /french/splash_inet.html?size=5.8KB HTTP/1.0 400 2024-03-20T12:45:30 5800",
+            false,
+            1710938730000L,
+            "247.37.0.0",
+            "DELETE",
+            5.8,
+            5800L },
+        {
+            "10.20.30.40 GET /path/to/resource?size=6.3KB HTTP/1.0 200 2024-03-20T13:10:15 6300",
+            true,
+            1710940215000L,
+            "10.20.30.40",
+            "GET",
+            6.3,
+            6300L },
+        {
+            "50.60.70.80 GET /path/to/resource?size=7.2KB HTTP/1.0 404 2024-03-20T14:20:50 7200",
+            false,
+            1710944450000L,
+            "50.60.70.80",
+            "GET",
+            7.2,
+            7200L },
+        {
+            "127.0.0.1 PUT /path/to/resource?size=8.9KB HTTP/1.0 500 2024-03-20T15:30:25 8900",
+            false,
+            1710948625000L,
+            "127.0.0.1",
+            "PUT",
+            8.9,
+            8900L },
+        {
+            "127.0.0.1 GET /path/to/resource?size=9.4KB HTTP/1.0 200 2024-03-20T16:40:15 9400",
+            true,
+            1710952815000L,
+            "127.0.0.1",
+            "GET",
+            9.4,
+            9400L },
+        {
+            "192.168.1.1 GET /path/to/resource?size=10.7KB HTTP/1.0 400 2024-03-20T17:50:40 10700",
+            false,
+            1710957040000L,
+            "192.168.1.1",
+            "GET",
+            10.7,
+            10700L } };
+
+    public void testAllPossibleQueriesOnDerivedFields() throws IOException {
+        MapperService mapperService = createMapperService(topMapping(b -> {
+            b.startObject("properties");
+            {
+                b.startObject("raw_message");
+                {
+                    b.field("type", "text");
+                }
+                b.endObject();
+            }
+            b.endObject();
+            b.startObject("derived");
+            {
+                b.startObject("request_succeeded");
+                {
+                    b.field("type", "boolean");
+                    b.field("script", "");
+                }
+                b.endObject();
+                b.startObject("@timestamp");
+                {
+                    b.field("type", "date");
+                    b.field("script", "");
+                }
+                b.endObject();
+                b.startObject("client_ip");
+                {
+                    b.field("type", "ip");
+                    b.field("script", "");
+                }
+                b.endObject();
+                b.startObject("method");
+                {
+                    b.field("type", "keyword");
+                    b.field("script", "");
+                }
+                b.endObject();
+                b.startObject("request_size");
+                {
+                    b.field("type", "double");
+                    b.field("script", "");
+                }
+                b.endObject();
+                b.startObject("duration");
+                {
+                    b.field("type", "long");
+                    b.field("script", "");
+                }
+                b.endObject();
+            }
+            b.endObject();
+        }));
+
+        List<Document> docs = new ArrayList<>();
+        for (Object[] request : raw_requests) {
+            Document document = new Document();
+            document.add(new TextField("raw_message", (String) request[0], Field.Store.YES));
+            docs.add(document);
+        }
+
+        int[] scriptIndex = { 1 };
+
+        // Mock DerivedFieldScript.Factory
+        DerivedFieldScript.Factory factory = (params, lookup) -> (DerivedFieldScript.LeafFactory) ctx -> new DerivedFieldScript(
+            params,
+            lookup,
+            ctx
+        ) {
+            int docId = 0;
+
+            @Override
+            public void setDocument(int docId) {
+                super.setDocument(docId);
+                this.docId = docId;
+            }
+
+            @Override
+            public Object execute() {
+                return raw_requests[docId][scriptIndex[0]];
+            }
+        };
+
+        QueryShardContext queryShardContext = createQueryShardContext(mapperService);
+        when(queryShardContext.compile(Mockito.any(), Mockito.any())).thenReturn(factory);
+        when(queryShardContext.sourcePath("raw_message")).thenReturn(Set.of("raw_message"));
+        when(queryShardContext.index()).thenReturn(new Index("test_index", "uuid"));
+
+        // Index and Search
+        try (Directory dir = newDirectory()) {
+            IndexWriter iw = new IndexWriter(dir, new IndexWriterConfig(Lucene.STANDARD_ANALYZER));
+            for (Document d : docs) {
+                iw.addDocument(d);
+            }
+            try (IndexReader reader = DirectoryReader.open(iw)) {
+                iw.close();
+
+                IndexSearcher searcher = new IndexSearcher(reader);
+                Query query = QueryBuilders.termQuery("request_succeeded", "true").toQuery(queryShardContext);
+                TopDocs topDocs = searcher.search(query, 10);
+                assertEquals(4, topDocs.totalHits.value);
+
+                // IP Field Term Query
+                scriptIndex[0] = 3;
+                query = QueryBuilders.termQuery("client_ip", "192.168.0.0/16").toQuery(queryShardContext);
+                topDocs = searcher.search(query, 10);
+                assertEquals(1, topDocs.totalHits.value);
+
+                scriptIndex[0] = 4;
+                query = QueryBuilders.termsQuery("method", "DELETE", "PUT").toQuery(queryShardContext);
+                topDocs = searcher.search(query, 10);
+                assertEquals(3, topDocs.totalHits.value);
+
+                query = QueryBuilders.termsQuery("method", "delete").toQuery(queryShardContext);
+                topDocs = searcher.search(query, 10);
+                assertEquals(0, topDocs.totalHits.value);
+
+                query = QueryBuilders.termQuery("method", "delete").caseInsensitive(true).toQuery(queryShardContext);
+                topDocs = searcher.search(query, 10);
+                assertEquals(2, topDocs.totalHits.value);
+
+                // Range queries of types - date, long and double
+                scriptIndex[0] = 2;
+                query = QueryBuilders.rangeQuery("@timestamp").from("2024-03-20T14:20:50").toQuery(queryShardContext);
+                topDocs = searcher.search(query, 10);
+                assertEquals(4, topDocs.totalHits.value);
+
+                scriptIndex[0] = 5;
+                query = QueryBuilders.rangeQuery("request_size").from("4.1").toQuery(queryShardContext);
+                topDocs = searcher.search(query, 10);
+                assertEquals(7, topDocs.totalHits.value);
+
+                scriptIndex[0] = 6;
+                query = QueryBuilders.rangeQuery("duration").from("5800").toQuery(queryShardContext);
+                topDocs = searcher.search(query, 10);
+                assertEquals(6, topDocs.totalHits.value);
+
+                scriptIndex[0] = 4;
+
+                // Prefix Query
+                query = QueryBuilders.prefixQuery("method", "DE").toQuery(queryShardContext);
+                topDocs = searcher.search(query, 10);
+                assertEquals(2, topDocs.totalHits.value);
+
+                scriptIndex[0] = 4;
+                query = QueryBuilders.wildcardQuery("method", "G*").toQuery(queryShardContext);
+                topDocs = searcher.search(query, 10);
+                assertEquals(7, topDocs.totalHits.value);
+
+                // Regexp Query
+                scriptIndex[0] = 4;
+                query = QueryBuilders.regexpQuery("method", ".*LET.*").toQuery(queryShardContext);
+                topDocs = searcher.search(query, 10);
+                assertEquals(2, topDocs.totalHits.value);
+            }
+        }
+    }
+}
diff --git a/server/src/test/java/org/opensearch/index/mapper/DerivedFieldMapperTests.java b/server/src/test/java/org/opensearch/index/mapper/DerivedFieldMapperTests.java
index 19c5ccdc8faf3..e3e2b978369ff 100644
--- a/server/src/test/java/org/opensearch/index/mapper/DerivedFieldMapperTests.java
+++ b/server/src/test/java/org/opensearch/index/mapper/DerivedFieldMapperTests.java
@@ -44,12 +44,12 @@ protected void minimalMapping(XContentBuilder b) throws IOException {
 
     protected void registerParameters(ParameterChecker checker) throws IOException {
         // TODO Any conflicts or updates to check for here? Parameters index, store, doc_values and boost are not
-        //  supported for DerivedFieldMapper (we explicitly set these values on initialization)
+        // supported for DerivedFieldMapper (we explicitly set these values on initialization)
     }
 
     // TODO: Can update this once the query implementation is completed
-    //  This is also being left blank because the super assertExistsQuery is trying to parse
-    //  an empty source and fails.
+    // This is also being left blank because the super assertExistsQuery is trying to parse
+    // an empty source and fails.
     @Override
     protected void assertExistsQuery(MapperService mapperService) {}
 
@@ -82,39 +82,36 @@ public void testParsesScript() throws IOException {
         XContentBuilder builder = XContentFactory.jsonBuilder().startObject();
         mapper.toXContent(builder, ToXContent.EMPTY_PARAMS);
         builder.endObject();
-        assertEquals("{\"field\":{\"type\":\"keyword\",\"script\":{\"source\":\"doc['test'].value\",\"lang\":\"painless\"}}}", builder.toString());
+        assertEquals(
+            "{\"field\":{\"type\":\"keyword\",\"script\":{\"source\":\"doc['test'].value\",\"lang\":\"painless\"}}}",
+            builder.toString()
+        );
     }
 
-    //  TODO: Is there a case where we want to allow the field to be defined in both 'derived' and 'properties'?
-    //    If the user wants to move a field they were testing as derived to be indexed, they should be able to update
-    //    the mappings in index template to move the field from 'derived' to 'properties' and it should take affect
-    //    during the next index rollover (so even in this case, the field is only defined in one or the other).
+    // TODO: Is there a case where we want to allow the field to be defined in both 'derived' and 'properties'?
+    // If the user wants to move a field they were testing as derived to be indexed, they should be able to update
+    // the mappings in index template to move the field from 'derived' to 'properties' and it should take affect
+    // during the next index rollover (so even in this case, the field is only defined in one or the other).
     public void testFieldInDerivedAndProperties() throws IOException {
-        MapperParsingException ex = expectThrows(
-            MapperParsingException.class,
-            () -> createDocumentMapper(topMapping(b -> {
-                b.startObject("derived");
-                b.startObject("field");
-                b.field("type", "keyword");
-                b.endObject();
-                b.endObject();
-                b.startObject("properties");
-                b.startObject("field");
-                b.field("type", "keyword");
-                b.endObject();
-                b.endObject();
-            }))
-        );
+        MapperParsingException ex = expectThrows(MapperParsingException.class, () -> createDocumentMapper(topMapping(b -> {
+            b.startObject("derived");
+            b.startObject("field");
+            b.field("type", "keyword");
+            b.endObject();
+            b.endObject();
+            b.startObject("properties");
+            b.startObject("field");
+            b.field("type", "keyword");
+            b.endObject();
+            b.endObject();
+        })));
         // TODO: Do we want to handle this as a different error? As it stands, it fails as a merge conflict which makes sense.
-        //  If it didn't fail here, it would hit the MapperParsingException for the field being defined more than once
-        //  when MappingLookup is initialized
-        assertEquals(
-            "Failed to parse mapping [_doc]: mapper [field] cannot be changed from type [derived] to [keyword]",
-            ex.getMessage()
-        );
+        // If it didn't fail here, it would hit the MapperParsingException for the field being defined more than once
+        // when MappingLookup is initialized
+        assertEquals("Failed to parse mapping [_doc]: mapper [field] cannot be changed from type [derived] to [keyword]", ex.getMessage());
     }
 
     // TODO TESTCASE: testWithFieldInSource() (derived field with that field present in source)
-    //  This is more checking search behavior so may need to revisit this after query implementation
+    // This is more checking search behavior so may need to revisit this after query implementation
 
 }
diff --git a/server/src/test/java/org/opensearch/index/mapper/ObjectMapperTests.java b/server/src/test/java/org/opensearch/index/mapper/ObjectMapperTests.java
index 66a2faa89aedc..abdfe577d225a 100644
--- a/server/src/test/java/org/opensearch/index/mapper/ObjectMapperTests.java
+++ b/server/src/test/java/org/opensearch/index/mapper/ObjectMapperTests.java
@@ -465,8 +465,7 @@ public void testDerivedFields() throws Exception {
             .endObject()
             .toString();
 
-        DocumentMapper documentMapper = createIndex("test")
-            .mapperService()
+        DocumentMapper documentMapper = createIndex("test").mapperService()
             .documentMapperParser()
             .parse("tweet", new CompressedXContent(mapping));