Skip to content

perf(misk-aws2-dynamodb): truncate tables on external fixture reset instead of drop/recreate #3864

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rnc505 wants to merge 2 commits into
base: master
Choose a base branch
from
Open

perf(misk-aws2-dynamodb): truncate tables on external fixture reset instead of drop/recreate #3864

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a6d03ab
perf(misk-aws2-dynamodb): truncate tables on external reset
rnc505 Jul 20, 2026
51c0498
feat(misk-aws2-dynamodb): allow drop-recreate resets
rnc505 Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions misk-aws2-dynamodb/api/misk-aws2-dynamodb.api
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,9 @@ public final class misk/aws2/dynamodb/testing/DynamoDbTable {

public final class misk/aws2/dynamodb/testing/ExternalTestDynamoDbClientModule : misk/inject/KAbstractModule {
public fun <init> (ILjava/util/List;)V
public fun <init> (ILjava/util/List;Lmisk/aws2/dynamodb/testing/ResetStrategy;)V
public synthetic fun <init> (ILjava/util/List;Lmisk/aws2/dynamodb/testing/ResetStrategy;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
public fun <init> (ILmisk/aws2/dynamodb/testing/ResetStrategy;[Lmisk/aws2/dynamodb/testing/DynamoDbTable;)V
public fun <init> (I[Lmisk/aws2/dynamodb/testing/DynamoDbTable;)V
public final fun provideRequiredTables ()Ljava/util/List;
public final fun providesAmazonDynamoDB (Lapp/cash/tempest2/testing/TestDynamoDbClient;)Lsoftware/amazon/awssdk/services/dynamodb/DynamoDbClient;
Expand All @@ -89,6 +92,14 @@ public final class misk/aws2/dynamodb/testing/ParallelTestsTableNameMapper : mis
public fun mapName (Ljava/lang/String;)Ljava/lang/String;
}

public final class misk/aws2/dynamodb/testing/ResetStrategy : java/lang/Enum {
public static final field DROP_RECREATE Lmisk/aws2/dynamodb/testing/ResetStrategy;
public static final field TRUNCATE Lmisk/aws2/dynamodb/testing/ResetStrategy;
public static fun getEntries ()Lkotlin/enums/EnumEntries;
public static fun valueOf (Ljava/lang/String;)Lmisk/aws2/dynamodb/testing/ResetStrategy;
public static fun values ()[Lmisk/aws2/dynamodb/testing/ResetStrategy;
}

public final class misk/aws2/dynamodb/testing/TestDynamoDb : com/google/common/util/concurrent/Service, misk/testing/TestFixture {
public fun <init> (Lapp/cash/tempest2/testing/internal/TestDynamoDbService;)V
public fun addListener (Lcom/google/common/util/concurrent/Service$Listener;Ljava/util/concurrent/Executor;)V
Expand Down
151 changes: 151 additions & 0 deletions ...namodb/src/test/kotlin/misk/aws2/dynamodb/testing/ExternalTestDynamoDbClientModuleTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,151 @@
package misk.aws2.dynamodb.testing

import app.cash.tempest2.testing.JvmDynamoDbServer
import app.cash.tempest2.testing.TestDynamoDbClient
import jakarta.inject.Inject
import java.net.ServerSocket
import java.time.LocalDate
import misk.MiskTestingServiceModule
import misk.inject.KAbstractModule
import misk.testing.MiskTest
import misk.testing.MiskTestModule
import misk.testing.TestFixture
import org.assertj.core.api.Assertions.assertThat
import org.junit.jupiter.api.AfterAll
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.TestInstance
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient
import software.amazon.awssdk.enhanced.dynamodb.model.EnhancedGlobalSecondaryIndex
import software.amazon.awssdk.enhanced.dynamodb.model.QueryConditional
import software.amazon.awssdk.services.dynamodb.model.ProjectionType

@MiskTest(startService = true)
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
class ExternalTestDynamoDbClientModuleTest {
private val server =
JvmDynamoDbServer.Factory.create(ServerSocket(0).use { it.localPort }).also { it.startAsync().awaitRunning() }

@MiskTestModule val module = ExternalDynamoTestModule(server.port)

@Inject lateinit var client: TestDynamoDbClient
@Inject lateinit var testFixtures: Set<TestFixture>

@AfterAll
fun stopServer() {
server.stopAsync().awaitTerminated()
}

@Test
fun `reset truncates tables and indexes`() {
val movieTable =
DynamoDbEnhancedClient.builder()
.dynamoDbClient(client.dynamoDb)
.build()
.table(client.tables.single().tableName, AbstractDynamoDbTest.MOVIE_TABLE_SCHEMA)
val tableCreatedAt =
client.dynamoDb.describeTable { it.tableName(client.tables.single().tableName) }.table().creationDateTime()

assertThat(movieTable.scan().items()).isEmpty()

repeat(30) { index ->
movieTable.putItem(
DyMovie().apply {
name = "Movie $index"
release_date = LocalDate.of(2000, 1, 1).plusDays(index.toLong())
directed_by = "Director"
}
)
}
assertThat(movieTable.scan().items()).hasSize(30)
assertThat(
movieTable
.index("movies.release_date_index")
.query(QueryConditional.keyEqualTo { it.partitionValue("Director") })
.flatMap { it.items() }
)
.hasSize(30)

testFixtures.single { it.javaClass.simpleName == "TestDynamoDbFixture" }.reset()

assertThat(
client.dynamoDb.describeTable { it.tableName(client.tables.single().tableName) }.table().creationDateTime()
)
.isEqualTo(tableCreatedAt)
assertThat(movieTable.scan().items()).isEmpty()
assertThat(
movieTable
.index("movies.release_date_index")
.query(QueryConditional.keyEqualTo { it.partitionValue("Director") })
.flatMap { it.items() }
)
.isEmpty()
}
}

@MiskTest(startService = true)
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
class DropRecreateExternalTestDynamoDbClientModuleTest {
private val server =
JvmDynamoDbServer.Factory.create(ServerSocket(0).use { it.localPort }).also { it.startAsync().awaitRunning() }

@MiskTestModule val module = ExternalDynamoTestModule(server.port, ResetStrategy.DROP_RECREATE)

@Inject lateinit var client: TestDynamoDbClient
@Inject lateinit var testFixtures: Set<TestFixture>

@AfterAll
fun stopServer() {
server.stopAsync().awaitTerminated()
}

@Test
fun `reset drops and recreates tables`() {
val tableName = client.tables.single().tableName
val movieTable =
DynamoDbEnhancedClient.builder()
.dynamoDbClient(client.dynamoDb)
.build()
.table(tableName, AbstractDynamoDbTest.MOVIE_TABLE_SCHEMA)
val tableCreatedAt = client.dynamoDb.describeTable { it.tableName(tableName) }.table().creationDateTime()

movieTable.putItem(
DyMovie().apply {
name = "Movie"
release_date = LocalDate.of(2000, 1, 1)
directed_by = "Director"
}
)

testFixtures.single { it.javaClass.simpleName == "TestDynamoDbFixture" }.reset()

assertThat(client.dynamoDb.describeTable { it.tableName(tableName) }.table().creationDateTime())
.isNotEqualTo(tableCreatedAt)
assertThat(movieTable.scan().items()).isEmpty()
}
}

class ExternalDynamoTestModule(
private val port: Int,
private val resetStrategy: ResetStrategy = ResetStrategy.TRUNCATE,
) : KAbstractModule() {
override fun configure() {
install(MiskTestingServiceModule())
install(
ExternalTestDynamoDbClientModule(
port,
listOf(
DynamoDbTable("movies", DyMovie::class) { createTableEnhancedRequest ->
createTableEnhancedRequest.globalSecondaryIndices(
EnhancedGlobalSecondaryIndex.builder()
.indexName("movies.release_date_index")
.projection { it.projectionType(ProjectionType.ALL) }
.provisionedThroughput { it.readCapacityUnits(40_000L).writeCapacityUnits(40_000L) }
.build()
)
}
),
resetStrategy,
)
)
}
}
112 changes: 101 additions & 11 deletions ...db/src/testFixtures/kotlin/misk/aws2/dynamodb/testing/ExternalTestDynamoDbClientModule.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import com.google.common.util.concurrent.AbstractService
import com.google.inject.Provides
import jakarta.inject.Inject
import jakarta.inject.Singleton
import java.util.concurrent.ConcurrentHashMap
import misk.ServiceModule
import misk.aws2.dynamodb.DynamoDbService
import misk.aws2.dynamodb.RequiredDynamoDbTable
Expand All @@ -17,25 +18,51 @@ import misk.inject.KAbstractModule
import misk.testing.TestFixture
import misk.testing.updateForParallelTests
import software.amazon.awssdk.services.dynamodb.DynamoDbClient
import software.amazon.awssdk.services.dynamodb.model.AttributeValue
import software.amazon.awssdk.services.dynamodb.model.DeleteRequest
import software.amazon.awssdk.services.dynamodb.model.DeleteTableRequest
import software.amazon.awssdk.services.dynamodb.model.ResourceNotFoundException
import software.amazon.awssdk.services.dynamodb.model.WriteRequest
import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient

enum class ResetStrategy {
TRUNCATE,
DROP_RECREATE,
}

/**
* Unlike the InProcessDynamoDbModule and DockerDynamoDbModule classes, this module does not internally start DynamoDB,
* and instead relies on an external DynamoDB server already running on the given port, e.g. using a Gradle task as a
* dependency of the test task, which starts DynamoDB in a Docker container.
*
* By default, tables are created once per injector and rows are truncated on later resets to avoid slow DynamoDB DDL
* calls before every test. Use [ResetStrategy.DROP_RECREATE] if tests use DynamoDB Streams or modify table definitions:
* truncation preserves streams (including delete records from the reset) and does not undo changes to TTL, billing
* mode, or indexes.
*/
class ExternalTestDynamoDbClientModule(private val port: Int, originalTables: List<DynamoDbTable>) : KAbstractModule() {
class ExternalTestDynamoDbClientModule
@JvmOverloads
constructor(
private val port: Int,
originalTables: List<DynamoDbTable>,
private val resetStrategy: ResetStrategy = ResetStrategy.TRUNCATE,
) : KAbstractModule() {

private val tables = originalTables.map { it.copy(tableName = ParallelTestsTableNameMapper.mapName(it.tableName)) }

constructor(port: Int, vararg tables: DynamoDbTable) : this(port, tables.toList())

constructor(
port: Int,
resetStrategy: ResetStrategy,
vararg tables: DynamoDbTable,
) : this(port, tables.toList(), resetStrategy)

override fun configure() {
for (table in tables) {
multibind<DynamoDbTable>().toInstance(table)
}
bind<ResetStrategy>().toInstance(resetStrategy)
bind<DynamoDbService>().to<TestDynamoDbService>()
install(ServiceModule<DynamoDbService>().dependsOn<TestDynamoDbFixture>())
install(ServiceModule<TestDynamoDbFixture>())
Expand Down Expand Up @@ -71,8 +98,14 @@ class ExternalTestDynamoDbClientModule(private val port: Int, originalTables: Li
@Singleton
private class TestDynamoDbFixture
@Inject
constructor(private val client: TestDynamoDbClient, private val tables: List<DynamoDbTable>) :
AbstractIdleService(), TestFixture {
constructor(
private val client: TestDynamoDbClient,
private val tables: List<DynamoDbTable>,
private val resetStrategy: ResetStrategy,
) : AbstractIdleService(), TestFixture {

private var tablesCreated = false
private val keyAttributeNames = ConcurrentHashMap<String, List<String>>()

override fun startUp() {
reset()
Expand All @@ -81,20 +114,77 @@ constructor(private val client: TestDynamoDbClient, private val tables: List<Dyn
override fun shutDown() {}

override fun reset() {
if (!tablesCreated || resetStrategy == ResetStrategy.DROP_RECREATE) {
recreateTables()
tablesCreated = true
return
}

client.tables.forEach(::truncate)
}

/**
* The first reset recreates tables to remove any state left by a previous process. Under [ResetStrategy.TRUNCATE],
* later resets only delete rows, avoiding slow DynamoDB DDL calls before every test.
*/
private fun recreateTables() {
for (tableName in tables.map { it.tableName }) {
try {
client.dynamoDb.deleteTable(DeleteTableRequest.builder().tableName(tableName).build())
} catch (e: ResourceNotFoundException) {
// Ignore if the table doesn't exist
}
} catch (_: ResourceNotFoundException) {}
}
for (table in
tables.map { table ->
TestTable.create(table.tableName, table.tableClass) { table.configureTable(it.toBuilder()).build() }
}) {
client.dynamoDb.createTable(table)

client.tables.forEach { client.dynamoDb.createTable(it) }
}

private fun truncate(table: TestTable) {
val tableName = table.tableName
val keyNames =
keyAttributeNames.getOrPut(tableName) {
client.dynamoDb.describeTable { it.tableName(tableName) }.table().keySchema().map { it.attributeName() }
}
val keyAliases = keyNames.withIndex().associate { (index, name) -> "#k$index" to name }
var exclusiveStartKey: Map<String, AttributeValue>? = null

do {
val startKey = exclusiveStartKey
val scan =
client.dynamoDb.scan {
it
.tableName(tableName)
.projectionExpression(keyAliases.keys.joinToString(", "))
.expressionAttributeNames(keyAliases)
.consistentRead(true)
if (startKey != null) it.exclusiveStartKey(startKey)
}
deleteItems(tableName, scan.items())
exclusiveStartKey = if (scan.hasLastEvaluatedKey()) scan.lastEvaluatedKey() else null
} while (exclusiveStartKey != null)
}

private fun deleteItems(tableName: String, keys: List<Map<String, AttributeValue>>) {
for (chunk in keys.chunked(MAX_BATCH_WRITE_ITEMS)) {
var writes =
chunk.map { key -> WriteRequest.builder().deleteRequest(DeleteRequest.builder().key(key).build()).build() }
var attempts = 0
while (writes.isNotEmpty()) {
if (attempts > 0) Thread.sleep(UNPROCESSED_RETRY_DELAY_MS)
check(attempts < MAX_BATCH_WRITE_ATTEMPTS) {
"BatchWriteItem left ${writes.size} unprocessed delete(s) for table $tableName after " +
"$MAX_BATCH_WRITE_ATTEMPTS attempts"
}
val response = client.dynamoDb.batchWriteItem { it.requestItems(mapOf(tableName to writes)) }
writes = response.unprocessedItems()[tableName].orEmpty()
attempts++
}
}
}

private companion object {
const val MAX_BATCH_WRITE_ITEMS = 25
const val MAX_BATCH_WRITE_ATTEMPTS = 5
const val UNPROCESSED_RETRY_DELAY_MS = 50L
}
}

@Singleton
Expand Down