Simplify transaction state modelling

Author: michaelbull

src/main/kotlin/com/github/michaelbull/jdbc/Transaction.kt

@@ -1,82 +1,105 @@
 package com.github.michaelbull.jdbc
 
-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
 import kotlin.contracts.contract
-import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.coroutineContext
 
 /**
  * Calls the specified suspending [block] in the context of a [CoroutineTransaction], suspends until it completes, and
  * returns the result.
  *
- * When there exists a [CoroutineTransaction] in the current [CoroutineContext], the [block] will be immediately invoked
- * if the [transaction is running][CoroutineTransaction.isRunning], otherwise an [IllegalStateException] will be thrown.
+ * When the [coroutineContext] has no [CoroutineTransaction], the specified suspending [block] will be
+ * [ran transactionally][runTransactionally] [with the context of a connection][withConnection].
  *
- * When no [CoroutineTransaction] exists in the current [CoroutineContext], the [block] will be invoked
- * [with the context][withContext] of a new [CoroutineTransaction].
+ * When the [coroutineContext] has an [incomplete][CoroutineTransaction.incomplete] [CoroutineTransaction], the
+ * specified suspending [block] will be called [with this context][withContext].
  *
- * The [block] will be invoked [with a connection][withConnection] in its [CoroutineContext]. The connection's
- * [autoCommit][Connection.setAutoCommit] is set to `false` before the invocation. If the [block] throws a [Throwable],
- * 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].
+ * When the [coroutineContext] has a [completed][CoroutineTransaction.completed] [CoroutineTransaction], an
+ * [IllegalStateException] will be thrown.
  */
 suspend inline fun <T> transaction(crossinline block: suspend CoroutineScope.() -> T): T {
     contract {
         callsInPlace(block, InvocationKind.AT_MOST_ONCE)
     }
 
-    val existingTransaction = coroutineContext.transaction
+    val existingTransaction = coroutineContext[CoroutineTransaction]
 
     return when {
-        existingTransaction == null -> withContext(CoroutineTransaction()) {
+        existingTransaction == null -> {
             withConnection {
-                execute(block)
+                runTransactionally {
+                    block()
+                }
             }
         }
 
-        existingTransaction.isRunning -> withContext(coroutineContext) {
-            block()
+        existingTransaction.incomplete -> {
+            withContext(coroutineContext) {
+                block()
+            }
         }
 
         else -> error("Attempted to start new transaction within: $existingTransaction")
     }
 }
 
 /**
- * [Starts][CoroutineTransaction.start] the current [CoroutineTransaction] and sets the
- * current [CoroutineConnection]'s [autoCommit][Connection.setAutoCommit] to `false`, calls the specified suspending
- * [block], suspends until it completes, then [commits][Connection.commit] and returns the result.
+ * Calls the specified suspending [block] [with the context][withContext] of a [CoroutineTransaction] and returns its
+ * result.
+ *
+ * If invocation of the suspending [block] was successful, [commit][Connection.commit] is then called on the
+ * [Connection] in the [coroutineContext].
  *
- * If the [block] throws a [Throwable], the connection will [rollback][Connection.rollback] and not
- * [commit][Connection.commit].
+ * If invocation of the suspending [block] throws a [Throwable] exception, [rollback][Connection.rollback] is then
+ * called on the [Connection] in the [coroutineContext] and the exception is thrown.
  */
 @PublishedApi
-internal suspend inline fun <T> execute(crossinline block: suspend CoroutineScope.() -> T): T {
+internal suspend inline fun <T> runTransactionally(crossinline block: suspend CoroutineScope.() -> T): T {
     contract {
-        callsInPlace(block, InvocationKind.AT_MOST_ONCE)
+        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
     }
 
-    val transaction = coroutineContext.transaction ?: error("No transaction in context")
-    transaction.start()
+    coroutineContext.connection.runWithManualCommit {
+        val transaction = CoroutineTransaction()
+
+        try {
+            val result = withContext(transaction) {
+                block()
+            }
+
+            commit()
+            return result
+        } catch (ex: Throwable) {
+            rollback()
+            throw ex
+        } finally {
+            transaction.complete()
+        }
+    }
+}
+
+/**
+ * Disables [autoCommit][Connection.getAutoCommit] mode on `this` [Connection], then calls a specific function [block]
+ * with `this` [Connection] as its receiver and returns its result, then sets the [autoCommit][Connection.getAutoCommit]
+ * mode on `this` [Connection] back to its original value.
+ */
+@PublishedApi
+internal inline fun <T> Connection.runWithManualCommit(block: Connection.() -> T): T {
+    contract {
+        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+    }
 
-    val connection = coroutineContext.connection
-    connection.autoCommit = false
+    val before = autoCommit
 
-    try {
-        val result = withContext(coroutineContext) { block() }
-        transaction.complete()
-        connection.commit()
-        return result
-    } catch (ex: Throwable) {
-        transaction.complete()
-        connection.rollback()
-        throw ex
+    return try {
+        autoCommit = false
+        this.run(block)
+    } finally {
+        autoCommit = before
     }
 }

src/main/kotlin/com/github/michaelbull/jdbc/TransactionState.kt

@@ -1,33 +1,21 @@
 package com.github.michaelbull.jdbc.context
 
-import com.github.michaelbull.jdbc.TransactionState
 import kotlin.coroutines.AbstractCoroutineContextElement
 import kotlin.coroutines.CoroutineContext
 
 @PublishedApi
-internal val CoroutineContext.transaction: CoroutineTransaction?
-    get() = get(CoroutineTransaction)
-
-@PublishedApi
-internal class CoroutineTransaction : AbstractCoroutineContextElement(CoroutineTransaction) {
+internal class CoroutineTransaction(
+    private var completed: Boolean = false
+) : AbstractCoroutineContextElement(CoroutineTransaction) {
 
     companion object Key : CoroutineContext.Key<CoroutineTransaction>
 
-    var state: TransactionState = TransactionState.Idle
-        private set
-
-    val isRunning: Boolean
-        get() = state == TransactionState.Running
-
-    fun start() {
-        check(state == TransactionState.Idle) { "cannot start: $this" }
-        state = TransactionState.Running
-    }
+    val incomplete: Boolean
+        get() = !completed
 
     fun complete() {
-        check(state == TransactionState.Running) { "cannot complete: $this" }
-        state = TransactionState.Completed
+        completed = true
     }
 
-    override fun toString(): String = "CoroutineTransaction(state=$state)"
+    override fun toString(): String = "CoroutineTransaction(completed=$completed)"
 }