Skip to content

Commit a2bf4de

Browse files
mridula-s109mridula-s109
committed
made parsing robust
1 parent 7433023 commit a2bf4de

File tree

2 files changed

+165
-25
lines changed

2 files changed

+165
-25
lines changed

‎x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverComponent.java‎

Lines changed: 58 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -77,38 +77,71 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContentObject.Para
7777
}
7878

7979
public static RRFRetrieverComponent fromXContent(XContentParser parser, RetrieverParserContext context) throws IOException {
80-
RetrieverBuilder innerRetriever = null;
81-
float weight = DEFAULT_WEIGHT;
82-
8380
if (parser.currentToken() != XContentParser.Token.START_OBJECT) {
84-
throw new ParsingException(parser.getTokenLocation(), "[{}] expected object", parser.currentToken());
81+
throw new ParsingException(parser.getTokenLocation(), "expected object but found [{}]", parser.currentToken());
8582
}
8683

87-
while ((parser.nextToken()) != XContentParser.Token.END_OBJECT) {
88-
var name = parser.currentName();
84+
// Peek at the first field to determine the format
85+
XContentParser.Token token = parser.nextToken();
86+
if (token == XContentParser.Token.END_OBJECT) {
87+
throw new ParsingException(parser.getTokenLocation(), "retriever component must contain a retriever");
88+
}
89+
if (token != XContentParser.Token.FIELD_NAME) {
90+
throw new ParsingException(parser.getTokenLocation(), "expected field name but found [{}]", token);
91+
}
8992

90-
if (name.equals(RETRIEVER_FIELD.getPreferredName())) {
91-
if (parser.nextToken() != XContentParser.Token.START_OBJECT) {
92-
throw new ParsingException(parser.getTokenLocation(), "[{}] expected object", parser.currentToken());
93-
}
94-
parser.nextToken();
95-
96-
name = parser.currentName();
97-
innerRetriever = parser.namedObject(RetrieverBuilder.class, name, context);
98-
parser.nextToken();
99-
} else if (name.equals(WEIGHT_FIELD.getPreferredName())) {
100-
if (parser.nextToken() != XContentParser.Token.VALUE_NUMBER) {
101-
throw new ParsingException(parser.getTokenLocation(), "[{}] expected number", parser.currentToken());
93+
String firstFieldName = parser.currentName();
94+
95+
// Check if this is a structured component (starts with "retriever" or "weight")
96+
if (RETRIEVER_FIELD.match(firstFieldName, parser.getDeprecationHandler())
97+
|| WEIGHT_FIELD.match(firstFieldName, parser.getDeprecationHandler())) {
98+
// This is a structured component - parse manually
99+
RetrieverBuilder retriever = null;
100+
Float weight = null;
101+
102+
do {
103+
String fieldName = parser.currentName();
104+
if (RETRIEVER_FIELD.match(fieldName, parser.getDeprecationHandler())) {
105+
if (retriever != null) {
106+
throw new ParsingException(parser.getTokenLocation(), "only one retriever can be specified");
107+
}
108+
parser.nextToken();
109+
parser.nextToken();
110+
String retrieverType = parser.currentName();
111+
retriever = parser.namedObject(RetrieverBuilder.class, retrieverType, context);
112+
context.trackRetrieverUsage(retriever.getName());
113+
parser.nextToken();
114+
} else if (WEIGHT_FIELD.match(fieldName, parser.getDeprecationHandler())) {
115+
if (weight != null) {
116+
throw new ParsingException(parser.getTokenLocation(), "[weight] field can only be specified once");
117+
}
118+
parser.nextToken();
119+
weight = parser.floatValue();
120+
} else {
121+
if (retriever != null) {
122+
throw new ParsingException(parser.getTokenLocation(), "only one retriever can be specified");
123+
}
124+
throw new ParsingException(
125+
parser.getTokenLocation(),
126+
"unknown field [{}], expected [{}] or [{}]",
127+
fieldName,
128+
RETRIEVER_FIELD.getPreferredName(),
129+
WEIGHT_FIELD.getPreferredName()
130+
);
102131
}
132+
} while (parser.nextToken() == XContentParser.Token.FIELD_NAME);
133+
134+
if (retriever == null) {
135+
throw new ParsingException(parser.getTokenLocation(), "retriever component must contain a retriever");
136+
}
103137

104-
weight = parser.floatValue();
105-
} else {
106-
innerRetriever = parser.namedObject(RetrieverBuilder.class, name, context);
107-
context.trackRetrieverUsage(innerRetriever.getName());
108-
parser.nextToken();
109-
break;
138+
return new RRFRetrieverComponent(retriever, weight);
139+
} else {
140+
RetrieverBuilder retriever = parser.namedObject(RetrieverBuilder.class, firstFieldName, context);
141+
context.trackRetrieverUsage(retriever.getName());
142+
while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
110143
}
144+
return new RRFRetrieverComponent(retriever, DEFAULT_WEIGHT);
111145
}
112-
return new RRFRetrieverComponent(innerRetriever, weight);
113146
}
114147
}

‎x-pack/plugin/rank-rrf/src/test/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilderParsingTests.java‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import java.util.ArrayList;
2828
import java.util.List;
2929

30+
import static org.hamcrest.Matchers.containsString;
3031
import static org.hamcrest.Matchers.equalTo;
3132
import static org.hamcrest.Matchers.instanceOf;
3233

@@ -251,4 +252,110 @@ public void testRRFRetrieverParsingWithDefaultWeights() throws IOException {
251252
""";
252253
checkRRFRetrieverParsing(restContent);
253254
}
255+
256+
public void testRRFRetrieverComponentErrorCases() throws IOException {
257+
// Test case 1: Multiple retrievers in same component
258+
String multipleRetrieversContent = """
259+
{
260+
"retriever": {
261+
"rrf": {
262+
"retrievers": [
263+
{
264+
"retriever": { "test": { "value": "first" } },
265+
"standard": { "query": { "match_all": {} } }
266+
}
267+
],
268+
"rank_window_size": 100,
269+
"rank_constant": 10,
270+
"min_score": 20.0,
271+
"_name": "foo_rrf"
272+
}
273+
}
274+
}
275+
""";
276+
277+
expectParsingException(multipleRetrieversContent, "only one retriever can be specified");
278+
279+
// Test case 2: Weight without retriever
280+
String weightOnlyContent = """
281+
{
282+
"retriever": {
283+
"rrf": {
284+
"retrievers": [
285+
{
286+
"weight": 2.0
287+
}
288+
],
289+
"rank_window_size": 100,
290+
"rank_constant": 10,
291+
"min_score": 20.0,
292+
"_name": "foo_rrf"
293+
}
294+
}
295+
}
296+
""";
297+
298+
expectParsingException(weightOnlyContent, "retriever component must contain a retriever");
299+
300+
// Test case 3: Empty retriever component
301+
String emptyComponentContent = """
302+
{
303+
"retriever": {
304+
"rrf": {
305+
"retrievers": [
306+
{}
307+
],
308+
"rank_window_size": 100,
309+
"rank_constant": 10,
310+
"min_score": 20.0,
311+
"_name": "foo_rrf"
312+
}
313+
}
314+
}
315+
""";
316+
317+
expectParsingException(emptyComponentContent, "retriever component must contain a retriever");
318+
319+
// Test case 4: Negative weight
320+
String negativeWeightContent = """
321+
{
322+
"retriever": {
323+
"rrf": {
324+
"retrievers": [
325+
{
326+
"retriever": { "test": { "value": "test" } },
327+
"weight": -1.0
328+
}
329+
],
330+
"rank_window_size": 100,
331+
"rank_constant": 10,
332+
"min_score": 20.0,
333+
"_name": "foo_rrf"
334+
}
335+
}
336+
}
337+
""";
338+
339+
expectParsingException(negativeWeightContent, "weight] must be non-negative");
340+
}
341+
342+
private void expectParsingException(String restContent, String expectedMessageFragment) throws IOException {
343+
SearchUsageHolder searchUsageHolder = new UsageService().getSearchUsageHolder();
344+
try (XContentParser jsonParser = createParser(JsonXContent.jsonXContent, restContent)) {
345+
Exception exception = expectThrows(Exception.class, () -> {
346+
new SearchSourceBuilder().parseXContent(jsonParser, true, searchUsageHolder, nf -> true);
347+
});
348+
349+
String message = exception.getMessage();
350+
if (exception.getCause() != null) {
351+
message = exception.getCause().getMessage();
352+
}
353+
354+
assertThat(
355+
"Expected error message to contain: " + expectedMessageFragment + ", but got: " + message,
356+
message,
357+
containsString(expectedMessageFragment)
358+
);
359+
}
360+
}
254361
}

0 commit comments

Comments
 (0)