Correct search_after handling (#50629)

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/action/EqlSearchRequest.java

@@ -19,11 +19,10 @@ import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.index.query.AbstractQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.search.searchafter.SearchAfterBuilder;
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Collections;
-import java.util.List;
 import java.util.Objects;
 import java.util.function.Supplier;
 
@@ -40,7 +39,7 @@ public class EqlSearchRequest extends ActionRequest implements IndicesRequest.Re
     private String eventTypeField = "event.category";
     private String implicitJoinKeyField = "agent.id";
     private int fetchSize = 50;
-    private List<String> searchAfter = Collections.emptyList();
+    private SearchAfterBuilder searchAfterBuilder;
     private String rule;
 
     static final String KEY_QUERY = "query";
@@ -74,23 +73,10 @@ public class EqlSearchRequest extends ActionRequest implements IndicesRequest.Re
         eventTypeField = in.readString();
         implicitJoinKeyField = in.readString();
         fetchSize = in.readVInt();
-        searchAfter = in.readList(StreamInput::readString);
+        searchAfterBuilder = in.readOptionalWriteable(SearchAfterBuilder::new);
         rule = in.readString();
     }
 
-    public EqlSearchRequest(String[] indices, QueryBuilder query,
-                            String timestampField, String eventTypeField, String implicitJoinKeyField,
-                            int fetchSize, List<String> searchAfter, String rule) {
-        this.indices = indices;
-        this.query = query;
-        this.timestampField = timestampField;
-        this.eventTypeField = eventTypeField;
-        this.implicitJoinKeyField = implicitJoinKeyField;
-        this.fetchSize = fetchSize;
-        this.searchAfter = searchAfter;
-        this.rule = rule;
-    }
-
     @Override
     public ActionRequestValidationException validate() {
         ActionRequestValidationException validationException = null;
@@ -122,10 +108,6 @@ public class EqlSearchRequest extends ActionRequest implements IndicesRequest.Re
             validationException = addValidationError("size must be more than 0.", validationException);
         }
 
-        if (searchAfter == null) {
-            validationException = addValidationError("search after is null", validationException);
-        }
-
         return validationException;
     }
 
@@ -141,13 +123,10 @@ public class EqlSearchRequest extends ActionRequest implements IndicesRequest.Re
         }
         builder.field(KEY_SIZE, fetchSize());
 
-        if (this.searchAfter != null && !this.searchAfter.isEmpty()) {
+        if (searchAfterBuilder != null) {
-            builder.startArray(KEY_SEARCH_AFTER);
+            builder.array(SEARCH_AFTER.getPreferredName(), searchAfterBuilder.getSortValues());
-            for (String val : this.searchAfter) {
-                builder.value(val);
-            }
-            builder.endArray();
         }
+
         builder.field(KEY_RULE, rule);
 
         return builder;
@@ -165,11 +144,13 @@ public class EqlSearchRequest extends ActionRequest implements IndicesRequest.Re
         parser.declareString(EqlSearchRequest::eventTypeField, EVENT_TYPE_FIELD);
         parser.declareString(EqlSearchRequest::implicitJoinKeyField, IMPLICIT_JOIN_KEY_FIELD);
         parser.declareInt(EqlSearchRequest::fetchSize, SIZE);
-        parser.declareStringArray(EqlSearchRequest::searchAfter, SEARCH_AFTER);
+        parser.declareField(EqlSearchRequest::setSearchAfter, SearchAfterBuilder::fromXContent, SEARCH_AFTER,
+            ObjectParser.ValueType.OBJECT_ARRAY);
         parser.declareString(EqlSearchRequest::rule, RULE);
         return parser;
     }
 
+    @Override
     public EqlSearchRequest indices(String... indices) {
         this.indices = indices;
         return this;
@@ -219,22 +200,26 @@ public class EqlSearchRequest extends ActionRequest implements IndicesRequest.Re
         return this;
     }
 
-    public List<String> searchAfter() {
+    public Object[] searchAfter() {
-        return searchAfter;
+        if (searchAfterBuilder == null) {
+            return null;
+        }
+        return searchAfterBuilder.getSortValues();
     }
 
-    public EqlSearchRequest searchAfter(List<String> searchAfter) {
+    public EqlSearchRequest searchAfter(Object[] values) {
-        if (searchAfter != null && !searchAfter.isEmpty()) {
+        this.searchAfterBuilder = new SearchAfterBuilder().setSortValues(values);
-            this.searchAfter = searchAfter;
-        }
         return this;
     }
 
+    private EqlSearchRequest setSearchAfter(SearchAfterBuilder builder) {
+        this.searchAfterBuilder = builder;
+        return this;
+    }
 
     public String rule() { return this.rule; }
 
     public EqlSearchRequest rule(String rule) {
-        // TODO: possibly attempt to parse the rule here
         this.rule = rule;
         return this;
     }
@@ -249,7 +234,7 @@ public class EqlSearchRequest extends ActionRequest implements IndicesRequest.Re
         out.writeString(eventTypeField);
         out.writeString(implicitJoinKeyField);
         out.writeVInt(fetchSize);
-        out.writeStringCollection(searchAfter);
+        out.writeOptionalWriteable(searchAfterBuilder);
         out.writeString(rule);
     }
 
@@ -270,7 +255,7 @@ public class EqlSearchRequest extends ActionRequest implements IndicesRequest.Re
             Objects.equals(timestampField, that.timestampField) &&
             Objects.equals(eventTypeField, that.eventTypeField) &&
             Objects.equals(implicitJoinKeyField, that.implicitJoinKeyField) &&
-            Objects.equals(searchAfter, that.searchAfter) &&
+            Objects.equals(searchAfterBuilder, that.searchAfterBuilder) &&
             Objects.equals(rule, that.rule);
     }
 
@@ -284,7 +269,7 @@ public class EqlSearchRequest extends ActionRequest implements IndicesRequest.Re
             timestampField,
             eventTypeField,
             implicitJoinKeyField,
-            searchAfter,
+            searchAfterBuilder,
             rule);
     }
 

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/plugin/RestEqlSearchAction.java

@@ -28,12 +28,11 @@ import static org.elasticsearch.rest.RestRequest.Method.GET;
 import static org.elasticsearch.rest.RestRequest.Method.POST;
 
 public class RestEqlSearchAction extends BaseRestHandler {
-    public RestEqlSearchAction(RestController controller) {
+    private static final String SEARCH_PATH = "/{index}/_eql/search";
-        // Not sure yet if we will always have index in the path or not
-        // TODO: finalize the endpoints
-        controller.registerHandler(GET, "/{index}/_eql/search", this);
-        controller.registerHandler(POST, "/{index}/_eql/search", this);
 
+    public RestEqlSearchAction(RestController controller) {
+        controller.registerHandler(GET, SEARCH_PATH, this);
+        controller.registerHandler(POST, SEARCH_PATH, this);
     }
 
     @Override

x-pack/plugin/eql/src/test/java/org/elasticsearch/xpack/eql/action/EqlRequestParserTests.java

@@ -16,7 +16,6 @@ import org.elasticsearch.search.SearchModule;
 import org.elasticsearch.test.ESTestCase;
 
 import java.io.IOException;
-import java.util.Arrays;
 import java.util.List;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -54,7 +53,7 @@ public class EqlRequestParserTests extends ESTestCase {
             + "\"timestamp_field\" : \"tsf\", "
             + "\"event_type_field\" : \"etf\","
             + "\"implicit_join_key_field\" : \"imjf\","
-            + "\"search_after\" : [ \"device-20184\", \"/user/local/foo.exe\", \"2019-11-26T00:45:43.542\" ],"
+            + "\"search_after\" : [ 12345678, \"device-20184\", \"/user/local/foo.exe\", \"2019-11-26T00:45:43.542\" ],"
             + "\"size\" : \"101\","
             + "\"rule\" : \"file where user != 'SYSTEM' by file_path\""
             + "}", EqlSearchRequest::fromXContent);
@@ -67,8 +66,7 @@ public class EqlRequestParserTests extends ESTestCase {
         assertEquals("tsf", request.timestampField());
         assertEquals("etf", request.eventTypeField());
         assertEquals("imjf", request.implicitJoinKeyField());
-        assertArrayEquals(Arrays.asList("device-20184", "/user/local/foo.exe", "2019-11-26T00:45:43.542").toArray(),
+        assertArrayEquals(new Object[]{12345678, "device-20184", "/user/local/foo.exe", "2019-11-26T00:45:43.542"}, request.searchAfter());
-            request.searchAfter().toArray());
         assertEquals(101, request.fetchSize());
         assertEquals("file where user != 'SYSTEM' by file_path", request.rule());
     }

x-pack/plugin/eql/src/test/java/org/elasticsearch/xpack/eql/action/EqlSearchRequestTests.java

@@ -5,21 +5,27 @@
  */
 package org.elasticsearch.xpack.eql.action;
 
+import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.text.Text;
 import org.elasticsearch.common.xcontent.NamedXContentRegistry;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.json.JsonXContent;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.search.SearchModule;
+import org.elasticsearch.search.searchafter.SearchAfterBuilder;
 import org.elasticsearch.test.AbstractSerializingTestCase;
+import org.elasticsearch.test.ESTestCase;
 import org.junit.Before;
 
 import java.io.IOException;
-import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
-import java.util.List;
+import java.util.function.Supplier;
 
 import static org.elasticsearch.index.query.AbstractQueryBuilder.parseInnerQueryBuilder;
 
@@ -52,15 +58,25 @@ public class EqlSearchRequestTests extends AbstractSerializingTestCase<EqlSearch
 
     @Override
     protected EqlSearchRequest createTestInstance() {
-        QueryBuilder query = null;
         try {
-            query = parseQuery(defaultTestQuery);
+            QueryBuilder query = parseQuery(defaultTestQuery);
+            EqlSearchRequest request = new EqlSearchRequest()
+                .indices(new String[]{defaultTestIndex})
+                .query(query)
+                .timestampField(randomAlphaOfLength(10))
+                .eventTypeField(randomAlphaOfLength(10))
+                .implicitJoinKeyField(randomAlphaOfLength(10))
+                .fetchSize(randomIntBetween(1, 50))
+                .rule(randomAlphaOfLength(10));
+
+            if (randomBoolean()) {
+                request.searchAfter(randomJsonSearchFromBuilder());
+            }
+            return request;
         } catch (IOException ex) {
             assertNotNull("unexpected IOException " + ex.getCause().getMessage(), ex);
         }
-        return new EqlSearchRequest(new String[]{defaultTestIndex}, query,
+        return null;
-                randomAlphaOfLength(10), randomAlphaOfLength(10), randomAlphaOfLength(10),
-                randomIntBetween(1, 50), randomSearchAfter(), randomAlphaOfLength(10));
     }
 
     protected QueryBuilder parseQuery(String queryAsString) throws IOException {
@@ -74,16 +90,36 @@ public class EqlSearchRequestTests extends AbstractSerializingTestCase<EqlSearch
         return parseInnerQueryBuilder;
     }
 
-    private List<String> randomSearchAfter() {
+    private Object randomValue() {
-        if (randomBoolean()) {
+        Supplier<Object> value = randomFrom(Arrays.asList(
-            return Collections.emptyList();
+            ESTestCase::randomInt,
-        } else {
+            ESTestCase::randomFloat,
-            int size = randomIntBetween(1, 50);
+            ESTestCase::randomLong,
-            List<String> arr = new ArrayList<>(size);
+            ESTestCase::randomDouble,
-            for (int i = 0; i < size; i++) {
+            () -> randomAlphaOfLengthBetween(5, 20),
-                arr.add(randomAlphaOfLength(randomIntBetween(1, 15)));
+            ESTestCase::randomBoolean,
-            }
+            ESTestCase::randomByte,
-            return Collections.unmodifiableList(arr);
+            ESTestCase::randomShort,
+            () -> new Text(randomAlphaOfLengthBetween(5, 20)),
+            () -> null));
+        return value.get();
+    }
+
+    private Object[] randomJsonSearchFromBuilder() throws IOException {
+        int numSearchAfter = randomIntBetween(1, 10);
+        XContentBuilder jsonBuilder = XContentFactory.jsonBuilder();
+        jsonBuilder.startObject();
+        jsonBuilder.startArray("search_after");
+        for (int i = 0; i < numSearchAfter; i++) {
+            jsonBuilder.value(randomValue());
+        }
+        jsonBuilder.endArray();
+        jsonBuilder.endObject();
+        try (XContentParser parser = createParser(JsonXContent.jsonXContent, BytesReference.bytes(jsonBuilder))) {
+            parser.nextToken();
+            parser.nextToken();
+            parser.nextToken();
+            return SearchAfterBuilder.fromXContent(parser).getSortValues();
         }
     }