feat: EXPOSED-396 Supports fetchBatchedResults with sorting order (#2102

)
JetBrains · Jun 4, 2024 · 72a8d35 · 72a8d35
1 parent b4b1336
commit 72a8d35
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 8 deletions.
diff --git a/exposed-core/api/exposed-core.api b/exposed-core/api/exposed-core.api
@@ -1773,8 +1773,8 @@ public class org/jetbrains/exposed/sql/Query : org/jetbrains/exposed/sql/Abstrac
 	public fun empty ()Z
 	public synthetic fun executeInternal (Lorg/jetbrains/exposed/sql/statements/api/PreparedStatementApi;Lorg/jetbrains/exposed/sql/Transaction;)Ljava/lang/Object;
 	public fun executeInternal (Lorg/jetbrains/exposed/sql/statements/api/PreparedStatementApi;Lorg/jetbrains/exposed/sql/Transaction;)Ljava/sql/ResultSet;
-	public final fun fetchBatchedResults (I)Ljava/lang/Iterable;
-	public static synthetic fun fetchBatchedResults$default (Lorg/jetbrains/exposed/sql/Query;IILjava/lang/Object;)Ljava/lang/Iterable;
+	public final fun fetchBatchedResults (ILorg/jetbrains/exposed/sql/SortOrder;)Ljava/lang/Iterable;
+	public static synthetic fun fetchBatchedResults$default (Lorg/jetbrains/exposed/sql/Query;ILorg/jetbrains/exposed/sql/SortOrder;ILjava/lang/Object;)Ljava/lang/Iterable;
 	public fun forUpdate (Lorg/jetbrains/exposed/sql/vendors/ForUpdateOption;)Lorg/jetbrains/exposed/sql/Query;
 	public synthetic fun forUpdate (Lorg/jetbrains/exposed/sql/vendors/ForUpdateOption;)Lorg/jetbrains/exposed/sql/SizedIterable;
 	public final fun getDistinct ()Z

diff --git a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/Query.kt b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/Query.kt
@@ -2,6 +2,7 @@ package org.jetbrains.exposed.sql
 
 import org.jetbrains.exposed.dao.id.EntityID
 import org.jetbrains.exposed.sql.SqlExpressionBuilder.greater
+import org.jetbrains.exposed.sql.SqlExpressionBuilder.less
 import org.jetbrains.exposed.sql.statements.Statement
 import org.jetbrains.exposed.sql.statements.api.PreparedStatementApi
 import org.jetbrains.exposed.sql.vendors.ForUpdateOption
@@ -244,10 +245,11 @@ open class Query(override var set: FieldSet, where: Op<Boolean>?) : AbstractQuer
      * This query's [FieldSet] will be ordered by the first auto-increment column.
      *
      * @param batchSize Size of each sub-collection to return.
+     * @param sortOrder Order in which the results should be retrieved.
      * @return Retrieved results as a collection of batched [ResultRow] sub-collections.
-     * @sample org.jetbrains.exposed.sql.tests.shared.dml.SelectBatchedTests.testFetchBatchedResultsWithWhereAndSetBatchSize
+     * @sample org.jetbrains.exposed.sql.tests.shared.dml.FetchBatchedResultsTests.testFetchBatchedResultsWithWhereAndSetBatchSize
      */
-    fun fetchBatchedResults(batchSize: Int = 1000): Iterable<Iterable<ResultRow>> {
+    fun fetchBatchedResults(batchSize: Int = 1000, sortOrder: SortOrder = SortOrder.ASC): Iterable<Iterable<ResultRow>> {
         require(batchSize > 0) { "Batch size should be greater than 0." }
         require(limit == null) { "A manual `LIMIT` clause should not be set. By default, `batchSize` will be used." }
         require(orderByExpressions.isEmpty()) {
@@ -260,16 +262,23 @@ open class Query(override var set: FieldSet, where: Op<Boolean>?) : AbstractQuer
             throw UnsupportedOperationException("Batched select only works on tables with an auto-incrementing column")
         }
         limit = batchSize
-        (orderByExpressions as MutableList).add(autoIncColumn to SortOrder.ASC)
+        (orderByExpressions as MutableList).add(autoIncColumn to sortOrder)
         val whereOp = where ?: Op.TRUE
+        val fetchInAscendingOrder = sortOrder in listOf(SortOrder.ASC, SortOrder.ASC_NULLS_FIRST, SortOrder.ASC_NULLS_LAST)
 
         return object : Iterable<Iterable<ResultRow>> {
             override fun iterator(): Iterator<Iterable<ResultRow>> {
                 return iterator {
-                    var lastOffset = 0L
+                    var lastOffset = if (fetchInAscendingOrder) 0L else null
                     while (true) {
                         val query = this@Query.copy().adjustWhere {
-                            whereOp and (autoIncColumn greater lastOffset)
+                            lastOffset?.let {
+                                whereOp and if (fetchInAscendingOrder) {
+                                    (autoIncColumn greater it)
+                                } else {
+                                    (autoIncColumn less it)
+                                }
+                            } ?: whereOp
                         }
 
                         val results = query.iterator().asSequence().toList()

diff --git a/...ql/tests/shared/dml/SelectBatchedTests.kt → ...ts/shared/dml/FetchBatchedResultsTests.kt b/...ql/tests/shared/dml/SelectBatchedTests.kt → ...ts/shared/dml/FetchBatchedResultsTests.kt
@@ -6,7 +6,7 @@ import org.jetbrains.exposed.sql.tests.shared.assertEqualLists
 import org.junit.Test
 import java.util.*
 
-class SelectBatchedTests : DatabaseTestsBase() {
+class FetchBatchedResultsTests : DatabaseTestsBase() {
     @Test
     fun testFetchBatchedResultsWithWhereAndSetBatchSize() {
         val cities = DMLTestsData.Cities
@@ -29,6 +29,28 @@ class SelectBatchedTests : DatabaseTestsBase() {
         }
     }
 
+    @Test
+    fun `when sortOrder is given, fetchBatchedResults should return batches in the given order`() {
+        val cities = DMLTestsData.Cities
+        withTables(cities) {
+            val names = List(100) { UUID.randomUUID().toString() }
+            cities.batchInsert(names) { name -> this[cities.name] = name }
+
+            val batches = cities.selectAll().where { cities.id less 51 }
+                .fetchBatchedResults(batchSize = 25, sortOrder = SortOrder.DESC)
+                .toList().map { it.toCityNameList() }
+
+            val expectedNames = names.take(50).reversed()
+            assertEqualLists(
+                listOf(
+                    expectedNames.take(25),
+                    expectedNames.takeLast(25)
+                ),
+                batches
+            )
+        }
+    }
+
     @Test
     fun `when batch size is greater than the amount of available items, fetchBatchedResults should return 1 batch`() {
         val cities = DMLTestsData.Cities