Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Feat: EXPOSED-396 Supports fetchBatchedResults with sorting order #2102

Merged
merged 5 commits into from
Jun 4, 2024
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions exposed-core/api/exposed-core.api
Original file line number Diff line number Diff line change
Expand Up @@ -1761,8 +1761,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
Expand Down
22 changes: 17 additions & 5 deletions exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/Query.kt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()) {
Expand All @@ -260,16 +262,26 @@ 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

return object : Iterable<Iterable<ResultRow>> {
override fun iterator(): Iterator<Iterable<ResultRow>> {
return iterator {
var lastOffset = 0L
var lastOffset: Long? = null
while (true) {
val query = this@Query.copy().adjustWhere {
whereOp and (autoIncColumn greater lastOffset)
return@adjustWhere lastOffset.let {
if (it == null) return@let whereOp

return@let if (listOf(SortOrder.ASC, SortOrder.ASC_NULLS_FIRST, SortOrder.ASC_NULLS_LAST)
.contains(sortOrder)
) {
whereOp and (autoIncColumn greater it)
} else {
whereOp and (autoIncColumn less it)
}
}
bog-walk marked this conversation as resolved.
Show resolved Hide resolved
}

val results = query.iterator().asSequence().toList()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
Loading