Make nested transactions no-ops on SQLite

Author: tulir

mautrix/util/async_db/aiosqlite.py

@@ -7,6 +7,7 @@
 
 from typing import Any, AsyncContextManager
 from contextlib import asynccontextmanager
+from contextvars import ContextVar
 import asyncio
 import logging
 import os
@@ -24,6 +25,9 @@
 POSITIONAL_PARAM_PATTERN = re.compile(r"\$(\d+)")
 
 
+in_transaction = ContextVar("in_transaction", default=False)
+
+
 class TxnConnection(aiosqlite.Connection):
     def __init__(self, path: str, **kwargs) -> None:
         def connector() -> sqlite3.Connection:
@@ -35,14 +39,20 @@ def connector() -> sqlite3.Connection:
 
     @asynccontextmanager
     async def transaction(self) -> None:
+        if in_transaction.get():
+            yield
+            return
         await self.execute("BEGIN TRANSACTION")
+        token = in_transaction.set(True)
         try:
             yield
         except Exception:
             await self.rollback()
             raise
         else:
             await self.commit()
+        finally:
+            in_transaction.reset(token)
 
     def __execute(self, query: str, *args: Any):
         query = POSITIONAL_PARAM_PATTERN.sub(r"?\1", query)