Ensure CompositeQueryLogic error codes are sent to metrics

core/query/src/main/java/datawave/core/query/logic/composite/CompositeLogicException.java

@@ -7,38 +7,88 @@
 import datawave.webservice.query.exception.QueryException;
 
 public class CompositeLogicException extends RuntimeException {
+
     public CompositeLogicException(String message, String logicName, Exception exception) {
-        super(getMessage(message, Collections.singletonMap(logicName, exception)), exception);
+        super(getMessage(message, Collections.singletonMap(logicName, exception)), getRaisedQueryException(exception));
     }
 
     public CompositeLogicException(String message, Map<String,Exception> exceptions) {
-        super(getMessage(message, exceptions), getQueryException(exceptions.values()));
+        super(getMessage(message, exceptions), getCause(exceptions.values()));
         if (exceptions.size() > 1) {
             exceptions.values().forEach(this::addSuppressed);
         }
     }
 
-    // looking for an exception that has a nested QueryException such that we may return an error code
-    private static Exception getQueryException(Collection<Exception> exceptions) {
+    /**
+     * Return the cause to use, prioritizing the first {@link QueryException} instance that we see. In the case where the {@link QueryException} is found to be
+     * the cause or further nested in the stack of an {@link Exception}, a {@link CompositeRaisedQueryException} will be returned with the query exception's
+     * error code, and the original exception as the cause. This is necessary to ensure the error code is passed to query metrics.
+     */
+    private static Exception getCause(Collection<Exception> exceptions) {
         if (exceptions.size() == 1) {
             return exceptions.iterator().next();
         }
-        Exception e = null;
-        for (Exception test : exceptions) {
-            if (e == null) {
-                e = test;
-            } else if (isQueryException(test)) {
-                e = test;
+        Exception cause = null;
+        for (Exception exception : exceptions) {
+            // Establish the initial cause as the first seen exception.
+            if (cause == null) {
+                cause = getRaisedQueryException(exception);
+                // If the first cause we see is a QueryException, there's nothing further to do.
+                if (cause instanceof QueryException) {
+                    return cause;
+                }
+                // If a subsequent exception is a or contains a QueryException in its stack, return it with the query exception error code available at the root
+                // exception.
+            } else if (hasQueryExceptionInStack(exception)) {
+                return getRaisedQueryException(exception);
             }
-            if (isQueryException(e)) {
-                break;
+        }
+        return cause;
+    }
+
+    /**
+     * Return whether the given throwable contains at least one {@link QueryException} in its stack trace (including itself).
+     */
+    private static boolean hasQueryExceptionInStack(Throwable throwable) {
+        return getFirstQueryExceptionInStack(throwable) != null;
+    }
+
+    /**
+     * Return the given exception with query exception's error code (if present) available at the root exception. This means one of the following cases will
+     * occur:
+     * <ul>
+     * <li>The exception is not a {@link QueryException} and no {@link QueryException} exists in the exception's stack: The exception will be returned.</li>
+     * <li>The exception is a {@link QueryException}: The exception will be returned.</li>
+     * <li>The exception is not a {@link QueryException}, but a {@link QueryException} exists in the exception's stack. A {@link CompositeRaisedQueryException}
+     * will be returned with the error code of the first {@link QueryException} found in the stack, and the original exception as its cause.</li>
+     * </ul>
+     */
+    private static Exception getRaisedQueryException(Exception exception) {
+        if (exception instanceof QueryException) {
+            return exception;
+        } else {
+            // TODO - should we fetch the top-most or bottom-most query exception in the stack?
+            QueryException queryException = getFirstQueryExceptionInStack(exception);
+            if (queryException != null) {
+                return new CompositeRaisedQueryException(exception, queryException.getErrorCode());
+            } else {
+                return exception;
             }
         }
-        return e;
     }
 
-    private static boolean isQueryException(Exception e) {
-        return new QueryException(e).getQueryExceptionsInStack().size() > 1;
+    /**
+     * Return the first {@link QueryException} found in the stack, or null if none were found.
+     */
+    private static QueryException getFirstQueryExceptionInStack(Throwable throwable) {
+        if (throwable != null) {
+            if (throwable instanceof QueryException) {
+                return (QueryException) throwable;
+            } else {
+                return getFirstQueryExceptionInStack(throwable.getCause());
+            }
+        }
+        return null;
     }
 
     private static String getMessage(String message, Map<String,Exception> exceptions) {

core/query/src/main/java/datawave/core/query/logic/composite/CompositeRaisedQueryException.java

@@ -0,0 +1,16 @@
+package datawave.core.query.logic.composite;
+
+import datawave.webservice.query.exception.QueryException;
+
+/**
+ * This class exists to be used when a {@link CompositeLogicException} has a cause that is not a {@link QueryException}, but contains a {@link QueryException}
+ * in its stack trace. In order for the error code to be properly passed to query metrics, the error code must be present as part of the
+ * {@link CompositeLogicException}'s cause. This exception is intended to be a wrapper for the original cause, with the error code of the identified query
+ * exception.
+ */
+public class CompositeRaisedQueryException extends QueryException {
+
+    public CompositeRaisedQueryException(Throwable cause, String errorCode) {
+        super(cause, errorCode);
+    }
+}

core/query/src/test/java/datawave/core/query/logic/composite/CompositeLogicExceptionTest.java

@@ -0,0 +1,100 @@
+package datawave.core.query.logic.composite;
+
+import static org.junit.Assert.assertEquals;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+import org.junit.Test;
+
+import datawave.webservice.query.exception.DatawaveErrorCode;
+import datawave.webservice.query.exception.QueryException;
+
+public class CompositeLogicExceptionTest {
+
+    @Test
+    public void testSingleNonQueryExceptionCause() {
+        IllegalArgumentException cause = new IllegalArgumentException("illegal argument");
+        CompositeLogicException exception = new CompositeLogicException("composite error occurred", "LogicName", cause);
+        assertEquals("composite error occurred:\nLogicName: illegal argument", exception.getMessage());
+        assertEquals(cause, exception.getCause());
+    }
+
+    @Test
+    public void testSingleQueryExceptionCause() {
+        QueryException cause = new QueryException(DatawaveErrorCode.MODEL_FETCH_ERROR, "connection failed");
+        CompositeLogicException exception = new CompositeLogicException("composite error occurred", "LogicName", cause);
+
+        assertEquals("composite error occurred:\nLogicName: Could not get model. connection failed", exception.getMessage());
+        assertEquals(cause, exception.getCause());
+        assertEquals(DatawaveErrorCode.MODEL_FETCH_ERROR.getErrorCode(), ((QueryException) exception.getCause()).getErrorCode());
+    }
+
+    @Test
+    public void testNestedSingleQueryExceptionCause() {
+        QueryException nestedCause = new QueryException(DatawaveErrorCode.MODEL_FETCH_ERROR, "connection failed");
+        IllegalArgumentException cause = new IllegalArgumentException("illegal argument", nestedCause);
+        CompositeLogicException exception = new CompositeLogicException("composite error occurred", "LogicName", cause);
+        assertEquals("composite error occurred:\nLogicName: illegal argument", exception.getMessage());
+        assertEquals(CompositeRaisedQueryException.class, exception.getCause().getClass());
+        assertEquals(DatawaveErrorCode.MODEL_FETCH_ERROR.getErrorCode(), ((CompositeRaisedQueryException) exception.getCause()).getErrorCode());
+    }
+
+    @Test
+    public void testMultipleNonQueryExceptionCauses() {
+        IllegalArgumentException expectedCause = new IllegalArgumentException("illegal name");
+        Map<String,Exception> exceptions = new LinkedHashMap<>();
+        exceptions.put("logic1", expectedCause);
+        exceptions.put("logic2", new NullPointerException("null value"));
+        exceptions.put("logic3", new IllegalStateException("bad state"));
+
+        CompositeLogicException exception = new CompositeLogicException("failed to complete", exceptions);
+        assertEquals("failed to complete:\nlogic1: illegal name\nlogic2: null value\nlogic3: bad state", exception.getMessage());
+        assertEquals(expectedCause, exception.getCause());
+    }
+
+    @Test
+    public void testMultipleExceptionWithOneTopLevelQueryException() {
+        QueryException expectedCause = new QueryException(DatawaveErrorCode.MODEL_FETCH_ERROR, "connection failed");
+        Map<String,Exception> exceptions = new LinkedHashMap<>();
+        exceptions.put("logic1", new IllegalArgumentException("illegal name"));
+        exceptions.put("logic2", new NullPointerException("null value"));
+        exceptions.put("logic3", expectedCause);
+        exceptions.put("logic4", new IllegalStateException("bad state"));
+
+        CompositeLogicException exception = new CompositeLogicException("failed to complete", exceptions);
+        assertEquals("failed to complete:\nlogic1: illegal name\nlogic2: null value\nlogic3: Could not get model. connection failed\nlogic4: bad state",
+                        exception.getMessage());
+        assertEquals(expectedCause, exception.getCause());
+    }
+
+    @Test
+    public void testMultipleExceptionWithOneNestedQueryException() {
+        QueryException nestedCause = new QueryException(DatawaveErrorCode.MODEL_FETCH_ERROR, "connection failed");
+        IllegalStateException topCause = new IllegalStateException("bad state", nestedCause);
+        Map<String,Exception> exceptions = new LinkedHashMap<>();
+        exceptions.put("logic1", new IllegalArgumentException("illegal name"));
+        exceptions.put("logic2", topCause);
+        exceptions.put("logic3", new NullPointerException("null value"));
+
+        CompositeLogicException exception = new CompositeLogicException("failed to complete", exceptions);
+        assertEquals("failed to complete:\nlogic1: illegal name\nlogic2: bad state\nlogic3: null value", exception.getMessage());
+        assertEquals(CompositeRaisedQueryException.class, exception.getCause().getClass());
+        assertEquals(DatawaveErrorCode.MODEL_FETCH_ERROR.getErrorCode(), ((CompositeRaisedQueryException) exception.getCause()).getErrorCode());
+    }
+
+    @Test
+    public void testMultipleExceptionWithNestedQueryExceptionSeenFirst() {
+        QueryException nestedCause = new QueryException(DatawaveErrorCode.MODEL_FETCH_ERROR, "connection failed");
+        IllegalStateException topCause = new IllegalStateException("bad state", nestedCause);
+        Map<String,Exception> exceptions = new LinkedHashMap<>();
+        exceptions.put("logic1", topCause);
+        exceptions.put("logic2", new IllegalArgumentException("illegal name"));
+        exceptions.put("logic3", new NullPointerException("null value"));
+
+        CompositeLogicException exception = new CompositeLogicException("failed to complete", exceptions);
+        assertEquals("failed to complete:\nlogic1: bad state\nlogic2: illegal name\nlogic3: null value", exception.getMessage());
+        assertEquals(CompositeRaisedQueryException.class, exception.getCause().getClass());
+        assertEquals(DatawaveErrorCode.MODEL_FETCH_ERROR.getErrorCode(), ((CompositeRaisedQueryException) exception.getCause()).getErrorCode());
+    }
+}

web-services/query/src/test/java/datawave/webservice/query/logic/composite/CompositeQueryLogicTest.java

@@ -805,7 +805,7 @@ public GenericQueryConfiguration initialize(AccumuloClient connection, Query set
 
             c.getTransformer(settings);
         } catch (CompositeLogicException e) {
-            Assert.assertEquals("query initialize failed", e.getCause().getCause().getMessage());
+            Assert.assertEquals("datawave.webservice.query.exception.QueryException: query initialize failed", e.getCause().getCause().getMessage());
         }
     }