Have FlowFrameworkException status recognized by ExceptionsHelper (#811)

* Have FlowFrameworkException status recognized by ExceptionsHelper Signed-off-by: Daniel Widdis <[email protected]> * Add tests Signed-off-by: Daniel Widdis <[email protected]> --------- Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Jul 26, 2024 · 7ec848a · 7ec848a
1 parent 834903f
commit 7ec848a
Show file tree

Hide file tree

Showing 4 changed files with 85 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -25,6 +25,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.1.0/)
 ### Bug Fixes
 - Handle Not Found exceptions as successful deletions for agents and models ([#805](https://github.com/opensearch-project/flow-framework/pull/805))
 - Wrap CreateIndexRequest mappings in _doc key as required ([#809](https://github.com/opensearch-project/flow-framework/pull/809))
+- Have FlowFrameworkException status recognized by ExceptionsHelper ([#811](https://github.com/opensearch-project/flow-framework/pull/811))
 
 ### Infrastructure
 ### Documentation

diff --git a/release-notes/opensearch-flow-framework.release-notes-2.16.0.0.md b/release-notes/opensearch-flow-framework.release-notes-2.16.0.0.md
@@ -11,3 +11,4 @@ Compatible with OpenSearch 2.16.0
 ### Bug Fixes
 - Handle Not Found deprovision exceptions as successful deletions ([#805](https://github.com/opensearch-project/flow-framework/pull/805))
 - Wrap CreateIndexRequest mappings in _doc key as required ([#809](https://github.com/opensearch-project/flow-framework/pull/809))
+- Have FlowFrameworkException status recognized by ExceptionsHelper ([#811](https://github.com/opensearch-project/flow-framework/pull/811))
diff --git a/src/main/java/org/opensearch/flowframework/exception/FlowFrameworkException.java b/src/main/java/org/opensearch/flowframework/exception/FlowFrameworkException.java
@@ -8,6 +8,9 @@
  */
 package org.opensearch.flowframework.exception;
 
+import org.opensearch.OpenSearchException;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.ToXContentObject;
 import org.opensearch.core.xcontent.XContentBuilder;
@@ -17,7 +20,7 @@
 /**
  * Representation of Flow Framework Exceptions
  */
-public class FlowFrameworkException extends RuntimeException implements ToXContentObject {
+public class FlowFrameworkException extends OpenSearchException implements ToXContentObject {
 
     private static final long serialVersionUID = 1L;
 
@@ -56,6 +59,16 @@ public FlowFrameworkException(String message, Throwable cause, RestStatus restSt
         this.restStatus = restStatus;
     }
 
+    /**
+     * Read from a stream.
+     * @param in THe input stream
+     * @throws IOException on stream reading failure
+     */
+    public FlowFrameworkException(StreamInput in) throws IOException {
+        super(in);
+        restStatus = RestStatus.readFrom(in);
+    }
+
     /**
      * Getter for restStatus.
      *
@@ -65,8 +78,26 @@ public RestStatus getRestStatus() {
         return restStatus;
     }
 
+    // Same getter but for superclass
+    @Override
+    public final RestStatus status() {
+        return restStatus;
+    }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         return builder.startObject().field("error", this.getMessage()).endObject();
     }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        RestStatus.writeTo(out, restStatus);
+    }
+
+    // Keeping toXContentObject for backwards compatibility but this is needed for overriding superclass fragment
+    @Override
+    public boolean isFragment() {
+        return false;
+    }
 }
diff --git a/src/test/java/org/opensearch/flowframework/exception/FlowFrameworkExceptionTests.java b/src/test/java/org/opensearch/flowframework/exception/FlowFrameworkExceptionTests.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+package org.opensearch.flowframework.exception;
+
+import org.opensearch.ExceptionsHelper;
+import org.opensearch.OpenSearchException;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.common.xcontent.json.JsonXContent;
+import org.opensearch.core.common.bytes.BytesReference;
+import org.opensearch.core.common.io.stream.BytesStreamInput;
+import org.opensearch.core.rest.RestStatus;
+import org.opensearch.core.xcontent.ToXContent;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.io.IOException;
+
+public class FlowFrameworkExceptionTests extends OpenSearchTestCase {
+
+    public void testExceptions() {
+        WorkflowStepException wse = new WorkflowStepException("WSE", RestStatus.OK);
+        assertTrue(wse instanceof FlowFrameworkException);
+        assertTrue(wse instanceof OpenSearchException);
+        assertEquals(RestStatus.OK, ExceptionsHelper.status(wse));
+        assertFalse(wse.isFragment());
+    }
+
+    public void testSerialize() throws IOException {
+        FlowFrameworkException ffe = new FlowFrameworkException("FFE", RestStatus.OK);
+        assertTrue(ffe instanceof OpenSearchException);
+        assertEquals(RestStatus.OK, ExceptionsHelper.status(ffe));
+
+        try (BytesStreamOutput out = new BytesStreamOutput()) {
+            ffe.writeTo(out);
+            try (BytesStreamInput in = new BytesStreamInput(BytesReference.toBytes(out.bytes()))) {
+                ffe = new FlowFrameworkException(in);
+                assertTrue(ffe instanceof OpenSearchException);
+                assertEquals(RestStatus.OK, ExceptionsHelper.status(ffe));
+            }
+        }
+
+        XContentBuilder builder = JsonXContent.contentBuilder();
+        assertEquals("{\"error\":\"FFE\"}", ffe.toXContent(builder, ToXContent.EMPTY_PARAMS).toString());
+    }
+}