Skip to content

Commit d08b78e

Browse files
ianbotsfianbotsf
authored
fix: correctly convert key fields during paginated Scan and Query operations (#1614)
1 parent 90d5f6a commit d08b78e

File tree

8 files changed

+165
-14
lines changed

8 files changed

+165
-14
lines changed

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
{
2+
"id": "02bbf5e7-68d7-480d-a94c-35cd39dfe0db",
3+
"type": "bugfix",
4+
"description": "Correctly convert key fields during paginated `Scan` and `Query` operations",
5+
"issues": [
6+
"awslabs/aws-sdk-kotlin#1596"
7+
]
8+
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ public object MapperTypes {
6262
}
6363

6464
public object Model {
65+
public val intersectKeys: TypeRef = TypeRef(MapperPkg.Hl.Model, "intersectKeys")
6566
public fun tablePartitionKey(objectType: TypeRef, pkType: TypeRef): TypeRef = TypeRef(
6667
MapperPkg.Hl.Model,
6768
"Table.PartitionKey",

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

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,10 @@ package aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model
77
import aws.sdk.kotlin.hll.codegen.model.*
88
import aws.sdk.kotlin.hll.dynamodbmapper.codegen.model.MapperPkg
99
import aws.sdk.kotlin.hll.dynamodbmapper.codegen.model.MapperTypes
10-
import aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model.ExpressionArgumentsType.*
11-
import aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model.ExpressionLiteralType.*
10+
import aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model.ExpressionArgumentsType.AttributeNames
11+
import aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model.ExpressionArgumentsType.AttributeValues
12+
import aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model.ExpressionLiteralType.Filter
13+
import aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model.ExpressionLiteralType.KeyCondition
1214
import aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.model.MemberCodegenBehavior.*
1315

1416
/**
@@ -156,6 +158,6 @@ private val rules = listOf(
156158

157159
// Mappable members
158160
Rule(".*".toRegex(), Types.Kotlin.list(MapperTypes.AttributeMap), ListMapAll),
159-
Rule("key", MapperTypes.AttributeMap, MapKeys),
161+
Rule("key|lastEvaluatedKey|exclusiveStartKey".toRegex(), MapperTypes.AttributeMap, MapKeys),
160162
Rule(".*".toRegex(), MapperTypes.AttributeMap, MapAll),
161163
)

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
*/
55
package aws.sdk.kotlin.hll.dynamodbmapper.codegen.operations.rendering
66

7-
import aws.sdk.kotlin.hll.codegen.core.*
7+
import aws.sdk.kotlin.hll.codegen.core.ImportDirective
88
import aws.sdk.kotlin.hll.codegen.model.Member
99
import aws.sdk.kotlin.hll.codegen.model.Operation
1010
import aws.sdk.kotlin.hll.codegen.model.Structure
@@ -106,9 +106,10 @@ internal class OperationRenderer(
106106
requestMembers(MemberCodegenBehavior.PassThrough) { write("#L = this@convert.#L", name, highLevel.name) }
107107
requestMembers(MemberCodegenBehavior.MapKeys) {
108108
write(
109-
"this@convert.#L?.let { #L = schema.converter.convertTo(it, schema.keyAttributeNames) }",
109+
"this@convert.#L?.let { #L = schema.converter.convertTo(it, schema.keyAttributeNames).#T(schema.keyAttributeNames) }",
110110
highLevel.name,
111111
name,
112+
MapperTypes.Model.intersectKeys,
112113
)
113114
}
114115
requestMembers(MemberCodegenBehavior.MapAll) {

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import aws.sdk.kotlin.hll.dynamodbmapper.model.internal.ItemImpl
88
import aws.sdk.kotlin.hll.dynamodbmapper.util.dynamicAttr
99
import aws.sdk.kotlin.services.dynamodb.model.AttributeValue
1010
import aws.smithy.kotlin.runtime.ExperimentalApi
11-
import kotlin.jvm.JvmName
1211

1312
/**
1413
* 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
5958
*/
6059
@JvmName("itemOfPairStringAny")
6160
internal fun itemOf(vararg pairs: Pair<String, Any?>): Item = mapOf(*pairs).toItem()
61+
62+
/**
63+
* Truncate this item to include only the specified key attributes, removing any non-key attributes
64+
* @param keys The collection of key names which should be preserved
65+
*/
66+
internal fun Item.intersectKeys(keys: Collection<String>): Item = filterKeys { keys.contains(it) }.toItem()

‎hll/dynamodb-mapper/dynamodb-mapper/common/test/aws/sdk/kotlin/hll/dynamodbmapper/items/SimpleItemConverterTest.kt‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import kotlin.test.assertTrue
1313

1414
class SimpleItemConverterTest {
1515
@Test
16-
fun testSomething() {
16+
fun testBasicConversion() {
1717
val converter = SimpleItemConverter(
1818
::ProductBuilder,
1919
ProductBuilder::build,
@@ -33,6 +33,24 @@ class SimpleItemConverterTest {
3333
val unconverted = converter.convertFrom(item)
3434
assertEquals(foo, unconverted)
3535
}
36+
37+
@Test
38+
fun testKeyOnlyConversion() {
39+
val converter = SimpleItemConverter(
40+
::ProductBuilder,
41+
ProductBuilder::build,
42+
AttributeDescriptor("id", Product::id, ProductBuilder::id::set, IntConverter),
43+
AttributeDescriptor("name", Product::name, ProductBuilder::name::set, StringConverter),
44+
AttributeDescriptor("in-stock", Product::inStock, ProductBuilder::inStock::set, BooleanConverter),
45+
)
46+
47+
val foo = Product(42, "Foo 2.0", inStock = true)
48+
val item = converter.convertTo(foo, setOf("id", "name"))
49+
50+
assertEquals(2, item.size)
51+
assertEquals(42, item.getValue("id").asN().toInt())
52+
assertEquals("Foo 2.0", item.getValue("name").asS())
53+
}
3654
}
3755

3856
private data class Product(

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

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
package aws.sdk.kotlin.hll.dynamodbmapper.operations
2+
3+
import aws.sdk.kotlin.hll.dynamodbmapper.items.*
4+
import aws.sdk.kotlin.hll.dynamodbmapper.model.Item
5+
import aws.sdk.kotlin.hll.dynamodbmapper.model.itemOf
6+
import aws.sdk.kotlin.hll.dynamodbmapper.testutils.DdbLocalTest
7+
import aws.sdk.kotlin.hll.dynamodbmapper.values.scalars.IntConverter
8+
import aws.sdk.kotlin.hll.dynamodbmapper.values.scalars.StringConverter
9+
import kotlinx.coroutines.flow.map
10+
import kotlinx.coroutines.flow.toSet
11+
import kotlinx.coroutines.test.runTest
12+
import kotlin.test.assertEquals
13+
import kotlin.test.assertTrue
14+
15+
class PaginatedScanTest : DdbLocalTest() {
16+
companion object {
17+
private const val TABLE_NAME = "paginated-scan-test"
18+
19+
private fun rankName(rank: Int) = when (rank) {
20+
2, 3, 4, 5, 6, 7, 8, 9, 10 -> rank.toString()
21+
11 -> "Jack"
22+
12 -> "Queen"
23+
13 -> "King"
24+
14 -> "Ace"
25+
else -> "Unknown ($rank)"
26+
}
27+
28+
private data class Card(
29+
var rank: Int = 14,
30+
var suit: String = "Spades",
31+
var description: String = "This is the Ace of Spades. ".repeat(2000),
32+
) {
33+
val rankName: String
34+
get() = rankName(rank)
35+
36+
override fun toString() = "$rankName of $suit" // Easier when debugging big outputs
37+
}
38+
39+
private val converter = object : ItemConverter<Card> {
40+
val delegate = SimpleItemConverter(
41+
::Card,
42+
{ this },
43+
AttributeDescriptor("rank", Card::rank, Card::rank::set, IntConverter),
44+
AttributeDescriptor("suit", Card::suit, Card::suit::set, StringConverter),
45+
AttributeDescriptor("description", Card::description, Card::description::set, StringConverter),
46+
)
47+
48+
override fun convertFrom(to: Item): Card = delegate.convertFrom(to)
49+
50+
override fun convertTo(from: Card, onlyAttributes: Set<String>?): Item =
51+
delegate.convertTo(from, null) // Ignore `onlyAttributes` arg to mock badly-behaved converter
52+
}
53+
54+
private val schema = ItemSchema(converter, KeySpec.String("suit"), KeySpec.Number("rank"))
55+
56+
private val allCards = listOf("Spades", "Clubs", "Hearts", "Diamonds").flatMap { suit ->
57+
(2..14).map { rank ->
58+
Card(rank, suit, "This is the ${rankName(rank)} of $suit. ".repeat(2000))
59+
}
60+
}.toSet()
61+
}
62+
63+
@BeforeAll
64+
fun setUp() = runTest {
65+
createTable(
66+
name = TABLE_NAME,
67+
schema = schema,
68+
items = allCards.map { card ->
69+
itemOf(
70+
"rank" to card.rank,
71+
"suit" to card.suit,
72+
"description" to card.description,
73+
)
74+
}.toTypedArray(),
75+
)
76+
}
77+
78+
@Test
79+
fun testPaginatedScan() = runTest {
80+
val mapper = mapper()
81+
val table = mapper.getTable(TABLE_NAME, schema)
82+
var pages = 0
83+
84+
val actual = table
85+
.scanPaginated { }
86+
.map {
87+
pages++
88+
it
89+
}
90+
.items()
91+
.toSet()
92+
assertEquals(allCards, actual)
93+
assertTrue(pages > 1, "Got only $pages pages but expected at least 2")
94+
}
95+
96+
@Test
97+
fun testPaginatedScanWithOffset() = runTest {
98+
val mapper = mapper()
99+
val table = mapper.getTable(TABLE_NAME, schema)
100+
val startKey = allCards.first()
101+
var pages = 0
102+
103+
val actual = table
104+
.scanPaginated {
105+
exclusiveStartKey = startKey
106+
}
107+
.map {
108+
pages++
109+
it
110+
}
111+
.items()
112+
.toSet()
113+
assertTrue(startKey !in actual)
114+
assertTrue(pages > 1, "Got only $pages pages but expected at least 2")
115+
}
116+
}

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -88,17 +88,17 @@ suspend fun DynamoDbClient.getItem(tableName: String, vararg keys: Pair<String,
8888
* @param items A collection of maps of strings to values to be mapped and persisted to the table
8989
*/
9090
suspend fun DynamoDbClient.putItems(tableName: String, items: List<Item>) {
91-
val writeRequests = items.map { item ->
92-
WriteRequest {
93-
putRequest {
94-
this.item = item
91+
val batches = items
92+
.map { item ->
93+
WriteRequest {
94+
putRequest { this.item = item }
9595
}
9696
}
97-
}
97+
.chunked(25) // Max batchWriteItem page size
9898

99-
if (writeRequests.isNotEmpty()) {
99+
batches.forEach { batch ->
100100
batchWriteItem {
101-
requestItems = mapOf(tableName to writeRequests)
101+
requestItems = mapOf(tableName to batch)
102102
}
103103
}
104104
}

0 commit comments

Comments
 (0)