Skip to content

Commit 327b2a0

Browse files
Mikep86Mikep86
committed
Update Test Framework To Handle Query Rewrites That Rely on Non-Null Searchers (elastic#129160)
1 parent 3e05896 commit 327b2a0

File tree

1 file changed

+144
-68
lines changed

1 file changed

+144
-68
lines changed

‎test/framework/src/main/java/org/elasticsearch/test/AbstractQueryTestCase.java‎

Lines changed: 144 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,14 @@
99

1010
package org.elasticsearch.test;
1111

12+
import org.apache.lucene.index.IndexReader;
1213
import org.apache.lucene.search.BoostQuery;
14+
import org.apache.lucene.search.IndexSearcher;
1315
import org.apache.lucene.search.MatchNoDocsQuery;
1416
import org.apache.lucene.search.Query;
1517
import org.apache.lucene.search.TermQuery;
18+
import org.apache.lucene.store.Directory;
19+
import org.apache.lucene.tests.index.RandomIndexWriter;
1620
import org.elasticsearch.ElasticsearchParseException;
1721
import org.elasticsearch.TransportVersion;
1822
import org.elasticsearch.action.support.PlainActionFuture;
@@ -45,6 +49,7 @@
4549
import org.elasticsearch.xcontent.json.JsonStringEncoder;
4650
import org.elasticsearch.xcontent.json.JsonXContent;
4751

52+
import java.io.Closeable;
4853
import java.io.IOException;
4954
import java.time.Instant;
5055
import java.time.ZoneOffset;
@@ -453,82 +458,87 @@ protected boolean builderGeneratesCacheableQueries() {
453458
return true;
454459
}
455460

461+
protected IndexReaderManager getIndexReaderManager() {
462+
return NullIndexReaderManager.INSTANCE;
463+
}
464+
456465
/**
457466
* Test creates the {@link Query} from the {@link QueryBuilder} under test and delegates the
458467
* assertions being made on the result to the implementing subclass.
459468
*/
460469
public void testToQuery() throws IOException {
461470
for (int runs = 0; runs < NUMBER_OF_TESTQUERIES; runs++) {
462-
SearchExecutionContext context = createSearchExecutionContext();
463-
assert context.isCacheable();
464-
context.setAllowUnmappedFields(true);
465-
QB firstQuery = createTestQueryBuilder();
466-
QB controlQuery = copyQuery(firstQuery);
467-
/* we use a private rewrite context here since we want the most realistic way of asserting that we are cacheable or not.
468-
* We do it this way in SearchService where
469-
* we first rewrite the query with a private context, then reset the context and then build the actual lucene query*/
470-
QueryBuilder rewritten = rewriteQuery(firstQuery, createQueryRewriteContext(), new SearchExecutionContext(context));
471-
Query firstLuceneQuery = rewritten.toQuery(context);
472-
assertNotNull("toQuery should not return null", firstLuceneQuery);
473-
assertLuceneQuery(firstQuery, firstLuceneQuery, context);
474-
// remove after assertLuceneQuery since the assertLuceneQuery impl might access the context as well
475-
assertEquals(
476-
"query is not equal to its copy after calling toQuery, firstQuery: " + firstQuery + ", secondQuery: " + controlQuery,
477-
firstQuery,
478-
controlQuery
479-
);
480-
assertEquals(
481-
"equals is not symmetric after calling toQuery, firstQuery: " + firstQuery + ", secondQuery: " + controlQuery,
482-
controlQuery,
483-
firstQuery
484-
);
485-
assertThat(
486-
"query copy's hashcode is different from original hashcode after calling toQuery, firstQuery: "
487-
+ firstQuery
488-
+ ", secondQuery: "
489-
+ controlQuery,
490-
controlQuery.hashCode(),
491-
equalTo(firstQuery.hashCode())
492-
);
493-
494-
QB secondQuery = copyQuery(firstQuery);
495-
// query _name never should affect the result of toQuery, we randomly set it to make sure
496-
if (randomBoolean()) {
497-
secondQuery.queryName(
498-
secondQuery.queryName() == null
499-
? randomAlphaOfLengthBetween(1, 30)
500-
: secondQuery.queryName() + randomAlphaOfLengthBetween(1, 10)
501-
);
502-
}
503-
context = new SearchExecutionContext(context);
504-
Query secondLuceneQuery = rewriteQuery(secondQuery, createQueryRewriteContext(), new SearchExecutionContext(context)).toQuery(
505-
context
506-
);
507-
assertNotNull("toQuery should not return null", secondLuceneQuery);
508-
assertLuceneQuery(secondQuery, secondLuceneQuery, context);
509-
510-
if (builderGeneratesCacheableQueries()) {
471+
try (IndexReaderManager irm = getIndexReaderManager()) {
472+
SearchExecutionContext context = createSearchExecutionContext(irm.getIndexSearcher());
473+
assert context.isCacheable();
474+
context.setAllowUnmappedFields(true);
475+
QB firstQuery = createTestQueryBuilder();
476+
QB controlQuery = copyQuery(firstQuery);
477+
/* we use a private rewrite context here since we want the most realistic way of asserting that we are cacheable or not.
478+
* We do it this way in SearchService where
479+
* we first rewrite the query with a private context, then reset the context and then build the actual lucene query*/
480+
QueryBuilder rewritten = rewriteQuery(firstQuery, createQueryRewriteContext(), new SearchExecutionContext(context));
481+
Query firstLuceneQuery = rewritten.toQuery(context);
482+
assertNotNull("toQuery should not return null", firstLuceneQuery);
483+
assertLuceneQuery(firstQuery, firstLuceneQuery, context);
484+
// remove after assertLuceneQuery since the assertLuceneQuery impl might access the context as well
511485
assertEquals(
512-
"two equivalent query builders lead to different lucene queries hashcode",
513-
secondLuceneQuery.hashCode(),
514-
firstLuceneQuery.hashCode()
486+
"query is not equal to its copy after calling toQuery, firstQuery: " + firstQuery + ", secondQuery: " + controlQuery,
487+
firstQuery,
488+
controlQuery
515489
);
516490
assertEquals(
517-
"two equivalent query builders lead to different lucene queries",
518-
rewrite(secondLuceneQuery),
519-
rewrite(firstLuceneQuery)
491+
"equals is not symmetric after calling toQuery, firstQuery: " + firstQuery + ", secondQuery: " + controlQuery,
492+
controlQuery,
493+
firstQuery
494+
);
495+
assertThat(
496+
"query copy's hashcode is different from original hashcode after calling toQuery, firstQuery: "
497+
+ firstQuery
498+
+ ", secondQuery: "
499+
+ controlQuery,
500+
controlQuery.hashCode(),
501+
equalTo(firstQuery.hashCode())
520502
);
521-
}
522503

523-
if (supportsBoost() && firstLuceneQuery instanceof MatchNoDocsQuery == false) {
524-
secondQuery.boost(firstQuery.boost() + 1f + randomFloat());
525-
Query thirdLuceneQuery = rewriteQuery(secondQuery, createQueryRewriteContext(), new SearchExecutionContext(context))
504+
QB secondQuery = copyQuery(firstQuery);
505+
// query _name never should affect the result of toQuery, we randomly set it to make sure
506+
if (randomBoolean()) {
507+
secondQuery.queryName(
508+
secondQuery.queryName() == null
509+
? randomAlphaOfLengthBetween(1, 30)
510+
: secondQuery.queryName() + randomAlphaOfLengthBetween(1, 10)
511+
);
512+
}
513+
context = new SearchExecutionContext(context);
514+
Query secondLuceneQuery = rewriteQuery(secondQuery, createQueryRewriteContext(), new SearchExecutionContext(context))
526515
.toQuery(context);
527-
assertNotEquals(
528-
"modifying the boost doesn't affect the corresponding lucene query",
529-
rewrite(firstLuceneQuery),
530-
rewrite(thirdLuceneQuery)
531-
);
516+
assertNotNull("toQuery should not return null", secondLuceneQuery);
517+
assertLuceneQuery(secondQuery, secondLuceneQuery, context);
518+
519+
if (builderGeneratesCacheableQueries()) {
520+
assertEquals(
521+
"two equivalent query builders lead to different lucene queries hashcode",
522+
secondLuceneQuery.hashCode(),
523+
firstLuceneQuery.hashCode()
524+
);
525+
assertEquals(
526+
"two equivalent query builders lead to different lucene queries",
527+
rewrite(secondLuceneQuery),
528+
rewrite(firstLuceneQuery)
529+
);
530+
}
531+
532+
if (supportsBoost() && firstLuceneQuery instanceof MatchNoDocsQuery == false) {
533+
secondQuery.boost(firstQuery.boost() + 1f + randomFloat());
534+
Query thirdLuceneQuery = rewriteQuery(secondQuery, createQueryRewriteContext(), new SearchExecutionContext(context))
535+
.toQuery(context);
536+
assertNotEquals(
537+
"modifying the boost doesn't affect the corresponding lucene query",
538+
rewrite(firstLuceneQuery),
539+
rewrite(thirdLuceneQuery)
540+
);
541+
}
532542
}
533543
}
534544
}
@@ -938,9 +948,75 @@ public boolean isTextField(String fieldName) {
938948
*/
939949
public void testCacheability() throws IOException {
940950
QB queryBuilder = createTestQueryBuilder();
941-
SearchExecutionContext context = createSearchExecutionContext();
942-
QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, createQueryRewriteContext(), new SearchExecutionContext(context));
943-
assertNotNull(rewriteQuery.toQuery(context));
944-
assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
951+
try (IndexReaderManager irm = getIndexReaderManager()) {
952+
SearchExecutionContext context = createSearchExecutionContext(irm.getIndexSearcher());
953+
QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, createQueryRewriteContext(), new SearchExecutionContext(context));
954+
assertNotNull(rewriteQuery.toQuery(context));
955+
assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
956+
}
957+
}
958+
959+
public static class IndexReaderManager implements Closeable {
960+
private final Directory directory;
961+
private RandomIndexWriter indexWriter;
962+
private IndexReader indexReader;
963+
private IndexSearcher indexSearcher;
964+
965+
public IndexReaderManager() {
966+
this.directory = newDirectory();
967+
}
968+
969+
private IndexReaderManager(Directory directory) {
970+
this.directory = directory;
971+
}
972+
973+
public IndexReader getIndexReader() throws IOException {
974+
if (indexReader == null) {
975+
indexWriter = new RandomIndexWriter(random(), directory);
976+
initIndexWriter(indexWriter);
977+
indexReader = indexWriter.getReader();
978+
}
979+
return indexReader;
980+
}
981+
982+
public IndexSearcher getIndexSearcher() throws IOException {
983+
if (indexSearcher == null) {
984+
indexSearcher = newSearcher(getIndexReader());
985+
}
986+
return indexSearcher;
987+
}
988+
989+
@Override
990+
public void close() throws IOException {
991+
if (indexReader != null) {
992+
indexReader.close();
993+
}
994+
if (indexWriter != null) {
995+
indexWriter.close();
996+
}
997+
if (directory != null) {
998+
directory.close();
999+
}
1000+
}
1001+
1002+
protected void initIndexWriter(RandomIndexWriter indexWriter) {}
1003+
}
1004+
1005+
public static class NullIndexReaderManager extends IndexReaderManager {
1006+
public static final NullIndexReaderManager INSTANCE = new NullIndexReaderManager();
1007+
1008+
public NullIndexReaderManager() {
1009+
super(null);
1010+
}
1011+
1012+
@Override
1013+
public IndexReader getIndexReader() {
1014+
return null;
1015+
}
1016+
1017+
@Override
1018+
public IndexSearcher getIndexSearcher() {
1019+
return null;
1020+
}
9451021
}
9461022
}

0 commit comments

Comments
 (0)