Skip to content

Commit 5290e0c

Browse files
ioanatiaioanatia
authored
ES|QL: Add license check for RRF (#125225)
1 parent d2c65f5 commit 5290e0c

File tree

4 files changed

+148
-27
lines changed

4 files changed

+148
-27
lines changed

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/ForkIT.java‎

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -350,32 +350,6 @@ public void testScoringKeepAndSort() {
350350
}
351351
}
352352

353-
public void testRrf() {
354-
assumeTrue("requires RRF capability", EsqlCapabilities.Cap.RRF.isEnabled());
355-
356-
var query = """
357-
FROM test METADATA _score, _id, _index
358-
| WHERE id > 2
359-
| FORK
360-
( WHERE content:"fox" | SORT _score, _id DESC )
361-
( WHERE content:"dog" | SORT _score, _id DESC )
362-
| RRF
363-
| EVAL _score = round(_score, 4)
364-
| KEEP id, content, _score, _fork
365-
""";
366-
try (var resp = run(query)) {
367-
assertColumnNames(resp.columns(), List.of("id", "content", "_score", "_fork"));
368-
assertColumnTypes(resp.columns(), List.of("integer", "keyword", "double", "keyword"));
369-
assertThat(getValuesList(resp.values()).size(), equalTo(3));
370-
Iterable<Iterable<Object>> expectedValues = List.of(
371-
List.of(6, "The quick brown fox jumps over the lazy dog", 0.0325, List.of("fork1", "fork2")),
372-
List.of(4, "The dog is brown but this document is very very long", 0.0164, "fork2"),
373-
List.of(3, "This dog is really brown", 0.0159, "fork2")
374-
);
375-
assertValues(resp.values(), expectedValues);
376-
}
377-
}
378-
379353
public void testThreeSubQueries() {
380354
var query = """
381355
FROM test

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/RrfIT.java‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.action;
9+
10+
import org.elasticsearch.action.index.IndexRequest;
11+
import org.elasticsearch.action.support.WriteRequest;
12+
import org.elasticsearch.common.settings.Settings;
13+
import org.elasticsearch.plugins.Plugin;
14+
import org.junit.Before;
15+
16+
import java.util.Collection;
17+
import java.util.List;
18+
19+
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
20+
import static org.elasticsearch.xpack.esql.EsqlTestUtils.getValuesList;
21+
import static org.hamcrest.Matchers.equalTo;
22+
23+
public class RrfIT extends AbstractEsqlIntegTestCase {
24+
@Override
25+
protected Collection<Class<? extends Plugin>> nodePlugins() {
26+
return List.of(EsqlPluginWithEnterpriseOrTrialLicense.class);
27+
}
28+
29+
@Before
30+
public void setupIndex() {
31+
assumeTrue("requires RRF capability", EsqlCapabilities.Cap.RRF.isEnabled());
32+
createAndPopulateIndex();
33+
}
34+
35+
public void testRrf() {
36+
var query = """
37+
FROM test METADATA _score, _id, _index
38+
| WHERE id > 2
39+
| FORK
40+
( WHERE content:"fox" | SORT _score, _id DESC )
41+
( WHERE content:"dog" | SORT _score, _id DESC )
42+
| RRF
43+
| EVAL _score = round(_score, 4)
44+
| KEEP id, content, _score, _fork
45+
""";
46+
try (var resp = run(query)) {
47+
assertColumnNames(resp.columns(), List.of("id", "content", "_score", "_fork"));
48+
assertColumnTypes(resp.columns(), List.of("integer", "keyword", "double", "keyword"));
49+
assertThat(getValuesList(resp.values()).size(), equalTo(3));
50+
Iterable<Iterable<Object>> expectedValues = List.of(
51+
List.of(6, "The quick brown fox jumps over the lazy dog", 0.0325, List.of("fork1", "fork2")),
52+
List.of(4, "The dog is brown but this document is very very long", 0.0164, "fork2"),
53+
List.of(3, "This dog is really brown", 0.0159, "fork2")
54+
);
55+
assertValues(resp.values(), expectedValues);
56+
}
57+
}
58+
59+
private void createAndPopulateIndex() {
60+
var indexName = "test";
61+
var client = client().admin().indices();
62+
var CreateRequest = client.prepareCreate(indexName)
63+
.setSettings(Settings.builder().put("index.number_of_shards", 1))
64+
.setMapping("id", "type=integer", "content", "type=text");
65+
assertAcked(CreateRequest);
66+
client().prepareBulk()
67+
.add(new IndexRequest(indexName).id("1").source("id", 1, "content", "This is a brown fox"))
68+
.add(new IndexRequest(indexName).id("2").source("id", 2, "content", "This is a brown dog"))
69+
.add(new IndexRequest(indexName).id("3").source("id", 3, "content", "This dog is really brown"))
70+
.add(new IndexRequest(indexName).id("4").source("id", 4, "content", "The dog is brown but this document is very very long"))
71+
.add(new IndexRequest(indexName).id("5").source("id", 5, "content", "There is also a white cat"))
72+
.add(new IndexRequest(indexName).id("6").source("id", 6, "content", "The quick brown fox jumps over the lazy dog"))
73+
.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
74+
.get();
75+
ensureYellow(indexName);
76+
}
77+
}

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/RrfWithInvalidLicenseIT.java‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.action;
9+
10+
import org.elasticsearch.ElasticsearchException;
11+
import org.elasticsearch.action.index.IndexRequestBuilder;
12+
import org.elasticsearch.action.support.WriteRequest;
13+
import org.elasticsearch.common.settings.Settings;
14+
import org.elasticsearch.plugins.Plugin;
15+
import org.elasticsearch.xpack.esql.VerificationException;
16+
import org.junit.Before;
17+
18+
import java.util.Collection;
19+
import java.util.List;
20+
21+
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
22+
import static org.hamcrest.Matchers.containsString;
23+
24+
public class RrfWithInvalidLicenseIT extends AbstractEsqlIntegTestCase {
25+
private static final String LICENSE_ERROR_MESSAGE = "current license is non-compliant for [RRF]";
26+
27+
@Override
28+
protected Collection<Class<? extends Plugin>> nodePlugins() {
29+
return List.of(EsqlPluginWithNonEnterpriseOrExpiredLicense.class);
30+
}
31+
32+
@Before
33+
public void setupIndex() {
34+
assumeTrue("requires RRF capability", EsqlCapabilities.Cap.RRF.isEnabled());
35+
var indexName = "test";
36+
var client = client().admin().indices();
37+
var CreateRequest = client.prepareCreate(indexName)
38+
.setSettings(Settings.builder().put("index.number_of_shards", 1))
39+
.setMapping("id", "type=integer", "content", "type=text");
40+
assertAcked(CreateRequest);
41+
client().prepareBulk()
42+
.add(new IndexRequestBuilder(client, indexName).setId("1").setSource("id", 1, "content", "This is a brown fox"))
43+
.add(new IndexRequestBuilder(client, indexName).setId("2").setSource("id", 2, "content", "This is a brown dog"))
44+
.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
45+
.get();
46+
ensureYellow(indexName);
47+
}
48+
49+
public void testRrf() {
50+
var query = """
51+
FROM test METADATA _score, _id, _index
52+
| FORK
53+
( WHERE content:"fox" )
54+
( WHERE content:"dog" )
55+
| RRF
56+
""";
57+
58+
ElasticsearchException e = expectThrows(VerificationException.class, () -> run(query));
59+
60+
assertThat(e.getMessage(), containsString(LICENSE_ERROR_MESSAGE));
61+
}
62+
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/RrfScoreEval.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@
88
package org.elasticsearch.xpack.esql.plan.logical;
99

1010
import org.elasticsearch.common.io.stream.StreamOutput;
11+
import org.elasticsearch.license.License;
12+
import org.elasticsearch.license.XPackLicenseState;
13+
import org.elasticsearch.xpack.esql.LicenseAware;
1114
import org.elasticsearch.xpack.esql.capabilities.PostAnalysisVerificationAware;
1215
import org.elasticsearch.xpack.esql.common.Failures;
1316
import org.elasticsearch.xpack.esql.core.expression.Attribute;
@@ -20,7 +23,7 @@
2023

2124
import static org.elasticsearch.xpack.esql.common.Failure.fail;
2225

23-
public class RrfScoreEval extends UnaryPlan implements PostAnalysisVerificationAware {
26+
public class RrfScoreEval extends UnaryPlan implements PostAnalysisVerificationAware, LicenseAware {
2427
private final Attribute forkAttr;
2528
private final Attribute scoreAttr;
2629

@@ -93,4 +96,9 @@ public boolean equals(Object obj) {
9396
RrfScoreEval rrf = (RrfScoreEval) obj;
9497
return child().equals(rrf.child()) && scoreAttr.equals(rrf.scoreAttribute()) && forkAttr.equals(forkAttribute());
9598
}
99+
100+
@Override
101+
public boolean licenseCheck(XPackLicenseState state) {
102+
return state.isAllowedByLicense(License.OperationMode.ENTERPRISE);
103+
}
96104
}

0 commit comments

Comments
 (0)