drop async session scope. Allow user to manage session connection scope and also manage closure when app has a response

Author: eadwinCode

ellar_sql/module.py

@@ -27,33 +27,32 @@ def _raise_exception():
     return _raise_exception
 
 
-async def _session_cleanup(
-    db_service: EllarSQLService, session: t.Union[Session, AsyncSession]
-) -> None:
-    res = session.close()
-    if isinstance(res, t.Coroutine):
-        await res
-
-    res = db_service.session_factory.remove()
-    if isinstance(res, t.Coroutine):
-        await res
-
-
 @as_middleware
 async def session_middleware(
     context: IHostContext, call_next: t.Callable[..., t.Coroutine]
 ):
     connection = context.switch_to_http_connection().get_client()
-
     db_service = context.get_service_provider().get(EllarSQLService)
-    session = db_service.session_factory()
 
+    # Create a NEW session for this request
+    session = db_service.session_factory()
     connection.state.session = session
 
     try:
         await call_next()
+    except Exception as ex:
+        # Only rollback if session is still active
+        if session.is_active and session.in_transaction():
+            res = session.rollback()
+            if isinstance(res, t.Coroutine):
+                await res
+        raise ex
     finally:
-        await _session_cleanup(db_service, session)
+        # Always clean up
+        if session.is_active:
+            res = session.close()
+            if isinstance(res, t.Coroutine):
+                await res
 
 
 @Module(

ellar_sql/pagination/base.py

@@ -286,7 +286,7 @@ async def _close_session(self) -> None:
     def _get_session(self) -> t.Union[sa_orm.Session, AsyncSession, t.Any]:
         self._created_session = True
         service = current_injector.get(EllarSQLService)
-        return service.get_scoped_session()()
+        return service.session_factory_maker()()
 
     def _query_items(self) -> t.List[t.Any]:
         if self._is_async:

ellar_sql/query/utils.py

@@ -19,7 +19,7 @@ async def get_or_404(
 ) -> _O:
     """ """
     db_service = current_injector.get(EllarSQLService)
-    session = db_service.get_scoped_session()()
+    session = db_service.session_factory_maker()()
 
     value = session.get(entity, ident, **kwargs)
 
@@ -39,7 +39,7 @@ async def get_or_none(
 ) -> t.Optional[_O]:
     """ """
     db_service = current_injector.get(EllarSQLService)
-    session = db_service.get_scoped_session()()
+    session = db_service.session_factory_maker()()
 
     value = session.get(entity, ident, **kwargs)
 

ellar_sql/services/base.py

@@ -1,6 +1,5 @@
 import os
 import typing as t
-from threading import get_ident
 from weakref import WeakKeyDictionary
 
 import sqlalchemy as sa
@@ -14,7 +13,6 @@
 )
 from sqlalchemy.ext.asyncio import (
     AsyncSession,
-    async_scoped_session,
     async_sessionmaker,
 )
 
@@ -32,6 +30,10 @@
 
 
 class EllarSQLService:
+    session_factory: t.Union[
+        sa_orm.sessionmaker[sa_orm.Session], async_sessionmaker[AsyncSession]
+    ]
+
     def __init__(
         self,
         databases: t.Union[str, t.Dict[str, t.Any]],
@@ -61,7 +63,7 @@ def __init__(
         self._has_async_engine_driver: bool = False
 
         self._setup(databases, models=models, echo=echo)
-        self.session_factory = self.get_scoped_session()
+        self.session_factory = self.session_factory_maker()
 
     @property
     def has_async_engine_driver(self) -> bool:
@@ -177,24 +179,16 @@ def reflect(self, *databases: str) -> None:
                 continue
             metadata_engine.reflect()
 
-    def get_scoped_session(
+    def session_factory_maker(
         self,
         **extra_options: t.Any,
-    ) -> t.Union[
-        sa_orm.scoped_session[sa_orm.Session],
-        async_scoped_session[t.Union[AsyncSession, t.Any]],
-    ]:
+    ) -> t.Union[sa_orm.sessionmaker[sa_orm.Session], async_sessionmaker[AsyncSession]]:
         options = self._session_options.copy()
         options.update(extra_options)
 
-        scope = options.pop("scopefunc", get_ident)
-
-        factory = self._make_session_factory(options)
-
-        if self.has_async_engine_driver:
-            return async_scoped_session(factory, scope)  # type:ignore[arg-type]
+        scope = options.pop("scopefunc", None)  # noqa: F841
 
-        return sa_orm.scoped_session(factory, scope)  # type:ignore[arg-type]
+        return self._make_session_factory(options)
 
     def _make_session_factory(
         self, options: t.Dict[str, t.Any]

tests/test_model_export.py

@@ -43,7 +43,7 @@ def test_model_export_without_filter(self, db_service, ignore_base):
             "id": 1,
             "name": "Ellar",
         }
-        db_service.session_factory.close()
+        # db_service.session_factory.close()
 
     def test_model_exclude_none(self, db_service, ignore_base):
         user_factory = get_model_factory(db_service)
@@ -59,7 +59,7 @@ def test_model_exclude_none(self, db_service, ignore_base):
             "id": 1,
             "name": "Ellar",
         }
-        db_service.session_factory.close()
+        # db_service.session_factory.close()
 
     def test_model_export_include(self, db_service, ignore_base):
         user_factory = get_model_factory(db_service)
@@ -73,7 +73,7 @@ def test_model_export_include(self, db_service, ignore_base):
             "id",
             "name",
         }
-        db_service.session_factory.close()
+        # db_service.session_factory.close()
 
     def test_model_export_exclude(self, db_service, ignore_base):
         user_factory = get_model_factory(db_service)
@@ -83,7 +83,7 @@ def test_model_export_exclude(self, db_service, ignore_base):
         user = user_factory()
 
         assert user.dict(exclude={"email", "name"}).keys() == {"address", "city", "id"}
-        db_service.session_factory.close()
+        # db_service.session_factory.close()
 
 
 @pytest.mark.asyncio
@@ -105,7 +105,6 @@ async def test_model_export_without_filter_async(
             "id": 1,
             "name": "Ellar",
         }
-        await db_service_async.session_factory.close()
 
     async def test_model_exclude_none_async(self, db_service_async, ignore_base):
         user_factory = get_model_factory(db_service_async)
@@ -121,7 +120,6 @@ async def test_model_exclude_none_async(self, db_service_async, ignore_base):
             "id": 1,
             "name": "Ellar",
         }
-        await db_service_async.session_factory.close()
 
     async def test_model_export_include_async(self, db_service_async, ignore_base):
         user_factory = get_model_factory(db_service_async)
@@ -135,7 +133,6 @@ async def test_model_export_include_async(self, db_service_async, ignore_base):
             "id",
             "name",
         }
-        await db_service_async.session_factory.close()
 
     async def test_model_export_exclude_async(self, db_service_async, ignore_base):
         user_factory = get_model_factory(db_service_async)
@@ -145,4 +142,3 @@ async def test_model_export_exclude_async(self, db_service_async, ignore_base):
         user = user_factory()
 
         assert user.dict(exclude={"email", "name"}).keys() == {"address", "city", "id"}
-        await db_service_async.session_factory.close()