Fix fallback error handling to show original Calcite error

core/src/main/java/org/opensearch/sql/executor/QueryService.java

@@ -64,6 +64,23 @@ public class QueryService {
   @Getter(lazy = true)
   private final CalciteRelNodeVisitor relNodeVisitor = new CalciteRelNodeVisitor(dataSourceService);
 
+  /** Helper: depending on the type of error, either re-raise or propagate to the listener. */
+  private void propagateCalciteError(Throwable t, ResponseListener<?> listener)
+      throws VirtualMachineError {
+    if (t instanceof Exception) {
+      listener.onFailure((Exception) t);
+    } else if (t instanceof ExceptionInInitializerError
+        && ((ExceptionInInitializerError) t).getException() instanceof Exception) {
+      listener.onFailure((Exception) ((ExceptionInInitializerError) t).getException());
+    } else if (t instanceof VirtualMachineError) {
+      // throw and fast fail the VM errors such as OOM (same with v2).
+      throw (VirtualMachineError) t;
+    } else {
+      // Calcite may throw AssertError during query execution.
+      listener.onFailure(new CalciteUnsupportedException(t.getMessage(), t));
+    }
+  }
+
   /** Execute the {@link UnresolvedPlan}, using {@link ResponseListener} to get response.<br> */
   public void execute(
       UnresolvedPlan plan,
@@ -112,18 +129,7 @@ public void executeWithCalcite(
               log.warn("Fallback to V2 query engine since got exception", t);
               executeWithLegacy(plan, queryType, listener, Optional.of(t));
             } else {
-              if (t instanceof Exception) {
-                listener.onFailure((Exception) t);
-              } else if (t instanceof ExceptionInInitializerError
-                  && ((ExceptionInInitializerError) t).getException() instanceof Exception) {
-                listener.onFailure((Exception) ((ExceptionInInitializerError) t).getException());
-              } else if (t instanceof VirtualMachineError) {
-                // throw and fast fail the VM errors such as OOM (same with v2).
-                throw t;
-              } else {
-                // Calcite may throw AssertError during query execution.
-                listener.onFailure(new CalciteUnsupportedException(t.getMessage(), t));
-              }
+              propagateCalciteError(t, listener);
             }
           }
         },
@@ -154,12 +160,7 @@ public void explainWithCalcite(
               log.warn("Fallback to V2 query engine since got exception", t);
               explainWithLegacy(plan, queryType, listener, mode, Optional.of(t));
             } else {
-              if (t instanceof Error) {
-                // Calcite may throw AssertError during query execution.
-                listener.onFailure(new CalciteUnsupportedException(t.getMessage(), t));
-              } else {
-                listener.onFailure((Exception) t);
-              }
+              propagateCalciteError(t, listener);
             }
           }
         },
@@ -174,11 +175,10 @@ public void executeWithLegacy(
     try {
       executePlan(analyze(plan, queryType), PlanContext.emptyPlanContext(), listener);
     } catch (Exception e) {
-      if (shouldUseCalcite(queryType) && isCalciteFallbackAllowed(null)) {
-        // if there is a failure thrown from Calcite and execution after fallback V2
-        // keeps failure, we should throw the failure from Calcite.
-        calciteFailure.ifPresentOrElse(
-            t -> listener.onFailure(new RuntimeException(t)), () -> listener.onFailure(e));
+      // if there is a failure thrown from Calcite and execution after fallback V2
+      // keeps failure, we should throw the failure from Calcite.
+      if (calciteFailure.isPresent()) {
+        propagateCalciteError(calciteFailure.get(), listener);
       } else {
         listener.onFailure(e);
       }
@@ -207,11 +207,10 @@ public void explainWithLegacy(
       }
       executionEngine.explain(plan(analyze(plan, queryType)), listener);
     } catch (Exception e) {
-      if (shouldUseCalcite(queryType) && isCalciteFallbackAllowed(null)) {
-        // if there is a failure thrown from Calcite and execution after fallback V2
-        // keeps failure, we should throw the failure from Calcite.
-        calciteFailure.ifPresentOrElse(
-            t -> listener.onFailure(new RuntimeException(t)), () -> listener.onFailure(e));
+      // if there is a failure thrown from Calcite and execution after fallback V2
+      // keeps failure, we should throw the failure from Calcite.
+      if (calciteFailure.isPresent()) {
+        propagateCalciteError(calciteFailure.get(), listener);
       } else {
         listener.onFailure(e);
       }

core/src/test/java/org/opensearch/sql/executor/QueryServiceTest.java

@@ -91,6 +91,124 @@ public void analyzeExceptionShouldBeCached() {
     queryService().analyzeFail().handledByOnFailure();
   }
 
+  @Test
+  public void testExecuteWithLegacyShouldReturnCalciteErrorWhenBothFail() {
+    UnsupportedOperationException calciteException =
+        new UnsupportedOperationException("Calcite error");
+    IllegalStateException v2Exception = new IllegalStateException("V2 error");
+
+    ResponseListener<ExecutionEngine.QueryResponse> responseListener =
+        new ResponseListener<>() {
+          @Override
+          public void onResponse(ExecutionEngine.QueryResponse pplQueryResponse) {
+            fail("Expected onFailure to be called");
+          }
+
+          @Override
+          public void onFailure(Exception e) {
+            // Should get the Calcite error directly (not wrapped), not the V2 error
+            assertNotNull(e);
+            assertTrue(e instanceof UnsupportedOperationException);
+            assertTrue(e.getMessage().contains("Calcite error"));
+          }
+        };
+
+    lenient().when(settings.getSettingValue(Key.CALCITE_ENGINE_ENABLED)).thenReturn(false);
+    lenient().when(analyzer.analyze(any(), any())).thenThrow(v2Exception);
+
+    QueryService service = new QueryService(analyzer, executionEngine, planner, null, settings);
+    service.executeWithLegacy(ast, queryType, responseListener, Optional.of(calciteException));
+  }
+
+  @Test
+  public void testExplainWithLegacyShouldReturnCalciteErrorWhenBothFail() {
+    UnsupportedOperationException calciteException =
+        new UnsupportedOperationException("Calcite error");
+    IllegalStateException v2Exception = new IllegalStateException("V2 error");
+
+    ResponseListener<ExecutionEngine.ExplainResponse> responseListener =
+        new ResponseListener<>() {
+          @Override
+          public void onResponse(ExecutionEngine.ExplainResponse explainResponse) {
+            fail("Expected onFailure to be called");
+          }
+
+          @Override
+          public void onFailure(Exception e) {
+            // Should get the Calcite error directly (not wrapped), not the V2 error
+            assertNotNull(e);
+            assertTrue(e instanceof UnsupportedOperationException);
+            assertTrue(e.getMessage().contains("Calcite error"));
+          }
+        };
+
+    lenient().when(settings.getSettingValue(Key.CALCITE_ENGINE_ENABLED)).thenReturn(false);
+    lenient().when(analyzer.analyze(any(), any())).thenThrow(v2Exception);
+
+    QueryService service = new QueryService(analyzer, executionEngine, planner, null, settings);
+    service.explainWithLegacy(
+        ast, queryType, responseListener, ExplainMode.STANDARD, Optional.of(calciteException));
+  }
+
+  @Test
+  public void testExecuteWithLegacyShouldWrapCalciteErrorInCalciteUnsupportedException() {
+    AssertionError calciteError = new AssertionError("Calcite assertion failed");
+    IllegalStateException v2Exception = new IllegalStateException("V2 error");
+
+    ResponseListener<ExecutionEngine.QueryResponse> responseListener =
+        new ResponseListener<>() {
+          @Override
+          public void onResponse(ExecutionEngine.QueryResponse pplQueryResponse) {
+            fail("Expected onFailure to be called");
+          }
+
+          @Override
+          public void onFailure(Exception e) {
+            // Errors should be wrapped in CalciteUnsupportedException
+            assertNotNull(e);
+            assertTrue(e instanceof org.opensearch.sql.exception.CalciteUnsupportedException);
+            assertTrue(e.getCause() instanceof AssertionError);
+            assertTrue(e.getMessage().contains("Calcite assertion failed"));
+          }
+        };
+
+    lenient().when(settings.getSettingValue(Key.CALCITE_ENGINE_ENABLED)).thenReturn(false);
+    lenient().when(analyzer.analyze(any(), any())).thenThrow(v2Exception);
+
+    QueryService service = new QueryService(analyzer, executionEngine, planner, null, settings);
+    service.executeWithLegacy(ast, queryType, responseListener, Optional.of(calciteError));
+  }
+
+  @Test
+  public void testExplainWithLegacyShouldWrapCalciteErrorInCalciteUnsupportedException() {
+    AssertionError calciteError = new AssertionError("Calcite assertion failed");
+    IllegalStateException v2Exception = new IllegalStateException("V2 error");
+
+    ResponseListener<ExecutionEngine.ExplainResponse> responseListener =
+        new ResponseListener<>() {
+          @Override
+          public void onResponse(ExecutionEngine.ExplainResponse explainResponse) {
+            fail("Expected onFailure to be called");
+          }
+
+          @Override
+          public void onFailure(Exception e) {
+            // Errors should be wrapped in CalciteUnsupportedException
+            assertNotNull(e);
+            assertTrue(e instanceof org.opensearch.sql.exception.CalciteUnsupportedException);
+            assertTrue(e.getCause() instanceof AssertionError);
+            assertTrue(e.getMessage().contains("Calcite assertion failed"));
+          }
+        };
+
+    lenient().when(settings.getSettingValue(Key.CALCITE_ENGINE_ENABLED)).thenReturn(false);
+    lenient().when(analyzer.analyze(any(), any())).thenThrow(v2Exception);
+
+    QueryService service = new QueryService(analyzer, executionEngine, planner, null, settings);
+    service.explainWithLegacy(
+        ast, queryType, responseListener, ExplainMode.STANDARD, Optional.of(calciteError));
+  }
+
   Helper queryService() {
     return new Helper();
   }