bugfix: Consume all result pages in DynamoDbLogicalDb.batchLoad

samples/musiclibrary2/src/main/kotlin/app/cash/tempest2/musiclibrary/MusicItem.kt

@@ -54,6 +54,7 @@ class MusicItem {
   var track_title: String? = null
   @get:DynamoDbConvertedBy(DurationTypeConverter::class)
   var run_length: Duration? = null
+  var track_description: String? = null
 
   // PlaylistInfo.
   var playlist_name: String? = null

samples/musiclibrary2/src/main/kotlin/app/cash/tempest2/musiclibrary/MusicTable.kt

@@ -94,14 +94,16 @@ data class AlbumTrack(
   @Attribute(name = "sort_key", prefix = "TRACK_")
   val track_token: String,
   val track_title: String,
-  val run_length: Duration
+  val run_length: Duration,
+  val track_description: String = "",
 ) {
   constructor(
     album_token: String,
     track_number: Long,
     track_title: String,
-    run_length: Duration
-  ) : this(album_token, "%016x".format(track_number), track_title, run_length)
+    run_length: Duration,
+    track_description: String = "",
+  ) : this(album_token, "%016x".format(track_number), track_title, run_length, track_description)
 
   @Transient
   val key = Key(album_token, track_token)

tempest2/src/main/kotlin/app/cash/tempest2/internal/DynamoDbLogicalDb.kt

@@ -101,8 +101,8 @@ internal class DynamoDbLogicalDb(
         returnConsumedCapacity
       )
 
-      val pages = batchRequests.map {
-        dynamoDbEnhancedClient.batchGetItem(it).iterator().next()
+      val pages = batchRequests.flatMap {
+        dynamoDbEnhancedClient.batchGetItem(it).iterator().asSequence().toList()
       }
 
       return toBatchLoadResponse(keysByTable, requestKeys, pages)

tempest2/src/test/kotlin/app/cash/tempest2/AsyncLogicalDbBatchTest.kt

@@ -23,7 +23,6 @@ import app.cash.tempest2.musiclibrary.testDb
 import app.cash.tempest2.testing.asyncLogicalDb
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.Test
-import org.junit.jupiter.api.Disabled
 import org.junit.jupiter.api.extension.RegisterExtension
 import java.time.Duration
 

tempest2/src/test/kotlin/app/cash/tempest2/LogicalDbBatchTest.kt

@@ -86,6 +86,39 @@ class LogicalDbBatchTest {
     assertThat(loadedItems.getItems<PlaylistInfo>()).containsExactly(playlistInfo)
   }
 
+  @Test
+  fun `batchLoad more than 16MB`() = runBlockingTest {
+    val twoHundredKbDescription = "a".repeat(200_000)
+
+    // Generate ~20MB of data. Dynamo only supports reading 16MB per batch request, so we need to handled retries to
+    // get the second page of data
+    val albumTracks = (1 until (MAX_BATCH_READ)).map {
+      AlbumTrack(
+        "ALBUM_1",
+        it.toLong(),
+        "track $it",
+        Duration.parse("PT3M28S"),
+        twoHundredKbDescription,
+      )
+    }
+    for (albumTrack in albumTracks) {
+      musicTable.albumTracks.save(albumTrack)
+    }
+    val playlistInfo = PlaylistInfo(
+      playlist_token = "PLAYLIST_1",
+      playlist_name = "WFH Music",
+      playlist_tracks = listOf(AlbumTrack.Key("ALBUM_1", 1))
+    )
+    musicTable.playlistInfo.save(playlistInfo)
+
+    val loadedItems = musicDb.batchLoad(
+      PlaylistInfo.Key("PLAYLIST_1"),
+      *(albumTracks.map { AlbumTrack.Key("ALBUM_1", track_number = it.track_number) }.toTypedArray())
+    )
+    assertThat(loadedItems.getItems<AlbumTrack>()).containsExactlyInAnyOrderElementsOf(albumTracks)
+    assertThat(loadedItems.getItems<PlaylistInfo>()).containsExactly(playlistInfo)
+  }
+
   @Test
   fun batchLoadMultipleTables() {
     val albumTracks = listOf(

tempest2/src/test/kotlin/app/cash/tempest2/WritingPagerTest.kt

@@ -31,8 +31,6 @@ import org.junit.jupiter.api.extension.RegisterExtension
 import software.amazon.awssdk.enhanced.dynamodb.Expression
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue
 import software.amazon.awssdk.services.dynamodb.model.TransactionCanceledException
-import java.util.function.BiFunction
-import kotlin.streams.toList
 
 class WritingPagerTest {
 

tempest2/src/test/kotlin/app/cash/tempest2/internal/SchemaTest.kt

@@ -86,7 +86,7 @@ class SchemaTest {
 
     assertThatIllegalArgumentException().isThrownBy {
       musicDb.music.inlineView(Any::class, BadItem5::class)
-    }.withMessage("Expect nonexistent_attribute, required by class app.cash.tempest2.internal.BadItem5, to be declared in class app.cash.tempest2.musiclibrary.MusicItem. But found [album_title, artist_name, genre_name, label_name, partition_key, playlist_name, playlist_size, playlist_tracks, playlist_version, release_date, run_length, sort_key, track_title, track_token]. Use @Transient to exclude it. You might see this error message if the property name starts with `is`. See: https://github.com/cashapp/tempest/issues/53.")
+    }.withMessage("Expect nonexistent_attribute, required by class app.cash.tempest2.internal.BadItem5, to be declared in class app.cash.tempest2.musiclibrary.MusicItem. But found [album_title, artist_name, genre_name, label_name, partition_key, playlist_name, playlist_size, playlist_tracks, playlist_version, release_date, run_length, sort_key, track_description, track_title, track_token]. Use @Transient to exclude it. You might see this error message if the property name starts with `is`. See: https://github.com/cashapp/tempest/issues/53.")
 
     assertThatIllegalArgumentException().isThrownBy {
       musicDb.music.inlineView(Any::class, BadItem6::class)