Restore original exception handling expectations

Signed-off-by: Daniel Widdis <[email protected]>
dbwiddis · Jun 5, 2024 · 6f20123 · 6f20123
1 parent c26c11d
commit 6f20123
Show file tree

Hide file tree

Showing 8 changed files with 85 additions and 44 deletions.
diff --git a/common/src/main/java/org/opensearch/sdk/SdkClient.java b/common/src/main/java/org/opensearch/sdk/SdkClient.java
@@ -80,7 +80,12 @@ default GetDataObjectResponse getDataObject(GetDataObjectRequest request) {
             if (cause instanceof InterruptedException) {
                 Thread.currentThread().interrupt();
             }
-            throw new OpenSearchException(cause);
+            // Rethrow unchecked Exceptions
+            if (cause instanceof RuntimeException) {
+                throw (RuntimeException) cause;
+            } else {
+                throw new OpenSearchException(cause);                
+            }
         }
     }
 
@@ -114,7 +119,12 @@ default DeleteDataObjectResponse deleteDataObject(DeleteDataObjectRequest reques
             if (cause instanceof InterruptedException) {
                 Thread.currentThread().interrupt();
             }
-            throw new OpenSearchException(cause);
+            // Rethrow unchecked Exceptions
+            if (cause instanceof RuntimeException) {
+                throw (RuntimeException) cause;
+            } else {
+                throw new OpenSearchException(cause);                
+            }
         }
     }
 }
diff --git a/plugin/src/main/java/org/opensearch/ml/action/connector/DeleteConnectorTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/connector/DeleteConnectorTransportAction.java
@@ -15,6 +15,7 @@
 import java.util.List;
 import java.util.Objects;
 
+import org.opensearch.OpenSearchException;
 import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.delete.DeleteRequest;
@@ -125,7 +126,7 @@ private void checkForModelsUsingConnector(String connectorId, String tenantId, A
                 sourceBuilder.query(QueryBuilders.matchQuery(TENANT_ID, tenantId));
             }
             searchRequest.source(sourceBuilder);
-            // TODO: User SDK client not client.
+            // TODO: Use SDK client not client.
             client.search(searchRequest, ActionListener.runBefore(ActionListener.wrap(searchResponse -> {
                 SearchHit[] searchHits = searchResponse.getHits().getHits();
                 if (searchHits.length == 0) {
@@ -213,8 +214,13 @@ private void handleDeleteResponse(
         ActionListener<DeleteResponse> actionListener
     ) {
         if (throwable != null) {
-            log.error("Failed to delete ML connector: {}", connectorId, throwable);
-            actionListener.onFailure(new RuntimeException(throwable));
+            Throwable cause = throwable.getCause() == null ? throwable : throwable.getCause();
+            log.error("Failed to delete ML connector: {}", connectorId, cause);
+            if (cause instanceof Exception) {
+                actionListener.onFailure((Exception) cause);
+            } else {
+                actionListener.onFailure(new OpenSearchException(cause));
+            }
         } else {
             log.info("Connector deletion result: {}, connector id: {}", response.deleted(), response.id());
             DeleteResponse deleteResponse = new DeleteResponse(response.shardId(), response.id(), 0, 0, 0, response.deleted());

diff --git a/plugin/src/main/java/org/opensearch/ml/action/connector/GetConnectorTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/connector/GetConnectorTransportAction.java
@@ -10,6 +10,7 @@
 import static org.opensearch.ml.plugin.MachineLearningPlugin.GENERAL_THREAD_POOL;
 import static org.opensearch.ml.utils.RestActionUtils.getFetchSourceContext;
 
+import org.opensearch.OpenSearchException;
 import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.support.ActionFilters;
@@ -95,8 +96,12 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLConn
                             log.error("Failed to get connector index", cause);
                             actionListener.onFailure(new OpenSearchStatusException("Failed to find connector", RestStatus.NOT_FOUND));
                         } else {
-                            log.error("Failed to get ML connector {}", connectorId, cause);
-                            actionListener.onFailure(new RuntimeException(cause));
+                            log.error("Failed to get ML connector " + connectorId, cause);
+                            if (cause instanceof Exception) {
+                                actionListener.onFailure((Exception) cause);
+                            } else {
+                                actionListener.onFailure(new OpenSearchException(cause));
+                            }
                         }
                     } else {
                         if (r != null && r.parser().isPresent()) {
@@ -140,6 +145,5 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLConn
             log.error("Failed to get ML connector " + connectorId, e);
             actionListener.onFailure(e);
         }
-
     }
 }
diff --git a/plugin/src/main/java/org/opensearch/ml/action/connector/TransportCreateConnectorAction.java b/plugin/src/main/java/org/opensearch/ml/action/connector/TransportCreateConnectorAction.java
@@ -12,6 +12,7 @@
 import java.util.HashSet;
 import java.util.List;
 
+import org.opensearch.OpenSearchException;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.action.support.HandledTransportAction;
@@ -147,7 +148,13 @@ private void indexConnector(Connector connector, ActionListener<MLCreateConnecto
                     .whenComplete((r, throwable) -> {
                         context.restore();
                         if (throwable != null) {
-                            listener.onFailure(new RuntimeException(throwable));
+                            Throwable cause = throwable.getCause() == null ? throwable : throwable.getCause();
+                            log.error("Failed to create ML connector", cause);
+                            if (cause instanceof Exception) {
+                                listener.onFailure((Exception) cause);
+                            } else {
+                                listener.onFailure(new OpenSearchException(cause));
+                            }
                         } else {
                             log.info("Connector creation result: {}, connector id: {}", r.created(), r.id());
                             MLCreateConnectorResponse response = new MLCreateConnectorResponse(r.id());

diff --git a/plugin/src/main/java/org/opensearch/ml/sdkclient/LocalClusterIndicesClient.java b/plugin/src/main/java/org/opensearch/ml/sdkclient/LocalClusterIndicesClient.java
@@ -14,14 +14,14 @@
 import static org.opensearch.common.xcontent.json.JsonXContent.jsonXContent;
 import static org.opensearch.core.xcontent.ToXContent.EMPTY_PARAMS;
 
+import java.io.IOException;
 import java.security.AccessController;
 import java.security.PrivilegedAction;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.Executor;
 
-import org.opensearch.OpenSearchException;
 import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.delete.DeleteRequest;
 import org.opensearch.action.delete.DeleteResponse;
@@ -32,10 +32,10 @@
 import org.opensearch.client.Client;
 import org.opensearch.common.xcontent.LoggingDeprecationHandler;
 import org.opensearch.common.xcontent.XContentFactory;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
-import org.opensearch.index.IndexNotFoundException;
 import org.opensearch.sdk.DeleteDataObjectRequest;
 import org.opensearch.sdk.DeleteDataObjectResponse;
 import org.opensearch.sdk.GetDataObjectRequest;
@@ -79,8 +79,12 @@ public CompletionStage<PutDataObjectResponse> putDataObjectAsync(PutDataObjectRe
                     .actionGet();
                 log.info("Creation status for id {}: {}", indexResponse.getId(), indexResponse.getResult());
                 return new PutDataObjectResponse.Builder().id(indexResponse.getId()).created(indexResponse.getResult() == CREATED).build();
-            } catch (Exception e) {
-                throw new OpenSearchException(e);
+            } catch (IOException e) {
+                // Rethrow unchecked exception on XContent parsing error
+                throw new OpenSearchStatusException(
+                    "Failed to parse data object to put in index " + request.index(),
+                    RestStatus.BAD_REQUEST
+                );
             }
         }), executor);
     }
@@ -90,38 +94,38 @@ public CompletionStage<GetDataObjectResponse> getDataObjectAsync(GetDataObjectRe
         return CompletableFuture.supplyAsync(() -> AccessController.doPrivileged((PrivilegedAction<GetDataObjectResponse>) () -> {
             try {
                 log.info("Getting {} from {}", request.id(), request.index());
-                GetResponse getResponse = client.get(new GetRequest(request.index(), request.id())).actionGet();
+                GetResponse getResponse = client
+                    .get(new GetRequest(request.index(), request.id()).fetchSourceContext(request.fetchSourceContext()))
+                    .actionGet();
                 if (getResponse == null || !getResponse.isExists()) {
                     return new GetDataObjectResponse.Builder().id(request.id()).build();
                 }
                 XContentParser parser = jsonXContent
                     .createParser(xContentRegistry, LoggingDeprecationHandler.INSTANCE, getResponse.getSourceAsString());
                 log.info("Retrieved data object");
                 return new GetDataObjectResponse.Builder().id(getResponse.getId()).parser(Optional.of(parser)).build();
-            } catch (OpenSearchStatusException | IndexNotFoundException notFound) {
-                throw notFound;
-            } catch (Exception e) {
-                throw new OpenSearchException(e);
+            } catch (IOException e) {
+                // Rethrow unchecked exception on XContent parser creation error
+                throw new OpenSearchStatusException(
+                    "Failed to create parser for data object retrieved from index " + request.index(),
+                    RestStatus.INTERNAL_SERVER_ERROR
+                );
             }
         }), executor);
     }
 
     @Override
     public CompletionStage<DeleteDataObjectResponse> deleteDataObjectAsync(DeleteDataObjectRequest request, Executor executor) {
         return CompletableFuture.supplyAsync(() -> AccessController.doPrivileged((PrivilegedAction<DeleteDataObjectResponse>) () -> {
-            try {
-                log.info("Deleting {} from {}", request.id(), request.index());
-                DeleteResponse deleteResponse = client.delete(new DeleteRequest(request.index(), request.id())).actionGet();
-                log.info("Deletion status for id {}: {}", deleteResponse.getId(), deleteResponse.getResult());
-                return new DeleteDataObjectResponse.Builder()
-                    .id(deleteResponse.getId())
-                    .shardId(deleteResponse.getShardId())
-                    .shardInfo(deleteResponse.getShardInfo())
-                    .deleted(deleteResponse.getResult() == DELETED)
-                    .build();
-            } catch (Exception e) {
-                throw new OpenSearchException(e);
-            }
+            log.info("Deleting {} from {}", request.id(), request.index());
+            DeleteResponse deleteResponse = client.delete(new DeleteRequest(request.index(), request.id())).actionGet();
+            log.info("Deletion status for id {}: {}", deleteResponse.getId(), deleteResponse.getResult());
+            return new DeleteDataObjectResponse.Builder()
+                .id(deleteResponse.getId())
+                .shardId(deleteResponse.getShardId())
+                .shardInfo(deleteResponse.getShardInfo())
+                .deleted(deleteResponse.getResult() == DELETED)
+                .build();
         }), executor);
     }
 }
diff --git a/plugin/src/main/java/org/opensearch/ml/sdkclient/RemoteClusterIndicesClient.java b/plugin/src/main/java/org/opensearch/ml/sdkclient/RemoteClusterIndicesClient.java
@@ -11,6 +11,7 @@
 import static org.opensearch.client.opensearch._types.Result.Created;
 import static org.opensearch.client.opensearch._types.Result.Deleted;
 
+import java.io.IOException;
 import java.security.AccessController;
 import java.security.PrivilegedAction;
 import java.util.Map;
@@ -19,7 +20,7 @@
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.Executor;
 
-import org.opensearch.OpenSearchException;
+import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.support.replication.ReplicationResponse.ShardInfo;
 import org.opensearch.client.opensearch.OpenSearchClient;
 import org.opensearch.client.opensearch.core.DeleteRequest;
@@ -30,6 +31,7 @@
 import org.opensearch.client.opensearch.core.IndexResponse;
 import org.opensearch.common.xcontent.LoggingDeprecationHandler;
 import org.opensearch.common.xcontent.json.JsonXContent;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.sdk.DeleteDataObjectRequest;
@@ -70,8 +72,12 @@ public CompletionStage<PutDataObjectResponse> putDataObjectAsync(PutDataObjectRe
                 IndexResponse indexResponse = openSearchClient.index(indexRequest);
                 log.info("Creation status for id {}: {}", indexResponse.id(), indexResponse.result());
                 return new PutDataObjectResponse.Builder().id(indexResponse.id()).created(indexResponse.result() == Created).build();
-            } catch (Exception e) {
-                throw new OpenSearchException("Error occurred while indexing data object", e);
+            } catch (IOException e) {
+                // Rethrow unchecked exception on XContent parsing error
+                throw new OpenSearchStatusException(
+                    "Failed to parse data object to put in index " + request.index(),
+                    RestStatus.BAD_REQUEST
+                );
             }
         }), executor);
     }
@@ -92,8 +98,12 @@ public CompletionStage<GetDataObjectResponse> getDataObjectAsync(GetDataObjectRe
                 XContentParser parser = JsonXContent.jsonXContent
                     .createParser(NamedXContentRegistry.EMPTY, LoggingDeprecationHandler.INSTANCE, json);
                 return new GetDataObjectResponse.Builder().id(getResponse.id()).parser(Optional.of(parser)).build();
-            } catch (Exception e) {
-                throw new OpenSearchException(e);
+            } catch (IOException e) {
+                // Rethrow unchecked exception on XContent parser creation error
+                throw new OpenSearchStatusException(
+                    "Failed to create parser for data object retrieved from index " + request.index(),
+                    RestStatus.INTERNAL_SERVER_ERROR
+                );
             }
         }), executor);
     }
@@ -116,8 +126,12 @@ public CompletionStage<DeleteDataObjectResponse> deleteDataObjectAsync(DeleteDat
                     .shardInfo(shardInfo)
                     .deleted(deleteResponse.result() == Deleted)
                     .build();
-            } catch (Exception e) {
-                throw new OpenSearchException("Error occurred while deleting data object", e);
+            } catch (IOException e) {
+                // Rethrow unchecked exception on deletion IOException
+                throw new OpenSearchStatusException(
+                    "IOException occurred while deleting data object " + request.id() + " from index " + request.index(),
+                    RestStatus.INTERNAL_SERVER_ERROR
+                );
             }
         }), executor);
     }

diff --git a/...src/test/java/org/opensearch/ml/action/connector/DeleteConnectorTransportActionTests.java b/...src/test/java/org/opensearch/ml/action/connector/DeleteConnectorTransportActionTests.java
@@ -310,9 +310,7 @@ public void testDeleteConnector_ResourceNotFoundException() throws IOException,
 
         ArgumentCaptor<Exception> argumentCaptor = ArgumentCaptor.forClass(RuntimeException.class);
         verify(actionListener).onFailure(argumentCaptor.capture());
-        // TODO: fix all this exception nesting
-        // java.util.concurrent.CompletionException: OpenSearchException[ResourceNotFoundException[errorMessage]]; nested: ResourceNotFoundException[errorMessage];
-        assertEquals("errorMessage", argumentCaptor.getValue().getCause().getCause().getCause().getMessage());
+        assertEquals("errorMessage", argumentCaptor.getValue().getMessage());
     }
 
     public void test_ValidationFailedException() throws IOException {

diff --git a/...in/src/test/java/org/opensearch/ml/action/connector/GetConnectorTransportActionTests.java b/...in/src/test/java/org/opensearch/ml/action/connector/GetConnectorTransportActionTests.java
@@ -226,9 +226,7 @@ public void testGetConnector_RuntimeException() throws InterruptedException {
 
         ArgumentCaptor<Exception> argumentCaptor = ArgumentCaptor.forClass(Exception.class);
         verify(actionListener).onFailure(argumentCaptor.capture());
-        // TODO: Fix this nesting
-        // [OpenSearchException[java.lang.RuntimeException: errorMessage]; nested: RuntimeException[errorMessage];
-        assertEquals("errorMessage", argumentCaptor.getValue().getCause().getCause().getMessage());
+        assertEquals("errorMessage", argumentCaptor.getValue().getMessage());
     }
 
     public void testGetConnector_MultiTenancyEnabled_Success() throws IOException, InterruptedException {