Batch exception fix when unwrapping exception

server/implementation/src/main/java/io/smallrye/graphql/bootstrap/Config.java

@@ -26,6 +26,10 @@ default Optional<List<String>> getShowErrorMessageList() {
         return Optional.empty();
     }
 
+    default Optional<List<String>> getUnwrapExceptions() {
+        return Optional.empty();
+    }
+
     default boolean isAllowGet() {
         return false;
     }
@@ -74,10 +78,6 @@ default String getFieldVisibility() {
         return FIELD_VISIBILITY_DEFAULT;
     }
 
-    default Optional<List<String>> getUnwrapExceptions() {
-        return Optional.empty();
-    }
-
     default <T> T getConfigValue(String key, Class<T> type, T defaultValue) {
         return defaultValue;
     }

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/AbstractDataFetcher.java

@@ -1,9 +1,5 @@
 package io.smallrye.graphql.execution.datafetcher;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.concurrent.CompletionException;
-
 import org.dataloader.BatchLoaderWithContext;
 import org.eclipse.microprofile.graphql.GraphQLException;
 
@@ -38,7 +34,6 @@ public abstract class AbstractDataFetcher<K, T> implements DataFetcher<T>, Batch
     protected ArgumentHelper argumentHelper;
     protected EventEmitter eventEmitter;
     protected BatchLoaderHelper batchLoaderHelper;
-    protected List<String> unwrapExceptions = new ArrayList<>();
 
     public AbstractDataFetcher(Operation operation, Config config) {
         this.operation = operation;
@@ -48,10 +43,6 @@ public AbstractDataFetcher(Operation operation, Config config) {
         this.argumentHelper = new ArgumentHelper(operation.getArguments());
         this.partialResultHelper = new PartialResultHelper();
         this.batchLoaderHelper = new BatchLoaderHelper();
-        if (config != null && config.getUnwrapExceptions().isPresent()) {
-            this.unwrapExceptions.addAll(config.getUnwrapExceptions().get());
-        }
-        this.unwrapExceptions.addAll(DEFAULT_EXCEPTION_UNWRAP);
     }
 
     @Override
@@ -92,23 +83,4 @@ protected abstract <T> T invokeAndTransform(DataFetchingEnvironment dfe, DataFet
 
     protected abstract <T> T invokeFailure(DataFetcherResult.Builder<Object> resultBuilder);
 
-    protected Throwable unwrapThrowable(Throwable t) {
-        if (shouldUnwrapThrowable(t)) {
-            t = t.getCause();
-            return unwrapThrowable(t);
-        }
-        return t;
-    }
-
-    private boolean shouldUnwrapThrowable(Throwable t) {
-        return unwrapExceptions.contains(t.getClass().getName()) && t.getCause() != null;
-    }
-
-    private static final List<String> DEFAULT_EXCEPTION_UNWRAP = new ArrayList<>();
-
-    static {
-        DEFAULT_EXCEPTION_UNWRAP.add(CompletionException.class.getName());
-        DEFAULT_EXCEPTION_UNWRAP.add("javax.ejb.EJBException");
-        DEFAULT_EXCEPTION_UNWRAP.add("jakarta.ejb.EJBException");
-    }
 }

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/CompletionStageDataFetcher.java

@@ -44,8 +44,6 @@ protected <T> T invokeAndTransform(DataFetchingEnvironment dfe, DataFetcherResul
             return (T) futureResultFromMethodCall.handle((result, throwable) -> {
 
                 if (throwable != null) {
-                    throwable = unwrapThrowable(throwable);
-
                     eventEmitter.fireOnDataFetchError(dfe.getExecutionId().toString(), throwable);
                     if (throwable instanceof GraphQLException) {
                         GraphQLException graphQLException = (GraphQLException) throwable;

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/DefaultDataFetcher.java

@@ -1,5 +1,7 @@
 package io.smallrye.graphql.execution.datafetcher;
 
+import static io.smallrye.graphql.SmallRyeGraphQLServerMessages.msg;
+
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
@@ -37,8 +39,6 @@ public <T> T invokeAndTransform(DataFetchingEnvironment dfe, DataFetcherResult.B
             Object resultFromTransform = fieldHelper.transformResponse(resultFromMethodCall);
             resultBuilder.data(resultFromTransform);
             return (T) resultBuilder.build();
-        } catch (Exception e) {
-            throw (Exception) unwrapThrowable(e);
         } finally {
             SmallRyeContext.remove();
         }
@@ -58,10 +58,23 @@ public CompletionStage<List<T>> load(List<K> keys, BatchLoaderEnvironment ble) {
         ThreadContext threadContext = ThreadContext.builder().build();
         try {
             SmallRyeContext.setContext(context);
+
+            CompletableFuture<List<T>> reflectionSupplier = CompletableFuture.supplyAsync(() -> {
+                try {
+                    return (List<T>) reflectionHelper.invokePrivileged(tccl, arguments);
+                } catch (Exception e) {
+                    if (e instanceof RuntimeException && e.getCause() != null && !(e.getCause() instanceof RuntimeException)) {
+                        throw msg.dataFetcherException(operation, e.getCause());
+                    } else if (e instanceof RuntimeException) {
+                        throw (RuntimeException) e;
+                    } else {
+                        throw msg.dataFetcherException(operation, e);
+                    }
+                }
+            }, threadContext.currentContextExecutor());
+
             return threadContext
-                    .withContextCapture(
-                            CompletableFuture.supplyAsync(() -> (List<T>) reflectionHelper.invokePrivileged(tccl, arguments),
-                                    threadContext.currentContextExecutor()));
+                    .withContextCapture(reflectionSupplier);
         } finally {
             SmallRyeContext.remove();
         }

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/UniDataFetcher.java

@@ -41,7 +41,6 @@ protected <T> T invokeAndTransform(DataFetchingEnvironment dfe, DataFetcherResul
                     .onItemOrFailure().transform((result, throwable) -> {
 
                         if (throwable != null) {
-                            throwable = unwrapThrowable(throwable);
                             eventEmitter.fireOnDataFetchError(dfe.getExecutionId().toString(), throwable);
                             if (throwable instanceof GraphQLException) {
                                 GraphQLException graphQLException = (GraphQLException) throwable;
@@ -61,8 +60,6 @@ protected <T> T invokeAndTransform(DataFetchingEnvironment dfe, DataFetcherResul
 
                         return resultBuilder.build();
                     }).runSubscriptionOn(Infrastructure.getDefaultExecutor()).subscribe().asCompletionStage();
-        } catch (Exception e) {
-            throw (Exception) unwrapThrowable(e);
         } finally {
             SmallRyeContext.remove();
         }

server/implementation/src/main/java/io/smallrye/graphql/execution/error/ExceptionHandler.java

@@ -2,13 +2,18 @@
 
 import static io.smallrye.graphql.SmallRyeGraphQLServerLogging.log;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletionException;
+
 import graphql.ExceptionWhileDataFetching;
 import graphql.execution.DataFetcherExceptionHandler;
 import graphql.execution.DataFetcherExceptionHandlerParameters;
 import graphql.execution.DataFetcherExceptionHandlerResult;
 import graphql.execution.ResultPath;
 import graphql.language.SourceLocation;
 import io.smallrye.graphql.bootstrap.Config;
+import io.smallrye.graphql.execution.datafetcher.DataFetcherException;
 
 /**
  * Here we have the ability to mask certain messages to the client (for security reasons)
@@ -19,15 +24,21 @@ public class ExceptionHandler implements DataFetcherExceptionHandler {
 
     private final Config config;
     private final ExceptionLists exceptionLists;
+    protected List<String> unwrapExceptions = new ArrayList<>();
 
     public ExceptionHandler(Config config) {
         this.config = config;
         this.exceptionLists = new ExceptionLists(config.getHideErrorMessageList(), config.getShowErrorMessageList());
+        if (config.getUnwrapExceptions().isPresent()) {
+            this.unwrapExceptions.addAll(config.getUnwrapExceptions().get());
+        }
+        this.unwrapExceptions.addAll(DEFAULT_EXCEPTION_UNWRAP);
     }
 
     @Override
     public DataFetcherExceptionHandlerResult onException(DataFetcherExceptionHandlerParameters handlerParameters) {
         Throwable throwable = handlerParameters.getException();
+        throwable = unwrapThrowable(throwable);
         SourceLocation sourceLocation = handlerParameters.getSourceLocation();
         ResultPath path = handlerParameters.getPath();
         ExceptionWhileDataFetching error = getExceptionWhileDataFetching(throwable, sourceLocation, path);
@@ -57,4 +68,25 @@ private ExceptionWhileDataFetching getExceptionWhileDataFetching(Throwable throw
             }
         }
     }
+
+    private Throwable unwrapThrowable(Throwable t) {
+        if (shouldUnwrapThrowable(t)) {
+            t = t.getCause();
+            return unwrapThrowable(t);
+        }
+        return t;
+    }
+
+    private boolean shouldUnwrapThrowable(Throwable t) {
+        return unwrapExceptions.contains(t.getClass().getName()) && t.getCause() != null;
+    }
+
+    private static final List<String> DEFAULT_EXCEPTION_UNWRAP = new ArrayList<>();
+
+    static {
+        DEFAULT_EXCEPTION_UNWRAP.add(CompletionException.class.getName());
+        DEFAULT_EXCEPTION_UNWRAP.add(DataFetcherException.class.getName());
+        DEFAULT_EXCEPTION_UNWRAP.add("javax.ejb.EJBException");
+        DEFAULT_EXCEPTION_UNWRAP.add("jakarta.ejb.EJBException");
+    }
 }

server/pom.xml

@@ -16,7 +16,7 @@
 
     <properties>
         <!-- We use WF to run the TCK and do IT and run the runner project -->
-        <version.wildfly>21.0.1.Final</version.wildfly>
+        <version.wildfly>22.0.0.Final</version.wildfly>
     </properties>
 
     <modules>

server/tck/src/test/java/io/smallrye/graphql/SmallRyeGraphQLArchiveProcessor.java

@@ -12,6 +12,7 @@
 import org.jboss.shrinkwrap.resolver.api.maven.Maven;
 
 import io.smallrye.graphql.test.apps.async.api.AsyncApi;
+import io.smallrye.graphql.test.apps.batch.api.BatchApi;
 import io.smallrye.graphql.test.apps.context.api.ContextApi;
 import io.smallrye.graphql.test.apps.defaultvalue.api.DefaultValueParrotAPI;
 import io.smallrye.graphql.test.apps.error.api.ErrorApi;
@@ -86,6 +87,7 @@ public void process(Archive<?> applicationArchive, TestClass testClass) {
             war.addPackage(MutinyApi.class.getPackage());
             war.addPackage(ContextApi.class.getPackage());
             war.addPackage(JsonPApi.class.getPackage());
+            war.addPackage(BatchApi.class.getPackage());
 
             System.out.println(war.toString(true));
         }

server/tck/src/test/java/io/smallrye/graphql/test/apps/batch/api/BatchApi.java

@@ -0,0 +1,60 @@
+package io.smallrye.graphql.test.apps.batch.api;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.UUID;
+
+import org.eclipse.microprofile.graphql.GraphQLApi;
+import org.eclipse.microprofile.graphql.Query;
+import org.eclipse.microprofile.graphql.Source;
+
+@GraphQLApi
+public class BatchApi {
+
+    // Normal List Query
+    @Query
+    public List<BatchPojo> batchPojos() {
+        List<BatchPojo> l = new ArrayList<>();
+        l.add(new BatchPojo(1));
+        l.add(new BatchPojo(2));
+        l.add(new BatchPojo(3));
+        l.add(new BatchPojo(4));
+
+        return l;
+    }
+
+    // Nornal Source
+    public String greeting(@Source BatchPojo batchPojo) {
+        return "hello";
+    }
+
+    // Normal Batch Source
+    public List<UUID> uuid(@Source List<BatchPojo> batchPojos) {
+        List<UUID> uuids = new ArrayList<>();
+        for (BatchPojo batchPojo : batchPojos) {
+            uuids.add(UUID.fromString("88aaea3a-a48c-46de-a675-d6f2a65a9b2" + batchPojo.id));
+        }
+        return uuids;
+    }
+
+    // Runtime Exception Source
+    public String runtime(@Source BatchPojo batchPojo) {
+        throw new RuntimeException("Some runtime exception");
+    }
+
+    // Runtime Exception Batch Source
+    public List<String> runtimes(@Source List<BatchPojo> batchPojos) {
+        throw new RuntimeException("Some runtimes exception");
+    }
+
+    // Business Exception Source
+    public String business(@Source BatchPojo batchPojo) throws BusinessException {
+        throw new BusinessException("Some business exception");
+    }
+
+    // Business Exception Batch Source
+    public List<String> businesses(@Source List<BatchPojo> batchPojos) throws BusinessException {
+        throw new BusinessException("Some businesses exception");
+    }
+
+}

server/tck/src/test/java/io/smallrye/graphql/test/apps/batch/api/BatchPojo.java

@@ -0,0 +1,17 @@
+package io.smallrye.graphql.test.apps.batch.api;
+
+public class BatchPojo {
+
+    public int id;
+
+    public BatchPojo() {
+    }
+
+    public BatchPojo(int id) {
+        this.id = id;
+    }
+
+    public String getSimpleField() {
+        return "Some String";
+    }
+}

server/tck/src/test/java/io/smallrye/graphql/test/apps/batch/api/BusinessException.java

@@ -0,0 +1,24 @@
+package io.smallrye.graphql.test.apps.batch.api;
+
+public class BusinessException extends Exception {
+
+    public BusinessException() {
+    }
+
+    public BusinessException(String message) {
+        super(message);
+    }
+
+    public BusinessException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+    public BusinessException(Throwable cause) {
+        super(cause);
+    }
+
+    public BusinessException(String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
+        super(message, cause, enableSuppression, writableStackTrace);
+    }
+
+}

server/tck/src/test/resources/tests/batch/error/input.graphql

@@ -0,0 +1,8 @@
+{
+  batchPojos{
+    business
+    runtime
+    businesses
+    runtimes
+  }
+}