Add batch select #642

merge conflict resolved

exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/Queries.kt

@@ -22,6 +22,19 @@ fun FieldSet.select(where: Op<Boolean>) : Query = Query(this, where)
  */
 fun FieldSet.selectAll() : Query = Query(this, null)
 
+fun FieldSet.selectBatched(
+        batchSize: Int = 1000,
+        where: SqlExpressionBuilder.() -> Op<Boolean>
+): Iterable<Iterable<ResultRow>> {
+    return selectBatched(batchSize, SqlExpressionBuilder.where())
+}
+
+fun FieldSet.selectAllBatched(
+        batchSize: Int = 1000
+): Iterable<Iterable<ResultRow>> {
+    return selectBatched(batchSize, Op.TRUE)
+}
+
 /**
  * @sample org.jetbrains.exposed.sql.tests.shared.DMLTests.testDelete01
  */
@@ -196,6 +209,50 @@ fun checkExcessiveIndices(vararg tables: Table) {
     }
 }
 
+private fun FieldSet.selectBatched(
+        batchSize: Int = 1000,
+        whereOp: Op<Boolean>
+): Iterable<Iterable<ResultRow>> {
+    require(batchSize > 0) { "Batch size should be greater than 0" }
+
+    val autoIncColumn = try {
+        source.columns.first { it.columnType.isAutoInc }
+    } catch (_: NoSuchElementException) {
+        throw UnsupportedOperationException("Batched select only works on tables with an autoincrementing column")
+    }
+
+    return object : Iterable<Iterable<ResultRow>> {
+        override fun iterator(): Iterator<Iterable<ResultRow>> {
+            return iterator {
+                var lastOffset = 0L
+                while (true) {
+                    val query =
+                            select { whereOp and (autoIncColumn greater lastOffset) }
+                                    .limit(batchSize)
+                                    .orderBy(autoIncColumn, SortOrder.ASC)
+
+                    // query.iterator() executes the query
+                    val results = query.iterator().asSequence().toList()
+
+                    if (results.isEmpty()) break
+
+                    yield(results)
+
+                    if (results.size < batchSize) break
+
+                    lastOffset = toLong(results.last()[autoIncColumn]!!)
+                }
+            }
+        }
+
+        private fun toLong(autoIncVal: Any): Long = when (autoIncVal) {
+            is EntityID<*> -> autoIncVal.value as Long
+            is Int -> autoIncVal.toLong()
+            else -> autoIncVal as Long
+        }
+    }
+}
+
 /** Returns list of indices missed in database **/
 private fun checkMissingIndices(vararg tables: Table): List<Index> {
     fun Collection<Index>.log(mainMessage: String) {

exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/DMLTests.kt

@@ -992,6 +992,71 @@ class DMLTests : DatabaseTestsBase() {
         }
     }
 
+    @Test
+    fun `selectBatched should respect 'where' expression and the provided batch size`() {
+        val Cities = DMLTestsData.Cities
+        withTables(Cities) {
+            val names = List(100) { UUID.randomUUID().toString() }
+            Cities.batchInsert(names) { name -> this[Cities.name] = name }
+
+            val batches = Cities.selectBatched(batchSize = 25) { Cities.id less 51 }
+                    .toList().map { it.toCityNameList() }
+
+            val expectedNames = names.take(50)
+            assertEqualLists(listOf(
+                    expectedNames.take(25),
+                    expectedNames.takeLast(25)
+            ), batches)
+        }
+    }
+
+    @Test
+    fun `when batch size is greater than the amount of available items, selectAllBatched should return 1 batch`() {
+        val Cities = DMLTestsData.Cities
+        withTables(Cities) {
+            val names = List(25) { UUID.randomUUID().toString() }
+            Cities.batchInsert(names) { name -> this[Cities.name] = name }
+
+            val batches = Cities.selectAllBatched(batchSize = 100).toList().map { it.toCityNameList() }
+
+            assertEqualLists(listOf(names), batches)
+        }
+    }
+
+    @Test
+    fun `when there are no items, selectAllBatched should return an empty iterable`() {
+        val Cities = DMLTestsData.Cities
+        withTables(Cities) {
+            val batches = Cities.selectAllBatched().toList()
+
+            assertEqualLists(batches, emptyList())
+        }
+    }
+
+    @Test
+    fun `when there are no items of the given condition, should return an empty iterable`() {
+        val Cities = DMLTestsData.Cities
+        withTables(Cities) {
+            val names = List(25) { UUID.randomUUID().toString() }
+            Cities.batchInsert(names) { name -> this[Cities.name] = name }
+
+            val batches = Cities.selectBatched(batchSize = 100) { Cities.id greater 50 }
+                    .toList().map { it.toCityNameList() }
+
+            assertEqualLists(emptyList(), batches)
+        }
+    }
+
+    @Test(expected = java.lang.UnsupportedOperationException::class)
+    fun `when the table doesn't have an autoinc column, selectAllBatched should throw an exception`() {
+        DMLTestsData.UserData.selectAllBatched()
+    }
+
+    @Test(expected = IllegalArgumentException::class)
+    fun `when batch size is 0 or less, should throw an exception`() {
+        DMLTestsData.Cities.selectAllBatched(batchSize = -1)
+    }
+
     @Test
     fun testJoinWithAlias01() {
         withCitiesAndUsers { cities, users, userData ->
@@ -1246,7 +1311,7 @@ class DMLTests : DatabaseTestsBase() {
 
     @Test fun testTRUEandFALSEOps() {
         withCitiesAndUsers { cities, _, _ ->
-            val allSities = cities.selectAll().map { it[cities.name] }
+            val allSities = cities.selectAll().toCityNameList()
             assertEquals(0, cities.select { Op.FALSE }.count())
             assertEquals(allSities.size, cities.select { Op.TRUE }.count())
         }
@@ -1357,6 +1422,7 @@ class DMLTests : DatabaseTestsBase() {
             )
         }
     }
+    private fun Iterable<ResultRow>.toCityNameList(): List<String> = map { it[DMLTestsData.Cities.name] }
 }
 
 object OrgMemberships : IntIdTable() {