Pass backfill_id in EmbeddedBackfila (#297)

More realistic
cashapp · Mar 15, 2023 · cb86ff2 · cb86ff2
1 parent d4da5c5
commit cb86ff2
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 0 deletions.
diff --git a/backfila-embedded/src/main/kotlin/app/cash/backfila/embedded/BackfillRun.kt b/backfila-embedded/src/main/kotlin/app/cash/backfila/embedded/BackfillRun.kt
@@ -16,6 +16,7 @@ interface BackfillRun<B : Backfill> {
   val parameters: Map<String, ByteString>
   val rangeStart: String?
   val rangeEnd: String?
+  val backfillId: String
 
   var batchSize: Long
   var scanSize: Long

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
@@ -82,6 +82,7 @@ internal class EmbeddedBackfila @Inject internal constructor(
       parameters = createRequest.parameter_map.toMutableMap(),
       rangeStart = createRequest.pkey_range_start?.utf8(),
       rangeEnd = createRequest.pkey_range_end?.utf8(),
+      backfillId = backfillId.toString(),
     )
 
     run.execute()
@@ -131,6 +132,7 @@ internal class EmbeddedBackfila @Inject internal constructor(
       parameters = parameters.toMutableMap(),
       rangeStart = rangeStart,
       rangeEnd = rangeEnd,
+      backfillId = backfillId,
     )
   }
 }
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
@@ -26,6 +26,7 @@ internal class EmbeddedBackfillRun<B : Backfill>(
   override val parameters: MutableMap<String, ByteString>,
   override val rangeStart: String?,
   override val rangeEnd: String?,
+  override val backfillId: String,
 
   override var batchSize: Long = 100L,
   override var scanSize: Long = 10_000L,
@@ -75,6 +76,7 @@ internal class EmbeddedBackfillRun<B : Backfill>(
     val response =
       operator.getNextBatchRange(
         GetNextBatchRangeRequest.Builder()
+          .backfill_id(backfillId)
           .partition_name(cursor.partitionName)
           .backfill_range(cursor.keyRange)
           .previous_end_key(cursor.previousEndKey)
@@ -113,6 +115,7 @@ internal class EmbeddedBackfillRun<B : Backfill>(
     val response =
       operator.getNextBatchRange(
         GetNextBatchRangeRequest.Builder()
+          .backfill_id(backfillId)
           .partition_name(cursor.partitionName)
           .backfill_range(cursor.keyRange)
           .previous_end_key(cursor.previousEndKey)
@@ -164,6 +167,7 @@ internal class EmbeddedBackfillRun<B : Backfill>(
       val remainingBatch = batch.copy(batchRange = remainingRange)
       response = operator.runBatch(
         RunBatchRequest.Builder()
+          .backfill_id(backfillId)
           .partition_name(remainingBatch.partitionName)
           .batch_range(remainingBatch.batchRange)
           .parameters(parameters)