Improve exception unwrapping flexibility for SdkClientUtils

Signed-off-by: Daniel Widdis <[email protected]>

CHANGELOG.md

@@ -3,8 +3,11 @@ All notable changes to this project are documented in this file.
 
 Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.1.0/)
 
+## [Unreleased 2.x](https://github.com/opensearch-project/flow-framework/compare/2.19...2.x)
 ### Features
 ### Enhancements
+- Improve exception unwrapping flexibility for SdkClientUtils ([#67](https://github.com/opensearch-project/opensearch-remote-metadata-sdk/pull/67))
+
 ### Bug Fixes
 ### Infrastructure
 ### Documentation

core/src/main/java/org/opensearch/remote/metadata/common/SdkClientUtils.java

@@ -13,8 +13,9 @@
 import org.opensearch.common.util.concurrent.FutureUtils;
 import org.opensearch.common.util.concurrent.UncategorizedExecutionException;
 
+import java.util.Arrays;
+import java.util.List;
 import java.util.Locale;
-import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.ExecutionException;
 import java.util.regex.Matcher;
@@ -30,12 +31,22 @@ private SdkClientUtils() {}
     /**
      * Unwraps the cause of a {@link CompletionException}. If the cause is an {@link Exception}, rethrows the exception.
      * Otherwise wraps it in an {@link OpenSearchException}. Properly re-interrupts the thread on {@link InterruptedException}.
-     * @param throwable a throwable, expected to be a {@link CompletionException} or {@link CancellationException}.
+     * @param throwable a throwable.
+     * @param exceptionTypesToUnwrap optional list of exception types to unwrap. Defaults to {@link CompletionException}.
      * @return the cause of the completion exception or the throwable, directly if an {@link Exception} or wrapped in an OpenSearchException otherwise.
      */
-    public static Exception unwrapAndConvertToException(Throwable throwable) {
-        // Unwrap completion exception or pass through other exceptions
-        Throwable cause = throwable instanceof CompletionException ? throwable.getCause() : throwable;
+    @SafeVarargs
+    public static Exception unwrapAndConvertToException(Throwable throwable, Class<? extends Throwable>... exceptionTypesToUnwrap) {
+        // Unwrap specified exception types or pass through other exceptions
+        List<Class<? extends Throwable>> unwrapTypes = (exceptionTypesToUnwrap.length > 0)
+            ? Arrays.asList(exceptionTypesToUnwrap)
+            : List.of(CompletionException.class);
+
+        Throwable cause = throwable;
+        while (cause != null && unwrapTypes.contains(cause.getClass()) && cause.getCause() != null) {
+            cause = cause.getCause();
+        }
+
         // Double-unwrap checked exceptions wrapped in ExecutionException
         cause = getRethrownExecutionExceptionRootCause(cause);
         if (cause instanceof InterruptedException) {

core/src/test/java/org/opensearch/remote/metadata/common/SdkClientUtilsTests.java

@@ -8,6 +8,7 @@
  */
 package org.opensearch.remote.metadata.common;
 
+import org.opensearch.OpenSearchException;
 import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.support.PlainActionFuture;
 import org.opensearch.core.rest.RestStatus;
@@ -69,6 +70,61 @@ public void testUnwrapAndConvertToException_Unwrapped() {
         assertSame(ioException, e);
     }
 
+    @Test
+    public void testUnwrapAndConvertToException_VarargsUnwrap() {
+        // Create nested exceptions
+        OpenSearchException openSearchException = new OpenSearchException("Custom exception");
+        CompletionException completionException = new CompletionException(openSearchException);
+        OpenSearchStatusException statusException = new OpenSearchStatusException(
+            "Status exception",
+            RestStatus.INTERNAL_SERVER_ERROR,
+            completionException
+        );
+
+        // Test unwrapping with multiple exception types
+        Exception result = SdkClientUtils.unwrapAndConvertToException(
+            statusException,
+            OpenSearchStatusException.class,
+            CompletionException.class,
+            OpenSearchException.class
+        );
+        assertSame(openSearchException, result, "Should unwrap to the OpenSearchException");
+
+        // Test with a different order of exception types (order shouldn't matter now)
+        result = SdkClientUtils.unwrapAndConvertToException(
+            statusException,
+            CompletionException.class,
+            OpenSearchException.class,
+            OpenSearchStatusException.class
+        );
+        assertSame(openSearchException, result, "Should still unwrap to the OpenSearchException regardless of order");
+
+        // Test with only one exception type
+        result = SdkClientUtils.unwrapAndConvertToException(statusException, OpenSearchStatusException.class);
+        assertSame(completionException, result, "Should unwrap to the CompletionException (cause of OpenSearchStatusException)");
+
+        // Test with no matching exception type
+        IOException ioException = new IOException("IO Exception");
+        result = SdkClientUtils.unwrapAndConvertToException(ioException, OpenSearchException.class, CompletionException.class);
+        assertSame(ioException, result, "Should return the original exception when no matching type is found");
+
+        // Test with default behavior (only CompletionException)
+        result = SdkClientUtils.unwrapAndConvertToException(completionException);
+        assertSame(openSearchException, result, "Should unwrap CompletionException by default");
+
+        // Test with InterruptedException
+        InterruptedException interruptedException = new InterruptedException("Interrupted");
+        result = SdkClientUtils.unwrapAndConvertToException(interruptedException);
+        assertSame(interruptedException, result, "Should return InterruptedException and set interrupt flag");
+        assertTrue(Thread.interrupted(), "Interrupt flag should be set");
+
+        // Test with a non-Exception Throwable
+        Error error = new Error("Some error");
+        result = SdkClientUtils.unwrapAndConvertToException(error);
+        assertTrue(result instanceof OpenSearchException, "Should wrap non-Exception Throwable in OpenSearchException");
+        assertSame(error, result.getCause(), "Wrapped OpenSearchException should have original Error as cause");
+    }
+
     @Test
     public void testGetRethrownExecutionException_Unwrapped() {
         PlainActionFuture<Object> future = PlainActionFuture.newFuture();