Skip to content

Commit 2abe636

Browse files
parkertimminsparkertimmins
committed
Move query generation into separate classes
1 parent f11085e commit 2abe636

File tree

3 files changed

+231
-229
lines changed

3 files changed

+231
-229
lines changed

‎test/framework/src/main/java/org/elasticsearch/datageneration/queries/LeafQueryGenerator.java‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
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", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.datageneration.queries;
11+
12+
import org.elasticsearch.datageneration.FieldType;
13+
import org.elasticsearch.index.query.QueryBuilder;
14+
import org.elasticsearch.index.query.QueryBuilders;
15+
import org.elasticsearch.test.ESTestCase;
16+
17+
import java.util.ArrayList;
18+
import java.util.Arrays;
19+
import java.util.List;
20+
import java.util.Map;
21+
22+
public interface LeafQueryGenerator {
23+
24+
List<QueryBuilder> generate(Map<String, Object> fieldMapping, String path, Object value);
25+
26+
/**
27+
* Build a query for a specific type. If the field is nested, this query will need to be wrapped in nested queries.
28+
* @param type the type to build a query for
29+
* @return a generator that can build queries for this type
30+
*/
31+
static LeafQueryGenerator buildForType(String type) {
32+
LeafQueryGenerator noQueries = (Map<String, Object> fieldMapping, String path, Object value) -> List.of();
33+
34+
FieldType fieldType = FieldType.tryParse(type);
35+
if (fieldType == null) {
36+
return noQueries;
37+
}
38+
39+
return switch (fieldType) {
40+
case KEYWORD -> new KeywordQueryGenerator();
41+
case TEXT -> new TextQueryGenerator();
42+
case WILDCARD -> new WildcardQueryGenerator();
43+
default -> noQueries;
44+
};
45+
}
46+
47+
class KeywordQueryGenerator implements LeafQueryGenerator {
48+
public List<QueryBuilder> generate(Map<String, Object> fieldMapping, String path, Object value) {
49+
var ignoreAbove = (Integer) fieldMapping.getOrDefault("ignore_above", Integer.MAX_VALUE);
50+
var s = (String) value;
51+
if (s.isEmpty() || ignoreAbove < s.length()) {
52+
return List.of();
53+
}
54+
return List.of(QueryBuilders.termQuery(path, value));
55+
}
56+
}
57+
58+
class WildcardQueryGenerator implements LeafQueryGenerator {
59+
public List<QueryBuilder> generate(Map<String, Object> fieldMapping, String path, Object value) {
60+
var ignoreAbove = (Integer) fieldMapping.getOrDefault("ignore_above", Integer.MAX_VALUE);
61+
var s = (String) value;
62+
if (s.isEmpty() || ignoreAbove < s.length()) {
63+
return List.of();
64+
}
65+
return List.of(
66+
QueryBuilders.termQuery(path, value),
67+
QueryBuilders.wildcardQuery(path, value + "*")
68+
);
69+
}
70+
}
71+
72+
class TextQueryGenerator implements LeafQueryGenerator {
73+
public List<QueryBuilder> generate(Map<String, Object> fieldMapping, String path, Object value) {
74+
if (((String) value).isEmpty()) {
75+
return List.of();
76+
}
77+
var results = new ArrayList<QueryBuilder>();
78+
results.add(QueryBuilders.matchQuery(path, value));
79+
var phraseQuery = buildPhraseQuery(path, value);
80+
if (phraseQuery != null) {
81+
results.add(phraseQuery);
82+
}
83+
return results;
84+
}
85+
86+
private static QueryBuilder buildPhraseQuery(String path, Object value) {
87+
String needle = (String) value;
88+
var tokens = Arrays.asList(needle.split("[^a-zA-Z0-9]"));
89+
if (tokens.isEmpty()) {
90+
return null;
91+
}
92+
93+
int low = ESTestCase.randomIntBetween(0, tokens.size() - 1);
94+
int hi = ESTestCase.randomIntBetween(low+1, tokens.size());
95+
var phrase = String.join(" ", tokens.subList(low, hi));
96+
return QueryBuilders.matchPhraseQuery(path, phrase);
97+
}
98+
}
99+
}

‎test/framework/src/main/java/org/elasticsearch/datageneration/queries/QueryGenerator.java‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
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", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.datageneration.queries;
11+
12+
import org.apache.lucene.search.join.ScoreMode;
13+
import org.elasticsearch.index.query.QueryBuilder;
14+
import org.elasticsearch.index.query.QueryBuilders;
15+
16+
import java.util.Arrays;
17+
import java.util.ArrayList;
18+
import java.util.List;
19+
import java.util.Map;
20+
21+
22+
public class QueryGenerator {
23+
24+
private final Map<String, Map<String, Object>> mappingLookup;
25+
private final Map<String, Object> mappingRaw;
26+
27+
public QueryGenerator(Map<String, Map<String, Object>> mappingLookup, Map<String, Object> mappingRaw) {
28+
this.mappingLookup = mappingLookup;
29+
this.mappingRaw = mappingRaw;
30+
}
31+
32+
public List<QueryBuilder> generateQueries(
33+
String path,
34+
Map<String, Object> mapping,
35+
Object value
36+
) {
37+
// This query generator cannot handle fields with periods in the name.
38+
if (path.equals("host.name")) {
39+
return List.of();
40+
}
41+
if (mapping == null || isEnabled(path) == false) {
42+
return List.of();
43+
}
44+
boolean isIndexed = (Boolean) mapping.getOrDefault("index", true);
45+
if (isIndexed == false) {
46+
return List.of();
47+
}
48+
var type = (String) mapping.get("type");
49+
var leafQueryGenerator = LeafQueryGenerator.buildForType(type);
50+
var leafQueries = leafQueryGenerator.generate(mapping, path, value);
51+
return leafQueries.stream().map(q -> wrapInNestedQuery(path, q)).toList();
52+
}
53+
54+
private QueryBuilder wrapInNestedQuery(String path, QueryBuilder leafQuery) {
55+
String[] parts = path.split("\\.");
56+
List<String> nestedPaths = getNestedPathPrefixes(parts);
57+
QueryBuilder query = leafQuery;
58+
for (String nestedPath : nestedPaths.reversed()) {
59+
query = QueryBuilders.nestedQuery(nestedPath, query, ScoreMode.Max);
60+
}
61+
return query;
62+
}
63+
64+
@SuppressWarnings("unchecked")
65+
private List<String> getNestedPathPrefixes(String[] path) {
66+
Map<String, Object> mapping = mappingRaw;
67+
mapping = (Map<String, Object>) mapping.get("_doc");
68+
mapping = (Map<String, Object>) mapping.get("properties");
69+
70+
var result = new ArrayList<String>();
71+
for (int i = 0; i < path.length - 1; i++) {
72+
var field = path[i];
73+
mapping = (Map<String, Object>) mapping.get(field);
74+
boolean nested = "nested".equals(mapping.get("type"));
75+
if (nested) {
76+
result.add(String.join(".", Arrays.copyOfRange(path, 0, i + 1)));
77+
}
78+
mapping = (Map<String, Object>) mapping.get("properties");
79+
}
80+
81+
mapping = (Map<String, Object>) mapping.get(path[path.length - 1]);
82+
assert mapping.containsKey("properties") == false;
83+
return result;
84+
}
85+
86+
/**
87+
* Traverse down mapping tree and check that all objects are enabled in path
88+
*/
89+
private boolean isEnabled(String path) {
90+
String[] parts = path.split("\\.");
91+
for (int i = 0; i < parts.length - 1; i++) {
92+
var pathToHere = String.join(".", Arrays.copyOfRange(parts, 0, i + 1));
93+
Map<String, Object> mapping = mappingLookup.get(pathToHere);
94+
95+
boolean enabled = true;
96+
if (mapping.containsKey("enabled") && mapping.get("enabled") instanceof Boolean) {
97+
enabled = (Boolean) mapping.get("enabled");
98+
}
99+
if (mapping.containsKey("enabled") && mapping.get("enabled") instanceof String) {
100+
enabled = Boolean.parseBoolean((String) mapping.get("enabled"));
101+
}
102+
103+
if (enabled == false) {
104+
return false;
105+
}
106+
}
107+
return true;
108+
}
109+
}

0 commit comments

Comments
 (0)