scoping transaction to engine

Author: dantownsend

piccolo/engine/postgres.py

@@ -17,11 +17,6 @@
 from piccolo.utils.sync import run_sync
 
 
-pg_transaction_connection = contextvars.ContextVar(
-    "pg_transaction_connection", default=None
-)
-
-
 @dataclass
 class AsyncBatch(Batch):
 
@@ -136,6 +131,8 @@ def run_sync(self):
 
 
 ###############################################################################
+
+
 class Transaction:
     """
     Used for wrapping queries in a transaction, using a context manager.
@@ -162,7 +159,7 @@ async def __aenter__(self):
 
         self.transaction = self.connection.transaction()
         await self.transaction.start()
-        self.context = pg_transaction_connection.set(self.connection)
+        self.context = self.engine.transaction_connection.set(self.connection)
 
     async def commit(self):
         await self.transaction.commit()
@@ -181,9 +178,9 @@ async def __aexit__(self, exception_type, exception, traceback):
         else:
             self.connection.close()
 
-        pg_transaction_connection.reset(self.context)
+        self.engine.transaction_connection.reset(self.context)
 
-        return exception is not None
+        return exception is None
 
 
 ###############################################################################
@@ -199,6 +196,10 @@ class PostgresEngine(Engine):
     def __init__(self, config: t.Dict[str, t.Any]) -> None:
         self.config = config
         self.pool: t.Optional[Pool] = None
+        database_name = config.get("database", "Unknown")
+        self.transaction_connection = contextvars.ContextVar(
+            f"pg_transaction_connection_{database_name}", default=None
+        )
         super().__init__()
 
     def get_version(self) -> float:
@@ -283,7 +284,7 @@ async def run_querystring(
         )
 
         # If running inside a transaction:
-        connection = pg_transaction_connection.get()
+        connection = self.transaction_connection.get()
         if connection:
             return await connection.fetch(query, *query_args)
         elif in_pool and self.pool:

piccolo/engine/sqlite.py

@@ -14,11 +14,6 @@
 from piccolo.utils.sync import run_sync
 
 
-sqlite_transaction_connection = contextvars.ContextVar(
-    "sqlite_transaction_connection", default=None
-)
-
-
 @dataclass
 class AsyncBatch(Batch):
 
@@ -135,7 +130,7 @@ def __init__(self, engine: SQLiteEngine):
     async def __aenter__(self):
         self.connection = await self.engine.get_connection()
         await self.connection.execute("BEGIN")
-        self.context = sqlite_transaction_connection.set(self.connection)
+        self.context = self.engine.transaction_connection.set(self.connection)
 
     async def __aexit__(self, exception_type, exception, traceback):
         if exception:
@@ -144,7 +139,9 @@ async def __aexit__(self, exception_type, exception, traceback):
             await self.connection.execute("COMMIT")
 
         await self.connection.close()
-        return exception is not None
+        self.engine.transaction_connection.reset(self.context)
+
+        return exception is None
 
 
 ###############################################################################
@@ -186,6 +183,11 @@ def __init__(
             }
         )
         self.connection_kwargs = connection_kwargs
+
+        self.transaction_connection = contextvars.ContextVar(
+            f"sqlite_transaction_connection_{path}", default=None
+        )
+
         super().__init__()
 
     @property
@@ -297,7 +299,7 @@ async def run_querystring(
         )
 
         # If running inside a transaction:
-        connection = sqlite_transaction_connection.get()
+        connection = self.transaction_connection.get()
         if connection:
             return await self._run_in_existing_connection(
                 connection=connection,
@@ -307,7 +309,7 @@ async def run_querystring(
             )
 
         return await self._run_in_new_connection(
-            query=query, args=query_args, query_type=querystring.query_type,
+            query=query, args=query_args, query_type=querystring.query_type
         )
 
     def atomic(self) -> Atomic:

tests/engine/test_nested_transaction.py

@@ -0,0 +1,55 @@
+import asyncio
+from unittest import TestCase
+
+from piccolo.columns.column_types import Varchar
+from piccolo.engine.sqlite import SQLiteEngine
+from piccolo.table import Table
+
+from ..base import sqlite_only
+
+
+ENGINE_1 = SQLiteEngine(path="engine1.sqlite")
+ENGINE_2 = SQLiteEngine(path="engine2.sqlite")
+
+
+class Musician(Table, db=ENGINE_1):
+    name = Varchar(length=100)
+
+
+class Roadie(Table, db=ENGINE_2):
+    name = Varchar(length=100)
+
+
+@sqlite_only
+class TestNestedTransaction(TestCase):
+    def setUp(self):
+        ENGINE_1.remove_db_file()
+        ENGINE_2.remove_db_file()
+
+    def tearDown(self):
+        ENGINE_1.remove_db_file()
+        ENGINE_2.remove_db_file()
+
+    async def run_nested(self):
+        """
+        Make sure nested transactions which reference different databases work
+        as expected.
+        """
+        async with Musician._meta.db.transaction():
+            await Musician.create_table().run()
+            await Musician(name="Bob").save().run()
+
+            async with Roadie._meta.db.transaction():
+                await Roadie.create_table().run()
+                await Roadie(name="Dave").save().run()
+
+        self.assertTrue(await Musician.table_exists().run())
+        musician = await Musician.select("name").first().run()
+        self.assertTrue(musician["name"] == "Bob")
+
+        self.assertTrue(await Roadie.table_exists().run())
+        roadie = await Roadie.select("name").first().run()
+        self.assertTrue(roadie["name"] == "Dave")
+
+    def test_nested(self):
+        asyncio.run(self.run_nested())

tests/engine/test_transaction.py

@@ -1,8 +1,10 @@
 import asyncio
 from unittest import TestCase
 
+from piccolo.engine.sqlite import SQLiteEngine
+
 from ..example_project.tables import Band, Manager
-from ..base import postgres_only
+from ..base import postgres_only, sqlite_only
 
 
 class TestAtomic(TestCase):