...21import com.alibaba.druid.sql.parser.Token;22import java.util.Arrays;23import java.util.List;24import java.util.function.Function;25import org.hamcrest.FeatureMatcher;26import org.hamcrest.Matcher;27import org.hamcrest.Matchers;28import org.junit.Test;29import org.mockito.Mockito;30import org.opensearch.action.search.SearchAction;31import org.opensearch.action.search.SearchRequest;32import org.opensearch.action.search.SearchRequestBuilder;33import org.opensearch.client.Client;34import org.opensearch.index.query.BoolQueryBuilder;35import org.opensearch.index.query.NestedQueryBuilder;36import org.opensearch.index.query.QueryBuilder;37import org.opensearch.search.builder.SearchSourceBuilder;38import org.opensearch.search.fetch.subphase.FetchSourceContext;39import org.opensearch.sql.legacy.domain.Select;40import org.opensearch.sql.legacy.exception.SqlParseException;41import org.opensearch.sql.legacy.parser.ElasticSqlExprParser;42import org.opensearch.sql.legacy.parser.SqlParser;43import org.opensearch.sql.legacy.query.DefaultQueryAction;44import org.opensearch.sql.legacy.query.maker.QueryMaker;45import org.opensearch.sql.legacy.rewriter.nestedfield.NestedFieldProjection;46import org.opensearch.sql.legacy.rewriter.nestedfield.NestedFieldRewriter;47import org.opensearch.sql.legacy.util.HasFieldWithValue;48public class NestedFieldProjectionTest {49 @Test50 public void regression() {51 assertThat(query("SELECT region FROM team"), is(anything()));52 assertThat(query("SELECT region FROM team WHERE nested(employees.age) = 30"), is(anything()));53 assertThat(query("SELECT * FROM team WHERE region = 'US'"), is(anything()));54 }55 @Test56 public void nestedFieldSelectAll() {57 assertThat(58 query("SELECT nested(employees.*) FROM team"),59 source(60 boolQuery(61 filter(62 boolQuery(63 must(64 nestedQuery(65 path("employees"),66 innerHits("employees.*")67 )68 )69 )70 )71 )72 )73 );74 }75 @Test76 public void nestedFieldInSelect() {77 assertThat(78 query("SELECT nested(employees.firstname) FROM team"),79 source(80 boolQuery(81 filter(82 boolQuery(83 must(84 nestedQuery(85 path("employees"),86 innerHits("employees.firstname")87 )88 )89 )90 )91 )92 )93 );94 }95 @Test96 public void regularAndNestedFieldInSelect() {97 assertThat(98 query("SELECT region, nested(employees.firstname) FROM team"),99 source(100 boolQuery(101 filter(102 boolQuery(103 must(104 nestedQuery(105 path("employees"),106 innerHits("employees.firstname")107 )108 )109 )110 )111 ),112 fetchSource("region")113 )114 );115 }116 /*117 // Should be integration test118 @Test119 public void nestedFieldInWhereSelectAll() {}120 */121 @Test122 public void nestedFieldInSelectAndWhere() {123 assertThat(124 query("SELECT nested(employees.firstname) " +125 " FROM team " +126 " WHERE nested(employees.age) = 30"),127 source(128 boolQuery(129 filter(130 boolQuery(131 must(132 nestedQuery(133 path("employees"),134 innerHits("employees.firstname")135 )136 )137 )138 )139 )140 )141 );142 }143 @Test144 public void regularAndNestedFieldInSelectAndWhere() {145 assertThat(146 query("SELECT region, nested(employees.firstname) " +147 " FROM team " +148 " WHERE nested(employees.age) = 30"),149 source(150 boolQuery(151 filter(152 boolQuery(153 must(154 nestedQuery(155 innerHits("employees.firstname")156 )157 )158 )159 )160 ),161 fetchSource("region")162 )163 );164 }165 @Test166 public void multipleSameNestedFields() {167 assertThat(168 query("SELECT nested(employees.firstname), nested(employees.lastname) " +169 " FROM team " +170 " WHERE nested(\"employees\", employees.age = 30 AND employees.firstname LIKE 'John')"),171 source(172 boolQuery(173 filter(174 boolQuery(175 must(176 nestedQuery(177 path("employees"),178 innerHits("employees.firstname", "employees.lastname")179 )180 )181 )182 )183 )184 )185 );186 }187 @Test188 public void multipleDifferentNestedFields() {189 assertThat(190 query("SELECT region, nested(employees.firstname), nested(manager.name) " +191 " FROM team " +192 " WHERE nested(employees.age) = 30 AND nested(manager.age) = 50"),193 source(194 boolQuery(195 filter(196 boolQuery(197 must(198 boolQuery(199 must(200 nestedQuery(201 path("employees"),202 innerHits("employees.firstname")203 ),204 nestedQuery(205 path("manager"),206 innerHits("manager.name")207 )208 )209 )210 )211 )212 )213 ),214 fetchSource("region")215 )216 );217 }218 @Test219 public void leftJoinWithSelectAll() {220 assertThat(221 query("SELECT * FROM team AS t LEFT JOIN t.projects AS p "),222 source(223 boolQuery(224 filter(225 boolQuery(226 should(227 boolQuery(228 mustNot(229 nestedQuery(230 path("projects")231 )232 )233 ),234 nestedQuery(235 path("projects"),236 innerHits("projects.*")237 )238 )239 )240 )241 )242 )243 );244 }245 @Test246 public void leftJoinWithSpecificFields() {247 assertThat(248 query("SELECT t.name, p.name, p.started_year FROM team AS t LEFT JOIN t.projects AS p "),249 source(250 boolQuery(251 filter(252 boolQuery(253 should(254 boolQuery(255 mustNot(256 nestedQuery(257 path("projects")258 )259 )260 ),261 nestedQuery(262 path("projects"),263 innerHits("projects.name", "projects.started_year")264 )265 )266 )267 )268 ),269 fetchSource("name")270 )271 );272 }273 private Matcher<SearchSourceBuilder> source(Matcher<QueryBuilder> queryMatcher) {274 return featureValueOf("query", queryMatcher, SearchSourceBuilder::query);275 }276 private Matcher<SearchSourceBuilder> source(Matcher<QueryBuilder> queryMatcher,277 Matcher<FetchSourceContext> fetchSourceMatcher) {278 return allOf(279 featureValueOf("query", queryMatcher, SearchSourceBuilder::query),280 featureValueOf("fetchSource", fetchSourceMatcher, SearchSourceBuilder::fetchSource)281 );282 }283 /** Asserting instanceOf and continue other chained matchers of subclass requires explicity cast */284 @SuppressWarnings("unchecked")285 private Matcher<QueryBuilder> boolQuery(Matcher<BoolQueryBuilder> matcher) {286 return (Matcher) allOf(instanceOf(BoolQueryBuilder.class), matcher);287 }288 @SafeVarargs289 @SuppressWarnings("unchecked")290 private final Matcher<QueryBuilder> nestedQuery(Matcher<NestedQueryBuilder>... matchers) {291 return (Matcher) both(is(Matchers.<NestedQueryBuilder>instanceOf(NestedQueryBuilder.class))).292 and(allOf(matchers));293 }294 @SafeVarargs295 private final FeatureMatcher<BoolQueryBuilder, List<QueryBuilder>> filter(Matcher<QueryBuilder>... matchers) {296 return hasClauses("filter", BoolQueryBuilder::filter, matchers);297 }298 @SafeVarargs299 private final FeatureMatcher<BoolQueryBuilder, List<QueryBuilder>> must(Matcher<QueryBuilder>... matchers) {300 return hasClauses("must", BoolQueryBuilder::must, matchers);301 }302 @SafeVarargs303 private final FeatureMatcher<BoolQueryBuilder, List<QueryBuilder>> mustNot(Matcher<QueryBuilder>... matchers) {304 return hasClauses("must_not", BoolQueryBuilder::mustNot, matchers);305 }306 @SafeVarargs307 private final FeatureMatcher<BoolQueryBuilder, List<QueryBuilder>> should(Matcher<QueryBuilder>... matchers) {308 return hasClauses("should", BoolQueryBuilder::should, matchers);309 }310 /** Hide contains() assertion to simplify */311 @SafeVarargs312 private final FeatureMatcher<BoolQueryBuilder, List<QueryBuilder>> hasClauses(String name,313 Function<BoolQueryBuilder, List<QueryBuilder>> func,314 Matcher<QueryBuilder>... matchers) {315 return new FeatureMatcher<BoolQueryBuilder, List<QueryBuilder>>(contains(matchers), name, name) {316 @Override317 protected List<QueryBuilder> featureValueOf(BoolQueryBuilder query) {318 return func.apply(query);319 }320 };321 }322 private Matcher<NestedQueryBuilder> path(String expected) {323 return HasFieldWithValue.hasFieldWithValue("path", "path", is(equalTo(expected)));324 }325 /** Skip intermediate property along the path. Hide arrayContaining assertion to simplify. */326 private FeatureMatcher<NestedQueryBuilder, String[]> innerHits(String... expected) {327 return featureValueOf("innerHits",328 arrayContaining(expected),329 (nestedQuery -> nestedQuery.innerHit().getFetchSourceContext().includes()));330 }331 @SuppressWarnings("unchecked")332 private Matcher<FetchSourceContext> fetchSource(String... expected) {333 if (expected.length == 0) {334 return anyOf(is(nullValue()),335 featureValueOf("includes", is(nullValue()), FetchSourceContext::includes),336 featureValueOf("includes", is(emptyArray()), FetchSourceContext::includes));337 }338 return featureValueOf("includes", contains(expected), fetchSource -> Arrays.asList(fetchSource.includes()));339 }340 private <T, U> FeatureMatcher<T, U> featureValueOf(String name, Matcher<U> subMatcher, Function<T, U> getter) {341 return new FeatureMatcher<T, U>(subMatcher, name, name) {342 @Override343 protected U featureValueOf(T actual) {344 return getter.apply(actual);345 }346 };347 }348 private SearchSourceBuilder query(String sql) {349 SQLQueryExpr expr = parseSql(sql);350 if (sql.contains("nested")) {351 return translate(expr).source();352 }353 expr = rewrite(expr);354 return translate(expr).source();355 }...