refactor: Extract query statement iterator logic to common class

- Add internal StatementIterator class with common functionality across ExplainQuery, AbstractQuery, and ReturningStatement
JetBrains · Nov 16, 2024 · ee1986c · ee1986c
1 parent 8463a8f
commit ee1986c
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 74 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -132,11 +132,11 @@ subprojects {
         }
     }
 
-//    testDb("sqlserver") {
-//        port = 3005
-//        dialects("SQLSERVER")
-//        dependencies {
-//            dependency(rootProject.libs.mssql)
-//        }
-//    }
+    testDb("sqlserver") {
+        port = 3005
+        dialects("SQLSERVER")
+        dependencies {
+            dependency(rootProject.libs.mssql)
+        }
+    }
 }
diff --git a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/AbstractQuery.kt b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/AbstractQuery.kt
@@ -1,6 +1,7 @@
 package org.jetbrains.exposed.sql
 
 import org.jetbrains.exposed.sql.statements.Statement
+import org.jetbrains.exposed.sql.statements.StatementIterator
 import org.jetbrains.exposed.sql.statements.StatementType
 import org.jetbrains.exposed.sql.transactions.TransactionManager
 import java.sql.ResultSet
@@ -98,33 +99,17 @@ abstract class AbstractQuery<T : AbstractQuery<T>>(
         }
     }
 
-    private inner class ResultIterator(val rs: ResultSet) : Iterator<ResultRow> {
-        private var hasNext = false
-            set(value) {
-                field = value
-                if (!field) {
-                    val statement = rs.statement
-                    rs.close()
-                    statement?.close()
-                    transaction.openResultSetsCount--
-                }
-            }
-
-        private val fieldsIndex = set.realFields.toSet().mapIndexed { index, expression -> expression to index }.toMap()
+    private inner class ResultIterator(rs: ResultSet) : StatementIterator<Expression<*>, ResultRow>(rs) {
+        override val fieldIndex = set.realFields.toSet()
+            .mapIndexed { index, expression -> expression to index }
+            .toMap()
 
         init {
-            hasNext = rs.next()
+            hasNext = result.next()
             if (hasNext) trackResultSet(transaction)
         }
 
-        override operator fun next(): ResultRow {
-            if (!hasNext) throw NoSuchElementException()
-            val result = ResultRow.create(rs, fieldsIndex)
-            hasNext = rs.next()
-            return result
-        }
-
-        override fun hasNext(): Boolean = hasNext
+        override fun createResultRow(): ResultRow = ResultRow.create(result, fieldIndex)
     }
 
     companion object {

diff --git a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/ExplainQuery.kt b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/ExplainQuery.kt
@@ -3,6 +3,7 @@ package org.jetbrains.exposed.sql
 import org.jetbrains.exposed.sql.statements.IStatementBuilder
 import org.jetbrains.exposed.sql.statements.Statement
 import org.jetbrains.exposed.sql.statements.StatementBuilder
+import org.jetbrains.exposed.sql.statements.StatementIterator
 import org.jetbrains.exposed.sql.statements.StatementType
 import org.jetbrains.exposed.sql.statements.api.PreparedStatementApi
 import org.jetbrains.exposed.sql.transactions.TransactionManager
@@ -36,34 +37,16 @@ open class ExplainQuery(
         return Iterable { resultIterator }.iterator()
     }
 
-    private inner class ResultIterator(private val rs: ResultSet) : Iterator<ExplainResultRow> {
-        private val fieldIndex: Map<String, Int> = List(rs.metaData.columnCount) { i ->
-            rs.metaData.getColumnName(i + 1) to i
+    private inner class ResultIterator(rs: ResultSet) : StatementIterator<String, ExplainResultRow>(rs) {
+        override val fieldIndex = List(result.metaData.columnCount) { i ->
+            result.metaData.getColumnName(i + 1) to i
         }.toMap()
 
-        private var hasNext = false
-            set(value) {
-                field = value
-                if (!field) {
-                    val statement = rs.statement
-                    rs.close()
-                    statement?.close()
-                    transaction.openResultSetsCount--
-                }
-            }
-
         init {
-            hasNext = rs.next()
+            hasNext = result.next()
         }
 
-        override fun hasNext(): Boolean = hasNext
-
-        override operator fun next(): ExplainResultRow {
-            if (!hasNext) throw NoSuchElementException()
-            val result = ExplainResultRow.create(rs, fieldIndex)
-            hasNext = rs.next()
-            return result
-        }
+        override fun createResultRow(): ExplainResultRow = ExplainResultRow.create(result, fieldIndex)
     }
 }
 

diff --git a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/statements/ReturningStatement.kt b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/statements/ReturningStatement.kt
@@ -38,31 +38,14 @@ open class ReturningStatement(
         return Iterable { resultIterator }.iterator()
     }
 
-    private inner class ResultIterator(val rs: ResultSet) : Iterator<ResultRow> {
-        val fieldIndex = returningExpressions.withIndex().associateBy({ it.value }, { it.index })
-
-        private var hasNext = false
-            set(value) {
-                field = value
-                if (!field) {
-                    val statement = rs.statement
-                    rs.close()
-                    statement?.close()
-                    transaction.openResultSetsCount--
-                }
-            }
+    private inner class ResultIterator(rs: ResultSet) : StatementIterator<Expression<*>, ResultRow>(rs) {
+        override val fieldIndex = returningExpressions.withIndex()
+            .associateBy({ it.value }, { it.index })
 
         init {
-            hasNext = rs.next()
+            hasNext = result.next()
         }
 
-        override fun hasNext(): Boolean = hasNext
-
-        override operator fun next(): ResultRow {
-            if (!hasNext) throw NoSuchElementException()
-            val result = ResultRow.create(rs, fieldIndex)
-            hasNext = rs.next()
-            return result
-        }
+        override fun createResultRow(): ResultRow = ResultRow.create(result, fieldIndex)
     }
 }
diff --git a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/statements/StatementIterator.kt b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/statements/StatementIterator.kt
@@ -0,0 +1,32 @@
+package org.jetbrains.exposed.sql.statements
+
+import org.jetbrains.exposed.sql.transactions.TransactionManager
+import java.sql.ResultSet
+
+internal abstract class StatementIterator<T, RR>(
+    protected val result: ResultSet
+) : Iterator<RR> {
+    protected abstract val fieldIndex: Map<T, Int>
+
+    protected abstract fun createResultRow(): RR
+
+    protected var hasNext = false
+        set(value) {
+            field = value
+            if (!field) {
+                val statement = result.statement
+                result.close()
+                statement?.close()
+                TransactionManager.current().openResultSetsCount--
+            }
+        }
+
+    override fun hasNext(): Boolean = hasNext
+
+    override operator fun next(): RR {
+        if (!hasNext) throw NoSuchElementException()
+        val resultRow = createResultRow()
+        hasNext = result.next()
+        return resultRow
+    }
+}