Add CoroutineScope receiver to transaction & withConnection

See: #1
michaelbull · Aug 25, 2020 · d5fd3a9 · d5fd3a9
1 parent 24c4b66
commit d5fd3a9
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 8 deletions.
diff --git a/README.md b/README.md
@@ -27,7 +27,7 @@ The primary higher-order function exposed by the library is the
 [`transaction`][transaction] function.
 
 ```kotlin
-suspend inline fun <T> transaction(crossinline block: suspend () -> T): T
+suspend inline fun <T> transaction(crossinline block: suspend CoroutineScope.() -> T): T
 ```
 
 Calling this function with a specific suspending block will run the block in

diff --git a/src/main/kotlin/com/github/michaelbull/jdbc/Connection.kt b/src/main/kotlin/com/github/michaelbull/jdbc/Connection.kt
@@ -4,6 +4,7 @@ import com.github.michaelbull.jdbc.context.CoroutineConnection
 import com.github.michaelbull.jdbc.context.CoroutineDataSource
 import com.github.michaelbull.jdbc.context.dataSource
 import com.github.michaelbull.logging.InlineLogger
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.withContext
 import java.sql.Connection
 import java.sql.SQLException
@@ -27,7 +28,7 @@ internal val logger = InlineLogger()
  * [with the context][withContext] of a new [CoroutineConnection] and an attempt will be made to [Connection.close]
  * it afterwards.
  */
-suspend inline fun <T> withConnection(crossinline block: suspend () -> T): T {
+suspend inline fun <T> withConnection(crossinline block: suspend CoroutineScope.() -> T): T {
     contract {
         callsInPlace(block, InvocationKind.AT_MOST_ONCE)
     }
@@ -37,7 +38,9 @@ suspend inline fun <T> withConnection(crossinline block: suspend () -> T): T {
     return if (connection.isNullOrClosed()) {
         newConnection(block)
     } else {
-        block()
+        withContext(coroutineContext) {
+            block()
+        }
     }
 }
 
@@ -49,7 +52,7 @@ suspend inline fun <T> withConnection(crossinline block: suspend () -> T): T {
  * [IllegalStateException] is thrown.
  */
 @PublishedApi
-internal suspend inline fun <T> newConnection(crossinline block: suspend () -> T): T {
+internal suspend inline fun <T> newConnection(crossinline block: suspend CoroutineScope.() -> T): T {
     contract {
         callsInPlace(block, InvocationKind.AT_MOST_ONCE)
     }

diff --git a/src/main/kotlin/com/github/michaelbull/jdbc/Transaction.kt b/src/main/kotlin/com/github/michaelbull/jdbc/Transaction.kt
@@ -4,6 +4,7 @@ import com.github.michaelbull.jdbc.context.CoroutineConnection
 import com.github.michaelbull.jdbc.context.CoroutineTransaction
 import com.github.michaelbull.jdbc.context.connection
 import com.github.michaelbull.jdbc.context.transaction
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.withContext
 import java.sql.Connection
 import kotlin.contracts.InvocationKind
@@ -26,7 +27,7 @@ import kotlin.coroutines.coroutineContext
  * the transaction will [rollback][Connection.rollback] and re-throw the [Throwable], otherwise the transaction will
  * [commit][Connection.commit] and return the result of type [T].
  */
-suspend inline fun <T> transaction(crossinline block: suspend () -> T): T {
+suspend inline fun <T> transaction(crossinline block: suspend CoroutineScope.() -> T): T {
     contract {
         callsInPlace(block, InvocationKind.AT_MOST_ONCE)
     }
@@ -39,7 +40,11 @@ suspend inline fun <T> transaction(crossinline block: suspend () -> T): T {
                 execute(block)
             }
         }
-        existingTransaction.isRunning -> block()
+
+        existingTransaction.isRunning -> withContext(coroutineContext) {
+            block()
+        }
+
         else -> error("Attempted to start new transaction within: $existingTransaction")
     }
 }
@@ -53,7 +58,7 @@ suspend inline fun <T> transaction(crossinline block: suspend () -> T): T {
  * [commit][Connection.commit].
  */
 @PublishedApi
-internal suspend inline fun <T> execute(crossinline block: suspend () -> T): T {
+internal suspend inline fun <T> execute(crossinline block: suspend CoroutineScope.() -> T): T {
     contract {
         callsInPlace(block, InvocationKind.AT_MOST_ONCE)
     }
@@ -65,7 +70,7 @@ internal suspend inline fun <T> execute(crossinline block: suspend () -> T): T {
     connection.autoCommit = false
 
     try {
-        val result = block()
+        val result = withContext(coroutineContext) { block() }
         transaction.complete()
         connection.commit()
         return result