diff --git a/api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java b/api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java index 74800db4a31..34872e2cd2e 100644 --- a/api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java +++ b/api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java @@ -14,6 +14,7 @@ import org.apache.calcite.rel.RelRoot; import org.apache.calcite.rel.core.Sort; import org.apache.calcite.rel.logical.LogicalSort; +import org.apache.calcite.sql.SqlKind; import org.apache.calcite.sql.SqlNode; import org.apache.calcite.tools.Frameworks; import org.apache.calcite.tools.Planner; @@ -60,8 +61,7 @@ public UnifiedQueryPlanner(UnifiedQueryContext context) { public RelNode plan(String query) { try { return context.measure(ANALYZE, () -> strategy.plan(query)); - } catch (SyntaxCheckException e) { - // Re-throw syntax error without wrapping + } catch (SyntaxCheckException | UnsupportedOperationException e) { throw e; } catch (Exception e) { throw new IllegalStateException("Failed to plan query", e); @@ -82,6 +82,11 @@ private static class CalciteNativeStrategy implements PlanningStrategy { public RelNode plan(String query) throws Exception { try (Planner planner = Frameworks.getPlanner(context.getPlanContext().config)) { SqlNode parsed = planner.parse(query); + if (!parsed.isA(SqlKind.QUERY)) { + throw new UnsupportedOperationException( + "Only query statements are supported. Got: " + parsed.getKind()); + } + SqlNode rewritten = parsed.accept(NamedArgRewriter.INSTANCE); SqlNode validated = planner.validate(rewritten); RelRoot relRoot = planner.rel(validated); diff --git a/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlTest.java b/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlTest.java index 53accd49715..70079cc9272 100644 --- a/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlTest.java +++ b/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlTest.java @@ -8,6 +8,7 @@ import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertThrows; +import java.util.List; import java.util.Map; import org.apache.calcite.schema.Schema; import org.apache.calcite.schema.impl.AbstractSchema; @@ -211,4 +212,50 @@ public void testSqlQueryPlanningWithMultipleCatalogs() { public void testInvalidSqlThrowsException() { assertThrows(IllegalStateException.class, () -> planner.plan("SELECT FROM")); } + + @Test + public void testNonQueryStatementsBlockedByWhitelist() { + List.of( + """ + INSERT INTO catalog.employees (id, name, age, department) + VALUES (99, 'injected', 0, 'hacked')\ + """, + """ + DELETE FROM catalog.employees + WHERE age > 30\ + """, + """ + UPDATE catalog.employees + SET department = 'Fired' + WHERE age > 50\ + """, + """ + EXPLAIN PLAN FOR + SELECT * FROM catalog.employees\ + """, + """ + MERGE INTO catalog.employees AS t + USING (SELECT 99 AS id) AS s ON t.id = s.id + WHEN MATCHED THEN UPDATE SET name = 'hacked'\ + """) + .forEach( + sql -> + givenInvalidQuery(sql).assertErrorMessage("Only query statements are supported")); + } + + @Test + public void testNonQueryStatementsBlockedByParser() { + List.of( + """ + CREATE MATERIALIZED VIEW mv AS + SELECT department, count(*) + FROM catalog.employees + GROUP BY department\ + """, + """ + SHOW TABLES\ + """) + .forEach( + sql -> givenInvalidQuery(sql).assertErrorMessage("Incorrect syntax near the keyword")); + } } diff --git a/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerTest.java b/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerTest.java index 9ad7aa42155..41ed12670f8 100644 --- a/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerTest.java +++ b/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerTest.java @@ -106,7 +106,7 @@ public void testPPLQueryPlanningWithMultipleCatalogsAndDefaultNamespace() { assertNotNull("Plan should be created with multiple catalogs", plan); } - @Test(expected = IllegalStateException.class) + @Test(expected = UnsupportedOperationException.class) public void testUnsupportedStatementType() { planner.plan("explain source = catalog.employees"); // explain statement }