Skip to content

Commit 4328f22

Browse files
karenyrxPeter Alfonsi
karenyrx
authored and
Peter Alfonsi
committed
[GRPC][Specialized queries] Implement GRPC Script query (opensearch-project#19455)
1 parent 7d50016 commit 4328f22

File tree

7 files changed

+345
-0
lines changed

7 files changed

+345
-0
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
3333
- Implement GRPC MatchPhrase, MultiMatch queries ([#19449](https://github.com/opensearch-project/OpenSearch/pull/19449))
3434
- Optimize gRPC transport thread management for improved throughput ([#19278](https://github.com/opensearch-project/OpenSearch/pull/19278))
3535
- Implement GRPC Boolean query and inject registry for all internal query converters ([#19391](https://github.com/opensearch-project/OpenSearch/pull/19391))
36+
- Implement GRPC Script query ([#19455](https://github.com/opensearch-project/OpenSearch/pull/19455))
3637

3738
### Changed
3839
- Refactor `if-else` chains to use `Java 17 pattern matching switch expressions`(([#18965](https://github.com/opensearch-project/OpenSearch/pull/18965))

‎modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/request/search/query/QueryBuilderProtoConverterRegistryImpl.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ protected void registerBuiltInConverters() {
5151
delegate.registerConverter(new MatchPhraseQueryBuilderProtoConverter());
5252
delegate.registerConverter(new MultiMatchQueryBuilderProtoConverter());
5353
delegate.registerConverter(new BoolQueryBuilderProtoConverter());
54+
delegate.registerConverter(new ScriptQueryBuilderProtoConverter());
5455

5556
// Set the registry on all converters so they can access each other
5657
delegate.setRegistryOnAllConverters(this);

‎modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/request/search/query/ScriptQueryBuilderProtoConverter.java‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
package org.opensearch.transport.grpc.proto.request.search.query;
9+
10+
import org.opensearch.index.query.QueryBuilder;
11+
import org.opensearch.protobufs.QueryContainer;
12+
import org.opensearch.transport.grpc.spi.QueryBuilderProtoConverter;
13+
14+
/**
15+
* Converter for Script queries.
16+
* This class implements the QueryBuilderProtoConverter interface to provide Script query support
17+
* for the gRPC transport module.
18+
*/
19+
public class ScriptQueryBuilderProtoConverter implements QueryBuilderProtoConverter {
20+
21+
/**
22+
* Constructs a new ScriptQueryBuilderProtoConverter.
23+
*/
24+
public ScriptQueryBuilderProtoConverter() {
25+
// Default constructor
26+
}
27+
28+
@Override
29+
public QueryContainer.QueryContainerCase getHandledQueryCase() {
30+
return QueryContainer.QueryContainerCase.SCRIPT;
31+
}
32+
33+
@Override
34+
public QueryBuilder fromProto(QueryContainer queryContainer) {
35+
if (queryContainer == null || !queryContainer.hasScript()) {
36+
throw new IllegalArgumentException("QueryContainer does not contain a Script query");
37+
}
38+
39+
return ScriptQueryBuilderProtoUtils.fromProto(queryContainer.getScript());
40+
}
41+
}

‎modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/request/search/query/ScriptQueryBuilderProtoUtils.java‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
package org.opensearch.transport.grpc.proto.request.search.query;
9+
10+
import org.opensearch.index.query.ScriptQueryBuilder;
11+
import org.opensearch.protobufs.ScriptQuery;
12+
import org.opensearch.script.Script;
13+
import org.opensearch.transport.grpc.proto.request.common.ScriptProtoUtils;
14+
15+
/**
16+
* Utility class for converting ScriptQuery Protocol Buffers to ScriptQueryBuilder objects.
17+
* This class handles the conversion of Protocol Buffer representations to their
18+
* corresponding OpenSearch query builder objects.
19+
*/
20+
class ScriptQueryBuilderProtoUtils {
21+
22+
private ScriptQueryBuilderProtoUtils() {
23+
// Utility class, no instances
24+
}
25+
26+
/**
27+
* Converts a ScriptQuery Protocol Buffer to a ScriptQueryBuilder object.
28+
* This method follows the same pattern as ScriptQueryBuilder.fromXContent().
29+
*
30+
* @param scriptQueryProto the Protocol Buffer ScriptQuery to convert
31+
* @return the converted ScriptQueryBuilder object
32+
* @throws IllegalArgumentException if the script query proto is null or invalid
33+
*/
34+
static ScriptQueryBuilder fromProto(ScriptQuery scriptQueryProto) {
35+
if (scriptQueryProto == null) {
36+
throw new IllegalArgumentException("ScriptQuery cannot be null");
37+
}
38+
39+
if (!scriptQueryProto.hasScript()) {
40+
throw new IllegalArgumentException("script must be provided with a [script] query");
41+
}
42+
43+
Script script = ScriptProtoUtils.parseFromProtoRequest(scriptQueryProto.getScript());
44+
45+
float boost = ScriptQueryBuilder.DEFAULT_BOOST;
46+
String queryName = null;
47+
48+
if (scriptQueryProto.hasBoost()) {
49+
boost = scriptQueryProto.getBoost();
50+
}
51+
52+
if (scriptQueryProto.hasXName()) {
53+
queryName = scriptQueryProto.getXName();
54+
}
55+
56+
return new ScriptQueryBuilder(script).boost(boost).queryName(queryName);
57+
}
58+
}

‎modules/transport-grpc/src/test/java/org/opensearch/transport/grpc/proto/request/search/query/QueryBuilderProtoConverterRegistryTests.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,14 @@
1010
import org.opensearch.index.query.QueryBuilder;
1111
import org.opensearch.protobufs.BoolQuery;
1212
import org.opensearch.protobufs.FieldValue;
13+
import org.opensearch.protobufs.InlineScript;
1314
import org.opensearch.protobufs.MatchAllQuery;
1415
import org.opensearch.protobufs.MatchPhraseQuery;
1516
import org.opensearch.protobufs.MinimumShouldMatch;
1617
import org.opensearch.protobufs.MultiMatchQuery;
1718
import org.opensearch.protobufs.QueryContainer;
19+
import org.opensearch.protobufs.Script;
20+
import org.opensearch.protobufs.ScriptQuery;
1821
import org.opensearch.protobufs.TermQuery;
1922
import org.opensearch.protobufs.TextQueryType;
2023
import org.opensearch.test.OpenSearchTestCase;
@@ -65,6 +68,22 @@ public void testTermQueryConversion() {
6568
assertEquals("Should be a TermQueryBuilder", "org.opensearch.index.query.TermQueryBuilder", queryBuilder.getClass().getName());
6669
}
6770

71+
public void testScriptQueryConversion() {
72+
// Create a Script query container with inline script
73+
Script script = Script.newBuilder().setInline(InlineScript.newBuilder().setSource("doc['field'].value > 100").build()).build();
74+
75+
ScriptQuery scriptQuery = ScriptQuery.newBuilder().setScript(script).setBoost(1.5f).setXName("test_script_query").build();
76+
77+
QueryContainer queryContainer = QueryContainer.newBuilder().setScript(scriptQuery).build();
78+
79+
// Convert using the registry
80+
QueryBuilder queryBuilder = registry.fromProto(queryContainer);
81+
82+
// Verify the result
83+
assertNotNull("QueryBuilder should not be null", queryBuilder);
84+
assertEquals("Should be a ScriptQueryBuilder", "org.opensearch.index.query.ScriptQueryBuilder", queryBuilder.getClass().getName());
85+
}
86+
6887
public void testNullQueryContainer() {
6988
expectThrows(IllegalArgumentException.class, () -> registry.fromProto(null));
7089
}

‎modules/transport-grpc/src/test/java/org/opensearch/transport/grpc/proto/request/search/query/ScriptQueryBuilderProtoConverterTests.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
package org.opensearch.transport.grpc.proto.request.search.query;
9+
10+
import org.opensearch.index.query.QueryBuilder;
11+
import org.opensearch.protobufs.InlineScript;
12+
import org.opensearch.protobufs.QueryContainer;
13+
import org.opensearch.protobufs.Script;
14+
import org.opensearch.protobufs.ScriptQuery;
15+
import org.opensearch.test.OpenSearchTestCase;
16+
17+
/**
18+
* Unit tests for ScriptQueryBuilderProtoConverter.
19+
* Tests only the converter-specific logic (QueryContainer handling).
20+
* The core conversion logic is tested in ScriptQueryBuilderProtoUtilsTests.
21+
*/
22+
public class ScriptQueryBuilderProtoConverterTests extends OpenSearchTestCase {
23+
24+
private ScriptQueryBuilderProtoConverter converter;
25+
26+
@Override
27+
public void setUp() throws Exception {
28+
super.setUp();
29+
converter = new ScriptQueryBuilderProtoConverter();
30+
}
31+
32+
public void testGetHandledQueryCase() {
33+
assertEquals(QueryContainer.QueryContainerCase.SCRIPT, converter.getHandledQueryCase());
34+
}
35+
36+
public void testFromProtoWithValidScriptQuery() {
37+
ScriptQuery scriptQuery = ScriptQuery.newBuilder()
38+
.setScript(Script.newBuilder().setInline(InlineScript.newBuilder().setSource("true").build()).build())
39+
.build();
40+
41+
QueryContainer container = QueryContainer.newBuilder().setScript(scriptQuery).build();
42+
43+
QueryBuilder result = converter.fromProto(container);
44+
45+
assertNotNull(result);
46+
assertTrue(result instanceof org.opensearch.index.query.ScriptQueryBuilder);
47+
}
48+
49+
public void testFromProtoWithNullContainer() {
50+
IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> converter.fromProto(null));
51+
assertEquals("QueryContainer does not contain a Script query", exception.getMessage());
52+
}
53+
54+
public void testFromProtoWithContainerWithoutScript() {
55+
QueryContainer container = QueryContainer.newBuilder().build();
56+
57+
IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> converter.fromProto(container));
58+
assertEquals("QueryContainer does not contain a Script query", exception.getMessage());
59+
}
60+
}

‎modules/transport-grpc/src/test/java/org/opensearch/transport/grpc/proto/request/search/query/ScriptQueryBuilderProtoUtilsTests.java‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
package org.opensearch.transport.grpc.proto.request.search.query;
9+
10+
import org.opensearch.index.query.ScriptQueryBuilder;
11+
import org.opensearch.protobufs.BuiltinScriptLanguage;
12+
import org.opensearch.protobufs.InlineScript;
13+
import org.opensearch.protobufs.Script;
14+
import org.opensearch.protobufs.ScriptLanguage;
15+
import org.opensearch.protobufs.ScriptQuery;
16+
import org.opensearch.test.OpenSearchTestCase;
17+
18+
/**
19+
* Unit tests for ScriptQueryBuilderProtoUtils.
20+
*/
21+
public class ScriptQueryBuilderProtoUtilsTests extends OpenSearchTestCase {
22+
23+
public void testFromProtoWithInlineScript() {
24+
// Test inline script with all features
25+
ScriptQuery scriptQuery = ScriptQuery.newBuilder()
26+
.setScript(
27+
Script.newBuilder()
28+
.setInline(
29+
InlineScript.newBuilder()
30+
.setSource("doc['field'].value > 0")
31+
.setLang(ScriptLanguage.newBuilder().setBuiltin(BuiltinScriptLanguage.BUILTIN_SCRIPT_LANGUAGE_PAINLESS))
32+
.setParams(
33+
org.opensearch.protobufs.ObjectMap.newBuilder()
34+
.putFields("multiplier", org.opensearch.protobufs.ObjectMap.Value.newBuilder().setDouble(2.5).build())
35+
.build()
36+
)
37+
.putOptions("content_type", "application/json")
38+
.build()
39+
)
40+
.build()
41+
)
42+
.setBoost(2.0f)
43+
.setXName("test_query")
44+
.build();
45+
46+
ScriptQueryBuilder result = ScriptQueryBuilderProtoUtils.fromProto(scriptQuery);
47+
48+
assertNotNull(result);
49+
assertEquals(2.0f, result.boost(), 0.001f);
50+
assertEquals("test_query", result.queryName());
51+
assertNotNull(result.script());
52+
assertEquals("doc['field'].value > 0", result.script().getIdOrCode());
53+
assertEquals("painless", result.script().getLang());
54+
assertNotNull(result.script().getParams());
55+
assertEquals(2.5, result.script().getParams().get("multiplier"));
56+
assertNotNull(result.script().getOptions());
57+
assertEquals("application/json", result.script().getOptions().get("content_type"));
58+
}
59+
60+
public void testFromProtoWithStoredScript() {
61+
ScriptQuery scriptQuery = ScriptQuery.newBuilder()
62+
.setScript(
63+
Script.newBuilder()
64+
.setStored(
65+
org.opensearch.protobufs.StoredScriptId.newBuilder()
66+
.setId("my_stored_script")
67+
.setParams(
68+
org.opensearch.protobufs.ObjectMap.newBuilder()
69+
.putFields(
70+
"param1",
71+
org.opensearch.protobufs.ObjectMap.Value.newBuilder().setString("test_value").build()
72+
)
73+
.build()
74+
)
75+
.build()
76+
)
77+
.build()
78+
)
79+
.build();
80+
81+
ScriptQueryBuilder result = ScriptQueryBuilderProtoUtils.fromProto(scriptQuery);
82+
83+
assertNotNull(result);
84+
assertEquals("my_stored_script", result.script().getIdOrCode());
85+
assertNull(result.script().getLang());
86+
assertNotNull(result.script().getParams());
87+
assertEquals("test_value", result.script().getParams().get("param1"));
88+
}
89+
90+
public void testFromProtoWithDifferentScriptLanguages() {
91+
// Test different script languages
92+
String[] languages = { "painless", "java", "mustache", "expression", "custom_lang" };
93+
BuiltinScriptLanguage[] builtinLangs = {
94+
BuiltinScriptLanguage.BUILTIN_SCRIPT_LANGUAGE_PAINLESS,
95+
BuiltinScriptLanguage.BUILTIN_SCRIPT_LANGUAGE_JAVA,
96+
BuiltinScriptLanguage.BUILTIN_SCRIPT_LANGUAGE_MUSTACHE,
97+
BuiltinScriptLanguage.BUILTIN_SCRIPT_LANGUAGE_EXPRESSION,
98+
null };
99+
100+
for (int i = 0; i < languages.length; i++) {
101+
ScriptQuery.Builder scriptQueryBuilder = ScriptQuery.newBuilder();
102+
InlineScript.Builder inlineScriptBuilder = InlineScript.newBuilder().setSource("test_script_" + i);
103+
104+
if (builtinLangs[i] != null) {
105+
inlineScriptBuilder.setLang(ScriptLanguage.newBuilder().setBuiltin(builtinLangs[i]));
106+
} else {
107+
inlineScriptBuilder.setLang(ScriptLanguage.newBuilder().setCustom(languages[i]));
108+
}
109+
110+
scriptQueryBuilder.setScript(Script.newBuilder().setInline(inlineScriptBuilder.build()).build());
111+
112+
ScriptQueryBuilder result = ScriptQueryBuilderProtoUtils.fromProto(scriptQueryBuilder.build());
113+
114+
assertNotNull(result);
115+
assertEquals(languages[i], result.script().getLang());
116+
assertEquals("test_script_" + i, result.script().getIdOrCode());
117+
}
118+
}
119+
120+
public void testFromProtoWithDefaults() {
121+
ScriptQuery scriptQuery = ScriptQuery.newBuilder()
122+
.setScript(Script.newBuilder().setInline(InlineScript.newBuilder().setSource("true").build()).build())
123+
.build();
124+
125+
ScriptQueryBuilder result = ScriptQueryBuilderProtoUtils.fromProto(scriptQuery);
126+
127+
assertNotNull(result);
128+
assertEquals(1.0f, result.boost(), 0.001f);
129+
assertNull(result.queryName());
130+
assertNotNull(result.script());
131+
assertEquals("true", result.script().getIdOrCode());
132+
}
133+
134+
public void testFromProtoWithEdgeCases() {
135+
ScriptQuery scriptQuery = ScriptQuery.newBuilder()
136+
.setScript(Script.newBuilder().setInline(InlineScript.newBuilder().setSource("true").build()).build())
137+
.setBoost(0.0f)
138+
.setXName("")
139+
.build();
140+
141+
ScriptQueryBuilder result = ScriptQueryBuilderProtoUtils.fromProto(scriptQuery);
142+
143+
assertNotNull(result);
144+
assertEquals(0.0f, result.boost(), 0.001f);
145+
assertEquals("", result.queryName());
146+
}
147+
148+
public void testFromProtoWithNullInput() {
149+
IllegalArgumentException exception = expectThrows(
150+
IllegalArgumentException.class,
151+
() -> ScriptQueryBuilderProtoUtils.fromProto(null)
152+
);
153+
assertEquals("ScriptQuery cannot be null", exception.getMessage());
154+
}
155+
156+
public void testFromProtoWithMissingScript() {
157+
ScriptQuery scriptQuery = ScriptQuery.newBuilder().build();
158+
159+
IllegalArgumentException exception = expectThrows(
160+
IllegalArgumentException.class,
161+
() -> ScriptQueryBuilderProtoUtils.fromProto(scriptQuery)
162+
);
163+
assertEquals("script must be provided with a [script] query", exception.getMessage());
164+
}
165+
}

0 commit comments

Comments
 (0)