diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java index 95ff51b0b9..8d07252c40 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java @@ -10,6 +10,7 @@ import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_WORKER_THREAD_POOL_NAME; import static org.opensearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY; +import java.util.Locale; import java.util.Map; import java.util.Optional; import org.apache.calcite.rel.RelNode; @@ -17,7 +18,9 @@ import org.apache.calcite.sql.SqlIdentifier; import org.apache.calcite.sql.SqlJoin; import org.apache.calcite.sql.SqlNode; +import org.apache.calcite.sql.SqlOrderBy; import org.apache.calcite.sql.SqlSelect; +import org.apache.calcite.sql.parser.SqlParser; import org.apache.calcite.sql.util.SqlBasicVisitor; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -214,13 +217,14 @@ private void forwardClusterSetting( * node. Uses the context's parser which supports both PPL and SQL. */ private static Optional extractIndexName( - String query, QueryType queryType, UnifiedQueryContext context) { + String query, QueryType queryType, UnifiedQueryContext context) throws Exception { if (queryType == QueryType.PPL) { UnresolvedPlan unresolvedPlan = (UnresolvedPlan) context.getParser().parse(query); return Optional.ofNullable(unresolvedPlan.accept(new IndexNameExtractor(), null)); } - SqlNode sqlNode = (SqlNode) context.getParser().parse(query); - return Optional.ofNullable(extractTableNameFromSqlNode(sqlNode)); + SqlNode sqlNode = SqlParser.create(query).parseQuery(); + String tableName = extractTableNameFromSqlNode(sqlNode); + return Optional.ofNullable(tableName != null ? tableName.toLowerCase(Locale.ROOT) : null); } /** AST visitor that extracts the source index name from a Relation node (PPL path). */ @@ -235,6 +239,9 @@ public String visitRelation(Relation node, Void context) { private static class SqlTableNameExtractor extends SqlBasicVisitor { @Override public String visit(SqlCall call) { + if (call instanceof SqlOrderBy orderBy) { + return orderBy.query.accept(this); + } if (call instanceof SqlSelect select) { return select.getFrom().accept(this); } diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryActionTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryActionTest.java index c801266036..cb5df14505 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryActionTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryActionTest.java @@ -65,6 +65,29 @@ public void pluggableDataformatIndexRoutesToAnalytics() { action.isAnalyticsIndex("source = opensearch.parquet_logs | fields ts", QueryType.PPL)); } + @Test + public void sqlQueryRoutesToAnalyticsForCompositeIndex() { + registerIndex( + "parquet_logs", + Settings.builder() + .put(IndexSettings.PLUGGABLE_DATAFORMAT_ENABLED_SETTING.getKey(), true) + .put(IndexSettings.PLUGGABLE_DATAFORMAT_VALUE_SETTING.getKey(), "composite") + .build()); + + assertTrue(action.isAnalyticsIndex("SELECT ts, msg FROM parquet_logs", QueryType.SQL)); + assertTrue(action.isAnalyticsIndex("SELECT ts FROM parquet_logs LIMIT 10", QueryType.SQL)); + assertTrue( + action.isAnalyticsIndex( + "SELECT score, name FROM parquet_logs WHERE score > 90", QueryType.SQL)); + } + + @Test + public void sqlQueryRoutesToLuceneForNonCompositeIndex() { + registerIndex("plain_logs", Settings.EMPTY); + + assertFalse(action.isAnalyticsIndex("SELECT ts FROM plain_logs", QueryType.SQL)); + } + @Test public void pluggableEnabledButLuceneFormatRoutesToLucene() { registerIndex(