Implementing Finalize Backfill Hook for All Backfila Clients

client-base/src/main/kotlin/app/cash/backfila/client/spi/BackfilaClientServiceHandler.kt

@@ -94,8 +94,14 @@ class BackfilaClientServiceHandler @Inject constructor(
     return loggingSetupProvider.withBackfillRunLogging(request.backfill_name, request.backfill_id) {
       logger.info { "Finalizing backfill `${request.backfill_name}::${request.backfill_id}`" }
 
-      // This is a stub for now
-      return@withBackfillRunLogging FinalizeBackfillResponse()
+      val operator = operatorFactory.create(request.backfill_name, request.backfill_id)
+      try {
+        return@withBackfillRunLogging operator.finalizeBackfill(request)
+      } catch (exception: Exception) {
+        return@withBackfillRunLogging FinalizeBackfillResponse.Builder()
+          .error_message(exception.message)
+          .build()
+      }
     }
   }
 

client-base/src/main/kotlin/app/cash/backfila/client/spi/BackfilaParameters.kt

@@ -4,7 +4,9 @@ import app.cash.backfila.client.BackfilaDefault
 import app.cash.backfila.client.BackfilaRequired
 import app.cash.backfila.client.BackfillConfig
 import app.cash.backfila.client.Description
+import app.cash.backfila.client.FinalizeBackfillConfig
 import app.cash.backfila.client.PrepareBackfillConfig
+import app.cash.backfila.protos.clientservice.FinalizeBackfillRequest
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeRequest
 import app.cash.backfila.protos.clientservice.PrepareBackfillRequest
 import app.cash.backfila.protos.clientservice.RunBatchRequest
@@ -60,6 +62,13 @@ class BackfilaParametersOperator<T : Any>(
     request.dry_run,
   )
 
+  fun constructBackfillConfig(request: FinalizeBackfillRequest): FinalizeBackfillConfig<T> =
+    FinalizeBackfillConfig(
+      constructParameters(request.parameters),
+      request.backfill_id,
+      request.dry_run,
+    )
+
   fun constructParameters(
     parameters: MutableMap<String, ByteString>,
   ): T {

client-base/src/main/kotlin/app/cash/backfila/client/spi/BackfillOperator.kt

@@ -1,6 +1,8 @@
 package app.cash.backfila.client.spi
 
 import app.cash.backfila.client.Backfill
+import app.cash.backfila.protos.clientservice.FinalizeBackfillRequest
+import app.cash.backfila.protos.clientservice.FinalizeBackfillResponse
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeRequest
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeResponse
 import app.cash.backfila.protos.clientservice.PrepareBackfillRequest
@@ -17,4 +19,5 @@ interface BackfillOperator {
   fun prepareBackfill(request: PrepareBackfillRequest): PrepareBackfillResponse
   fun getNextBatchRange(request: GetNextBatchRangeRequest): GetNextBatchRangeResponse
   fun runBatch(request: RunBatchRequest): RunBatchResponse
+  fun finalizeBackfill(request: FinalizeBackfillRequest): FinalizeBackfillResponse
 }

client-base/src/test/kotlin/app/cash/backfila/client/fixedset/FixedSetBackfill.kt

@@ -2,6 +2,7 @@ package app.cash.backfila.client.fixedset
 
 import app.cash.backfila.client.Backfill
 import app.cash.backfila.client.BackfillConfig
+import app.cash.backfila.client.FinalizeBackfillConfig
 import app.cash.backfila.client.ValidateResult
 
 abstract class FixedSetBackfill<Param : Any> : Backfill {
@@ -10,4 +11,9 @@ abstract class FixedSetBackfill<Param : Any> : Backfill {
   }
 
   abstract fun runOne(row: FixedSetRow, backfillConfig: BackfillConfig<Param>)
+
+  /**
+   * Override this to do any work after the backfill completes.
+   */
+  open fun finalize(config: FinalizeBackfillConfig<Param>) {}
 }

client-base/src/test/kotlin/app/cash/backfila/client/fixedset/FixedSetBackfillOperator.kt

@@ -3,6 +3,8 @@ 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.FinalizeBackfillRequest
+import app.cash.backfila.protos.clientservice.FinalizeBackfillResponse
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeRequest
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeResponse
 import app.cash.backfila.protos.clientservice.KeyRange
@@ -89,4 +91,12 @@ class FixedSetBackfillOperator<Param : Any>(
       .end(end.toString().encodeUtf8())
       .build()
   }
+
+  override fun finalizeBackfill(request: FinalizeBackfillRequest): FinalizeBackfillResponse {
+    val config = parametersOperator.constructBackfillConfig(request)
+    backfill.finalize(config)
+
+    return FinalizeBackfillResponse.Builder()
+      .build()
+  }
 }

client-dynamodb-v2/src/main/kotlin/app/cash/backfila/client/dynamodbv2/DynamoDbBackfill.kt

@@ -2,6 +2,7 @@ package app.cash.backfila.client.dynamodbv2
 
 import app.cash.backfila.client.Backfill
 import app.cash.backfila.client.BackfillConfig
+import app.cash.backfila.client.FinalizeBackfillConfig
 import app.cash.backfila.client.PrepareBackfillConfig
 import com.google.inject.TypeLiteral
 import com.squareup.moshi.Types
@@ -87,4 +88,9 @@ abstract class DynamoDbBackfill<I : Any, P : Any> : Backfill {
 
   /** See [ScanRequest.setExpressionAttributeNames]. */
   open fun expressionAttributeNames(config: BackfillConfig<P>): Map<String, String>? = null
+
+  /**
+   * Override this to do any work after the backfill completes.
+   */
+  open fun finalize(config: FinalizeBackfillConfig<P>) {}
 }

client-dynamodb-v2/src/main/kotlin/app/cash/backfila/client/dynamodbv2/internal/DynamoDbBackfillOperator.kt

@@ -3,6 +3,8 @@ package app.cash.backfila.client.dynamodbv2.internal
 import app.cash.backfila.client.dynamodbv2.DynamoDbBackfill
 import app.cash.backfila.client.spi.BackfilaParametersOperator
 import app.cash.backfila.client.spi.BackfillOperator
+import app.cash.backfila.protos.clientservice.FinalizeBackfillRequest
+import app.cash.backfila.protos.clientservice.FinalizeBackfillResponse
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeRequest
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeResponse
 import app.cash.backfila.protos.clientservice.KeyRange
@@ -160,6 +162,14 @@ class DynamoDbBackfillOperator<I : Any, P : Any>(
       .build()
   }
 
+  override fun finalizeBackfill(request: FinalizeBackfillRequest): FinalizeBackfillResponse {
+    val config = parametersOperator.constructBackfillConfig(request)
+    backfill.finalize(config)
+
+    return FinalizeBackfillResponse.Builder()
+      .build()
+  }
+
   private fun Map<String, AttributeValue>.toKeyRange(originalRange: DynamoDbKeyRange): KeyRange {
     require(originalRange.start + 1 == originalRange.end)
     return keyRangeCodec.encodeKeyRange(

client-dynamodb/src/main/kotlin/app/cash/backfila/client/dynamodb/DynamoDbBackfill.kt

@@ -2,6 +2,7 @@ package app.cash.backfila.client.dynamodb
 
 import app.cash.backfila.client.Backfill
 import app.cash.backfila.client.BackfillConfig
+import app.cash.backfila.client.FinalizeBackfillConfig
 import app.cash.backfila.client.PrepareBackfillConfig
 import com.amazonaws.services.dynamodbv2.model.AttributeValue
 import com.amazonaws.services.dynamodbv2.model.ScanRequest
@@ -79,4 +80,9 @@ abstract class DynamoDbBackfill<I : Any, P : Any> : Backfill {
 
   /** See [ScanRequest.setIndexName]. */
   open fun indexName(config: BackfillConfig<P>): String? = null
+
+  /**
+   * Override this to do any work after the backfill completes.
+   */
+  open fun finalize(config: FinalizeBackfillConfig<P>) {}
 }

client-dynamodb/src/main/kotlin/app/cash/backfila/client/dynamodb/internal/DynamoDbBackfillOperator.kt

@@ -3,6 +3,8 @@ package app.cash.backfila.client.dynamodb.internal
 import app.cash.backfila.client.dynamodb.DynamoDbBackfill
 import app.cash.backfila.client.spi.BackfilaParametersOperator
 import app.cash.backfila.client.spi.BackfillOperator
+import app.cash.backfila.protos.clientservice.FinalizeBackfillRequest
+import app.cash.backfila.protos.clientservice.FinalizeBackfillResponse
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeRequest
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeResponse
 import app.cash.backfila.protos.clientservice.KeyRange
@@ -134,6 +136,14 @@ class DynamoDbBackfillOperator<I : Any, P : Any>(
       .build()
   }
 
+  override fun finalizeBackfill(request: FinalizeBackfillRequest): FinalizeBackfillResponse {
+    val config = parametersOperator.constructBackfillConfig(request)
+    backfill.finalize(config)
+
+    return FinalizeBackfillResponse.Builder()
+      .build()
+  }
+
   private fun Map<String, AttributeValue>.toKeyRange(originalRange: DynamoDbKeyRange): KeyRange {
     require(originalRange.start + 1 == originalRange.end)
     return keyRangeCodec.encodeKeyRange(originalRange.start, originalRange.end, originalRange.count, this)

client-jooq/src/main/kotlin/app/cash/backfila/client/jooq/JooqBackfill.kt

@@ -2,6 +2,7 @@ package app.cash.backfila.client.jooq
 
 import app.cash.backfila.client.Backfill
 import app.cash.backfila.client.BackfillConfig
+import app.cash.backfila.client.FinalizeBackfillConfig
 import app.cash.backfila.client.PrepareBackfillConfig
 import app.cash.backfila.protos.clientservice.KeyRange
 import okio.ByteString
@@ -151,4 +152,9 @@ abstract class JooqBackfill<K, Param : Any> : Backfill {
   fun compoundKeyComparer(): CompoundKeyComparer<K> {
     return CompoundKeyComparer(compoundKeyFields)
   }
+
+  /**
+   * Override this to do any work after the backfill completes.
+   */
+  open fun finalize(config: FinalizeBackfillConfig<Param>) {}
 }

client-jooq/src/main/kotlin/app/cash/backfila/client/jooq/internal/JooqBackfillOperator.kt

@@ -5,6 +5,8 @@ import app.cash.backfila.client.jooq.CompoundKeyComparer
 import app.cash.backfila.client.jooq.JooqBackfill
 import app.cash.backfila.client.spi.BackfilaParametersOperator
 import app.cash.backfila.client.spi.BackfillOperator
+import app.cash.backfila.protos.clientservice.FinalizeBackfillRequest
+import app.cash.backfila.protos.clientservice.FinalizeBackfillResponse
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeRequest
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeResponse
 import app.cash.backfila.protos.clientservice.KeyRange
@@ -144,4 +146,12 @@ class JooqBackfillOperator<K, Param : Any> internal constructor(
       .end(rangeEnd)
       .build()
   }
+
+  override fun finalizeBackfill(request: FinalizeBackfillRequest): FinalizeBackfillResponse {
+    val config = parametersOperator.constructBackfillConfig(request)
+    backfill.finalize(config)
+
+    return FinalizeBackfillResponse.Builder()
+      .build()
+  }
 }

client-misk-hibernate/src/main/kotlin/app/cash/backfila/client/misk/hibernate/HibernateBackfill.kt

@@ -2,6 +2,7 @@ package app.cash.backfila.client.misk.hibernate
 
 import app.cash.backfila.client.Backfill
 import app.cash.backfila.client.BackfillConfig
+import app.cash.backfila.client.FinalizeBackfillConfig
 import app.cash.backfila.client.PrepareBackfillConfig
 import com.squareup.moshi.Types
 import java.lang.reflect.ParameterizedType
@@ -83,4 +84,9 @@ abstract class HibernateBackfill<E : DbEntity<E>, Pkey : Any, Param : Any> : Bac
    */
   open fun runOne(pkey: Pkey, config: BackfillConfig<Param>) {
   }
+
+  /**
+   * Override this to do any work after the backfill completes.
+   */
+  open fun finalizeBackfill(config: FinalizeBackfillConfig<Param>) {}
 }

client-misk-hibernate/src/main/kotlin/app/cash/backfila/client/misk/hibernate/internal/HibernateBackfillOperator.kt

@@ -6,6 +6,8 @@ import app.cash.backfila.client.misk.hibernate.HibernateBackfill
 import app.cash.backfila.client.misk.hibernate.PrimaryKeyCursorMapper
 import app.cash.backfila.client.spi.BackfilaParametersOperator
 import app.cash.backfila.client.spi.BackfillOperator
+import app.cash.backfila.protos.clientservice.FinalizeBackfillRequest
+import app.cash.backfila.protos.clientservice.FinalizeBackfillResponse
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeRequest
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeResponse
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeResponse.Batch
@@ -307,6 +309,14 @@ internal class HibernateBackfillOperator<E : DbEntity<E>, Pkey : Any, Param : An
     return RunBatchResponse.Builder().build()
   }
 
+  override fun finalizeBackfill(request: FinalizeBackfillRequest): FinalizeBackfillResponse {
+    val config = parametersOperator.constructBackfillConfig(request)
+    backfill.finalizeBackfill(config)
+
+    return FinalizeBackfillResponse.Builder()
+      .build()
+  }
+
   private fun HibernateBackfill<*, *, *>.getPrimaryKeyPath(queryRoot: Root<*>): Path<Number> {
     val fields = primaryKeyHibernateName().split('.')
     var path = queryRoot as Path<Number>

client-s3/src/main/kotlin/app/cash/backfila/client/s3/S3DatasourceBackfill.kt

@@ -2,6 +2,7 @@ package app.cash.backfila.client.s3
 
 import app.cash.backfila.client.Backfill
 import app.cash.backfila.client.BackfillConfig
+import app.cash.backfila.client.FinalizeBackfillConfig
 import app.cash.backfila.client.PrepareBackfillConfig
 import app.cash.backfila.client.s3.record.RecordStrategy
 
@@ -48,4 +49,9 @@ abstract class S3DatasourceBackfill<R : Any, P : Any> : Backfill {
    * Produces records from the S3 file.
    */
   abstract val recordStrategy: RecordStrategy<R>
+
+  /**
+   * Override this to do any work after the backfill completes.
+   */
+  open fun finalize(config: FinalizeBackfillConfig<P>) {}
 }

client-s3/src/main/kotlin/app/cash/backfila/client/s3/internal/S3DatasourceBackfillOperator.kt

@@ -4,6 +4,8 @@ import app.cash.backfila.client.s3.S3DatasourceBackfill
 import app.cash.backfila.client.s3.shim.S3Service
 import app.cash.backfila.client.spi.BackfilaParametersOperator
 import app.cash.backfila.client.spi.BackfillOperator
+import app.cash.backfila.protos.clientservice.FinalizeBackfillRequest
+import app.cash.backfila.protos.clientservice.FinalizeBackfillResponse
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeRequest
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeResponse
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeResponse.Batch
@@ -181,6 +183,14 @@ class S3DatasourceBackfillOperator<R : Any, P : Any>(
       .build()
   }
 
+  override fun finalizeBackfill(request: FinalizeBackfillRequest): FinalizeBackfillResponse {
+    val config = parametersOperator.constructBackfillConfig(request)
+    backfill.finalize(config)
+
+    return FinalizeBackfillResponse.Builder()
+      .build()
+  }
+
   data class DecodedRange(
     val start: Long,
     val end: Long,

client-static/src/main/kotlin/app/cash/backfila/client/stat/StaticDatasourceBackfill.kt

@@ -2,6 +2,7 @@ package app.cash.backfila.client.stat
 
 import app.cash.backfila.client.Backfill
 import app.cash.backfila.client.BackfillConfig
+import app.cash.backfila.client.FinalizeBackfillConfig
 import app.cash.backfila.client.PrepareBackfillConfig
 import com.google.inject.TypeLiteral
 import com.squareup.moshi.Types
@@ -54,4 +55,9 @@ abstract class StaticDatasourceBackfill<I : Any, P : Any> : Backfill {
    * This invokes the static list of items that the backfill will iterate over.
    */
   abstract val staticDatasource: List<I>
+
+  /**
+   * Override this to do any work after the backfill completes.
+   */
+  open fun finalize(config: FinalizeBackfillConfig<P>) {}
 }

client-static/src/main/kotlin/app/cash/backfila/client/stat/internal/StaticDatasourceBackfillOperator.kt

@@ -3,6 +3,8 @@ package app.cash.backfila.client.stat.internal
 import app.cash.backfila.client.spi.BackfilaParametersOperator
 import app.cash.backfila.client.spi.BackfillOperator
 import app.cash.backfila.client.stat.StaticDatasourceBackfill
+import app.cash.backfila.protos.clientservice.FinalizeBackfillRequest
+import app.cash.backfila.protos.clientservice.FinalizeBackfillResponse
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeRequest
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeResponse
 import app.cash.backfila.protos.clientservice.KeyRange
@@ -112,4 +114,12 @@ class StaticDatasourceBackfillOperator<I : Any, P : Any>(
   companion object {
     private const val PARTITION = "only"
   }
+
+  override fun finalizeBackfill(request: FinalizeBackfillRequest): FinalizeBackfillResponse {
+    val config = parametersOperator.constructBackfillConfig(request)
+    backfill.finalize(config)
+
+    return FinalizeBackfillResponse.Builder()
+      .build()
+  }
 }

client/src/main/kotlin/app/cash/backfila/client/BackfillConfig.kt

@@ -13,3 +13,9 @@ data class PrepareBackfillConfig<Param : Any>(
   val parameters: Param,
   val dryRun: Boolean,
 )
+
+data class FinalizeBackfillConfig<Param : Any>(
+  val parameters: Param,
+  val backfillId: String,
+  val dryRun: Boolean,
+)

client/src/main/proto/app/cash/backfila/client_service.proto

@@ -154,4 +154,5 @@ message FinalizeBackfillRequest {
 }
 
 message FinalizeBackfillResponse {
+  optional string error_message = 1;
 }

service-self-backfill/src/main/kotlin/app/cash/backfila/service/selfbackfill/BackfillRegisteredParameters.kt

@@ -21,9 +21,9 @@ internal class BackfillRegisteredParameters @Inject constructor(
   private val queryFactory: Query.Factory,
 ) : HibernateBackfill<DbRegisteredBackfill, Id<DbRegisteredBackfill>, NoParameters>() {
 
-  override fun runOne(id: Id<DbRegisteredBackfill>, config: BackfillConfig<NoParameters>) {
+  override fun runOne(pkey: Id<DbRegisteredBackfill>, config: BackfillConfig<NoParameters>) {
     transacter.transaction { session ->
-      val backfill = session.load(id)
+      val backfill = session.load(pkey)
       if (backfill.parameters.size == backfill.parameterNames().size) {
         return@transaction
       }

service-self-backfill/src/main/kotlin/app/cash/backfila/service/selfbackfill/LocalClientServiceClient.kt

@@ -2,6 +2,8 @@ package app.cash.backfila.service.selfbackfill
 
 import app.cash.backfila.client.BackfilaClientServiceClient
 import app.cash.backfila.client.spi.BackfilaClientServiceHandler
+import app.cash.backfila.protos.clientservice.FinalizeBackfillRequest
+import app.cash.backfila.protos.clientservice.FinalizeBackfillResponse
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeRequest
 import app.cash.backfila.protos.clientservice.GetNextBatchRangeResponse
 import app.cash.backfila.protos.clientservice.PrepareBackfillRequest
@@ -24,4 +26,8 @@ internal class LocalClientServiceClient @Inject internal constructor(
   override suspend fun runBatch(request: RunBatchRequest): RunBatchResponse {
     return backfilaClientServiceHandler.runBatch(request)
   }
+
+  override suspend fun finalizeBackfill(request: FinalizeBackfillRequest): FinalizeBackfillResponse {
+    return backfilaClientServiceHandler.finalizeBackfill(request)
+  }
 }