Skip to content

Commit 4945d32

Browse files
ahkcsahkcs
committed
Integrate SQL REST endpoint with analytics engine path
Add SQL query routing through the analytics engine for Parquet-backed indices. SQL queries targeting "parquet_*" indices are routed to RestUnifiedQueryAction via the unified query pipeline (Calcite SQL parser -> UnifiedQueryPlanner -> AnalyticsExecutionEngine). Changes: - Add SqlNode table extraction in RestUnifiedQueryAction.extractIndexName() to support SQL query routing (handles SqlSelect -> SqlIdentifier) - Add executeSql() and explainSql() methods in RestUnifiedQueryAction for SQL queries (parallel to existing PPL execute/explain) - Add analytics routing in RestSqlAction via optional BiFunction router that checks isAnalyticsIndex before delegating to SQLService - Wire the router in SQLPlugin.createSqlAnalyticsRouter() - Non-analytics SQL queries fall through to the existing V2 engine - Add AnalyticsSQLIT integration tests: SELECT *, column projection, explain, non-parquet fallback, syntax error handling Resolves: #5248 Signed-off-by: Kai Huang <ahkcs@amazon.com>
1 parent 24dd79c commit 4945d32

9 files changed

Lines changed: 415 additions & 3 deletions

File tree

‎integ-test/src/test/java/org/opensearch/sql/sql/AnalyticsSQLExplainIT.java‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.sql;
7+
8+
import static org.opensearch.sql.legacy.TestUtils.getResponseBody;
9+
import static org.opensearch.sql.legacy.TestUtils.isIndexExist;
10+
import static org.opensearch.sql.util.MatcherUtils.assertJsonEqualsIgnoreId;
11+
12+
import com.google.common.io.Resources;
13+
import java.io.IOException;
14+
import java.net.URI;
15+
import java.nio.file.Files;
16+
import java.nio.file.Paths;
17+
import org.junit.Test;
18+
import org.opensearch.client.Request;
19+
import org.opensearch.client.RequestOptions;
20+
import org.opensearch.client.Response;
21+
import org.opensearch.sql.legacy.SQLIntegTestCase;
22+
23+
/**
24+
* Explain integration tests for SQL queries routed through the analytics engine path (Project
25+
* Mustang). Validates that SQL queries targeting "parquet_*" indices produce correct logical plans
26+
* via the _plugins/_sql/_explain endpoint.
27+
*
28+
* <p>Expected output files are in resources/expectedOutput/analytics_sql/. Each test compares the
29+
* explain JSON output against its expected file.
30+
*/
31+
@SuppressWarnings("deprecation") // assertJsonEqualsIgnoreId is correct for JSON explain response
32+
public class AnalyticsSQLExplainIT extends SQLIntegTestCase {
33+
34+
@Override
35+
protected void init() throws Exception {
36+
if (!isIndexExist(client(), "parquet_logs")) {
37+
Request request = new Request("PUT", "/parquet_logs");
38+
request.setJsonEntity(
39+
"{"
40+
+ "\"mappings\": {"
41+
+ " \"properties\": {"
42+
+ " \"ts\": {\"type\": \"date\"},"
43+
+ " \"status\": {\"type\": \"integer\"},"
44+
+ " \"message\": {\"type\": \"keyword\"},"
45+
+ " \"ip_addr\": {\"type\": \"keyword\"}"
46+
+ " }"
47+
+ "}"
48+
+ "}");
49+
client().performRequest(request);
50+
}
51+
}
52+
53+
private String explainSqlQuery(String sql) throws IOException {
54+
Request request = new Request("POST", "/_plugins/_sql/_explain");
55+
request.setJsonEntity("{\"query\": \"" + sql + "\"}");
56+
request.setOptions(RequestOptions.DEFAULT);
57+
Response response = client().performRequest(request);
58+
return getResponseBody(response);
59+
}
60+
61+
private static String loadExpectedJson(String fileName) {
62+
return loadFromFile("expectedOutput/analytics_sql/" + fileName);
63+
}
64+
65+
private static String loadFromFile(String filename) {
66+
try {
67+
URI uri = Resources.getResource(filename).toURI();
68+
return new String(Files.readAllBytes(Paths.get(uri)));
69+
} catch (Exception e) {
70+
throw new RuntimeException(e);
71+
}
72+
}
73+
74+
/**
75+
* Normalize \r\n to \n for Windows compatibility. Calcite's RelOptUtil.toString() uses the system
76+
* line separator, which is \r\n on Windows.
77+
*/
78+
private static String normalizeLf(String s) {
79+
return s.replace("\r\n", "\n");
80+
}
81+
82+
@Test
83+
public void testExplainSelectStar() throws IOException {
84+
assertJsonEqualsIgnoreId(
85+
loadExpectedJson("explain_select_star.json"),
86+
normalizeLf(explainSqlQuery("SELECT * FROM parquet_logs")));
87+
}
88+
89+
@Test
90+
public void testExplainSelectColumns() throws IOException {
91+
assertJsonEqualsIgnoreId(
92+
loadExpectedJson("explain_select_columns.json"),
93+
normalizeLf(explainSqlQuery("SELECT ts, status FROM parquet_logs")));
94+
}
95+
96+
@Test
97+
public void testExplainSelectWithWhere() throws IOException {
98+
assertJsonEqualsIgnoreId(
99+
loadExpectedJson("explain_select_where.json"),
100+
normalizeLf(explainSqlQuery("SELECT ts, message FROM parquet_logs WHERE status = 200")));
101+
}
102+
}

‎integ-test/src/test/java/org/opensearch/sql/sql/AnalyticsSQLIT.java‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.sql;
7+
8+
import static org.opensearch.sql.legacy.TestUtils.getResponseBody;
9+
import static org.opensearch.sql.legacy.TestUtils.isIndexExist;
10+
import static org.opensearch.sql.util.MatcherUtils.rows;
11+
import static org.opensearch.sql.util.MatcherUtils.schema;
12+
import static org.opensearch.sql.util.MatcherUtils.verifyDataRows;
13+
import static org.opensearch.sql.util.MatcherUtils.verifySchema;
14+
15+
import java.io.IOException;
16+
import org.json.JSONObject;
17+
import org.junit.Test;
18+
import org.opensearch.client.Request;
19+
import org.opensearch.client.RequestOptions;
20+
import org.opensearch.client.Response;
21+
import org.opensearch.client.ResponseException;
22+
import org.opensearch.sql.legacy.SQLIntegTestCase;
23+
24+
/**
25+
* Integration tests for SQL queries routed through the analytics engine path (Project Mustang).
26+
* Queries targeting "parquet_*" indices are routed to {@code RestUnifiedQueryAction} which uses
27+
* {@code AnalyticsExecutionEngine} with a stub {@code QueryPlanExecutor}.
28+
*
29+
* <p>The stub executor returns rows in a fixed order [ts, status, message, ip_addr] regardless of
30+
* the plan. The schema from OpenSearchSchemaBuilder is alphabetical [ip_addr, message, status, ts].
31+
* AnalyticsExecutionEngine maps values by position, so the data values appear mismatched. This is
32+
* expected; the real analytics engine will evaluate the plan correctly.
33+
*/
34+
public class AnalyticsSQLIT extends SQLIntegTestCase {
35+
36+
@Override
37+
protected void init() throws Exception {
38+
createParquetLogsIndex();
39+
}
40+
41+
private void createParquetLogsIndex() throws IOException {
42+
if (isIndexExist(client(), "parquet_logs")) {
43+
return;
44+
}
45+
Request request = new Request("PUT", "/parquet_logs");
46+
request.setJsonEntity(
47+
"{"
48+
+ "\"mappings\": {"
49+
+ " \"properties\": {"
50+
+ " \"ts\": {\"type\": \"date\"},"
51+
+ " \"status\": {\"type\": \"integer\"},"
52+
+ " \"message\": {\"type\": \"keyword\"},"
53+
+ " \"ip_addr\": {\"type\": \"keyword\"}"
54+
+ " }"
55+
+ "}"
56+
+ "}");
57+
client().performRequest(request);
58+
}
59+
60+
private JSONObject executeSqlQuery(String sql) throws IOException {
61+
Request request = new Request("POST", "/_plugins/_sql");
62+
request.setJsonEntity("{\"query\": \"" + sql + "\"}");
63+
request.setOptions(RequestOptions.DEFAULT);
64+
Response response = client().performRequest(request);
65+
return new JSONObject(getResponseBody(response));
66+
}
67+
68+
@Test
69+
public void testSelectStarSchemaAndData() throws IOException {
70+
JSONObject result = executeSqlQuery("SELECT * FROM parquet_logs");
71+
verifySchema(
72+
result,
73+
schema("ip_addr", "string"),
74+
schema("message", "string"),
75+
schema("status", "integer"),
76+
schema("ts", "timestamp"));
77+
// Stub returns [ts, status, message, ip_addr] per row, mapped by position to
78+
// [ip_addr, message, status, ts] schema. Values appear mismatched — expected with stub.
79+
verifyDataRows(
80+
result,
81+
rows("2024-01-15 10:30:00", 200, "Request completed", "192.168.1.1"),
82+
rows("2024-01-15 10:31:00", 200, "Health check OK", "192.168.1.2"),
83+
rows("2024-01-15 10:32:00", 500, "Internal server error", "192.168.1.3"));
84+
}
85+
86+
@Test
87+
public void testSelectSpecificColumns() throws IOException {
88+
JSONObject result = executeSqlQuery("SELECT status, message FROM parquet_logs");
89+
verifySchema(result, schema("status", "integer"), schema("message", "string"));
90+
verifyDataRows(
91+
result,
92+
rows("2024-01-15 10:30:00", 200),
93+
rows("2024-01-15 10:31:00", 200),
94+
rows("2024-01-15 10:32:00", 500));
95+
}
96+
97+
@Test(expected = ResponseException.class)
98+
public void testSyntaxError() throws IOException {
99+
executeSqlQuery("SELEC * FROM parquet_logs");
100+
}
101+
}

‎integ-test/src/test/resources/expectedOutput/analytics_sql/explain_select_columns.json‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
{
2+
"calcite": {
3+
"logical": "LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n LogicalProject(ts=[$3], status=[$2])\n LogicalTableScan(table=[[opensearch, parquet_logs]])\n"
4+
}
5+
}

‎integ-test/src/test/resources/expectedOutput/analytics_sql/explain_select_star.json‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
{
2+
"calcite": {
3+
"logical": "LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n LogicalProject(ip_addr=[$0], message=[$1], status=[$2], ts=[$3])\n LogicalTableScan(table=[[opensearch, parquet_logs]])\n"
4+
}
5+
}

‎integ-test/src/test/resources/expectedOutput/analytics_sql/explain_select_where.json‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
{
2+
"calcite": {
3+
"logical": "LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n LogicalProject(ts=[$3], message=[$1])\n LogicalFilter(condition=[=($2, 200)])\n LogicalTableScan(table=[[opensearch, parquet_logs]])\n"
4+
}
5+
}

‎legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java‎

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.util.Map;
2020
import java.util.Optional;
2121
import java.util.Set;
22+
import java.util.function.BiFunction;
2223
import java.util.function.Predicate;
2324
import java.util.regex.Pattern;
2425
import org.apache.logging.log4j.LogManager;
@@ -83,10 +84,25 @@ public class RestSqlAction extends BaseRestHandler {
8384
/** New SQL query request handler. */
8485
private final RestSQLQueryAction newSqlQueryHandler;
8586

87+
/**
88+
* Optional analytics router. If set, it's called before the normal SQL engine. Accepts the
89+
* request and channel, returns {@code true} if it handled the request, {@code false} to fall
90+
* through to normal SQL engine.
91+
*/
92+
private final BiFunction<SQLQueryRequest, RestChannel, Boolean> analyticsRouter;
93+
8694
public RestSqlAction(Settings settings, Injector injector) {
95+
this(settings, injector, null);
96+
}
97+
98+
public RestSqlAction(
99+
Settings settings,
100+
Injector injector,
101+
BiFunction<SQLQueryRequest, RestChannel, Boolean> analyticsRouter) {
87102
super();
88103
this.allowExplicitIndex = MULTI_ALLOW_EXPLICIT_INDEX.get(settings);
89104
this.newSqlQueryHandler = new RestSQLQueryAction(injector);
105+
this.analyticsRouter = analyticsRouter;
90106
}
91107

92108
@Override
@@ -134,14 +150,44 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
134150

135151
Format format = SqlRequestParam.getFormat(request.params());
136152

137-
// Route request to new query engine if it's supported already
138153
SQLQueryRequest newSqlRequest =
139154
new SQLQueryRequest(
140155
sqlRequest.getJsonContent(),
141156
sqlRequest.getSql(),
142157
request.path(),
143158
request.params(),
144159
sqlRequest.cursor());
160+
161+
// Route to analytics engine for non-Lucene (e.g., Parquet-backed) indices.
162+
// The router returns true and sends the response directly if it handled the request.
163+
if (analyticsRouter != null) {
164+
final SQLQueryRequest finalRequest = newSqlRequest;
165+
return channel -> {
166+
if (!analyticsRouter.apply(finalRequest, channel)) {
167+
// Not an analytics query — delegate to normal SQL engine
168+
try {
169+
newSqlQueryHandler
170+
.prepareRequest(
171+
finalRequest,
172+
(ch, ex) -> {
173+
try {
174+
Format fmt = SqlRequestParam.getFormat(request.params());
175+
QueryAction qa = explainRequest(client, sqlRequest, fmt);
176+
executeSqlRequest(request, qa, client, ch);
177+
} catch (Exception e) {
178+
handleException(ch, e);
179+
}
180+
},
181+
this::handleException)
182+
.accept(channel);
183+
} catch (Exception e) {
184+
handleException(channel, e);
185+
}
186+
}
187+
};
188+
}
189+
190+
// Route request to new query engine if it's supported already
145191
return newSqlQueryHandler.prepareRequest(
146192
newSqlRequest,
147193
(restChannel, exception) -> {

‎plugin/build.gradle‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -348,6 +348,7 @@ def getJobSchedulerPlugin() {
348348

349349
testClusters.integTest {
350350
plugin(getJobSchedulerPlugin())
351+
plugin provider { (RegularFile) (() -> file("${rootDir}/libs/analytics-engine-3.6.0-SNAPSHOT.zip")) }
351352
plugin(project.tasks.bundlePlugin.archiveFile)
352353
testDistribution = "ARCHIVE"
353354

0 commit comments

Comments
 (0)