From d9038021cd221331b484bc7cef233f800f974157 Mon Sep 17 00:00:00 2001
From: Pwootage <pwootage@gmail.com>
Date: Sun, 22 Jul 2018 07:43:04 -0600
Subject: [PATCH] Allow overriding the fetch size for a Query (#327)

* Allow overriding the fetch size for a Query
* Allow setting default fetch size in Database
---
 .../kotlin/org/jetbrains/exposed/sql/Database.kt  |  8 ++++++++
 .../kotlin/org/jetbrains/exposed/sql/Query.kt     | 15 ++++++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/src/main/kotlin/org/jetbrains/exposed/sql/Database.kt b/src/main/kotlin/org/jetbrains/exposed/sql/Database.kt
index 4445e660d3..6187d2ae62 100644
--- a/src/main/kotlin/org/jetbrains/exposed/sql/Database.kt
+++ b/src/main/kotlin/org/jetbrains/exposed/sql/Database.kt
@@ -57,6 +57,14 @@ class Database private constructor(val connector: () -> Connection) {
         }
     }
 
+    var defaultFetchSize: Int? = null
+        private set
+
+    fun defaultFetchSize(size: Int): Database {
+        defaultFetchSize = size
+        return this
+    }
+
     private fun String.isIdentifier() = !isEmpty() && first().isIdentifierStart() && all { it.isIdentifierStart() || it in '0'..'9' }
     private fun Char.isIdentifierStart(): Boolean = this in 'a'..'z' || this in 'A'..'Z' || this == '_' || this in extraNameCharacters
 
diff --git a/src/main/kotlin/org/jetbrains/exposed/sql/Query.kt b/src/main/kotlin/org/jetbrains/exposed/sql/Query.kt
index 6ebbe1f110..72a1fd884b 100644
--- a/src/main/kotlin/org/jetbrains/exposed/sql/Query.kt
+++ b/src/main/kotlin/org/jetbrains/exposed/sql/Query.kt
@@ -93,6 +93,8 @@ open class Query(set: FieldSet, where: Op<Boolean>?): SizedIterable<ResultRow>,
         private set
     var offset: Int = 0
         private set
+    var fetchSize: Int? = null
+        private set
 
     /**
      * Changes [set.fields] field of a Query, [set.source] will be preserved
@@ -121,7 +123,13 @@ open class Query(set: FieldSet, where: Op<Boolean>?): SizedIterable<ResultRow>,
     fun hasCustomForUpdateState() = forUpdate != null
     fun isForUpdate() = (forUpdate ?: false) && currentDialect.supportsSelectForUpdate()
 
-    override fun PreparedStatement.executeInternal(transaction: Transaction): ResultSet? = executeQuery()
+    override fun PreparedStatement.executeInternal(transaction: Transaction): ResultSet? {
+        val fetchSize = this@Query.fetchSize ?: transaction.db.defaultFetchSize
+        if (fetchSize != null) {
+            this.fetchSize = fetchSize
+        }
+        return executeQuery()
+    }
 
     override fun arguments() = QueryBuilder(true).let {
         prepareSQL(it)
@@ -232,6 +240,11 @@ open class Query(set: FieldSet, where: Op<Boolean>?): SizedIterable<ResultRow>,
         return this
     }
 
+    fun fetchSize(n: Int): Query {
+        this.fetchSize = n
+        return this
+    }
+
     private inner class ResultIterator(val rs: ResultSet): Iterator<ResultRow> {
         private var hasNext: Boolean? = null
         private val fieldsIndex = HashMap<Expression<*>, Int>()