Skip to content

Commit b95d48d

Browse files
ahkcsahkcs
committed
Unify execute and explain methods for both PPL and SQL paths
Signed-off-by: Kai Huang <ahkcs@amazon.com>
1 parent b4f2500 commit b95d48d

3 files changed

Lines changed: 38 additions & 96 deletions

File tree

‎plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
import org.opensearch.script.ScriptContext;
6161
import org.opensearch.script.ScriptEngine;
6262
import org.opensearch.script.ScriptService;
63+
import org.opensearch.sql.ast.statement.ExplainMode;
6364
import org.opensearch.sql.common.response.ResponseListener;
6465
import org.opensearch.sql.datasource.DataSourceService;
6566
import org.opensearch.sql.datasources.auth.DataSourceUserAuthorizationHelper;
@@ -200,9 +201,10 @@ private BiFunction<SQLQueryRequest, RestChannel, Boolean> createSqlAnalyticsRout
200201
return false;
201202
}
202203
if (sqlRequest.isExplainRequest()) {
203-
unifiedQueryHandler.explainSql(
204+
unifiedQueryHandler.explain(
204205
sqlRequest.getQuery(),
205206
QueryType.SQL,
207+
ExplainMode.STANDARD,
206208
new ResponseListener<>() {
207209
@Override
208210
public void onResponse(ExplainResponse response) {
@@ -227,9 +229,10 @@ public void onFailure(Exception e) {
227229
}
228230
});
229231
} else {
230-
unifiedQueryHandler.executeSql(
232+
unifiedQueryHandler.execute(
231233
sqlRequest.getQuery(),
232234
QueryType.SQL,
235+
false,
233236
new ActionListener<>() {
234237
@Override
235238
public void onResponse(TransportPPLQueryResponse response) {

‎plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java‎

Lines changed: 28 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@
3939
import org.opensearch.sql.executor.analytics.AnalyticsExecutionEngine;
4040
import org.opensearch.sql.lang.LangSpec;
4141
import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse;
42-
import org.opensearch.sql.ppl.domain.PPLQueryRequest;
4342
import org.opensearch.sql.protocol.response.QueryResult;
4443
import org.opensearch.sql.protocol.response.format.ResponseFormatter;
4544
import org.opensearch.sql.protocol.response.format.SimpleJsonResponseFormatter;
@@ -99,70 +98,56 @@ public boolean isAnalyticsIndex(String query, QueryType queryType) {
9998
public void execute(
10099
String query,
101100
QueryType queryType,
102-
PPLQueryRequest pplRequest,
101+
boolean profiling,
103102
ActionListener<TransportPPLQueryResponse> listener) {
104103
client
105104
.threadPool()
106105
.schedule(
107-
withCurrentContext(() -> doExecute(query, queryType, pplRequest, listener)),
106+
withCurrentContext(
107+
() -> {
108+
try (UnifiedQueryContext context = buildContext(queryType, profiling)) {
109+
UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context);
110+
RelNode plan = planner.plan(query);
111+
CalcitePlanContext planContext = context.getPlanContext();
112+
plan = addQuerySizeLimit(plan, planContext);
113+
analyticsEngine.execute(
114+
plan, planContext, createQueryListener(queryType, listener));
115+
} catch (Exception e) {
116+
listener.onFailure(e);
117+
}
118+
}),
108119
new TimeValue(0),
109120
SQL_WORKER_THREAD_POOL_NAME);
110121
}
111122

112123
/**
113124
* Explain a query through the unified query pipeline on the sql-worker thread pool. Returns
114-
* ExplainResponse via ResponseListener so the caller (TransportPPLQueryAction) can format it
115-
* using its own createExplainResponseListener.
125+
* ExplainResponse via ResponseListener so the caller can format it.
116126
*/
117127
public void explain(
118128
String query,
119129
QueryType queryType,
120-
PPLQueryRequest pplRequest,
130+
ExplainMode mode,
121131
ResponseListener<ExplainResponse> listener) {
122132
client
123133
.threadPool()
124134
.schedule(
125-
withCurrentContext(() -> doExplain(query, queryType, pplRequest, listener)),
135+
withCurrentContext(
136+
() -> {
137+
try (UnifiedQueryContext context = buildContext(queryType, false)) {
138+
UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context);
139+
RelNode plan = planner.plan(query);
140+
CalcitePlanContext planContext = context.getPlanContext();
141+
plan = addQuerySizeLimit(plan, planContext);
142+
analyticsEngine.explain(plan, mode, planContext, listener);
143+
} catch (Exception e) {
144+
listener.onFailure(e);
145+
}
146+
}),
126147
new TimeValue(0),
127148
SQL_WORKER_THREAD_POOL_NAME);
128149
}
129150

130-
private void doExecute(
131-
String query,
132-
QueryType queryType,
133-
PPLQueryRequest pplRequest,
134-
ActionListener<TransportPPLQueryResponse> listener) {
135-
try (UnifiedQueryContext context = buildContext(queryType, pplRequest.profile())) {
136-
UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context);
137-
RelNode plan = planner.plan(query);
138-
139-
CalcitePlanContext planContext = context.getPlanContext();
140-
plan = addQuerySizeLimit(plan, planContext);
141-
142-
analyticsEngine.execute(plan, planContext, createQueryListener(queryType, listener));
143-
} catch (Exception e) {
144-
listener.onFailure(e);
145-
}
146-
}
147-
148-
private void doExplain(
149-
String query,
150-
QueryType queryType,
151-
PPLQueryRequest pplRequest,
152-
ResponseListener<ExplainResponse> listener) {
153-
try (UnifiedQueryContext context = buildContext(queryType, pplRequest.profile())) {
154-
UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context);
155-
RelNode plan = planner.plan(query);
156-
157-
CalcitePlanContext planContext = context.getPlanContext();
158-
plan = addQuerySizeLimit(plan, planContext);
159-
160-
analyticsEngine.explain(plan, pplRequest.mode(), planContext, listener);
161-
} catch (Exception e) {
162-
listener.onFailure(e);
163-
}
164-
}
165-
166151
/**
167152
* Build a lightweight context for parsing only (index name extraction). Does not require cluster
168153
* state or catalog schema.
@@ -194,55 +179,6 @@ private static Optional<String> extractIndexName(
194179
return Optional.ofNullable(extractTableNameFromSqlNode(sqlNode));
195180
}
196181

197-
/**
198-
* Execute a SQL query through the unified query pipeline. Uses {@link
199-
* org.opensearch.sql.plugin.transport.TransportPPLQueryResponse} as the transport response type
200-
* since both PPL and SQL share the same JSON response format.
201-
*/
202-
public void executeSql(
203-
String query, QueryType queryType, ActionListener<TransportPPLQueryResponse> listener) {
204-
client
205-
.threadPool()
206-
.schedule(
207-
withCurrentContext(
208-
() -> {
209-
try (UnifiedQueryContext context = buildContext(queryType, false)) {
210-
UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context);
211-
RelNode plan = planner.plan(query);
212-
CalcitePlanContext planContext = context.getPlanContext();
213-
plan = addQuerySizeLimit(plan, planContext);
214-
analyticsEngine.execute(
215-
plan, planContext, createQueryListener(queryType, listener));
216-
} catch (Exception e) {
217-
listener.onFailure(e);
218-
}
219-
}),
220-
new TimeValue(0),
221-
SQL_WORKER_THREAD_POOL_NAME);
222-
}
223-
224-
/** Explain a SQL query through the unified query pipeline. */
225-
public void explainSql(
226-
String query, QueryType queryType, ResponseListener<ExplainResponse> listener) {
227-
client
228-
.threadPool()
229-
.schedule(
230-
withCurrentContext(
231-
() -> {
232-
try (UnifiedQueryContext context = buildContext(queryType, false)) {
233-
UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context);
234-
RelNode plan = planner.plan(query);
235-
CalcitePlanContext planContext = context.getPlanContext();
236-
plan = addQuerySizeLimit(plan, planContext);
237-
analyticsEngine.explain(plan, ExplainMode.STANDARD, planContext, listener);
238-
} catch (Exception e) {
239-
listener.onFailure(e);
240-
}
241-
}),
242-
new TimeValue(0),
243-
SQL_WORKER_THREAD_POOL_NAME);
244-
}
245-
246182
/** Extracts the table name from a Calcite SqlNode parse tree. */
247183
private static String extractTableNameFromSqlNode(SqlNode sqlNode) {
248184
if (sqlNode instanceof SqlSelect select) {

‎plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -136,11 +136,14 @@ protected void doExecute(
136136
unifiedQueryHandler.explain(
137137
transformedRequest.getRequest(),
138138
QueryType.PPL,
139-
transformedRequest,
139+
transformedRequest.mode(),
140140
createExplainResponseListener(transformedRequest, clearingListener));
141141
} else {
142142
unifiedQueryHandler.execute(
143-
transformedRequest.getRequest(), QueryType.PPL, transformedRequest, clearingListener);
143+
transformedRequest.getRequest(),
144+
QueryType.PPL,
145+
transformedRequest.profile(),
146+
clearingListener);
144147
}
145148
return;
146149
}

0 commit comments

Comments
 (0)