fix: EXPOSED-400 ClassCastException when using fetchBatchedResults

JetBrains · Jun 10, 2024 · ad834d2 · ad834d2
1 parent 7bce640
commit ad834d2
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 11 deletions.
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
@@ -8,7 +8,6 @@ import org.jetbrains.exposed.sql.statements.api.PreparedStatementApi
 import org.jetbrains.exposed.sql.vendors.ForUpdateOption
 import org.jetbrains.exposed.sql.vendors.currentDialect
 import java.sql.ResultSet
-import java.util.*
 
 enum class SortOrder(val code: String) {
     ASC(code = "ASC"),
@@ -39,14 +38,15 @@ open class Query(override var set: FieldSet, where: Op<Boolean>?) : AbstractQuer
     var where: Op<Boolean>? = where
         private set
 
-    override val queryToExecute: Statement<ResultSet> get() {
-        val distinctExpressions = set.fields.distinct()
-        return if (distinctExpressions.size < set.fields.size) {
-            copy().adjustSelect { select(distinctExpressions) }
-        } else {
-            this
+    override val queryToExecute: Statement<ResultSet>
+        get() {
+            val distinctExpressions = set.fields.distinct()
+            return if (distinctExpressions.size < set.fields.size) {
+                copy().adjustSelect { select(distinctExpressions) }
+            } else {
+                this
+            }
         }
-    }
 
     /** Creates a new [Query] instance using all stored properties of this `SELECT` query. */
     override fun copy(): Query = Query(set, where).also { copy ->
@@ -272,11 +272,29 @@ open class Query(override var set: FieldSet, where: Op<Boolean>?) : AbstractQuer
                     var lastOffset = if (fetchInAscendingOrder) 0L else null
                     while (true) {
                         val query = this@Query.copy().adjustWhere {
-                            lastOffset?.let {
+                            lastOffset?.let { lastOffset ->
                                 whereOp and if (fetchInAscendingOrder) {
-                                    (autoIncColumn greater it)
+                                    when (autoIncColumn.columnType) {
+                                        is EntityIDColumnType<*> -> {
+                                            (autoIncColumn as? Column<EntityID<Long>>)?.let {
+                                                (it greater lastOffset)
+                                            } ?: (autoIncColumn as? Column<EntityID<Int>>)?.let {
+                                                (it greater lastOffset.toInt())
+                                            } ?: (autoIncColumn greater lastOffset)
+                                        }
+                                        else -> (autoIncColumn greater lastOffset)
+                                    }
                                 } else {
-                                    (autoIncColumn less it)
+                                    when (autoIncColumn.columnType) {
+                                        is EntityIDColumnType<*> -> {
+                                            (autoIncColumn as? Column<EntityID<Long>>)?.let {
+                                                (it less lastOffset)
+                                            } ?: (autoIncColumn as? Column<EntityID<Int>>)?.let {
+                                                (it less lastOffset.toInt())
+                                            } ?: (autoIncColumn less lastOffset)
+                                        }
+                                        else -> (autoIncColumn less lastOffset)
+                                    }
                                 }
                             } ?: whereOp
                         }

diff --git a/...ts/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/dml/FetchBatchedResultsTests.kt b/...ts/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/dml/FetchBatchedResultsTests.kt
@@ -1,5 +1,6 @@
 package org.jetbrains.exposed.sql.tests.shared.dml
 
+import org.jetbrains.exposed.dao.id.IntIdTable
 import org.jetbrains.exposed.sql.*
 import org.jetbrains.exposed.sql.tests.DatabaseTestsBase
 import org.jetbrains.exposed.sql.tests.shared.assertEqualLists
@@ -100,4 +101,20 @@ class FetchBatchedResultsTests : DatabaseTestsBase() {
     fun `when batch size is 0 or less, should throw an exception`() {
         DMLTestsData.Cities.selectAll().fetchBatchedResults(batchSize = -1)
     }
+
+    @Test
+    fun testFetchBatchedResultsWithAutoIncrementEntityId() {
+        val tester1 = object : IntIdTable("table_1") {
+            val data = varchar("data", 100)
+        }
+
+        val tester2 = object : IntIdTable("table_2") {
+            val moreData = varchar("more_data", 100)
+            val prevData = reference("prev_data", tester1, onUpdate = ReferenceOption.CASCADE)
+        }
+
+        withTables(tester1, tester2) {
+            (tester2 innerJoin tester1).selectAll().fetchBatchedResults(10_000).flatten()
+        }
+    }
 }