Improve exception handling in UnifiedQueryPlanner

api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java

@@ -19,10 +19,14 @@
 import org.apache.calcite.sql.util.SqlVisitor;
 import org.apache.calcite.tools.Frameworks;
 import org.apache.calcite.tools.Planner;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.opensearch.sql.api.parser.UnifiedQueryParser;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.calcite.CalciteRelNodeVisitor;
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
+import org.opensearch.sql.exception.QueryEngineException;
+import org.opensearch.sql.exception.SemanticCheckException;
 
 /**
  * {@code UnifiedQueryPlanner} provides a high-level API for parsing and analyzing queries using the
@@ -31,6 +35,8 @@
  */
 public class UnifiedQueryPlanner {
 
+  private static final Logger LOG = LogManager.getLogger(UnifiedQueryPlanner.class);
+
   /** Planning strategy selected at construction time based on query type. */
   private final PlanningStrategy strategy;
 
@@ -65,10 +71,21 @@ public RelNode plan(String query) {
             }
             return plan;
           });
-    } catch (SyntaxCheckException | UnsupportedOperationException e) {
+    } catch (SyntaxCheckException
+        | QueryEngineException
+        | UnsupportedOperationException
+        | IllegalArgumentException e) {
+      LOG.error("Failed to plan query: {}", e.getMessage());
       throw e;
+    } catch (AssertionError e) {
+      // Calcite throws assertion error directly when building bad RelNode
+      String message = "Failed to plan query: invalid plan structure";
+      LOG.error(message, e);
+      throw new SemanticCheckException(message, e);
     } catch (Exception e) {
-      throw new IllegalStateException("Failed to plan query", e);
+      String message = "Failed to plan query: unexpected error";
+      LOG.error(message, e);
+      throw new IllegalStateException(message, e);
     }
   }
 

api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerTest.java

@@ -14,6 +14,7 @@
 import org.apache.calcite.schema.impl.AbstractSchema;
 import org.junit.Test;
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
+import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.executor.QueryType;
 
 public class UnifiedQueryPlannerTest extends UnifiedQueryTestBase {
@@ -115,4 +116,32 @@ public void testUnsupportedStatementType() {
   public void testPlanPropagatingSyntaxCheckException() {
     planner.plan("source = catalog.employees | eval"); // Trigger syntax error from parser
   }
+
+  @Test
+  public void syntaxErrorIsRethrownAsSyntaxCheckException() {
+    givenInvalidQuery("INVALID +++")
+        .assertErrorType(SyntaxCheckException.class)
+        .assertErrorMessageContains("is not a valid term");
+  }
+
+  @Test
+  public void semanticErrorIsRethrownAsSemanticCheckException() {
+    givenInvalidQuery("source = catalog.employees | rename id* as x*y*")
+        .assertErrorType(SemanticCheckException.class)
+        .assertErrorMessageEquals("Source and target patterns have different wildcard counts");
+  }
+
+  @Test
+  public void assertionErrorIsWrappedAsSemanticCheckException() {
+    // Remove when the underlying Calcite assertion is fixed.
+    givenInvalidQuery(
+            """
+            source = catalog.employees
+            | eval ts = timestamp('2024-01-01')
+            | stats max(ts)
+            """)
+        .assertErrorType(SemanticCheckException.class)
+        .assertErrorMessageEquals("Failed to plan query: invalid plan structure")
+        .assertCauseType(AssertionError.class);
+  }
 }

api/src/testFixtures/java/org/opensearch/sql/api/QueryPlanAssertion.java

@@ -118,5 +118,40 @@ public QueryErrorAssert assertErrorMessage(String expected) {
           "Expected error to contain: " + expected + "\nActual: " + msg, msg.contains(expected));
       return this;
     }
+
+    /** Assert the outer error type matches the expected class. */
+    public QueryErrorAssert assertErrorType(Class<? extends Throwable> expected) {
+      assertTrue(
+          "Expected " + expected.getSimpleName() + " but got " + error.getClass().getSimpleName(),
+          expected.isInstance(error));
+      return this;
+    }
+
+    /** Assert the outer error message exactly matches expected. */
+    public QueryErrorAssert assertErrorMessageEquals(String expected) {
+      assertEquals(expected, error.getMessage());
+      return this;
+    }
+
+    /** Assert the outer error message contains the expected substring. */
+    public QueryErrorAssert assertErrorMessageContains(String substring) {
+      String msg = error.getMessage();
+      assertTrue(
+          "Expected message to contain '" + substring + "' but was: " + msg,
+          msg != null && msg.contains(substring));
+      return this;
+    }
+
+    /** Assert the immediate cause is of the expected type. */
+    public QueryErrorAssert assertCauseType(Class<? extends Throwable> expected) {
+      assertNotNull("Expected a cause but was null", error.getCause());
+      assertTrue(
+          "Expected cause "
+              + expected.getSimpleName()
+              + " but got "
+              + error.getCause().getClass().getSimpleName(),
+          expected.isInstance(error.getCause()));
+      return this;
+    }
   }
 }