From e18086e0d9844b15d7018b2a329e6ec96a5c7660 Mon Sep 17 00:00:00 2001
From: "Kamat, Trivikram" <16024985+trivikr@users.noreply.github.com>
Date: Wed, 28 Aug 2024 15:49:20 +0000
Subject: [PATCH] Process OperationContextParams in RuleSetParameterFinder

---
 .../endpointsV2/RuleSetParameterFinder.java   | 49 +++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/endpointsV2/RuleSetParameterFinder.java b/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/endpointsV2/RuleSetParameterFinder.java
index 26b00c8cd6d..cfaaff9a27a 100644
--- a/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/endpointsV2/RuleSetParameterFinder.java
+++ b/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/endpointsV2/RuleSetParameterFinder.java
@@ -32,6 +32,7 @@
 import software.amazon.smithy.rulesengine.traits.ClientContextParamsTrait;
 import software.amazon.smithy.rulesengine.traits.ContextParamTrait;
 import software.amazon.smithy.rulesengine.traits.EndpointRuleSetTrait;
+import software.amazon.smithy.rulesengine.traits.OperationContextParamsTrait;
 import software.amazon.smithy.rulesengine.traits.StaticContextParamsTrait;
 import software.amazon.smithy.utils.SmithyInternalApi;
 
@@ -152,6 +153,54 @@ public Map<String, String> getContextParams(Shape operationInput) {
         return map;
     }
 
+    /**
+     * Get map of params to JavaScript equivalent of provided JMESPath expressions.
+     */
+    public Map<String, String> getOperationContextParamValues(Shape operationInput) {
+        Map<String, String> map = new HashMap<>();
+
+        Optional<OperationContextParamsTrait> trait = operationInput.getTrait(OperationContextParamsTrait.class);
+        if (trait.isPresent()) {
+            OperationContextParamsTrait operationContextParamsTrait = trait.get();
+            operationContextParamsTrait.getParameters().forEach((name, definition) -> {
+                String separator = ".";
+                String value = "this" + separator + "input";
+                String path = definition.getPath();
+
+                // Split JMESPath expression string on separator and add JavaScript equivalent.
+                for (String part : path.split("[" + separator + "]")) {
+                    // Process keys https://jmespath.org/specification.html#keys
+                    if (part.startsWith("keys(")) {
+                        // Get provided object for which keys are to be extracted.
+                        String object = part.substring(5, part.length() - 1);
+                        value = "Object.keys(" + value + separator + object + ")";
+                        continue;
+                    }
+
+                    // Process list wildcard expression https://jmespath.org/specification.html#wildcard-expressions
+                    if (part.equals("*") || part.equals("[*]")) {
+                        value = "Object.values(" + value + ")";
+                        continue;
+                    }
+
+                    // Process hash wildcard expression https://jmespath.org/specification.html#wildcard-expressions
+                    if (part.endsWith("[*]")) {
+                        // Get key to run hash wildcard on.
+                        String key = part.substring(0, part.length() - 3);
+                        value = "Object.values(" + value + separator + key + ")";
+                        continue;
+                    }
+
+                    // Treat remaining part as identifier without spaces https://jmespath.org/specification.html#identifiers
+                    value += separator + part;
+                }
+
+            });
+        }
+
+        return map;
+    }
+
     private static class RuleSetParameterFinderVisitor extends NodeVisitor.Default<Void> {
         private final Map<String, String> map;