Allow optional varargs for exception types to unwrap

Signed-off-by: Daniel Widdis <[email protected]>
dbwiddis · Feb 11, 2025 · bb56f28 · bb56f28
1 parent 49d54d5
commit bb56f28
Showing 1 changed file with 44 additions and 14 deletions.
diff --git a/core/src/main/java/org/opensearch/remote/metadata/common/SdkClientUtils.java b/core/src/main/java/org/opensearch/remote/metadata/common/SdkClientUtils.java
@@ -50,9 +50,14 @@ private SdkClientUtils() {}
      * Wraps the completion of a PUT operation from the SdkClient into a format compatible with an ActionListener.
      *
      * @param listener The ActionListener that will receive the parsed IndexResponse or any errors
+     * @param exceptionTypesToUnwrap optional list of exception types to unwrap. Defaults to {@link OpenSearchStatusException} and {@link CompletionException}.
      * @return A BiConsumer that can be used directly with CompletionStage's whenComplete method
      */
-    public static BiConsumer<PutDataObjectResponse, Throwable> wrapPutCompletion(ActionListener<IndexResponse> listener) {
+    @SafeVarargs
+    public static BiConsumer<PutDataObjectResponse, Throwable> wrapPutCompletion(
+        ActionListener<IndexResponse> listener,
+        Class<? extends Throwable>... exceptionTypesToUnwrap
+    ) {
         return (r, throwable) -> {
             if (throwable == null) {
                 try {
@@ -62,7 +67,7 @@ public static BiConsumer<PutDataObjectResponse, Throwable> wrapPutCompletion(Act
                     handleParseFailure(listener, "put");
                 }
             } else {
-                handleThrowable(listener, throwable);
+                handleThrowable(listener, throwable, exceptionTypesToUnwrap);
             }
         };
     }
@@ -73,7 +78,11 @@ public static BiConsumer<PutDataObjectResponse, Throwable> wrapPutCompletion(Act
      * @param listener The ActionListener that will receive the parsed GetResponse or any errors
      * @return A BiConsumer that can be used directly with CompletionStage's whenComplete method
      */
-    public static BiConsumer<GetDataObjectResponse, Throwable> wrapGetCompletion(ActionListener<GetResponse> listener) {
+    @SafeVarargs
+    public static BiConsumer<GetDataObjectResponse, Throwable> wrapGetCompletion(
+        ActionListener<GetResponse> listener,
+        Class<? extends Throwable>... exceptionTypesToUnwrap
+    ) {
         return (r, throwable) -> {
             if (throwable == null) {
                 try {
@@ -83,7 +92,7 @@ public static BiConsumer<GetDataObjectResponse, Throwable> wrapGetCompletion(Act
                     handleParseFailure(listener, "get");
                 }
             } else {
-                handleThrowable(listener, throwable);
+                handleThrowable(listener, throwable, exceptionTypesToUnwrap);
             }
         };
     }
@@ -94,7 +103,11 @@ public static BiConsumer<GetDataObjectResponse, Throwable> wrapGetCompletion(Act
      * @param listener The ActionListener that will receive the parsed UpdateResponse or any errors
      * @return A BiConsumer that can be used directly with CompletionStage's whenComplete method
      */
-    public static BiConsumer<UpdateDataObjectResponse, Throwable> wrapUpdateCompletion(ActionListener<UpdateResponse> listener) {
+    @SafeVarargs
+    public static BiConsumer<UpdateDataObjectResponse, Throwable> wrapUpdateCompletion(
+        ActionListener<UpdateResponse> listener,
+        Class<? extends Throwable>... exceptionTypesToUnwrap
+    ) {
         return (r, throwable) -> {
             if (throwable == null) {
                 try {
@@ -104,7 +117,7 @@ public static BiConsumer<UpdateDataObjectResponse, Throwable> wrapUpdateCompleti
                     handleParseFailure(listener, "update");
                 }
             } else {
-                handleThrowable(listener, throwable);
+                handleThrowable(listener, throwable, exceptionTypesToUnwrap);
             }
         };
     }
@@ -115,7 +128,11 @@ public static BiConsumer<UpdateDataObjectResponse, Throwable> wrapUpdateCompleti
      * @param listener The ActionListener that will receive the parsed DeleteResponse or any errors
      * @return A BiConsumer that can be used directly with CompletionStage's whenComplete method
      */
-    public static BiConsumer<DeleteDataObjectResponse, Throwable> wrapDeleteCompletion(ActionListener<DeleteResponse> listener) {
+    @SafeVarargs
+    public static BiConsumer<DeleteDataObjectResponse, Throwable> wrapDeleteCompletion(
+        ActionListener<DeleteResponse> listener,
+        Class<? extends Throwable>... exceptionTypesToUnwrap
+    ) {
         return (r, throwable) -> {
             if (throwable == null) {
                 try {
@@ -125,7 +142,7 @@ public static BiConsumer<DeleteDataObjectResponse, Throwable> wrapDeleteCompleti
                     handleParseFailure(listener, "delete");
                 }
             } else {
-                handleThrowable(listener, throwable);
+                handleThrowable(listener, throwable, exceptionTypesToUnwrap);
             }
         };
     }
@@ -136,7 +153,11 @@ public static BiConsumer<DeleteDataObjectResponse, Throwable> wrapDeleteCompleti
      * @param listener The ActionListener that will receive the parsed BulkResponse or any errors
      * @return A BiConsumer that can be used directly with CompletionStage's whenComplete method
      */
-    public static BiConsumer<BulkDataObjectResponse, Throwable> wrapBulkCompletion(ActionListener<BulkResponse> listener) {
+    @SafeVarargs
+    public static BiConsumer<BulkDataObjectResponse, Throwable> wrapBulkCompletion(
+        ActionListener<BulkResponse> listener,
+        Class<? extends Throwable>... exceptionTypesToUnwrap
+    ) {
         return (r, throwable) -> {
             if (throwable == null) {
                 try {
@@ -146,7 +167,7 @@ public static BiConsumer<BulkDataObjectResponse, Throwable> wrapBulkCompletion(A
                     handleParseFailure(listener, "bulk");
                 }
             } else {
-                handleThrowable(listener, throwable);
+                handleThrowable(listener, throwable, exceptionTypesToUnwrap);
             }
         };
     }
@@ -157,7 +178,11 @@ public static BiConsumer<BulkDataObjectResponse, Throwable> wrapBulkCompletion(A
      * @param listener The ActionListener that will receive the parsed SearchResponse or any errors
      * @return A BiConsumer that can be used directly with CompletionStage's whenComplete method
      */
-    public static BiConsumer<SearchDataObjectResponse, Throwable> wrapSearchCompletion(ActionListener<SearchResponse> listener) {
+    @SafeVarargs
+    public static BiConsumer<SearchDataObjectResponse, Throwable> wrapSearchCompletion(
+        ActionListener<SearchResponse> listener,
+        Class<? extends Throwable>... exceptionTypesToUnwrap
+    ) {
         return (r, throwable) -> {
             if (throwable == null) {
                 try {
@@ -167,7 +192,7 @@ public static BiConsumer<SearchDataObjectResponse, Throwable> wrapSearchCompleti
                     handleParseFailure(listener, "search");
                 }
             } else {
-                handleThrowable(listener, throwable);
+                handleThrowable(listener, throwable, exceptionTypesToUnwrap);
             }
         };
     }
@@ -176,8 +201,13 @@ private static void handleParseFailure(ActionListener<?> listener, String operat
         listener.onFailure(new OpenSearchStatusException("Failed to parse " + operation + " response", INTERNAL_SERVER_ERROR));
     }
 
-    private static void handleThrowable(ActionListener<?> listener, Throwable throwable) {
-        Exception exception = unwrapAndConvertToException(throwable, CompletionException.class, OpenSearchStatusException.class);
+    @SafeVarargs
+    private static void handleThrowable(
+        ActionListener<?> listener,
+        Throwable throwable,
+        Class<? extends Throwable>... exceptionTypesToUnwrap
+    ) {
+        Exception exception = unwrapAndConvertToException(throwable, exceptionTypesToUnwrap);
         listener.onFailure(exception);
     }