Get schedules from automation store in batches of 999 ids (#1571)

urbanairship · Nov 16, 2024 · d2a2aec · d2a2aec
1 parent 11bbea2
commit d2a2aec
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 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
@@ -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<String>): List<AutomationScheduleData> {
-        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<String>, closure: (String, AutomationScheduleData?) -> AutomationScheduleData
     ): List<AutomationScheduleData> {
         val current = getSchedules(ids)
-            ?.mapNotNull { it.toScheduleData() }
-            ?.associateBy { it.schedule.identifier }
-            ?: mapOf()
+            .mapNotNull { it.toScheduleData() }
+            .associateBy { it.schedule.identifier }
 
         val result = mutableListOf<AutomationScheduleData>()
 
@@ -235,8 +236,16 @@ internal interface AutomationDao {
     @Query("SELECT * FROM schedules WHERE (`group` = :group)")
     suspend fun getSchedules(group: String): List<ScheduleEntity>?
 
+    @Transaction
+    suspend fun getSchedules(ids: List<String>): List<ScheduleEntity> =
+        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<String>): List<ScheduleEntity>?
+    suspend fun getSchedulesBatchInternal(ids: List<String>): List<ScheduleEntity>?
 
     @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
@@ -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<String, AutomationScheduleData>()
+        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
@@ -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 <T> The list type to split into batches.
+     * @param <R> The result type of the batched query.
+     */
+    public suspend fun <T, R> collectBatched(
+        items: List<T>,
+        callback: suspend (List<T>) -> List<R>?
+    ): List<R> {
+        val result = mutableListOf<R>()
+
+        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.
      *