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 e1bb84778d..79abf5143e 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 @@ -13,12 +13,6 @@ import java.util.Map; import java.util.Optional; import org.apache.calcite.rel.RelNode; -import org.apache.calcite.sql.SqlCall; -import org.apache.calcite.sql.SqlIdentifier; -import org.apache.calcite.sql.SqlJoin; -import org.apache.calcite.sql.SqlNode; -import org.apache.calcite.sql.SqlSelect; -import org.apache.calcite.sql.util.SqlBasicVisitor; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.apache.logging.log4j.ThreadContext; @@ -213,12 +207,8 @@ private UnifiedQueryContext.Builder applyClusterOverrides(UnifiedQueryContext.Bu */ private static Optional extractIndexName( String query, QueryType queryType, UnifiedQueryContext context) { - 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)); + UnresolvedPlan unresolvedPlan = (UnresolvedPlan) context.getParser().parse(query); + return Optional.ofNullable(unresolvedPlan.accept(new IndexNameExtractor(), null)); } /** AST visitor that extracts the source index name from a Relation node (PPL path). */ @@ -229,29 +219,6 @@ public String visitRelation(Relation node, Void context) { } } - /** SqlNode visitor that extracts the source table name from a SQL parse tree. */ - private static class SqlTableNameExtractor extends SqlBasicVisitor { - @Override - public String visit(SqlCall call) { - if (call instanceof SqlSelect select) { - return select.getFrom().accept(this); - } - if (call instanceof SqlJoin join) { - return join.getLeft().accept(this); - } - return null; - } - - @Override - public String visit(SqlIdentifier id) { - return id.toString(); - } - } - - private static String extractTableNameFromSqlNode(SqlNode sqlNode) { - return sqlNode.accept(new SqlTableNameExtractor()); - } - private static RelNode addQuerySizeLimit(RelNode plan, CalcitePlanContext context) { return LogicalSystemLimit.create( LogicalSystemLimit.SystemLimitType.QUERY_SIZE_LIMIT, 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..a8ccf34fc2 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 @@ -89,10 +89,75 @@ public void missingIndexRoutesToLucene() { assertFalse(action.isAnalyticsIndex("source = does_not_exist | fields ts", QueryType.PPL)); } + @Test + public void sqlQueryRoutesToAnalyticsForPluggableIndex() { + 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 * FROM parquet_logs", QueryType.SQL)); + assertTrue( + action.isAnalyticsIndex( + "SELECT ts, level FROM parquet_logs WHERE level = 'ERROR'", QueryType.SQL)); + } + + @Test + public void sqlQueryWithSchemaRoutesToAnalytics() { + 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 * FROM opensearch.parquet_logs", QueryType.SQL)); + } + + @Test + public void sqlQueryRoutesToLuceneForNonPluggableIndex() { + registerIndex("plain_logs", Settings.EMPTY); + + assertFalse(action.isAnalyticsIndex("SELECT * FROM plain_logs", QueryType.SQL)); + } + + @Test + public void sqlQueryRoutesToLuceneForMissingIndex() { + assertFalse(action.isAnalyticsIndex("SELECT * FROM does_not_exist", QueryType.SQL)); + } + @Test public void nullAndEmptyQueriesRouteToLucene() { assertFalse(action.isAnalyticsIndex(null, QueryType.PPL)); assertFalse(action.isAnalyticsIndex("", QueryType.PPL)); + assertFalse(action.isAnalyticsIndex(null, QueryType.SQL)); + assertFalse(action.isAnalyticsIndex("", QueryType.SQL)); + } + + @Test + public void showStatementRoutesToLucene() { + registerIndex( + "parquet_logs", + Settings.builder() + .put(IndexSettings.PLUGGABLE_DATAFORMAT_ENABLED_SETTING.getKey(), true) + .put(IndexSettings.PLUGGABLE_DATAFORMAT_VALUE_SETTING.getKey(), "composite") + .build()); + + assertFalse(action.isAnalyticsIndex("SHOW TABLES LIKE 'parquet_logs'", QueryType.SQL)); + } + + @Test + public void describeStatementRoutesToLucene() { + registerIndex( + "parquet_logs", + Settings.builder() + .put(IndexSettings.PLUGGABLE_DATAFORMAT_ENABLED_SETTING.getKey(), true) + .put(IndexSettings.PLUGGABLE_DATAFORMAT_VALUE_SETTING.getKey(), "composite") + .build()); + + assertFalse(action.isAnalyticsIndex("DESCRIBE TABLES LIKE 'parquet_logs'", QueryType.SQL)); } private void registerIndex(String name, Settings settings) {