feat: add functions (keys, values, merge) to JMESPath visitor (#948)

0marperez · web-flow · commit ca206c8f518b · 2023-09-13T12:29:05.000-04:00
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/AcceptorGenerator.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/AcceptorGenerator.kt
@@ -89,12 +89,12 @@ private fun KotlinWriter.renderPathAcceptor(wi: WaiterInfo, directive: String, i
         val comparison = when (matcher.comparator!!) {
             PathComparator.STRING_EQUALS -> "${actual.identifier} == ${expected.dq()}"
             PathComparator.BOOLEAN_EQUALS -> "${actual.identifier} == ${expected.toBoolean()}"
-            PathComparator.ANY_STRING_EQUALS -> "(${actual.identifier} as List<String>?)?.any { it == ${expected.dq()} } ?: false"
+            PathComparator.ANY_STRING_EQUALS -> "(${actual.identifier} as List<String?>?)?.any { it == ${expected.dq()} } ?: false"
 
             // NOTE: the isNotEmpty check is necessary because the waiter spec says that `allStringEquals` requires
             // at least one value unlike Kotlin's `all` which returns true if the collection is empty
             PathComparator.ALL_STRING_EQUALS ->
-                "!(${actual.identifier} as List<String>).isNullOrEmpty() && ${actual.identifier}.all { it == ${expected.dq()} }"
+                "!(${actual.identifier} as List<String?>).isNullOrEmpty() && ${actual.identifier}.all { it == ${expected.dq()} }"
         }
         write(comparison)
     }
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/KotlinJmespathExpressionVisitor.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/KotlinJmespathExpressionVisitor.kt
@@ -105,11 +105,12 @@ class KotlinJmespathExpressionVisitor(
         return VisitedExpression(outerName, leftShape, innerResult.shape)
     }
 
-    private fun subfield(expression: FieldExpression, parentName: String): VisitedExpression {
+    private fun subfield(expression: FieldExpression, parentName: String, isObject: Boolean = false): VisitedExpression {
         val member = currentShape.targetOrSelf(ctx.model).getMember(expression.name).getOrNull()
 
         val name = expression.name.toCamelCase()
-        val nameExpr = ensureNullGuard(currentShape, name)
+        // User created objects are represented as hash maps in code-gen and are marked by `isObject`
+        val nameExpr = if (isObject) "[\"$name\"]" else ensureNullGuard(currentShape, name)
 
         val unwrapExpr = member?.let {
             val memberTarget = ctx.model.expectShape(member.target)
@@ -210,17 +211,20 @@ class KotlinJmespathExpressionVisitor(
         codegenReq(arguments.size == 1) { "Unexpected number of arguments to $this" }
         return acceptSubexpression(this.arguments[0])
     }
-
     private fun FunctionExpression.twoArgs(): Pair<VisitedExpression, VisitedExpression> {
         codegenReq(arguments.size == 2) { "Unexpected number of arguments to $this" }
         return acceptSubexpression(this.arguments[0]) to acceptSubexpression(this.arguments[1])
     }
 
-    private fun VisitedExpression.dotFunction(expression: FunctionExpression, expr: String, elvisExpr: String? = null): VisitedExpression {
+    private fun FunctionExpression.args(): List<VisitedExpression> =
+        this.arguments.map { acceptSubexpression(it) }
+
+    private fun VisitedExpression.dotFunction(expression: FunctionExpression, expr: String, elvisExpr: String? = null, isObject: Boolean = false): VisitedExpression {
         val dotFunctionExpr = ensureNullGuard(shape, expr, elvisExpr)
-        val ident = addTempVar(expression.name, "$identifier$dotFunctionExpr")
+        val ident = addTempVar(expression.name.toCamelCase(), "$identifier$dotFunctionExpr")
 
-        return VisitedExpression(ident, shape)
+        shape?.let { shapeCursor.addLast(shape) }
+        return VisitedExpression(ident, shape, isObject = isObject)
     }
 
     override fun visitFunction(expression: FunctionExpression): VisitedExpression = when (expression.name) {
@@ -267,6 +271,21 @@ class KotlinJmespathExpressionVisitor(
             subject.dotFunction(expression, "endsWith(${suffix.identifier})")
         }
 
+        "keys" -> {
+            val obj = expression.singleArg()
+            VisitedExpression(addTempVar("keys", obj.getKeys()))
+        }
+
+        "values" -> {
+            val obj = expression.singleArg()
+            VisitedExpression(addTempVar("values", obj.getValues()))
+        }
+
+        "merge" -> {
+            val objects = expression.args()
+            VisitedExpression(addTempVar("merge", objects.mergeProperties()), isObject = true)
+        }
+
         else -> throw CodegenException("Unknown function type in $expression")
     }
 
@@ -391,22 +410,22 @@ class KotlinJmespathExpressionVisitor(
     }
 
     override fun visitSubexpression(expression: Subexpression): VisitedExpression {
-        val leftName = expression.left.accept(this).identifier
-        return processRightSubexpression(expression.right, leftName)
+        val left = expression.left.accept(this)
+        return processRightSubexpression(expression.right, left.identifier, left.isObject)
     }
 
     private fun subexpression(expression: Subexpression, parentName: String): VisitedExpression {
-        val leftName = when (val left = expression.left) {
-            is FieldExpression -> subfield(left, parentName).identifier
-            is Subexpression -> subexpression(left, parentName).identifier
+        val left = when (val left = expression.left) {
+            is FieldExpression -> subfield(left, parentName)
+            is Subexpression -> subexpression(left, parentName)
             else -> throw CodegenException("Subexpression type $left is unsupported")
         }
-        return processRightSubexpression(expression.right, leftName)
+        return processRightSubexpression(expression.right, left.identifier, left.isObject)
     }
 
-    private fun processRightSubexpression(expression: JmespathExpression, leftName: String): VisitedExpression =
+    private fun processRightSubexpression(expression: JmespathExpression, leftName: String, isObject: Boolean = false): VisitedExpression =
         when (expression) {
-            is FieldExpression -> subfield(expression, leftName)
+            is FieldExpression -> subfield(expression, leftName, isObject)
             is IndexExpression -> index(expression, leftName)
             is Subexpression -> subexpression(expression, leftName)
             is ProjectionExpression -> projection(expression, leftName)
@@ -415,10 +434,7 @@ class KotlinJmespathExpressionVisitor(
 
     private fun index(expression: IndexExpression, parentName: String): VisitedExpression {
         val index = if (expression.index < 0) "$parentName.size${expression.index}" else expression.index
-
-        shapeCursor.addLast(currentShape.targetOrSelf(ctx.model).targetMemberOrSelf)
-        val indexExpr = ensureNullGuard(currentShape, "get($index)")
-        shapeCursor.removeLast()
+        val indexExpr = ensureNullGuard(currentShape.targetMemberOrSelf, "get($index)")
 
         return VisitedExpression(addTempVar("index", "$parentName$indexExpr"), currentShape)
     }
@@ -439,6 +455,33 @@ class KotlinJmespathExpressionVisitor(
             ".$expr"
         }
 
+    private fun VisitedExpression.getKeys(): String {
+        val keys = this.shape?.targetOrSelf(ctx.model)?.allMembers
+            ?.keys?.joinToString(", ", "listOf(", ")") { "\"$it\"" }
+        return keys ?: "listOf<String>()"
+    }
+
+    private fun VisitedExpression.getValues(): String {
+        val values = this.shape?.targetOrSelf(ctx.model)?.allMembers?.keys
+            ?.joinToString(", ", "listOf(", ")") { "${this.identifier}${ensureNullGuard(this.shape, it)}" }
+        return values ?: "listOf<String>()"
+    }
+
+    private fun List<VisitedExpression>.mergeProperties(): String {
+        val union = addTempVar("union", "HashMap<String, Any?>()")
+
+        forEach { obj ->
+            val keys = addTempVar("keys", obj.getKeys())
+            val values = addTempVar("values", obj.getValues())
+
+            writer.withBlock("for(i in $keys.indices){", "}") {
+                write("union[$keys[i]] = $values[i]")
+            }
+        }
+
+        return union
+    }
+
     private val Shape.isNullable: Boolean
         get() = this is MemberShape &&
             ctx.model.expectShape(target).let { !it.hasTrait<OperationInput>() && !it.hasTrait<OperationOutput>() }
@@ -461,5 +504,9 @@ class KotlinJmespathExpressionVisitor(
  *                  `foo[].bar[].baz.qux`, the shape that backs the identifier (and therefore determines overall nullability)
  *                  is `foo`, but the shape that needs carried through to subfield expressions in the following projection
  *                  is the target of `bar`, such that its subfields `baz` and `qux` can be recognized.
+ * @param nullable Boolean to indicate that a visited expression is nullable. Shape is used for this mostly but sometimes an
+ *                 expression is nullable for reasons that are not shape related
+ * @param isObject Boolean to indicate that a visited expression results in an object. Objects are represented as hash maps
+ *                 because it is not possible to construct a class at runtime
  */
-data class VisitedExpression(val identifier: String, val shape: Shape? = null, val projected: Shape? = null, val nullable: Boolean = false)
+data class VisitedExpression(val identifier: String, val shape: Shape? = null, val projected: Shape? = null, val nullable: Boolean = false, val isObject: Boolean = false)
diff --git a/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/AcceptorGeneratorTest.kt b/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/AcceptorGeneratorTest.kt
@@ -77,12 +77,12 @@ class AcceptorGeneratorTest {
                 InputOutputAcceptor(RetryDirective.TerminateAndSucceed) {
                     val output = it.output
                     val tags = output.tags
-                    !(tags as List<String>).isNullOrEmpty() && tags.all { it == "foo" }
+                    !(tags as List<String?>).isNullOrEmpty() && tags.all { it == "foo" }
                 },
                 InputOutputAcceptor(RetryDirective.TerminateAndSucceed) {
                     val output = it.output
                     val tags = output.tags
-                    (tags as List<String>?)?.any { it == "foo" } ?: false
+                    (tags as List<String?>?)?.any { it == "foo" } ?: false
                 },
             )
         """.trimIndent()
diff --git a/runtime/runtime-core/common/src/aws/smithy/kotlin/runtime/util/Convenience.kt b/runtime/runtime-core/common/src/aws/smithy/kotlin/runtime/util/Convenience.kt
@@ -6,7 +6,6 @@
 package aws.smithy.kotlin.runtime.util
 
 import aws.smithy.kotlin.runtime.InternalApi
-import kotlin.jvm.JvmName
 
 /**
  * Determines the length of a collection. This is a synonym for [Collection.size].
diff --git a/tests/codegen/waiter-tests/model/function-keys.smithy b/tests/codegen/waiter-tests/model/function-keys.smithy
@@ -0,0 +1,42 @@
+$version: "2"
+namespace com.test
+
+use smithy.waiters#waitable
+
+@waitable(
+    KeysFunctionPrimitivesStringEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "keys(primitives)",
+                        expected: "string",
+                        comparator: "anyStringEquals"
+                    }
+                }
+            }
+        ]
+    },
+    KeysFunctionPrimitivesIntegerEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "keys(primitives)",
+                        expected: "integer",
+                        comparator: "anyStringEquals"
+                    }
+                }
+            }
+        ]
+    },
+)
+@readonly
+@http(method: "GET", uri: "/keys/{name}", code: 200)
+operation GetFunctionKeysEquals {
+    input: GetEntityRequest,
+    output: GetEntityResponse,
+    errors: [NotFound],
+}
diff --git a/tests/codegen/waiter-tests/model/function-merge.smithy b/tests/codegen/waiter-tests/model/function-merge.smithy
@@ -0,0 +1,57 @@
+$version: "2"
+namespace com.test
+
+use smithy.waiters#waitable
+
+@suppress(["WaitableTraitJmespathProblem"])
+@waitable(
+    MergeFunctionOverrideObjectsOneEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "merge(objectOne, objectTwo).valueOne",
+                        expected: "foo",
+                        comparator: "stringEquals"
+                    }
+                }
+            }
+        ]
+    },
+    MergeFunctionOverrideObjectsTwoEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "merge(objectOne, objectTwo).valueTwo",
+                        expected: "bar",
+                        comparator: "stringEquals"
+                    }
+                }
+            }
+        ]
+    },
+    MergeFunctionOverrideObjectsThreeEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "merge(objectOne, objectTwo).valueThree",
+                        expected: "baz",
+                        comparator: "stringEquals"
+                    }
+                }
+            }
+        ]
+    },
+)
+@readonly
+@http(method: "GET", uri: "/merge/{name}", code: 200)
+operation GetFunctionMergeEquals {
+    input: GetEntityRequest,
+    output: GetEntityResponse,
+    errors: [NotFound],
+}
diff --git a/tests/codegen/waiter-tests/model/function-values.smithy b/tests/codegen/waiter-tests/model/function-values.smithy
@@ -0,0 +1,42 @@
+$version: "2"
+namespace com.test
+
+use smithy.waiters#waitable
+
+@waitable(
+    ValuesFunctionSampleValuesEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "values(sampleValues)",
+                        expected: "foo",
+                        comparator: "allStringEquals"
+                    }
+                }
+            }
+        ]
+    },
+    ValuesFunctionAnySampleValuesEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "values(sampleValues)",
+                        expected: "foo",
+                        comparator: "anyStringEquals"
+                    }
+                }
+            }
+        ]
+    },
+)
+@readonly
+@http(method: "GET", uri: "/values/{name}", code: 200)
+operation GetFunctionValuesEquals {
+    input: GetEntityRequest,
+    output: GetEntityResponse,
+    errors: [NotFound],
+}
diff --git a/tests/codegen/waiter-tests/model/structures.smithy b/tests/codegen/waiter-tests/model/structures.smithy
@@ -18,6 +18,9 @@ structure GetEntityResponse {
     lists: EntityLists,
     twoDimensionalLists: TwoDimensionalEntityLists,
     maps: EntityMaps,
+    sampleValues: Values,
+    objectOne: Values,
+    objectTwo: Values,
 }
 
 structure EntityPrimitives {
@@ -163,3 +166,9 @@ list SubStructList {
 structure SubStruct {
     subStructPrimitives: EntityPrimitives,
 }
+
+structure Values {
+    valueOne: String
+    valueTwo: String
+    valueThree: String
+}
diff --git a/tests/codegen/waiter-tests/model/waiter-operations.smithy b/tests/codegen/waiter-tests/model/waiter-operations.smithy
@@ -20,5 +20,8 @@ service WaitersTestService {
         GetFunctionJoinEquals,
         GetFunctionStartsWithEquals,
         GetFunctionEndsWithEquals,
+        GetFunctionKeysEquals,
+        GetFunctionValuesEquals,
+        GetFunctionMergeEquals,
     ]
 }
diff --git a/tests/codegen/waiter-tests/src/main/kotlin/com/test/DefaultWaitersTestClient.kt b/tests/codegen/waiter-tests/src/main/kotlin/com/test/DefaultWaitersTestClient.kt
@@ -52,6 +52,12 @@ class DefaultWaitersTestClient<T>(resultList: List<Result<T>>) : WaitersTestClie
 
     override suspend fun getFunctionEndsWithEquals(input: GetFunctionEndsWithEqualsRequest): GetFunctionEndsWithEqualsResponse = findSuccess()
 
+    override suspend fun getFunctionKeysEquals(input: GetFunctionKeysEqualsRequest): GetFunctionKeysEqualsResponse = findSuccess()
+
+    override suspend fun getFunctionValuesEquals(input: GetFunctionValuesEqualsRequest): GetFunctionValuesEqualsResponse = findSuccess()
+
+    override suspend fun getFunctionMergeEquals(input: GetFunctionMergeEqualsRequest): GetFunctionMergeEqualsResponse = findSuccess()
+
     private fun <Response> findSuccess(): Response {
         val nextResult = results.next()
         @Suppress("UNCHECKED_CAST")
diff --git a/tests/codegen/waiter-tests/src/test/kotlin/com/test/FunctionKeysTest.kt b/tests/codegen/waiter-tests/src/test/kotlin/com/test/FunctionKeysTest.kt
diff --git a/tests/codegen/waiter-tests/src/test/kotlin/com/test/FunctionMergeTest.kt b/tests/codegen/waiter-tests/src/test/kotlin/com/test/FunctionMergeTest.kt
diff --git a/tests/codegen/waiter-tests/src/test/kotlin/com/test/FunctionValuesTest.kt b/tests/codegen/waiter-tests/src/test/kotlin/com/test/FunctionValuesTest.kt

Original file line number	Diff line number	Diff line change
`@@ -89,12 +89,12 @@ private fun KotlinWriter.renderPathAcceptor(wi: WaiterInfo, directive: String, i`
`89`	`89`	`val comparison = when (matcher.comparator!!) {`
`90`	`90`	`PathComparator.STRING_EQUALS -> "${actual.identifier} == ${expected.dq()}"`
`91`	`91`	`PathComparator.BOOLEAN_EQUALS -> "${actual.identifier} == ${expected.toBoolean()}"`
`92`		`- PathComparator.ANY_STRING_EQUALS -> "(${actual.identifier} as List<String>?)?.any { it == ${expected.dq()} } ?: false"`
	`92`	`+ PathComparator.ANY_STRING_EQUALS -> "(${actual.identifier} as List<String?>?)?.any { it == ${expected.dq()} } ?: false"`
`93`	`93`
`94`	`94`	// NOTE: the isNotEmpty check is necessary because the waiter spec says that `allStringEquals` requires
`95`	`95`	// at least one value unlike Kotlin's `all` which returns true if the collection is empty
`96`	`96`	`PathComparator.ALL_STRING_EQUALS ->`
`97`		`- "!(${actual.identifier} as List<String>).isNullOrEmpty() && ${actual.identifier}.all { it == ${expected.dq()} }"`
	`97`	`+ "!(${actual.identifier} as List<String?>).isNullOrEmpty() && ${actual.identifier}.all { it == ${expected.dq()} }"`
`98`	`98`	`}`
`99`	`99`	`write(comparison)`
`100`	`100`	`}`
Original file line number	Diff line number	Diff line change
`@@ -20,5 +20,8 @@ service WaitersTestService {`
`20`	`20`	`GetFunctionJoinEquals,`
`21`	`21`	`GetFunctionStartsWithEquals,`
`22`	`22`	`GetFunctionEndsWithEquals,`
	`23`	`+ GetFunctionKeysEquals,`
	`24`	`+ GetFunctionValuesEquals,`
	`25`	`+ GetFunctionMergeEquals,`
`23`	`26`	`]`
`24`	`27`	`}`