From 511fb35d3bb8e9349d8de9499a251bd65844cb56 Mon Sep 17 00:00:00 2001 From: Michael Pawliszyn Date: Thu, 12 Sep 2024 10:27:34 -0400 Subject: [PATCH] Fixing Hibernate bug for @When annotations. I do this by always using raw SQL for queries that are expected to deal with a full view of the table. --- .../internal/BoundingRangeStrategy.kt | 271 ++++++++++++++++-- .../internal/HibernateBackfillOperator.kt | 62 ++-- .../client/misk/ClientMiskTestingModule.kt | 3 + .../backfila/client/misk/DbActiveCoupon.kt | 36 +++ .../SinglePartitionHibernateBackfillTest.kt | 34 ++- .../UnshardedOnlyHibernateBackfillTest.kt | 110 +++++++ .../schema/backfila_main/v0008__backfila.sql | 5 + 7 files changed, 444 insertions(+), 77 deletions(-) create mode 100644 client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/DbActiveCoupon.kt create mode 100644 client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/hibernate/UnshardedOnlyHibernateBackfillTest.kt create mode 100644 client-misk-hibernate/src/test/resources/schema/backfila_main/v0008__backfila.sql diff --git a/client-misk-hibernate/src/main/kotlin/app/cash/backfila/client/misk/hibernate/internal/BoundingRangeStrategy.kt b/client-misk-hibernate/src/main/kotlin/app/cash/backfila/client/misk/hibernate/internal/BoundingRangeStrategy.kt index ba2facb8d..34a8069f2 100644 --- a/client-misk-hibernate/src/main/kotlin/app/cash/backfila/client/misk/hibernate/internal/BoundingRangeStrategy.kt +++ b/client-misk-hibernate/src/main/kotlin/app/cash/backfila/client/misk/hibernate/internal/BoundingRangeStrategy.kt @@ -4,7 +4,6 @@ import app.cash.backfila.client.misk.hibernate.HibernateBackfill import app.cash.backfila.client.misk.hibernate.PartitionProvider import com.google.common.collect.Ordering import javax.persistence.Table -import kotlin.streams.toList import misk.hibernate.DbEntity import misk.hibernate.Session import misk.hibernate.Transacter @@ -13,11 +12,30 @@ import misk.hibernate.transaction import misk.vitess.Keyspace import org.hibernate.internal.SessionImpl +/** + * The queries that are provided by the strategy are used to establish a primary key slice of + * the table off which the backfill criteria is applied. + */ interface BoundingRangeStrategy, Pkey : Any> { + /** + * Computes the raw table min and max based on the primary key. Returns null if the table is empty. + */ + fun computeAbsoluteMinMax( + backfill: HibernateBackfill, + partitionName: String, + ): MinMax? + /** * Computes a bound of size request.scan_size, to get a set of records that can be scanned for * records that match the criteria. + * * Returns null if there is are no more records left in the table. + * The return value must be greater than or equal to [backfillRangeStart] and less than or equal + * to [backfillRangeEnd] and greater than [previousEndKey]. + * + * @param backfillRangeStart this is [MinMax.min] unless a specific range was specified. + * @param backfillRangeEnd this is [MinMax.max] unless a specific range was specified. + * @param previousEndKey is the null at the start or the result of a previous call to this function. */ fun computeBoundingRangeMax( backfill: HibernateBackfill, @@ -27,11 +45,46 @@ interface BoundingRangeStrategy, Pkey : Any> { backfillRangeEnd: Pkey, scanSize: Long?, ): Pkey? + + /** + * Gets the min and count for the range of records. + * + * The returned [MinCount.min] value must be greater than [previousEndKey] and greater than or equal to [backfillRangeStart]. + * If [previousEndKey] is null: + * The returned [MinCount.scannedCount] counts items in [backfillRangeStart] (inclusive) until [end] (inclusive) + * + * If [previousEndKey] is non-null: + * The returned [MinCount.scannedCount] counts items in [previousEndKey] (exclusive) until [end] (inclusive). + * + * @param backfillRangeStart this is [MinMax.min] unless a specific range was specified. + * @param end is the batch slice and is greater than or equal to [backfillRangeStart]. + * @param previousEndKey is the null at the start or the result of a previous call [computeBoundingRangeMax]. + */ + fun computeMinAndCountForRange( + backfill: HibernateBackfill, + session: Session, + previousEndKey: Pkey?, + backfillRangeStart: Pkey, + end: Pkey, + ): MinCount } class UnshardedHibernateBoundingRangeStrategy, Pkey : Any>( private val partitionProvider: PartitionProvider, ) : BoundingRangeStrategy { + override fun computeAbsoluteMinMax( + backfill: HibernateBackfill, + partitionName: String, + ): MinMax? { + return partitionProvider.transaction(partitionName) { session -> + selectMinAndMax( + backfill, + session, + schemaAndTable(backfill), + ) + } + } + override fun computeBoundingRangeMax( backfill: HibernateBackfill, partitionName: String, @@ -42,21 +95,51 @@ class UnshardedHibernateBoundingRangeStrategy, Pkey : Any>( ): Pkey? { return partitionProvider.transaction(partitionName) { session -> selectMaxBound( - backfill, - session, - schemaAndTable(backfill), - previousEndKey, - backfillRangeStart, - backfillRangeEnd, - scanSize, + backfill = backfill, + session = session, + schemaAndTable = schemaAndTable(backfill), + previousEndKey = previousEndKey, + backfillRangeStart = backfillRangeStart, + backfillRangeEnd = backfillRangeEnd, + scanSize = scanSize, ) } } + + override fun computeMinAndCountForRange( + backfill: HibernateBackfill, + session: Session, + previousEndKey: Pkey?, + backfillRangeStart: Pkey, + end: Pkey, + ): MinCount { + return selectMinAndCount( + backfill = backfill, + session = session, + schemaAndTable = schemaAndTable(backfill), + previousEndKey = previousEndKey, + backfillRangeStart = backfillRangeStart, + end = end, + ) + } } class VitessShardedBoundingRangeStrategy, Pkey : Any>( private val partitionProvider: PartitionProvider, ) : BoundingRangeStrategy { + override fun computeAbsoluteMinMax( + backfill: HibernateBackfill, + partitionName: String, + ): MinMax? { + return partitionProvider.transaction(partitionName) { session -> + selectMinAndMax( + backfill, + session, + onlyTable(backfill), + ) + } + } + override fun computeBoundingRangeMax( backfill: HibernateBackfill, partitionName: String, @@ -68,22 +151,51 @@ class VitessShardedBoundingRangeStrategy, Pkey : Any>( return partitionProvider.transaction(partitionName) { session -> // We don't provide a schema when pinned to a shard. selectMaxBound( - backfill, - session, - onlyTable(backfill), - previousEndKey, - backfillRangeStart, - backfillRangeEnd, - scanSize, + backfill = backfill, + session = session, + schemaAndTable = onlyTable(backfill), + previousEndKey = previousEndKey, + backfillRangeStart = backfillRangeStart, + backfillRangeEnd = backfillRangeEnd, + scanSize = scanSize, ) } } + + override fun computeMinAndCountForRange( + backfill: HibernateBackfill, + session: Session, + previousEndKey: Pkey?, + backfillRangeStart: Pkey, + end: Pkey, + ): MinCount { + return selectMinAndCount( + backfill = backfill, + session = session, + schemaAndTable = onlyTable(backfill), + previousEndKey = previousEndKey, + backfillRangeStart = backfillRangeStart, + end = end, + ) + } } class VitessSingleCursorBoundingRangeStrategy, Pkey : Any>( private val transacter: Transacter, private val keyspace: Keyspace, ) : BoundingRangeStrategy { + override fun computeAbsoluteMinMax( + backfill: HibernateBackfill, + partitionName: String, + ): MinMax? { + return transacter.transaction { session -> + selectMinAndMax( + backfill, + session, + onlyTable(backfill), + ) + } + } /** * Computes a bounding range by scanning all shards and returning the minimum of MAX(pkey). @@ -108,13 +220,13 @@ class VitessSingleCursorBoundingRangeStrategy, Pkey : Any>( transacter.transaction(it) { session -> // We don't provide a schema when pinned to a shard. selectMaxBound( - backfill, - session, - onlyTable(backfill), - previousEndKey, - backfillRangeStart, - backfillRangeEnd, - scanSize, + backfill = backfill, + session = session, + schemaAndTable = onlyTable(backfill), + previousEndKey = previousEndKey, + backfillRangeStart = backfillRangeStart, + backfillRangeEnd = backfillRangeEnd, + scanSize = scanSize, ) } }.toList() @@ -122,10 +234,67 @@ class VitessSingleCursorBoundingRangeStrategy, Pkey : Any>( // Pkey must have a natural ordering .minWithOrNull(Ordering.natural>() as Comparator) } + + override fun computeMinAndCountForRange( + backfill: HibernateBackfill, + session: Session, + previousEndKey: Pkey?, + backfillRangeStart: Pkey, + end: Pkey, + ): MinCount { + return selectMinAndCount( + backfill = backfill, + session = session, + schemaAndTable = onlyTable(backfill), + previousEndKey = previousEndKey, + backfillRangeStart = backfillRangeStart, + end = end, + ) + } } class SingleCursorVitess +private fun , Pkey : Any> selectMinAndMax( + backfill: HibernateBackfill, + session: Session, + schemaAndTable: String, +): MinMax? { + // This query uses raw sql to avoid bumping into hibernate features such as @Where and + // @SQLRestriction. + // All of [selectMaxBound], [selectMinAndMax] and [selectMinAndCount] must be raw SQL since + // they depend on each other having the same view of the table. + val pkeyName = backfill.primaryKeyName() + val sql = """ + |SELECT MIN($pkeyName) as min, MAX($pkeyName) as max + |FROM $schemaAndTable + """.trimMargin() + val minMax = session.useConnection { connection -> + connection.prepareStatement(sql).use { ps -> + val pkeyType = session.hibernateSession.typeHelper.basic(backfill.pkeyClass.java)!! + + val rs = ps.executeQuery() + rs.next() + val min = pkeyType.nullSafeGet(rs, "min", session.hibernateSession as SessionImpl, null) + val max = pkeyType.nullSafeGet(rs, "max", session.hibernateSession as SessionImpl, null) + if (min == null) { + // Empty table, no work to do for this partition. + return@use null + } else { + checkNotNull(max) { "Table max was null but min wasn't, this shouldn't happen" } + @Suppress("UNCHECKED_CAST") // Return type from the query should always be Pkey. + MinMax(min as Pkey, max as Pkey) + } + } + } + return minMax +} + +data class MinMax( + val min: Pkey, + val max: Pkey, +) + private fun , Pkey : Any> selectMaxBound( backfill: HibernateBackfill, session: Session, @@ -137,6 +306,8 @@ private fun , Pkey : Any> selectMaxBound( ): Pkey? { // Hibernate doesn't support subqueries in FROM, and we don't want to read in 100k+ records, // so we use raw SQL here. + // All of [selectMaxBound], [selectMinAndMax] and [selectMinAndCount] must be raw SQL since + // they depend on each other having the same view of the table. val pkeyName = backfill.primaryKeyName() val params = mutableListOf() var where = when { @@ -178,6 +349,62 @@ private fun , Pkey : Any> selectMaxBound( return max as Pkey? } +private fun , Pkey : Any> selectMinAndCount( + backfill: HibernateBackfill, + session: Session, + schemaAndTable: String, + previousEndKey: Pkey?, + backfillRangeStart: Pkey, + end: Pkey, +): MinCount { + // This query uses raw sql to avoid bumping into hibernate features such as @Where and + // @SQLRestriction. + // All of [selectMaxBound], [selectMinAndMax] and [selectMinAndCount] must be raw SQL since + // they depend on each other having the same view of the table. + val pkeyName = backfill.primaryKeyName() + val params = mutableListOf() + var where = when { + previousEndKey != null -> { + params.add(previousEndKey) + "WHERE $pkeyName > ?" + } + else -> { + params.add(backfillRangeStart) + "WHERE $pkeyName >= ?" + } + } + params.add(end) + where += " AND $pkeyName <= ?" + val sql = """ + |SELECT MIN($pkeyName) as start, COUNT(*) as scannedCount + |FROM $schemaAndTable + |$where + """.trimMargin() + val minCount = session.useConnection { connection -> + connection.prepareStatement(sql).use { ps -> + val pkeyType = session.hibernateSession.typeHelper.basic(backfill.pkeyClass.java)!! + + params.forEachIndexed { index, pkey -> + pkeyType.nullSafeSet(ps, pkey, index + 1, session.hibernateSession as SessionImpl) + } + + val rs = ps.executeQuery() + rs.next() + @Suppress("UNCHECKED_CAST") // Return type from the query should always be a Pkey and Long. + MinCount( + pkeyType.nullSafeGet(rs, "start", session.hibernateSession as SessionImpl, null) as Pkey, + rs.getLong("scannedCount"), + ) + } + } + return minCount +} + +data class MinCount( + val min: Pkey, + val scannedCount: Long, +) + private fun , Pkey : Any> schemaAndTable(backfill: HibernateBackfill): String { val tableAnnotation = backfill.entityClass.java.getAnnotation(Table::class.java) val schema = tableAnnotation.schema diff --git a/client-misk-hibernate/src/main/kotlin/app/cash/backfila/client/misk/hibernate/internal/HibernateBackfillOperator.kt b/client-misk-hibernate/src/main/kotlin/app/cash/backfila/client/misk/hibernate/internal/HibernateBackfillOperator.kt index 705ff5e70..26a394863 100644 --- a/client-misk-hibernate/src/main/kotlin/app/cash/backfila/client/misk/hibernate/internal/HibernateBackfillOperator.kt +++ b/client-misk-hibernate/src/main/kotlin/app/cash/backfila/client/misk/hibernate/internal/HibernateBackfillOperator.kt @@ -89,27 +89,15 @@ internal class HibernateBackfillOperator, Pkey : Any, Param : An .estimated_record_count(null) .build() } - val keyRange: KeyRange = partitionProvider.transaction(partitionName) { session -> - val minmax = queryFactory.dynamicQuery(backfill.entityClass) - .dynamicUniqueResult(session) { criteriaBuilder, queryRoot -> - criteriaBuilder.tuple( - criteriaBuilder.min(backfill.getPrimaryKeyPath(queryRoot)), - criteriaBuilder.max(backfill.getPrimaryKeyPath(queryRoot)), - ) - }!! - - val min = minmax[0] as Pkey? - val max = minmax[1] as Pkey? - if (min == null) { - // Empty table, no work to do for this partition. - KeyRange.Builder().build() - } else { - checkNotNull(max) { "Table max was null but min wasn't, this shouldn't happen" } - KeyRange.Builder() - .start(requestedRange?.start ?: primaryKeyCursorMapper.toByteString(min)) - .end(requestedRange?.end ?: primaryKeyCursorMapper.toByteString(max)) - .build() - } + val minMax = boundingRangeStrategy.computeAbsoluteMinMax(backfill, partitionName) + val keyRange: KeyRange = if (minMax == null) { + // Empty table, no work to do for this partition. + KeyRange.Builder().build() + } else { + KeyRange.Builder() + .start(requestedRange?.start ?: primaryKeyCursorMapper.toByteString(minMax.min)) + .end(requestedRange?.end ?: primaryKeyCursorMapper.toByteString(minMax.max)) + .build() } return Partition.Builder() .partition_name(partitionName) @@ -185,12 +173,12 @@ internal class HibernateBackfillOperator, Pkey : Any, Param : An val stopwatch = Stopwatch.createStarted() boundingMax = boundingRangeStrategy .computeBoundingRangeMax( - backfill, - partitionName, - previousEndKey, - primaryKeyCursorMapper.fromByteString(backfill.pkeyClass.java, backfillRange.start).getOrThrow(), - primaryKeyCursorMapper.fromByteString(backfill.pkeyClass.java, backfillRange.end).getOrThrow(), - scanSize, + backfill = backfill, + partitionName = partitionName, + previousEndKey = previousEndKey, + backfillRangeStart = primaryKeyCursorMapper.fromByteString(backfill.pkeyClass.java, backfillRange.start).getOrThrow(), + backfillRangeEnd = primaryKeyCursorMapper.fromByteString(backfill.pkeyClass.java, backfillRange.end).getOrThrow(), + scanSize = scanSize, ) if (boundingMax == null) { logger.info("Bounding range returned no records, done computing batches") @@ -247,19 +235,13 @@ internal class HibernateBackfillOperator, Pkey : Any, Param : An } // Get start pkey and scanned record count for this batch. - val result = queryFactory.dynamicQuery(backfill.entityClass).apply { - addBoundingMin(this) - dynamicAddConstraint(pkeyProperty, LE, end) - }.dynamicUniqueResult(session) { criteriaBuilder, queryRoot -> - criteriaBuilder.tuple( - criteriaBuilder.min(backfill.getPrimaryKeyPath(queryRoot)), - criteriaBuilder.count(queryRoot), - ) - }!! - - @Suppress("UNCHECKED_CAST") // Return type from the query should always match. - val start = result[0] as Pkey - val scannedCount = result[1] as Long + val (start, scannedCount) = boundingRangeStrategy.computeMinAndCountForRange( + backfill = backfill, + session = session, + previousEndKey = previousEndKey, + backfillRangeStart = primaryKeyCursorMapper.fromByteString(backfill.pkeyClass.java, backfillRange.start).getOrThrow(), + end = end, + ) TxResult( end, diff --git a/client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/ClientMiskTestingModule.kt b/client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/ClientMiskTestingModule.kt index 7b34d3dc6..438928e42 100644 --- a/client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/ClientMiskTestingModule.kt +++ b/client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/ClientMiskTestingModule.kt @@ -3,6 +3,7 @@ package app.cash.backfila.client.misk import app.cash.backfila.client.BackfilaClientLoggingSetupProvider import app.cash.backfila.client.BackfilaClientNoLoggingSetupProvider import app.cash.backfila.client.BackfilaHttpClientConfig +import app.cash.backfila.client.misk.hibernate.ActiveCouponBackfill import app.cash.backfila.client.misk.hibernate.ChickenToBeefBackfill import app.cash.backfila.client.misk.hibernate.HibernateBackfillModule import app.cash.backfila.client.misk.hibernate.NullableParameterBackfill @@ -46,6 +47,7 @@ internal class ClientMiskTestingModule( DbMenu::class, DbOrder::class, DbRestaurant::class, + DbActiveCoupon::class, ) } }, @@ -73,6 +75,7 @@ internal class ClientMiskTestingModule( install(HibernateBackfillModule.create()) install(HibernateBackfillModule.create()) install(HibernateBackfillModule.create()) + install(HibernateBackfillModule.create()) install(MenuStackModule()) } diff --git a/client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/DbActiveCoupon.kt b/client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/DbActiveCoupon.kt new file mode 100644 index 000000000..00b491661 --- /dev/null +++ b/client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/DbActiveCoupon.kt @@ -0,0 +1,36 @@ +package app.cash.backfila.client.misk + +import java.time.Instant +import javax.persistence.Column +import javax.persistence.Entity +import javax.persistence.GeneratedValue +import javax.persistence.Table +import misk.hibernate.Constraint +import misk.hibernate.DbUnsharded +import misk.hibernate.Id +import misk.hibernate.Operator +import misk.hibernate.Query +import org.hibernate.annotations.Where + +@Entity +@Table(name = "coupons") +@Where(clause = "expired_at IS NULL") +class DbActiveCoupon() : DbUnsharded { + @javax.persistence.Id + @GeneratedValue + override lateinit var id: Id + + @Column + lateinit var expired_at: Instant + + fun coupon() = Coupon(id) +} + +interface CouponQuery : Query { + @Constraint(path = "id", operator = Operator.EQ) + fun id(id: Id): MenuQuery +} + +data class Coupon( + val id: Id, +) diff --git a/client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/hibernate/SinglePartitionHibernateBackfillTest.kt b/client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/hibernate/SinglePartitionHibernateBackfillTest.kt index a5bb93714..0667b9898 100644 --- a/client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/hibernate/SinglePartitionHibernateBackfillTest.kt +++ b/client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/hibernate/SinglePartitionHibernateBackfillTest.kt @@ -313,27 +313,31 @@ abstract class SinglePartitionHibernateBackfillTest { } private fun createSome(): List> { - return transacter.transaction { session: Session -> - val expected = mutableListOf>() - repeat((0..9).count()) { - val id = session.save(DbMenu("chicken")) - expected.add(id) - } - - // Intersperse these to make sure we test skipping non matching records. - repeat((0..4).count()) { session.save(DbMenu("beef")) } + val expected = mutableListOf>() + expected += createMatching(10) + // Intersperse these to make sure we test skipping non matching records. + createNoMatching(5) + expected += createMatching(10) + return expected + } - repeat((0..9).count()) { - val id = session.save(DbMenu("chicken")) - expected.add(id) + private fun createMatching(times: Int): List> { + val expected = mutableListOf>() + transacter.transaction { session -> + repeat(times) { + expected.add(session.save(DbMenu("chicken"))) } - expected } + return expected } - private fun createNoMatching() { + private fun createNoMatching(times: Int = 5): List> { + val expected = mutableListOf>() transacter.transaction { session: Session -> - repeat((0..4).count()) { session.save(DbMenu("beef")) } + repeat(times) { + expected.add(session.save(DbMenu("beef"))) + } } + return expected } } diff --git a/client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/hibernate/UnshardedOnlyHibernateBackfillTest.kt b/client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/hibernate/UnshardedOnlyHibernateBackfillTest.kt new file mode 100644 index 000000000..2bafbff44 --- /dev/null +++ b/client-misk-hibernate/src/test/kotlin/app/cash/backfila/client/misk/hibernate/UnshardedOnlyHibernateBackfillTest.kt @@ -0,0 +1,110 @@ +package app.cash.backfila.client.misk.hibernate + +import app.cash.backfila.client.BackfillConfig +import app.cash.backfila.client.Description +import app.cash.backfila.client.NoParameters +import app.cash.backfila.client.misk.ClientMiskService +import app.cash.backfila.client.misk.ClientMiskTestingModule +import app.cash.backfila.client.misk.DbActiveCoupon +import app.cash.backfila.embedded.Backfila +import app.cash.backfila.embedded.BackfillRun +import app.cash.backfila.embedded.createWetRun +import com.google.inject.Module +import java.time.Instant +import javax.inject.Inject +import misk.hibernate.Id +import misk.hibernate.Query +import misk.hibernate.Transacter +import misk.hibernate.load +import misk.testing.MiskTest +import misk.testing.MiskTestModule +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.Test + +/** + * This holds edge case hibernate tests. + */ +@MiskTest(startService = true) +class UnshardedOnlyHibernateBackfillTest { + @Suppress("unused") + @MiskTestModule + val module: Module = ClientMiskTestingModule(false) + + @Inject @ClientMiskService + lateinit var transacter: Transacter + + @Inject lateinit var backfila: Backfila + + @Test fun `where annotation hibernate class counts scanned ids correctly`() { + val couponIds = createCoupons(100) + couponIds.filterIndexed { index, _ -> index % 2 == 1 }.forEach { expireCoupon(it) } + val run = backfila.createWetRun() + .apply { configureForTest() } + run.execute() + assertThat(run.backfill.couponsToProcess).size().isEqualTo(50) + assertThat(run.precomputeMatchingCount).isEqualTo(50L) + assertThat(run.precomputeScannedCount).isEqualTo(100) + } + + @Test fun `where annotation hibernate class works with projection gaps`() { + // We had a problem with a @Where annotation was limiting the view on all but one query that + // needs to be manually written and caused a null exception in a later query. We fixed this by + // manually writing both queries that require a complete view of the table. This test checks + // this case going forward. + createCoupons(20) + val couponIdsToExpire = createCoupons(50) + createCoupons(20) + couponIdsToExpire.forEach { expireCoupon(it) } + val run = backfila.createWetRun() + .apply { configureForTest() } + run.execute() + assertThat(run.backfill.couponsToProcess).size().isEqualTo(40L) + assertThat(run.precomputeMatchingCount).isEqualTo(40L) + assertThat(run.precomputeScannedCount).isEqualTo(90L) + } + + private fun BackfillRun<*>.configureForTest() { + this.batchSize = 5L + this.scanSize = 10L + this.computeCountLimit = 1L + } + + private fun createCoupons(times: Int): List> { + val expected = mutableListOf>() + transacter.transaction { session -> + repeat(times) { + expected.add(session.save(DbActiveCoupon())) + } + } + return expected + } + + private fun expireCoupon(couponId: Id) { + transacter.transaction { session -> + session.load(couponId).apply { + this.expired_at = Instant.now() + session.save(this) + } + } + } +} + +@Description("To process active coupons.") +class ActiveCouponBackfill @Inject constructor( + @ClientMiskService private val transacter: Transacter, + private val queryFactory: Query.Factory, +) : HibernateBackfill, NoParameters>() { + val couponsToProcess = mutableSetOf>() + + override fun backfillCriteria(config: BackfillConfig): Query { + return queryFactory.dynamicQuery(DbActiveCoupon::class) + } + + override fun runBatch(pkeys: List>, config: BackfillConfig) { + if (!config.dryRun) { + couponsToProcess.addAll(pkeys) + } + } + + override fun partitionProvider() = UnshardedPartitionProvider(transacter) +} diff --git a/client-misk-hibernate/src/test/resources/schema/backfila_main/v0008__backfila.sql b/client-misk-hibernate/src/test/resources/schema/backfila_main/v0008__backfila.sql new file mode 100644 index 000000000..8497658d5 --- /dev/null +++ b/client-misk-hibernate/src/test/resources/schema/backfila_main/v0008__backfila.sql @@ -0,0 +1,5 @@ +CREATE TABLE `coupons` ( + `id` bigint(20) NOT NULL AUTO_INCREMENT, + `expired_at` TIMESTAMP(3) NULL DEFAULT NULL, + PRIMARY KEY (`id`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8; \ No newline at end of file