From d2a2aecf20f28850ccd7972f88068ccb99961502 Mon Sep 17 00:00:00 2001 From: Josh Yaganeh <319444+jyaganeh@users.noreply.github.com> Date: Fri, 15 Nov 2024 16:40:50 -0800 Subject: [PATCH] Get schedules from automation store in batches of 999 ids (#1571) --- .../automation/engine/AutomationStore.kt | 19 +++++++++++---- .../automation/AutomationStoreTest.kt | 23 ++++++++++++++++++ .../db/SuspendingBatchedQueryHelper.kt | 24 +++++++++++++++++++ 3 files changed, 61 insertions(+), 5 deletions(-) diff --git a/urbanairship-automation/src/main/java/com/urbanairship/automation/engine/AutomationStore.kt b/urbanairship-automation/src/main/java/com/urbanairship/automation/engine/AutomationStore.kt index 7b692edfe..a3c9d889e 100644 --- a/urbanairship-automation/src/main/java/com/urbanairship/automation/engine/AutomationStore.kt +++ b/urbanairship-automation/src/main/java/com/urbanairship/automation/engine/AutomationStore.kt @@ -25,6 +25,7 @@ import com.urbanairship.UALog import com.urbanairship.automation.AutomationSchedule import com.urbanairship.automation.engine.triggerprocessor.TriggerData import com.urbanairship.config.AirshipRuntimeConfig +import com.urbanairship.db.SuspendingBatchedQueryHelper.collectBatched import com.urbanairship.db.SuspendingBatchedQueryHelper.runBatched import com.urbanairship.json.JsonException import com.urbanairship.json.JsonTypeConverters @@ -145,8 +146,9 @@ internal abstract class AutomationStore : RoomDatabase(), AutomationStoreInterfa return dao.getSchedules(group)?.mapNotNull { it.toScheduleData() } ?: listOf() } + @Transaction override suspend fun getSchedules(ids: List): List { - return dao.getSchedules(ids)?.mapNotNull { it.toScheduleData() } ?: listOf() + return dao.getSchedules(ids).mapNotNull { it.toScheduleData() } } override suspend fun updateSchedule( @@ -207,9 +209,8 @@ internal interface AutomationDao { ids: List, closure: (String, AutomationScheduleData?) -> AutomationScheduleData ): List { val current = getSchedules(ids) - ?.mapNotNull { it.toScheduleData() } - ?.associateBy { it.schedule.identifier } - ?: mapOf() + .mapNotNull { it.toScheduleData() } + .associateBy { it.schedule.identifier } val result = mutableListOf() @@ -235,8 +236,16 @@ internal interface AutomationDao { @Query("SELECT * FROM schedules WHERE (`group` = :group)") suspend fun getSchedules(group: String): List? + @Transaction + suspend fun getSchedules(ids: List): List = + collectBatched(ids) { batch -> getSchedulesBatchInternal(batch) } + + /** + * This query is only for internal use, with batched queries + * to avoid the max query params limit of 999. + */ @Query("SELECT * FROM schedules WHERE (scheduleId IN (:ids))") - suspend fun getSchedules(ids: List): List? + suspend fun getSchedulesBatchInternal(ids: List): List? @Transaction suspend fun updateSchedule( diff --git a/urbanairship-automation/src/test/java/com/urbanairship/automation/AutomationStoreTest.kt b/urbanairship-automation/src/test/java/com/urbanairship/automation/AutomationStoreTest.kt index 1c60e35e4..283257602 100644 --- a/urbanairship-automation/src/test/java/com/urbanairship/automation/AutomationStoreTest.kt +++ b/urbanairship-automation/src/test/java/com/urbanairship/automation/AutomationStoreTest.kt @@ -174,6 +174,29 @@ public class AutomationStoreTest { assertNull(store.getSchedule("non-existing")) } + @Test + public fun testUpsertAndGetTooManySchedules(): TestResult = runTest { + // Create a ton of schedules (the SQL lite variable limit is 999) + val schedulesById = mutableMapOf() + for (i in 0 until 2000) { + val schedule = makeSchedule(i.toString()) + schedulesById[schedule.schedule.identifier] = schedule + } + + val schedules = schedulesById.values.toList() + val scheduleIds = schedulesById.keys.toList() + + // Initial insert + store.upsertSchedules(scheduleIds) { id, _ -> requireNotNull(schedulesById[id]) } + + // Upsert the schedules again + store.upsertSchedules(scheduleIds) { id, _ -> requireNotNull(schedulesById[id]) } + + // Verify the inserted schedules + val result = store.getSchedules(scheduleIds) + assertEquals(schedules.toSet(), result.toSet()) + } + @Test public fun testGetSchedulesByGroup(): TestResult = runTest { val original = listOf( diff --git a/urbanairship-core/src/main/java/com/urbanairship/db/SuspendingBatchedQueryHelper.kt b/urbanairship-core/src/main/java/com/urbanairship/db/SuspendingBatchedQueryHelper.kt index 3fbc79066..2891d13ef 100644 --- a/urbanairship-core/src/main/java/com/urbanairship/db/SuspendingBatchedQueryHelper.kt +++ b/urbanairship-core/src/main/java/com/urbanairship/db/SuspendingBatchedQueryHelper.kt @@ -35,6 +35,30 @@ public object SuspendingBatchedQueryHelper { runBatched(MAX_STATEMENT_PARAMETERS, items, callback) } + /** + * Collects a list of `items` resulting from running batched queries. + * + * Note that the batch lists are backed by the full items list, so non-structural changes in the + * batches are reflected in the items list, and vice-versa. + * + * @param items The list of items to be split into batches. + * @param callback The callback that will receive batched sub-lists and return the result of the query. + * @param The list type to split into batches. + * @param The result type of the batched query. + */ + public suspend fun collectBatched( + items: List, + callback: suspend (List) -> List? + ): List { + val result = mutableListOf() + + runBatched(MAX_STATEMENT_PARAMETERS, items) { batch -> + callback(batch)?.let(result::addAll) + } + + return result.toList() + } + /** * Splits up the given `items` into batches of the given size and invokes the `callback` for each batch. *