fix: correctly convert key fields during paginated Scan and Query operations

.changes/02bbf5e7-68d7-480d-a94c-35cd39dfe0db.json

@@ -0,0 +1,8 @@
+{
+    "id": "02bbf5e7-68d7-480d-a94c-35cd39dfe0db",
+    "type": "bugfix",
+    "description": "Correctly convert key fields during paginated `Scan` and `Query` operations",
+    "issues": [
+        "awslabs/aws-sdk-kotlin#1596"
+    ]
+}

hll/dynamodb-mapper/dynamodb-mapper-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/model/MapperTypes.kt

@@ -62,6 +62,7 @@ public object MapperTypes {
     }
 
     public object Model {
+        public val intersectKeys: TypeRef = TypeRef(MapperPkg.Hl.Model, "intersectKeys")
         public fun tablePartitionKey(objectType: TypeRef, pkType: TypeRef): TypeRef = TypeRef(
             MapperPkg.Hl.Model,
             "Table.PartitionKey",

hll/dynamodb-mapper/dynamodb-mapper-ops-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/operations/model/MemberCodegenBehavior.kt

@@ -7,8 +7,10 @@ package aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model
 import aws.sdk.kotlin.hll.codegen.model.*
 import aws.sdk.kotlin.hll.dynamodbmapper.codegen.model.MapperPkg
 import aws.sdk.kotlin.hll.dynamodbmapper.codegen.model.MapperTypes
-import aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model.ExpressionArgumentsType.*
-import aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model.ExpressionLiteralType.*
+import aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model.ExpressionArgumentsType.AttributeNames
+import aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model.ExpressionArgumentsType.AttributeValues
+import aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model.ExpressionLiteralType.Filter
+import aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model.ExpressionLiteralType.KeyCondition
 import aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model.MemberCodegenBehavior.*
 
 /**
@@ -156,6 +158,6 @@ private val rules = listOf(
 
     // Mappable members
     Rule(".*".toRegex(), Types.Kotlin.list(MapperTypes.AttributeMap), ListMapAll),
-    Rule("key", MapperTypes.AttributeMap, MapKeys),
+    Rule("key|lastEvaluatedKey|exclusiveStartKey".toRegex(), MapperTypes.AttributeMap, MapKeys),
     Rule(".*".toRegex(), MapperTypes.AttributeMap, MapAll),
 )

hll/dynamodb-mapper/dynamodb-mapper-ops-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/operations/rendering/OperationRenderer.kt

@@ -4,7 +4,7 @@
  */
 package aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.rendering
 
-import aws.sdk.kotlin.hll.codegen.core.*
+import aws.sdk.kotlin.hll.codegen.core.ImportDirective
 import aws.sdk.kotlin.hll.codegen.model.Member
 import aws.sdk.kotlin.hll.codegen.model.Operation
 import aws.sdk.kotlin.hll.codegen.model.Structure
@@ -106,9 +106,10 @@ internal class OperationRenderer(
         requestMembers(MemberCodegenBehavior.PassThrough) { write("#L = this@convert.#L", name, highLevel.name) }
         requestMembers(MemberCodegenBehavior.MapKeys) {
             write(
-                "this@convert.#L?.let { #L = schema.converter.convertTo(it, schema.keyAttributeNames) }",
+                "this@convert.#L?.let { #L = schema.converter.convertTo(it, schema.keyAttributeNames).#T(schema.keyAttributeNames) }",
                 highLevel.name,
                 name,
+                MapperTypes.Model.intersectKeys,
             )
         }
         requestMembers(MemberCodegenBehavior.MapAll) {

hll/dynamodb-mapper/dynamodb-mapper/common/src/aws/sdk/kotlin/hll/dynamodbmapper/model/Item.kt

@@ -8,7 +8,6 @@ import aws.sdk.kotlin.hll.dynamodbmapper.model.internal.ItemImpl
 import aws.sdk.kotlin.hll.dynamodbmapper.util.dynamicAttr
 import aws.sdk.kotlin.services.dynamodb.model.AttributeValue
 import aws.smithy.kotlin.runtime.ExperimentalApi
-import kotlin.jvm.JvmName
 
 /**
  * An immutable representation of a low-level item in a DynamoDB table. Items consist of attributes, each of which have
@@ -59,3 +58,9 @@ public fun itemOf(vararg pairs: Pair<String, AttributeValue>): Item = mapOf(*pai
  */
 @JvmName("itemOfPairStringAny")
 internal fun itemOf(vararg pairs: Pair<String, Any?>): Item = mapOf(*pairs).toItem()
+
+/**
+ * Truncate this item to include only the specified key attributes, removing any non-key attributes
+ * @param keys The collection of key names which should be preserved
+ */
+internal fun Item.intersectKeys(keys: Collection<String>): Item = filterKeys { keys.contains(it) }.toItem()

hll/dynamodb-mapper/dynamodb-mapper/jvm/test/aws/sdk/kotlin/hll/dynamodbmapper/operations/PaginatedScanTest.kt

@@ -0,0 +1,116 @@
+package aws.sdk.kotlin.hll.dynamodbmapper.operations
+
+import aws.sdk.kotlin.hll.dynamodbmapper.items.*
+import aws.sdk.kotlin.hll.dynamodbmapper.model.Item
+import aws.sdk.kotlin.hll.dynamodbmapper.model.itemOf
+import aws.sdk.kotlin.hll.dynamodbmapper.testutils.DdbLocalTest
+import aws.sdk.kotlin.hll.dynamodbmapper.values.scalars.IntConverter
+import aws.sdk.kotlin.hll.dynamodbmapper.values.scalars.StringConverter
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.toSet
+import kotlinx.coroutines.test.runTest
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+class PaginatedScanTest : DdbLocalTest() {
+    companion object {
+        private const val TABLE_NAME = "paginated-scan-test"
+
+        private fun rankName(rank: Int) = when (rank) {
+            2, 3, 4, 5, 6, 7, 8, 9, 10 -> rank.toString()
+            11 -> "Jack"
+            12 -> "Queen"
+            13 -> "King"
+            14 -> "Ace"
+            else -> "Unknown ($rank)"
+        }
+
+        private data class Card(
+            var rank: Int = 14,
+            var suit: String = "Spades",
+            var description: String = "This is the Ace of Spades. ".repeat(2000),
+        ) {
+            val rankName: String
+                get() = rankName(rank)
+
+            override fun toString() = "$rankName of $suit" // Easier when debugging big outputs
+        }
+
+        private val converter = object : ItemConverter<Card> {
+            val delegate = SimpleItemConverter(
+                ::Card,
+                { this },
+                AttributeDescriptor("rank", Card::rank, Card::rank::set, IntConverter),
+                AttributeDescriptor("suit", Card::suit, Card::suit::set, StringConverter),
+                AttributeDescriptor("description", Card::description, Card::description::set, StringConverter),
+            )
+
+            override fun convertFrom(to: Item): Card = delegate.convertFrom(to)
+
+            override fun convertTo(from: Card, onlyAttributes: Set<String>?): Item =
+                delegate.convertTo(from, null) // Ignore `onlyAttributes` arg to mock badly-behaved converter
+        }
+
+        private val schema = ItemSchema(converter, KeySpec.String("suit"), KeySpec.Number("rank"))
+
+        private val allCards = listOf("Spades", "Clubs", "Hearts", "Diamonds").flatMap { suit ->
+            (2..14).map { rank ->
+                Card(rank, suit, "This is the ${rankName(rank)} of $suit. ".repeat(2000))
+            }
+        }.toSet()
+    }
+
+    @BeforeAll
+    fun setUp() = runTest {
+        createTable(
+            name = TABLE_NAME,
+            schema = schema,
+            items = allCards.map { card ->
+                itemOf(
+                    "rank" to card.rank,
+                    "suit" to card.suit,
+                    "description" to card.description,
+                )
+            }.toTypedArray(),
+        )
+    }
+
+    @Test
+    fun testPaginatedScan() = runTest {
+        val mapper = mapper()
+        val table = mapper.getTable(TABLE_NAME, schema)
+        var pages = 0
+
+        val actual = table
+            .scanPaginated { }
+            .map {
+                pages++
+                it
+            }
+            .items()
+            .toSet()
+        assertEquals(allCards, actual)
+        assertTrue(pages > 1, "Got only $pages pages but expected at least 2")
+    }
+
+    @Test
+    fun testPaginatedScanWithOffset() = runTest {
+        val mapper = mapper()
+        val table = mapper.getTable(TABLE_NAME, schema)
+        val startKey = allCards.first()
+        var pages = 0
+
+        val actual = table
+            .scanPaginated {
+                exclusiveStartKey = startKey
+            }
+            .map {
+                pages++
+                it
+            }
+            .items()
+            .toSet()
+        assertTrue(startKey !in actual)
+        assertTrue(pages > 1, "Got only $pages pages but expected at least 2")
+    }
+}

hll/dynamodb-mapper/dynamodb-mapper/jvm/test/aws/sdk/kotlin/hll/dynamodbmapper/testutils/DdbClientExtensions.kt

@@ -88,17 +88,17 @@ suspend fun DynamoDbClient.getItem(tableName: String, vararg keys: Pair<String,
  * @param items A collection of maps of strings to values to be mapped and persisted to the table
  */
 suspend fun DynamoDbClient.putItems(tableName: String, items: List<Item>) {
-    val writeRequests = items.map { item ->
-        WriteRequest {
-            putRequest {
-                this.item = item
+    val batches = items
+        .map { item ->
+            WriteRequest {
+                putRequest { this.item = item }
             }
         }
-    }
+        .chunked(25) // Max batchWriteItem page size
 
-    if (writeRequests.isNotEmpty()) {
+    batches.forEach { batch ->
         batchWriteItem {
-            requestItems = mapOf(tableName to writeRequests)
+            requestItems = mapOf(tableName to batch)
         }
     }
 }