From 12d970b1c1aa796ec4124716b2af3c83b2b27d7f Mon Sep 17 00:00:00 2001
From: Tapac <fantocci@gmail.com>
Date: Fri, 16 Aug 2019 14:52:14 +0200
Subject: [PATCH] Using coroutines with multiple databases #624

---
 .../transactions/experimental/Suspended.kt    | 11 ++++-
 .../exposed/sql/tests/h2/MultiDatabaseTest.kt | 48 +++++++++++++++++++
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/exposed/src/main/kotlin/org/jetbrains/exposed/sql/transactions/experimental/Suspended.kt b/exposed/src/main/kotlin/org/jetbrains/exposed/sql/transactions/experimental/Suspended.kt
index 7711899eb0..563736d128 100644
--- a/exposed/src/main/kotlin/org/jetbrains/exposed/sql/transactions/experimental/Suspended.kt
+++ b/exposed/src/main/kotlin/org/jetbrains/exposed/sql/transactions/experimental/Suspended.kt
@@ -115,7 +115,8 @@ private suspend fun <T> withTransactionScope(context: CoroutineContext?,
                                              currentTransaction: Transaction?,
                                              db: Database? = null,
                                              body: suspend TransactionScope.() -> T) : T {
-    return coroutineContext[TransactionScope]?.body() ?: run {
+    val currentScope = coroutineContext[TransactionScope]
+    suspend fun newScope() : T {
         val manager = (currentTransaction?.db ?: db)?.transactionManager ?: TransactionManager.manager
 
         val tx = currentTransaction ?: manager.newTransaction(manager.defaultIsolationLevel)
@@ -124,7 +125,13 @@ private suspend fun <T> withTransactionScope(context: CoroutineContext?,
 
         val newContext = context ?: coroutineContext
 
-        TransactionScope(tx, newContext + element).body()
+       return TransactionScope(tx, newContext + element).body()
+    }
+    return when {
+        currentScope == null -> newScope()
+        currentTransaction != null && currentScope.tx != currentTransaction -> newScope()
+        db != null && currentScope.tx.db != db -> newScope()
+        else -> currentScope.body()
     }
 }
 
diff --git a/exposed/src/test/kotlin/org/jetbrains/exposed/sql/tests/h2/MultiDatabaseTest.kt b/exposed/src/test/kotlin/org/jetbrains/exposed/sql/tests/h2/MultiDatabaseTest.kt
index 2f50e5d84d..16b6fa5932 100644
--- a/exposed/src/test/kotlin/org/jetbrains/exposed/sql/tests/h2/MultiDatabaseTest.kt
+++ b/exposed/src/test/kotlin/org/jetbrains/exposed/sql/tests/h2/MultiDatabaseTest.kt
@@ -1,9 +1,13 @@
 package org.jetbrains.exposed.sql.tests.h2
 
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.runBlocking
 import org.jetbrains.exposed.sql.*
 import org.jetbrains.exposed.sql.tests.shared.DMLTestsData
 import org.jetbrains.exposed.sql.tests.shared.assertEqualLists
 import org.jetbrains.exposed.sql.transactions.TransactionManager
+import org.jetbrains.exposed.sql.transactions.experimental.newSuspendedTransaction
+import org.jetbrains.exposed.sql.transactions.experimental.suspendedTransaction
 import org.jetbrains.exposed.sql.transactions.transaction
 import org.jetbrains.exposed.sql.transactions.transactionManager
 import org.junit.After
@@ -150,4 +154,48 @@ class MultiDatabaseTest {
             SchemaUtils.drop(DMLTestsData.Cities)
         }
     }
+
+    @Test
+    fun testCoroutinesWithMultiDb() = runBlocking {
+        newSuspendedTransaction(Dispatchers.IO, db1) {
+            val tr1 = this
+            SchemaUtils.create(DMLTestsData.Cities)
+            assertTrue(DMLTestsData.Cities.selectAll().empty())
+            DMLTestsData.Cities.insert {
+                it[DMLTestsData.Cities.name] = "city1"
+            }
+
+            newSuspendedTransaction(Dispatchers.IO, db2) {
+                assertFalse(DMLTestsData.Cities.exists())
+                SchemaUtils.create(DMLTestsData.Cities)
+                DMLTestsData.Cities.insert {
+                    it[DMLTestsData.Cities.name] = "city2"
+                }
+                DMLTestsData.Cities.insert {
+                    it[DMLTestsData.Cities.name] = "city3"
+                }
+                assertEquals(2, DMLTestsData.Cities.selectAll().count())
+                assertEquals("city3", DMLTestsData.Cities.selectAll().last()[DMLTestsData.Cities.name])
+
+                tr1.suspendedTransaction {
+                    assertEquals(1, DMLTestsData.Cities.selectAll().count())
+                    DMLTestsData.Cities.insert {
+                        it[DMLTestsData.Cities.name] = "city4"
+                    }
+                    DMLTestsData.Cities.insert {
+                        it[DMLTestsData.Cities.name] = "city5"
+                    }
+                    assertEquals(3, DMLTestsData.Cities.selectAll().count())
+                }
+
+                assertEquals(2, DMLTestsData.Cities.selectAll().count())
+                assertEquals("city3", DMLTestsData.Cities.selectAll().last()[DMLTestsData.Cities.name])
+                SchemaUtils.drop(DMLTestsData.Cities)
+            }
+
+            assertEquals(3, DMLTestsData.Cities.selectAll().count())
+            assertEqualLists(listOf("city1", "city4", "city5"), DMLTestsData.Cities.selectAll().map { it[DMLTestsData.Cities.name]})
+            SchemaUtils.drop(DMLTestsData.Cities)
+        }
+    }
 }
\ No newline at end of file