Skip to content

Instantly share code, notes, and snippets.

@msfroh
Created September 29, 2023 18:42
Show Gist options
  • Save msfroh/74aa3fee52f4074c5e7b8d85f76e88ab to your computer and use it in GitHub Desktop.
Save msfroh/74aa3fee52f4074c5e7b8d85f76e88ab to your computer and use it in GitHub Desktop.
public static class QueryShapeVisitor implements QueryBuilderVisitor {
private final SetOnce<String> queryType = new SetOnce<>();
private final Map<BooleanClause.Occur, List<QueryShapeVisitor>> childVisitors = new EnumMap<>(BooleanClause.Occur.class);
@Override
public void accept(QueryBuilder qb) {
queryType.set(qb.getName());
}
@Override
public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
// Should get called once per Occur value
if (childVisitors.containsKey(occur)) {
throw new IllegalStateException("getChildVisitor already called for " + occur);
}
final List<QueryShapeVisitor> childVisitorList = new ArrayList<>();
QueryBuilderVisitor childVisitorWrapper = new QueryBuilderVisitor() {
QueryShapeVisitor currentChild;
@Override
public void accept(QueryBuilder qb) {
currentChild = new QueryShapeVisitor();
childVisitorList.add(currentChild);
currentChild.accept(qb);
}
@Override
public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
return currentChild.getChildVisitor(occur);
}
};
childVisitors.put(occur, childVisitorList);
return childVisitorWrapper;
}
public String toJson() {
StringBuilder outputBuilder = new StringBuilder("{\"type\":\"").append(queryType.get()).append("\"");
for (Map.Entry<BooleanClause.Occur, List<QueryShapeVisitor>> entry : childVisitors.entrySet()) {
outputBuilder.append(",\"").append(entry.getKey().name().toLowerCase(Locale.ROOT)).append("\"[");
boolean first = true;
for (QueryShapeVisitor child : entry.getValue()) {
if (!first) {
outputBuilder.append(",");
}
outputBuilder.append(child.toJson());
first = false;
}
outputBuilder.append("]");
}
outputBuilder.append("}");
return outputBuilder.toString();
}
}
public void testQueryShapeVisitor() {
QueryBuilder builder = new BoolQueryBuilder()
.must(new TermQueryBuilder("foo", "bar"))
.filter(new ConstantScoreQueryBuilder(new RangeQueryBuilder("timestamp").from("12345677").to("2345678")))
.should(new BoolQueryBuilder()
.must(new MatchQueryBuilder("text", "this is some text"))
.mustNot(new RegexpQueryBuilder("color", "red.*"))
)
.must(new TermsQueryBuilder("genre", "action", "drama", "romance"));
QueryShapeVisitor shapeVisitor = new QueryShapeVisitor();
builder.visit(shapeVisitor);
// {"type":"bool","must"[{"type":"term"},{"type":"terms"}],"filter"[{"type":"constant_score","filter"[{"type":"range"}]}],"should"[{"type":"bool","must"[{"type":"match"}],"must_not"[{"type":"regexp"}]}]}
assertEquals(
"{\"type\":\"bool\",\"must\"[{\"type\":\"term\"},{\"type\":\"terms\"}],\"filter\"[{\"type\":\"constant_score\",\"filter\"[{\"type\":\"range\"}]}],\"should\"[{\"type\":\"bool\",\"must\"[{\"type\":\"match\"}],\"must_not\"[{\"type\":\"regexp\"}]}]}",
shapeVisitor.toJson());
}
@vibrantvarun
Copy link

Great presentation of the solution for the problem.

@deshsidd
Copy link

deshsidd commented Oct 4, 2023

Thanks @msfroh this is super helpful

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment