Consume REST params consumed by extension (#4415)

Signed-off-by: Daniel Widdis <[email protected]> Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Sep 9, 2022 · 8283ede · 8283ede
1 parent 14b899b
commit 8283ede
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/server/src/main/java/org/opensearch/extensions/rest/RestSendToExtensionAction.java b/server/src/main/java/org/opensearch/extensions/rest/RestSendToExtensionAction.java
@@ -28,9 +28,11 @@
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Map.Entry;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.unmodifiableList;
@@ -42,6 +44,7 @@ public class RestSendToExtensionAction extends BaseRestHandler {
 
     private static final String SEND_TO_EXTENSION_ACTION = "send_to_extension_action";
     private static final Logger logger = LogManager.getLogger(RestSendToExtensionAction.class);
+    private static final String CONSUMED_PARAMS_KEY = "extension.consumed.parameters";
 
     private final List<Route> routes;
     private final String uriPrefix;
@@ -121,7 +124,17 @@ public void handleResponse(RestExecuteOnExtensionResponse response) {
                 restExecuteOnExtensionResponse.setStatus(response.getStatus());
                 restExecuteOnExtensionResponse.setContentType(response.getContentType());
                 restExecuteOnExtensionResponse.setContent(response.getContent());
-                restExecuteOnExtensionResponse.setHeaders(response.getHeaders());
+                // Extract the consumed parameters from the header
+                Map<String, List<String>> headers = response.getHeaders();
+                List<String> consumedParams = headers.get(CONSUMED_PARAMS_KEY);
+                if (consumedParams != null) {
+                    consumedParams.stream().forEach(p -> request.param(p));
+                }
+                Map<String, List<String>> headersWithoutConsumedParams = headers.entrySet()
+                    .stream()
+                    .filter(e -> !e.getKey().equals(CONSUMED_PARAMS_KEY))
+                    .collect(Collectors.toMap(e -> e.getKey(), e -> e.getValue()));
+                restExecuteOnExtensionResponse.setHeaders(headersWithoutConsumedParams);
                 inProgressLatch.countDown();
             }