Add support for EQL samples queries (#91312)

Author: luigidellaquila

docs/changelog/91312.yaml

@@ -0,0 +1,5 @@
+pr: 91312
+summary: EQL samples
+area: EQL
+type: feature
+issues: []

x-pack/plugin/eql/qa/common/src/main/java/org/elasticsearch/test/eql/BaseEqlSpecTestCase.java

@@ -11,6 +11,7 @@
 import org.apache.http.client.config.RequestConfig;
 import org.elasticsearch.client.Request;
 import org.elasticsearch.client.RequestOptions;
+import org.elasticsearch.client.Response;
 import org.elasticsearch.client.ResponseException;
 import org.elasticsearch.client.RestClient;
 import org.elasticsearch.common.Strings;
@@ -29,6 +30,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.StringJoiner;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 
 public abstract class BaseEqlSpecTestCase extends RemoteClusterAwareEqlRestTestCase {
 
@@ -37,21 +40,43 @@ public abstract class BaseEqlSpecTestCase extends RemoteClusterAwareEqlRestTestC
     private final String index;
     private final String query;
     private final String name;
-    private final long[] eventIds;
+    private final List<long[]> eventIds;
     /**
      * Join keys can be of multiple types, but toml is very restrictive and doesn't allow mixed types values in the same array of values
      * For now, every value will be converted to a String.
      */
     private final String[] joinKeys;
 
+    /**
+     * any negative value means undefined (ie. no "size" will be passed to the query)
+     */
+    private final int size;
+
     @Before
     public void setup() throws Exception {
         RestClient provisioningClient = provisioningClient();
-        if (provisioningClient.performRequest(new Request("HEAD", "/" + unqualifiedIndexName())).getStatusLine().getStatusCode() == 404) {
+        boolean dataLoaded = Arrays.stream(index.split(","))
+            .anyMatch(
+                indexName -> doWithRequest(
+                    new Request("HEAD", "/" + unqualifiedIndexName(indexName)),
+                    provisioningClient,
+                    response -> response.getStatusLine().getStatusCode() == 200
+                )
+            );
+
+        if (dataLoaded == false) {
             DataLoader.loadDatasetIntoEs(highLevelClient(provisioningClient), this::createParser);
         }
     }
 
+    private boolean doWithRequest(Request request, RestClient client, Function<Response, Boolean> consumer) {
+        try {
+            return consumer.apply(client.performRequest(request));
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     @AfterClass
     public static void wipeTestData() throws IOException {
         try {
@@ -77,19 +102,20 @@ protected static List<Object[]> asArray(List<EqlSpec> specs) {
                 name = "" + (counter);
             }
 
-            results.add(new Object[] { spec.query(), name, spec.expectedEventIds(), spec.joinKeys() });
+            results.add(new Object[] { spec.query(), name, spec.expectedEventIds(), spec.joinKeys(), spec.size() });
         }
 
         return results;
     }
 
-    BaseEqlSpecTestCase(String index, String query, String name, long[] eventIds, String[] joinKeys) {
+    BaseEqlSpecTestCase(String index, String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
         this.index = index;
 
         this.query = query;
         this.name = name;
         this.eventIds = eventIds;
         this.joinKeys = joinKeys;
+        this.size = size == null ? -1 : size;
     }
 
     public void test() throws Exception {
@@ -109,7 +135,7 @@ private void assertResponse(ObjectPath response) throws Exception {
         }
     }
 
-    private ObjectPath runQuery(String index, String query) throws Exception {
+    protected ObjectPath runQuery(String index, String query) throws Exception {
         XContentBuilder builder = JsonXContent.contentBuilder();
         builder.startObject();
         builder.field("query", query);
@@ -119,7 +145,7 @@ private ObjectPath runQuery(String index, String query) throws Exception {
         if (tiebreaker != null) {
             builder.field("tiebreaker_field", tiebreaker);
         }
-        builder.field("size", requestSize());
+        builder.field("size", this.size < 0 ? requestSize() : this.size);
         builder.field("fetch_size", requestFetchSize());
         builder.field("result_position", requestResultPosition());
         builder.endObject();
@@ -149,20 +175,42 @@ public String toString() {
             }
         });
 
-        long[] expected = eventIds;
         long[] actual = extractIds(events);
-        assertArrayEquals(
-            LoggerMessageFormat.format(
-                null,
-                "unexpected result for spec[{}] [{}] -> {} vs {}",
-                name,
-                query,
-                Arrays.toString(expected),
-                Arrays.toString(actual)
-            ),
-            expected,
-            actual
-        );
+        if (eventIds.size() == 1) {
+            long[] expected = eventIds.get(0);
+            assertArrayEquals(
+                LoggerMessageFormat.format(
+                    null,
+                    "unexpected result for spec[{}] [{}] -> {} vs {}",
+                    name,
+                    query,
+                    Arrays.toString(expected),
+                    Arrays.toString(actual)
+                ),
+                expected,
+                actual
+            );
+        } else {
+            boolean succeeded = false;
+            for (long[] expected : eventIds) {
+                if (Arrays.equals(expected, actual)) {
+                    succeeded = true;
+                    break;
+                }
+            }
+            if (succeeded == false) {
+                String msg = LoggerMessageFormat.format(
+                    null,
+                    "unexpected result for spec[{}] [{}]. Found: {} - Expected one of the following: {}",
+                    name,
+                    query,
+                    Arrays.toString(actual),
+                    eventIds.stream().map(Arrays::toString).collect(Collectors.joining(", "))
+                );
+                fail(msg);
+            }
+        }
+
     }
 
     private String eventsToString(List<Map<String, Object>> events) {
@@ -182,7 +230,7 @@ private long[] extractIds(List<Map<String, Object>> events) {
         for (int i = 0; i < len; i++) {
             Map<String, Object> event = events.get(i);
             Map<String, Object> source = (Map<String, Object>) event.get("_source");
-            Object field = source.get(tiebreaker());
+            Object field = source.get(idField());
             ids[i] = ((Number) field).longValue();
         }
         return ids;
@@ -262,6 +310,10 @@ protected String eventCategory() {
         return "event.category";
     }
 
+    protected String idField() {
+        return tiebreaker();
+    }
+
     protected abstract String tiebreaker();
 
     protected int requestSize() {
@@ -278,8 +330,8 @@ protected String requestResultPosition() {
     }
 
     // strip any qualification from the received index string
-    private String unqualifiedIndexName() {
-        int offset = index.indexOf(':');
-        return offset >= 0 ? index.substring(offset + 1) : index;
+    private static String unqualifiedIndexName(String indexName) {
+        int offset = indexName.indexOf(':');
+        return offset >= 0 ? indexName.substring(offset + 1) : indexName;
     }
 }

x-pack/plugin/eql/qa/common/src/main/java/org/elasticsearch/test/eql/DataLoader.java

@@ -48,6 +48,7 @@
  * - endgame-140       - for existing data
  * - endgame-140-nanos - same as endgame-140, but with nano-precision timestamps
  * - extra             - additional data
+ * - sample*         - data for "sample" functionality
  *
  * While the loader could be made generic, the queries are bound to each index and generalizing that would make things way too complicated.
  */
@@ -56,6 +57,7 @@ public class DataLoader {
     public static final String TEST_INDEX = "endgame-140";
     public static final String TEST_EXTRA_INDEX = "extra";
     public static final String TEST_NANOS_INDEX = "endgame-140-nanos";
+    public static final String TEST_SAMPLE = "sample1,sample2,sample3";
 
     private static final Map<String, String[]> replacementPatterns = Collections.unmodifiableMap(getReplacementPatterns());
 
@@ -99,27 +101,31 @@ public static void loadDatasetIntoEs(
         // chosen Windows filetime timestamps (2017+) can coincidentally also be readily used as nano-resolution unix timestamps (1973+).
         // There are mixed values with and without nanos precision so that the filtering is properly tested for both cases.
         load(client, TEST_NANOS_INDEX, TEST_INDEX, DataLoader::timestampToUnixNanos, p);
+        load(client, TEST_SAMPLE, null, null, p);
     }
 
     private static void load(
         RestHighLevelClient client,
-        String indexName,
+        String indexNames,
         String dataName,
         Consumer<Map<String, Object>> datasetTransform,
         CheckedBiFunction<XContent, InputStream, XContentParser, IOException> p
     ) throws IOException {
-        String name = "/data/" + indexName + ".mapping";
-        URL mapping = DataLoader.class.getResource(name);
-        if (mapping == null) {
-            throw new IllegalArgumentException("Cannot find resource " + name);
-        }
-        name = "/data/" + (dataName != null ? dataName : indexName) + ".data";
-        URL data = DataLoader.class.getResource(name);
-        if (data == null) {
-            throw new IllegalArgumentException("Cannot find resource " + name);
+        String[] splitNames = indexNames.split(",");
+        for (String indexName : splitNames) {
+            String name = "/data/" + indexName + ".mapping";
+            URL mapping = DataLoader.class.getResource(name);
+            if (mapping == null) {
+                throw new IllegalArgumentException("Cannot find resource " + name);
+            }
+            name = "/data/" + (dataName != null ? dataName : indexName) + ".data";
+            URL data = DataLoader.class.getResource(name);
+            if (data == null) {
+                throw new IllegalArgumentException("Cannot find resource " + name);
+            }
+            createTestIndex(client, indexName, readMapping(mapping));
+            loadData(client, indexName, datasetTransform, data, p);
         }
-        createTestIndex(client, indexName, readMapping(mapping));
-        loadData(client, indexName, datasetTransform, data, p);
     }
 
     private static void createTestIndex(RestHighLevelClient client, String indexName, String mapping) throws IOException {

x-pack/plugin/eql/qa/common/src/main/java/org/elasticsearch/test/eql/EqlDateNanosSpecTestCase.java

@@ -21,13 +21,13 @@ public static List<Object[]> readTestSpecs() throws Exception {
     }
 
     // constructor for "local" rest tests
-    public EqlDateNanosSpecTestCase(String query, String name, long[] eventIds, String[] joinKeys) {
-        this(TEST_NANOS_INDEX, query, name, eventIds, joinKeys);
+    public EqlDateNanosSpecTestCase(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        this(TEST_NANOS_INDEX, query, name, eventIds, joinKeys, size);
     }
 
     // constructor for multi-cluster tests
-    public EqlDateNanosSpecTestCase(String index, String query, String name, long[] eventIds, String[] joinKeys) {
-        super(index, query, name, eventIds, joinKeys);
+    public EqlDateNanosSpecTestCase(String index, String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(index, query, name, eventIds, joinKeys, size);
     }
 
     @Override

x-pack/plugin/eql/qa/common/src/main/java/org/elasticsearch/test/eql/EqlExtraSpecTestCase.java

@@ -21,13 +21,13 @@ public static List<Object[]> readTestSpecs() throws Exception {
     }
 
     // constructor for "local" rest tests
-    public EqlExtraSpecTestCase(String query, String name, long[] eventIds, String[] joinKeys) {
-        this(TEST_EXTRA_INDEX, query, name, eventIds, joinKeys);
+    public EqlExtraSpecTestCase(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        this(TEST_EXTRA_INDEX, query, name, eventIds, joinKeys, size);
     }
 
     // constructor for multi-cluster tests
-    public EqlExtraSpecTestCase(String index, String query, String name, long[] eventIds, String[] joinKeys) {
-        super(index, query, name, eventIds, joinKeys);
+    public EqlExtraSpecTestCase(String index, String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(index, query, name, eventIds, joinKeys, size);
     }
 
     @Override

x-pack/plugin/eql/qa/common/src/main/java/org/elasticsearch/test/eql/EqlRestTestCase.java

@@ -50,6 +50,9 @@ public void checkSearchContent() throws Exception {
         { String.format(Locale.ROOT, """
             {"query": "%s", "size": -1}
             """, validQuery), "size must be greater than or equal to 0" },
+        { String.format(Locale.ROOT, """
+            {"query": "%s", "fetch_size": 1}
+            """, validQuery), "fetch size must be greater than 1" },
         { String.format(Locale.ROOT, """
             {"query": "%s", "filter": null}
             """, validQuery), "filter doesn't support values of type: VALUE_NULL" },
@@ -60,23 +63,35 @@ public void checkSearchContent() throws Exception {
     public void testBadRequests() throws Exception {
         createIndex(defaultValidationIndexName, (String) null);
 
-        final String contentType = "application/json";
         for (String[] test : testBadRequests) {
-            final String endpoint = "/" + indexPattern(defaultValidationIndexName) + "/_eql/search";
-            Request request = new Request("GET", endpoint);
-            request.setJsonEntity(test[0]);
-
-            ResponseException e = expectThrows(ResponseException.class, () -> client().performRequest(request));
-            Response response = e.getResponse();
-
-            assertThat(response.getHeader("Content-Type"), containsString(contentType));
-            assertThat(EntityUtils.toString(response.getEntity()), containsString(test[1]));
-            assertThat(response.getStatusLine().getStatusCode(), is(400));
+            assertBadRequest(test[0], test[1], 400);
         }
 
+        bulkIndex("""
+            {"index": {"_index": "%s", "_id": 1}}
+            {"event":{"category":"process"},"@timestamp":"2020-01-01T12:34:56Z"}
+            """.formatted(defaultValidationIndexName));
+        assertBadRequest("""
+            {"query": "sample by event.category [any where true] [any where true]",
+             "fetch_size": 1001}
+            """, "Fetch size cannot be greater than [1000]", 500);
+
         deleteIndexWithProvisioningClient(defaultValidationIndexName);
     }
 
+    private void assertBadRequest(String query, String errorMessage, int errorCode) throws IOException {
+        final String endpoint = "/" + indexPattern(defaultValidationIndexName) + "/_eql/search";
+        Request request = new Request("GET", endpoint);
+        request.setJsonEntity(query);
+
+        ResponseException e = expectThrows(ResponseException.class, () -> client().performRequest(request));
+        Response response = e.getResponse();
+
+        assertThat(response.getHeader("Content-Type"), containsString("application/json"));
+        assertThat(EntityUtils.toString(response.getEntity()), containsString(errorMessage));
+        assertThat(response.getStatusLine().getStatusCode(), is(errorCode));
+    }
+
     @SuppressWarnings("unchecked")
     public void testIndexWildcardPatterns() throws Exception {
         createIndex("test1", """

x-pack/plugin/eql/qa/common/src/main/java/org/elasticsearch/test/eql/EqlSampleTestCase.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.test.eql;
+
+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
+
+import java.util.List;
+
+import static org.elasticsearch.test.eql.DataLoader.TEST_SAMPLE;
+
+public abstract class EqlSampleTestCase extends BaseEqlSpecTestCase {
+
+    public EqlSampleTestCase(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        this(TEST_SAMPLE, query, name, eventIds, joinKeys, size);
+    }
+
+    public EqlSampleTestCase(String index, String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(index, query, name, eventIds, joinKeys, size);
+    }
+
+    @ParametersFactory(shuffle = false, argumentFormatting = PARAM_FORMATTING)
+    public static List<Object[]> readTestSpecs() throws Exception {
+        return asArray(EqlSpecLoader.load("/test_sample.toml"));
+    }
+
+    @Override
+    protected String tiebreaker() {
+        return null;
+    }
+
+    @Override
+    protected String idField() {
+        return "id";
+    }
+
+    @Override
+    protected int requestFetchSize() {
+        // a more relevant fetch_size value for Samples, from algorithm point of view, so we'll mostly test this value
+        return frequently() ? 2 : super.requestFetchSize();
+    }
+}