Fix async ORM connection ownership race

fastapi_startkit/src/fastapi_startkit/masoniteorm/connections/connection.py

@@ -1,31 +1,104 @@
-from typing import List
+from __future__ import annotations
+
+from contextvars import ContextVar, Token
+from types import TracebackType
+from typing import TYPE_CHECKING
 
 from sqlalchemy import text
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, AsyncTransaction
 
 from fastapi_startkit.masoniteorm.models.builder import QueryBuilder
 
+if TYPE_CHECKING:
+    from typing import Self
+
+
+class Transaction:
+    def __init__(self, owner: Connection):
+        self.owner = owner
+        self.connection: AsyncConnection | None = None
+        self.transaction: AsyncTransaction | None = None
+        self._token: Token[AsyncConnection | None] | None = None
+        self._owns_connection = False
+
+    async def __aenter__(self) -> Self:
+        connection = self.owner.connection
+        if connection is None:
+            connection = await self.owner.engine.connect()
+            self._owns_connection = True
+        self.connection = connection
+        self._token = self.owner._connection_context.set(connection)
+
+        try:
+            if connection.in_transaction():
+                self.transaction = await connection.begin_nested()
+            else:
+                self.transaction = await connection.begin()
+        except BaseException:
+            self.owner._connection_context.reset(self._token)
+            if self._owns_connection:
+                await connection.close()
+            raise
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        assert self.connection is not None
+        assert self.transaction is not None
+        assert self._token is not None
+        try:
+            await self.transaction.__aexit__(exc_type, exc_value, traceback)
+        finally:
+            self.owner._connection_context.reset(self._token)
+            if self._owns_connection:
+                await self.connection.close()
+
+    async def commit(self) -> None:
+        assert self.transaction is not None
+        await self.transaction.commit()
+
+    async def rollback(self) -> None:
+        assert self.transaction is not None
+        await self.transaction.rollback()
+
 
 class Connection:
     def __init__(self, engine: AsyncEngine, config: dict):
         self.config = config
         self.engine: AsyncEngine = engine
-        self.connection: AsyncConnection | None = None
-        self.transactions: List[AsyncTransaction] = []
+        self._connection_context: ContextVar[AsyncConnection | None] = ContextVar(
+            f"masoniteorm_connection_{id(self)}", default=None
+        )
+
+    @property
+    def connection(self) -> AsyncConnection | None:
+        return self._connection_context.get()
+
+    @property
+    def transactions(self) -> list[AsyncTransaction]:
+        connection = self.connection
+        if connection is None:
+            return []
+        nested = connection.get_nested_transaction()
+        root = connection.get_transaction()
+        return [transaction for transaction in (root, nested) if transaction is not None]
+
+    def transaction(self) -> Transaction:
+        return Transaction(self)
 
-    def query(self) -> "QueryBuilder":
+    def query(self) -> QueryBuilder:
         return QueryBuilder(
             connection=self,
             grammar=self.get_query_grammar(),
             processor=self.get_post_processor(),
         )
 
     async def get_connection(self) -> AsyncConnection:
-        if self.connection is None:
-            self.connection = await self.engine.connect()
-
-        assert self.connection is not None
-        return self.connection
+        return self.connection or await self.engine.connect()
 
     def get_query_grammar(cls):
         pass
@@ -34,38 +107,51 @@ def get_post_processor(self):
         pass
 
     async def begin_transaction(self) -> None:
-        connection = await self.get_connection()
-
-        if not self.transactions:
-            transaction = await connection.begin()
+        connection = self.connection
+        if connection is None:
+            connection = await self.engine.connect()
+            self._connection_context.set(connection)
+        if connection.in_transaction():
+            await connection.begin_nested()
         else:
-            transaction = await connection.begin_nested()
-
-        self.transactions.append(transaction)
+            await connection.begin()
 
     async def commit_transaction(self) -> None:
-        if not self.transactions:
+        connection = self.connection
+        if connection is None or not connection.in_transaction():
             raise RuntimeError("No active transaction to commit")
-
-        transaction = self.transactions.pop()
-        await transaction.commit()
-
-        await self._maybe_cleanup()
+        nested = connection.get_nested_transaction()
+        if nested is not None:
+            await nested.commit()
+        else:
+            transaction = connection.get_transaction()
+            assert transaction is not None
+            await transaction.commit()
+            await self._release_connection(connection)
 
     async def rollback(self) -> None:
-        if not self.transactions:
+        connection = self.connection
+        if connection is None or not connection.in_transaction():
             raise RuntimeError("No active transaction to rollback")
+        nested = connection.get_nested_transaction()
+        if nested is not None:
+            await nested.rollback()
+        else:
+            transaction = connection.get_transaction()
+            assert transaction is not None
+            await transaction.rollback()
+            await self._release_connection(connection)
 
-        transaction = self.transactions.pop()
-        await transaction.rollback()
-
-        await self._maybe_cleanup()
+    async def _release_connection(self, connection: AsyncConnection) -> None:
+        await connection.close()
+        if self.connection is connection:
+            self._connection_context.set(None)
 
     async def close(self) -> None:
-        if self.connection is not None:
-            await self.connection.close()
-            self.connection = None
-        self.transactions = []
+        connection = self.connection
+        if connection is not None:
+            await connection.close()
+            self._connection_context.set(None)
 
     async def reconnect(self) -> None:
         await self.close()
@@ -81,30 +167,22 @@ def sql_alchemy_bindings(query: str, bindings: list | None = None):
         return (query, params)
 
     async def run(self, query: str, bindings: list | None = None):
-        query, bindings = self.sql_alchemy_bindings(query, bindings)
-
-        conn = await self.get_connection()
-        result = await conn.execute(text(query), bindings or {})
-
-        if not self.transactions:
-            await conn.commit()
-
-        return result
+        query, params = self.sql_alchemy_bindings(query, bindings)
+        return await self._execute(query, params)
 
     async def execute(self, query: str, bindings: list | None = None):
-        query, bindings = self.sql_alchemy_bindings(query, bindings)
-
-        conn = await self.get_connection()
-        result = await conn.execute(text(query), bindings or {})
+        query, params = self.sql_alchemy_bindings(query, bindings)
+        return await self._execute(query, params)
 
-        if not self.transactions:
-            await conn.commit()
-
-        return result
+    async def _execute(self, query: str, bindings: dict):
+        connection = self.connection
+        if connection is not None:
+            return await connection.execute(text(query), bindings or {})
+        async with self.engine.begin() as operation_connection:
+            return await operation_connection.execute(text(query), bindings or {})
 
     async def insert(self, query: str, bindings: list | None = None) -> int | None:
         result = await self.execute(query, bindings)
-
         return getattr(result, "lastrowid", None)
 
     async def insert_get_id(self, query: str, bindings: list | None = None) -> int | None:
@@ -113,7 +191,6 @@ async def insert_get_id(self, query: str, bindings: list | None = None) -> int |
 
     async def update(self, query: str, bindings: list | None = None) -> int:
         result = await self.execute(query, bindings)
-
         return result.rowcount  # type: ignore[return-value]
 
     async def delete(self, query: str, bindings: list | None = None) -> int:
@@ -122,30 +199,14 @@ async def delete(self, query: str, bindings: list | None = None) -> int:
 
     async def select(self, query: str, bindings: list | None = None) -> list[dict]:
         result = await self.run(query, bindings)
-
         return result.mappings().all()
 
     async def select_one(self, query: str, bindings: list | None = None) -> dict | None:
         result = await self.run(query, bindings)
         row = result.fetchone()
-        result_dict = dict(zip(result.keys(), row)) if row else None
-        if not self.transactions and self.connection is not None:
-            await self.connection.commit()
-        return result_dict
+        return dict(zip(result.keys(), row)) if row else None
 
     async def statement(self, query: str, bindings: list | None = None) -> bool:
-        query, bindings = self.sql_alchemy_bindings(query, bindings)
-
-        conn = await self.get_connection()
-        await conn.execute(text(query), bindings or {})
-
-        # Only commit if NOT inside a transaction
-        if not self.transactions:
-            await conn.commit()
-
+        query, params = self.sql_alchemy_bindings(query, bindings)
+        await self._execute(query, params)
         return True
-
-    async def _maybe_cleanup(self):
-        if not self.transactions and self.connection:
-            await self.connection.close()
-            self.connection = None

fastapi_startkit/src/fastapi_startkit/masoniteorm/connections/postgres_connection.py

@@ -10,9 +10,6 @@ class PostgresConnection(Connection):
     async def insert_get_id(self, query: str, bindings: list | None = None) -> int | None:
         result = await self.run(query, bindings)
         row = result.fetchone()
-        if not self.transactions:
-            conn = await self.get_connection()
-            await conn.commit()
         return row[0] if row is not None else None
 
     @classmethod

fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/Blueprint.py

@@ -736,8 +736,9 @@ async def __aexit__(self, exc_type, exc_value, exc_traceback):
 
         sql = await self.to_sql()
         if isinstance(sql, list):
-            for q in sql:
-                await self.connection.statement(q, ())
+            async with self.connection.transaction():
+                for q in sql:
+                    await self.connection.statement(q, ())
             return
         return await self.connection.statement(sql, ())
 

fastapi_startkit/src/fastapi_startkit/masoniteorm/testing/transaction.py

@@ -3,10 +3,12 @@ async def asyncStartTestRun(self):
         from fastapi_startkit.masoniteorm.models import Model
 
         self.connection = Model.db_manager.connection(None)
-        await self.connection.begin_transaction()
+        self.transaction = self.connection.transaction()
+        await self.transaction.__aenter__()
 
     async def asyncStopTestRun(self):
-        await self.connection.rollback()
+        await self.transaction.rollback()
+        await self.transaction.__aexit__(None, None, None)
 
 
 class RefreshDatabase(DatabaseTransaction):

fastapi_startkit/tests/masoniteorm/connections/test_postgres_connection.py

@@ -0,0 +1,32 @@
+from unittest import IsolatedAsyncioTestCase
+from unittest.mock import AsyncMock, Mock
+
+from fastapi_startkit.masoniteorm.connections.postgres_connection import PostgresConnection
+from fastapi_startkit.masoniteorm.query.grammars import PostgresGrammar
+from fastapi_startkit.masoniteorm.query.processors import PostgresPostProcessor
+from fastapi_startkit.masoniteorm.schema.platforms import PostgresPlatform
+
+
+class TestPostgresConnection(IsolatedAsyncioTestCase):
+    @staticmethod
+    def _connection_with_result(row):
+        connection = PostgresConnection(engine=None, config={"driver": "postgres"})
+        result = Mock()
+        result.fetchone.return_value = row
+        connection.run = AsyncMock(return_value=result)
+        return connection
+
+    async def test_insert_get_id_returns_first_column_of_returning_row(self):
+        connection = self._connection_with_result((42,))
+        inserted_id = await connection.insert_get_id("INSERT INTO users (name) VALUES (?) RETURNING id", ["Joe"])
+        self.assertEqual(inserted_id, 42)
+
+    async def test_insert_get_id_returns_none_when_no_row(self):
+        connection = self._connection_with_result(None)
+        inserted_id = await connection.insert_get_id("INSERT INTO users (name) VALUES (?) RETURNING id", ["Joe"])
+        self.assertIsNone(inserted_id)
+
+    def test_grammar_platform_and_processor_classes(self):
+        self.assertIs(PostgresConnection.get_query_grammar(), PostgresGrammar)
+        self.assertIs(PostgresConnection.get_default_platform(), PostgresPlatform)
+        self.assertIs(PostgresConnection.get_post_processor(), PostgresPostProcessor)

fastapi_startkit/tests/masoniteorm/sqlite/builder/test_sqlite_builder_pagination.py

@@ -38,3 +38,11 @@ async def test_simple_paginate(self):
             self.assertIsInstance(user, User)
 
         self.assertIsInstance(paginator.to_json(), str)
+
+    async def test_concurrent_paginate_calls(self):
+        import asyncio
+
+        paginators = await asyncio.gather(*(User.query().paginate(1) for _ in range(8)))
+
+        self.assertTrue(all(paginator.total for paginator in paginators))
+        self.assertTrue(all(paginator.count == 1 for paginator in paginators))