update changelog and narrow type for TransportAction

Signed-off-by: mloufra <mloufra@amazon.com>
opensearch-project · Oct 17, 2022 · 60ad8b9 · 60ad8b9
1 parent 47d7276
commit 60ad8b9
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -176,6 +176,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
  - Return consumed params and content from extensions ([#4705](https://github.com/opensearch-project/OpenSearch/pull/4705))
  - Modified EnvironmentSettingsRequest to pass entire Settings object ([#4731](https://github.com/opensearch-project/OpenSearch/pull/4731))
  - Added contentParser method to ExtensionRestRequest ([#4760](https://github.com/opensearch-project/OpenSearch/pull/4760))
+ - Enforce type safety for RegisterTransportActionsRequest([#4796](https://github.com/opensearch-project/OpenSearch/pull/4796))
 
 ## [2.x]
 

diff --git a/server/src/main/java/org/opensearch/extensions/RegisterTransportActionsRequest.java b/server/src/main/java/org/opensearch/extensions/RegisterTransportActionsRequest.java
@@ -8,6 +8,9 @@
 
 package org.opensearch.extensions;
 
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.ActionResponse;
+import org.opensearch.action.support.TransportAction;
 import org.opensearch.common.io.stream.StreamInput;
 import org.opensearch.common.io.stream.StreamOutput;
 import org.opensearch.transport.TransportRequest;
@@ -16,6 +19,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Map.Entry;
 
 /**
  * Request to register extension Transport actions
@@ -24,22 +28,23 @@
  */
 public class RegisterTransportActionsRequest extends TransportRequest {
     private String uniqueId;
-    private Map<String, Class<?>> transportActions;
+    private Map<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> transportActions;
 
-    public RegisterTransportActionsRequest(String uniqueId, Map<String, Class<?>> transportActions) {
+    public RegisterTransportActionsRequest(String uniqueId, Map<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> transportActions) {
         this.uniqueId = uniqueId;
         this.transportActions = new HashMap<>(transportActions);
     }
 
     public RegisterTransportActionsRequest(StreamInput in) throws IOException {
         super(in);
         this.uniqueId = in.readString();
-        Map<String, Class<?>> actions = new HashMap<>();
+        Map<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> actions = new HashMap<>();
         int actionCount = in.readVInt();
         for (int i = 0; i < actionCount; i++) {
             try {
                 String actionName = in.readString();
-                Class<?> transportAction = Class.forName(in.readString());
+                @SuppressWarnings("unchecked")
+                Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>> transportAction = (Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>) Class.forName(in.readString());
                 actions.put(actionName, transportAction);
             } catch (ClassNotFoundException e) {
                 throw new IllegalArgumentException("Could not read transport action");
@@ -52,7 +57,7 @@ public String getUniqueId() {
         return uniqueId;
     }
 
-    public Map<String, Class<?>> getTransportActions() {
+    public Map<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> getTransportActions() {
         return transportActions;
     }
 
@@ -61,7 +66,7 @@ public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         out.writeString(uniqueId);
         out.writeVInt(this.transportActions.size());
-        for (Map.Entry<String, Class<?>> action : transportActions.entrySet()) {
+        for (Entry<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> action : transportActions.entrySet()) {
             out.writeString(action.getKey());
             out.writeString(action.getValue().getName());
         }