Skip to content

Commit

Permalink
fix getScriptFieldFromQuery can't get requestBuilder from QueryPlan
Browse files Browse the repository at this point in the history
Signed-off-by: Sean Kao <seankao@amazon.com>
  • Loading branch information
seankao-az committed Feb 28, 2025
1 parent fcd64f5 commit d96e854
Show file tree
Hide file tree
Showing 4 changed files with 12 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,6 @@ public Client getClient() {
*/
public QueryAction explain(QueryActionRequest queryActionRequest)
throws SqlParseException, SQLFeatureNotSupportedException, SQLFeatureDisabledException {
return OpenSearchActionFactory.create(client, queryActionRequest);
return OpenSearchActionFactory.create(client, queryActionRequest, false);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -63,16 +63,22 @@ public class OpenSearchActionFactory {

public static QueryAction create(Client client, String sql)
throws SqlParseException, SQLFeatureNotSupportedException, SQLFeatureDisabledException {
return create(client, new QueryActionRequest(sql, new ColumnTypeProvider(), Format.JDBC));
return create(client, sql, false);
}

public static QueryAction create(Client client, String sql, boolean bypassMigrateToQueryPlan)
throws SqlParseException, SQLFeatureNotSupportedException, SQLFeatureDisabledException {
return create(client, new QueryActionRequest(sql, new ColumnTypeProvider(), Format.JDBC), bypassMigrateToQueryPlan);
}

/**
* Create the compatible Query object based on the SQL query.
*
* @param request The SQL query.
* @param bypassMigrateToQueryPlan Avoid using QueryPlan.
* @return Query object.
*/
public static QueryAction create(Client client, QueryActionRequest request)
public static QueryAction create(Client client, QueryActionRequest request, boolean bypassMigrateToQueryPlan)
throws SqlParseException, SQLFeatureNotSupportedException, SQLFeatureDisabledException {
String sql = request.getSql();
// Remove line breaker anywhere and semicolon at the end
Expand Down Expand Up @@ -116,7 +122,7 @@ public static QueryAction create(Client client, QueryActionRequest request)
} else {
sqlExpr.accept(new TermFieldRewriter());
// migrate aggregation to query planner framework.
if (shouldMigrateToQueryPlan(sqlExpr)) {
if (!bypassMigrateToQueryPlan && shouldMigrateToQueryPlan(sqlExpr)) {
return new QueryPlanQueryAction(
new QueryPlanRequestBuilder(
new BindingTupleQueryPlanner(client, sqlExpr, request.getTypeProvider())));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -435,7 +435,7 @@ private String translate(String sql, JSONObject jsonRequest)
CheckScriptContents.stubMockClient(mockClient);
QueryAction queryAction =
OpenSearchActionFactory.create(
mockClient, new QueryActionRequest(sql, columnTypeProvider, Format.JDBC));
mockClient, new QueryActionRequest(sql, columnTypeProvider, Format.JDBC), false);

SqlRequest sqlRequest = new SqlRequest(sql, jsonRequest);
queryAction.setSqlRequest(sqlRequest);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ public static ScriptField getScriptFieldFromQuery(String query) {
try {
Client mockClient = mock(Client.class);
stubMockClient(mockClient);
QueryAction queryAction = OpenSearchActionFactory.create(mockClient, query);
QueryAction queryAction = OpenSearchActionFactory.create(mockClient, query, true);
SqlElasticRequestBuilder requestBuilder = queryAction.explain();

SearchRequestBuilder request = (SearchRequestBuilder) requestBuilder.getBuilder();
Expand Down

0 comments on commit d96e854

Please sign in to comment.