diff --git a/.changes/02bbf5e7-68d7-480d-a94c-35cd39dfe0db.json b/.changes/02bbf5e7-68d7-480d-a94c-35cd39dfe0db.json new file mode 100644 index 00000000000..4ab9159d961 --- /dev/null +++ b/.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" + ] +} \ No newline at end of file diff --git a/hll/dynamodb-mapper/dynamodb-mapper-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/model/MapperTypes.kt b/hll/dynamodb-mapper/dynamodb-mapper-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/model/MapperTypes.kt index 6c93e1001c4..6b8d38c2fca 100644 --- a/hll/dynamodb-mapper/dynamodb-mapper-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/model/MapperTypes.kt +++ b/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", diff --git a/hll/dynamodb-mapper/dynamodb-mapper-ops-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/operations/model/MemberCodegenBehavior.kt b/hll/dynamodb-mapper/dynamodb-mapper-ops-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/operations/model/MemberCodegenBehavior.kt index 6295f09637a..e0905466c1c 100644 --- a/hll/dynamodb-mapper/dynamodb-mapper-ops-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/operations/model/MemberCodegenBehavior.kt +++ b/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), ) diff --git a/hll/dynamodb-mapper/dynamodb-mapper-ops-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/operations/rendering/OperationRenderer.kt b/hll/dynamodb-mapper/dynamodb-mapper-ops-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/operations/rendering/OperationRenderer.kt index d6a74ff0b97..900aee48df4 100644 --- a/hll/dynamodb-mapper/dynamodb-mapper-ops-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/operations/rendering/OperationRenderer.kt +++ b/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) { diff --git a/hll/dynamodb-mapper/dynamodb-mapper/common/src/aws/sdk/kotlin/hll/dynamodbmapper/model/Item.kt b/hll/dynamodb-mapper/dynamodb-mapper/common/src/aws/sdk/kotlin/hll/dynamodbmapper/model/Item.kt index 3341b58ada1..779e93ad057 100644 --- a/hll/dynamodb-mapper/dynamodb-mapper/common/src/aws/sdk/kotlin/hll/dynamodbmapper/model/Item.kt +++ b/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): Item = mapOf(*pai */ @JvmName("itemOfPairStringAny") internal fun itemOf(vararg pairs: Pair): 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): Item = filterKeys { keys.contains(it) }.toItem() diff --git a/hll/dynamodb-mapper/dynamodb-mapper/common/test/aws/sdk/kotlin/hll/dynamodbmapper/items/SimpleItemConverterTest.kt b/hll/dynamodb-mapper/dynamodb-mapper/common/test/aws/sdk/kotlin/hll/dynamodbmapper/items/SimpleItemConverterTest.kt index dc9e33d22c4..1b580fbc3f2 100644 --- a/hll/dynamodb-mapper/dynamodb-mapper/common/test/aws/sdk/kotlin/hll/dynamodbmapper/items/SimpleItemConverterTest.kt +++ b/hll/dynamodb-mapper/dynamodb-mapper/common/test/aws/sdk/kotlin/hll/dynamodbmapper/items/SimpleItemConverterTest.kt @@ -13,7 +13,7 @@ import kotlin.test.assertTrue class SimpleItemConverterTest { @Test - fun testSomething() { + fun testBasicConversion() { val converter = SimpleItemConverter( ::ProductBuilder, ProductBuilder::build, @@ -33,6 +33,24 @@ class SimpleItemConverterTest { val unconverted = converter.convertFrom(item) assertEquals(foo, unconverted) } + + @Test + fun testKeyOnlyConversion() { + val converter = SimpleItemConverter( + ::ProductBuilder, + ProductBuilder::build, + AttributeDescriptor("id", Product::id, ProductBuilder::id::set, IntConverter), + AttributeDescriptor("name", Product::name, ProductBuilder::name::set, StringConverter), + AttributeDescriptor("in-stock", Product::inStock, ProductBuilder::inStock::set, BooleanConverter), + ) + + val foo = Product(42, "Foo 2.0", inStock = true) + val item = converter.convertTo(foo, setOf("id", "name")) + + assertEquals(2, item.size) + assertEquals(42, item.getValue("id").asN().toInt()) + assertEquals("Foo 2.0", item.getValue("name").asS()) + } } private data class Product( diff --git a/hll/dynamodb-mapper/dynamodb-mapper/jvm/test/aws/sdk/kotlin/hll/dynamodbmapper/operations/PaginatedScanTest.kt b/hll/dynamodb-mapper/dynamodb-mapper/jvm/test/aws/sdk/kotlin/hll/dynamodbmapper/operations/PaginatedScanTest.kt new file mode 100644 index 00000000000..17870798ee9 --- /dev/null +++ b/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 { + 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?): 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") + } +} diff --git a/hll/dynamodb-mapper/dynamodb-mapper/jvm/test/aws/sdk/kotlin/hll/dynamodbmapper/testutils/DdbClientExtensions.kt b/hll/dynamodb-mapper/dynamodb-mapper/jvm/test/aws/sdk/kotlin/hll/dynamodbmapper/testutils/DdbClientExtensions.kt index 76bb72717e0..a37c8ca7ca4 100644 --- a/hll/dynamodb-mapper/dynamodb-mapper/jvm/test/aws/sdk/kotlin/hll/dynamodbmapper/testutils/DdbClientExtensions.kt +++ b/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) { - 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) } } }