Embedded backfila support changing parameters in validate (#236)

Added ages ago and never added to embedded
cashapp · Apr 6, 2022 · 2dba0e1 · 2dba0e1
1 parent 56aaed4
commit 2dba0e1
Show file tree

Hide file tree

Showing 9 changed files with 108 additions and 29 deletions.
diff --git a/backfila-embedded/src/main/kotlin/app/cash/backfila/embedded/internal/EmbeddedBackfila.kt b/backfila-embedded/src/main/kotlin/app/cash/backfila/embedded/internal/EmbeddedBackfila.kt
@@ -20,10 +20,10 @@ import java.util.concurrent.atomic.AtomicInteger
 import javax.inject.Inject
 import javax.inject.Singleton
 import kotlin.reflect.KClass
+import kotlin.reflect.jvm.jvmName
 import okio.ByteString
 import retrofit2.Call
 import retrofit2.mock.Calls
-import kotlin.reflect.jvm.jvmName
 
 /**
  * A small implementation of Backfila suitable for use in test cases and development mode. Unlike
@@ -79,7 +79,7 @@ internal class EmbeddedBackfila @Inject internal constructor(
     val run = EmbeddedBackfillRun<Backfill>(
       operator = operator,
       dryRun = createRequest.dry_run,
-      parameters = createRequest.parameter_map,
+      parameters = createRequest.parameter_map.toMutableMap(),
       rangeStart = createRequest.pkey_range_start?.utf8(),
       rangeEnd = createRequest.pkey_range_end?.utf8()
     )
@@ -128,7 +128,7 @@ internal class EmbeddedBackfila @Inject internal constructor(
     return EmbeddedBackfillRun(
       operator = operator,
       dryRun = dryRun,
-      parameters = parameters,
+      parameters = parameters.toMutableMap(),
       rangeStart = rangeStart,
       rangeEnd = rangeEnd
     )

diff --git a/backfila-embedded/src/main/kotlin/app/cash/backfila/embedded/internal/EmbeddedBackfillRun.kt b/backfila-embedded/src/main/kotlin/app/cash/backfila/embedded/internal/EmbeddedBackfillRun.kt
@@ -1,8 +1,8 @@
 package app.cash.backfila.embedded.internal
 
 import app.cash.backfila.client.Backfill
-import app.cash.backfila.embedded.BackfillRun
 import app.cash.backfila.client.spi.BackfillOperator
+import app.cash.backfila.embedded.BackfillRun
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeRequest
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeResponse
 import app.cash.backfila.protos.clientservice.KeyRange
@@ -23,13 +23,13 @@ import okio.ByteString.Companion.encodeUtf8
 internal class EmbeddedBackfillRun<B : Backfill>(
   private val operator: BackfillOperator,
   override val dryRun: Boolean,
-  override val parameters: Map<String, ByteString>,
+  override val parameters: MutableMap<String, ByteString>,
   override val rangeStart: String?,
   override val rangeEnd: String?,
 
   override var batchSize: Long = 100L,
   override var scanSize: Long = 10_000L,
-  override var computeCountLimit: Long = 1L
+  override var computeCountLimit: Long = 1L,
 ) : BackfillRun<B> {
   override val backfill: B
     // This api on the operator returns the same backfill it was initialized with.
@@ -59,6 +59,7 @@ internal class EmbeddedBackfillRun<B : Backfill>(
         .dry_run(dryRun)
         .build()
     )
+    parameters += prepareBackfillResponse.parameters
     precomputeProgress = prepareBackfillResponse.partitions.associate {
       it.partition_name to MutablePartitionCursor(it.partition_name, it.backfill_range)
     }
@@ -194,7 +195,7 @@ internal class EmbeddedBackfillRun<B : Backfill>(
 
 private class MutablePartitionCursor(
   val partitionName: String,
-  val keyRange: KeyRange
+  val keyRange: KeyRange,
 ) {
   var previousEndKey: ByteString? = null
   var done: Boolean = false
@@ -209,7 +210,7 @@ data class PartitionCursor(
   val partitionName: String,
   val keyRange: KeyRange,
   val previousEndKey: ByteString?,
-  val done: Boolean
+  val done: Boolean,
 ) {
   fun utf8RangeStart() = keyRange.start?.utf8()
   fun utf8RangeEnd() = keyRange.end?.utf8()
@@ -220,7 +221,7 @@ data class BatchSnapshot(
   val partitionName: String,
   val batchRange: KeyRange,
   val scannedRecordCount: Long,
-  val matchingRecordCount: Long
+  val matchingRecordCount: Long,
 ) {
   fun utf8RangeStart() = batchRange.start.utf8()
   fun utf8RangeEnd() = batchRange.end.utf8()

diff --git a/client-base/src/test/kotlin/app/cash/backfila/client/BackfillsModule.kt b/client-base/src/test/kotlin/app/cash/backfila/client/BackfillsModule.kt
@@ -16,5 +16,6 @@ class BackfillsModule : KAbstractModule() {
     install(FixedSetBackfillModule.create<FixedSetTest.ToUpperCaseBackfill>())
     install(FixedSetBackfillModule.create<DeleteBackfillByTest.TenYearBackfill>())
     install(FixedSetBackfillModule.create<DeleteBackfillByTest.DeprecatedBackfill>())
+    install(FixedSetBackfillModule.create<PrepareWithParametersTest.PrepareWithParametersBackfill>())
   }
 }
diff --git a/client-base/src/test/kotlin/app/cash/backfila/client/DeleteBackfillByTest.kt b/client-base/src/test/kotlin/app/cash/backfila/client/DeleteBackfillByTest.kt
@@ -1,18 +1,18 @@
 package app.cash.backfila.client
 
-import app.cash.backfila.embedded.Backfila
 import app.cash.backfila.client.fixedset.FixedSetBackfill
 import app.cash.backfila.client.fixedset.FixedSetDatastore
 import app.cash.backfila.client.fixedset.FixedSetRow
+import app.cash.backfila.embedded.Backfila
 import com.google.inject.Module
+import java.time.LocalDate
+import java.time.ZoneOffset.UTC
 import javax.inject.Inject
 import misk.testing.MiskTest
 import misk.testing.MiskTestModule
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.Assertions.assertNull
 import org.junit.jupiter.api.Test
-import java.time.LocalDate
-import java.time.ZoneOffset.UTC
 
 @MiskTest(startService = true)
 class DeleteBackfillByTest {
@@ -28,20 +28,12 @@ class DeleteBackfillByTest {
     override fun runOne(row: FixedSetRow) {
       // We won't be running this one
     }
-
-    override fun checkBackfillConfig(backfillConfig: BackfillConfig<NoParameters>) {
-      // No parameters to check
-    }
   }
 
   class DeprecatedBackfill @Inject constructor() : FixedSetBackfill<NoParameters>() {
     override fun runOne(row: FixedSetRow) {
       // We won't be running this one
     }
-
-    override fun checkBackfillConfig(backfillConfig: BackfillConfig<NoParameters>) {
-      // No parameters to check
-    }
   }
 
   @Test fun `check that the delete_by was set`() {

diff --git a/client-base/src/test/kotlin/app/cash/backfila/client/FixedSetTest.kt b/client-base/src/test/kotlin/app/cash/backfila/client/FixedSetTest.kt
@@ -1,10 +1,10 @@
 package app.cash.backfila.client
 
-import app.cash.backfila.embedded.Backfila
-import app.cash.backfila.embedded.createWetRun
 import app.cash.backfila.client.fixedset.FixedSetBackfill
 import app.cash.backfila.client.fixedset.FixedSetDatastore
 import app.cash.backfila.client.fixedset.FixedSetRow
+import app.cash.backfila.embedded.Backfila
+import app.cash.backfila.embedded.createWetRun
 import com.google.inject.Module
 import java.util.Locale
 import javax.inject.Inject
@@ -28,10 +28,6 @@ class FixedSetTest {
       runOrder += row.value
       row.value = row.value.toUpperCase(Locale.ROOT)
     }
-
-    override fun checkBackfillConfig(backfillConfig: BackfillConfig<NoParameters>) {
-      // No parameters to check
-    }
   }
 
   @Test fun `happy path`() {

diff --git a/client-base/src/test/kotlin/app/cash/backfila/client/PrepareWithParametersTest.kt b/client-base/src/test/kotlin/app/cash/backfila/client/PrepareWithParametersTest.kt
@@ -0,0 +1,67 @@
+package app.cash.backfila.client
+
+import app.cash.backfila.client.fixedset.FixedSetBackfill
+import app.cash.backfila.client.fixedset.FixedSetDatastore
+import app.cash.backfila.client.fixedset.FixedSetRow
+import app.cash.backfila.embedded.Backfila
+import app.cash.backfila.embedded.createWetRun
+import com.google.inject.Module
+import javax.inject.Inject
+import misk.testing.MiskTest
+import misk.testing.MiskTestModule
+import okio.ByteString.Companion.encodeUtf8
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.Test
+
+/**
+ * Tests modifying parameters by returning them in the PrepareBackfillResponse.
+ */
+@MiskTest(startService = true)
+class PrepareWithParametersTest {
+  @Suppress("unused")
+  @MiskTestModule
+  val module: Module = TestingModule()
+
+  @Inject lateinit var backfila: Backfila
+  @Inject lateinit var datastore: FixedSetDatastore
+
+  data class PrepareParameters(
+    val favoriteNumber: Int? = null
+  )
+
+  class PrepareWithParametersBackfill @Inject constructor() : FixedSetBackfill<PrepareParameters>() {
+    override fun checkBackfillConfig(
+      backfillConfig: BackfillConfig<PrepareParameters>
+    ): ValidateResult<PrepareParameters> {
+      if (backfillConfig.parameters.favoriteNumber != null) {
+        return ValidateResult(backfillConfig.parameters)
+      }
+      return ValidateResult(
+        PrepareParameters(favoriteNumber = 42)
+      )
+    }
+
+    override fun runOne(row: FixedSetRow) {
+    }
+  }
+
+  @Test fun `add parameters from prepare response`() {
+    datastore.put("instance", "a", "b", "c")
+
+    val backfillRun = backfila.createWetRun<PrepareWithParametersBackfill>()
+    backfillRun.execute()
+    assertThat(backfillRun.parameters["favoriteNumber"]).isEqualTo("42".encodeUtf8())
+  }
+
+  @Test fun `keep original parameters`() {
+    datastore.put("instance", "a", "b", "c")
+
+    val backfillRun = backfila.createWetRun<PrepareWithParametersBackfill>(
+      parameters = PrepareParameters(
+        favoriteNumber = 12
+      )
+    )
+    backfillRun.execute()
+    assertThat(backfillRun.parameters["favoriteNumber"]).isEqualTo("12".encodeUtf8())
+  }
+}
diff --git a/client-base/src/test/kotlin/app/cash/backfila/client/fixedset/FixedSetBackfill.kt b/client-base/src/test/kotlin/app/cash/backfila/client/fixedset/FixedSetBackfill.kt
@@ -2,8 +2,12 @@ package app.cash.backfila.client.fixedset
 
 import app.cash.backfila.client.Backfill
 import app.cash.backfila.client.BackfillConfig
+import app.cash.backfila.client.ValidateResult
 
 abstract class FixedSetBackfill<Param : Any> : Backfill {
-  abstract fun checkBackfillConfig(backfillConfig: BackfillConfig<Param>)
+  open fun checkBackfillConfig(backfillConfig: BackfillConfig<Param>): ValidateResult<Param> {
+    return ValidateResult(backfillConfig.parameters)
+  }
+
   abstract fun runOne(row: FixedSetRow)
 }
diff --git a/client-base/src/test/kotlin/app/cash/backfila/client/fixedset/FixedSetBackfillOperator.kt b/client-base/src/test/kotlin/app/cash/backfila/client/fixedset/FixedSetBackfillOperator.kt
@@ -2,6 +2,7 @@ package app.cash.backfila.client.fixedset
 
 import app.cash.backfila.client.spi.BackfilaParametersOperator
 import app.cash.backfila.client.spi.BackfillOperator
+import app.cash.backfila.client.spi.parametersToBytes
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeRequest
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeResponse
 import app.cash.backfila.protos.clientservice.KeyRange
@@ -19,8 +20,10 @@ class FixedSetBackfillOperator<Param : Any>(
   override fun name() = backfill.javaClass.toString()
 
   override fun prepareBackfill(request: PrepareBackfillRequest): PrepareBackfillResponse {
-    backfill.checkBackfillConfig(
-      parametersOperator.constructBackfillConfig(request.parameters, request.dry_run)
+    val backfillConfig = parametersOperator.constructBackfillConfig(request.parameters, request.dry_run)
+
+    val validateResult = backfill.checkBackfillConfig(
+      backfillConfig
     )
 
     val partitions = mutableListOf<PrepareBackfillResponse.Partition>()
@@ -36,6 +39,11 @@ class FixedSetBackfillOperator<Param : Any>(
 
     return PrepareBackfillResponse.Builder()
       .partitions(partitions)
+      .apply {
+        if (backfillConfig.parameters != validateResult.parameters) {
+          parameters(parametersToBytes(validateResult.parameters))
+        }
+      }
       .build()
   }
 

diff --git a/client/src/main/kotlin/app/cash/backfila/client/ValidateResult.kt b/client/src/main/kotlin/app/cash/backfila/client/ValidateResult.kt
@@ -0,0 +1,10 @@
+package app.cash.backfila.client
+
+data class ValidateResult<Param : Any>(
+  /**
+   * Set to override parameters used for the backfill from the original parameters provided by the user.
+   *
+   * Example usage: pick a random token to use for the rest of the backfill.
+   */
+  val parameters: Param,
+)