feat(sqlalchemy): manage session lifecycle when the session is created implicitly

polyfactory/factories/sqlalchemy_factory.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import enum
-from collections.abc import Collection, Mapping
+from collections.abc import Awaitable, Collection, Mapping
 from dataclasses import is_dataclass
 from datetime import date, datetime
 from typing import (
@@ -31,16 +31,22 @@
     from sqlalchemy.dialects import mssql, mysql, postgresql, sqlite
     from sqlalchemy.exc import NoInspectionAvailable
     from sqlalchemy.ext.associationproxy import AssociationProxy
-    from sqlalchemy.orm import InstanceState, Mapper, RelationshipProperty
+    from sqlalchemy.ext.asyncio import async_scoped_session
+    from sqlalchemy.orm import InstanceState, Mapper, RelationshipProperty, scoped_session, sessionmaker
 except ImportError as e:
     msg = "sqlalchemy is not installed"
     raise MissingDependencyException(msg) from e
 
 if TYPE_CHECKING:
-    from sqlalchemy.ext.asyncio import AsyncSession, async_scoped_session
-    from sqlalchemy.orm import Session, scoped_session
+    from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+    from sqlalchemy.orm import Session
     from sqlalchemy.sql.type_api import TypeEngine
     from typing_extensions import NotRequired, TypeGuard
+else:
+    try:
+        from sqlalchemy.ext.asyncio import async_sessionmaker
+    except ImportError:
+        async_sessionmaker = sessionmaker
 
 
 T = TypeVar("T")
@@ -64,10 +70,12 @@ def __init__(
         self,
         session: Session,
         persistence_method: SQLAlchemyPersistenceMethod = SQLAlchemyPersistenceMethod.COMMIT,
+        cleanup_callback: Callable[[], None] | None = None,
     ) -> None:
         """Sync persistence handler for SQLAFactory."""
         self.session = session
         self.persistence_method = persistence_method
+        self.cleanup_callback = cleanup_callback
 
     def _flush_or_commit(self) -> None:
         if self.persistence_method == SQLAlchemyPersistenceMethod.FLUSH:
@@ -76,25 +84,35 @@ def _flush_or_commit(self) -> None:
             self.session.commit()
 
     def save(self, data: T) -> T:
-        self.session.add(data)
-        self._flush_or_commit()
-        return data
+        try:
+            self.session.add(data)
+            self._flush_or_commit()
+            return data
+        finally:
+            if self.cleanup_callback is not None:
+                self.cleanup_callback()
 
     def save_many(self, data: list[T]) -> list[T]:
-        self.session.add_all(data)
-        self._flush_or_commit()
-        return data
+        try:
+            self.session.add_all(data)
+            self._flush_or_commit()
+            return data
+        finally:
+            if self.cleanup_callback is not None:
+                self.cleanup_callback()
 
 
 class SQLAASyncPersistence(AsyncPersistenceProtocol[T]):
     def __init__(
         self,
         session: AsyncSession,
         persistence_method: SQLAlchemyPersistenceMethod = SQLAlchemyPersistenceMethod.COMMIT,
+        cleanup_callback: Callable[[], Awaitable[None]] | None = None,
     ) -> None:
         """Async persistence handler for SQLAFactory."""
         self.session = session
         self.persistence_method = persistence_method
+        self.cleanup_callback = cleanup_callback
 
     async def _flush_or_commit(self, session: AsyncSession) -> None:
         if self.persistence_method == SQLAlchemyPersistenceMethod.FLUSH:
@@ -103,25 +121,32 @@ async def _flush_or_commit(self, session: AsyncSession) -> None:
             await session.commit()
 
     async def save(self, data: T) -> T:
-        self.session.add(data)
-        await self._flush_or_commit(self.session)
-        await self.session.refresh(data)
-        return data
+        try:
+            self.session.add(data)
+            await self._flush_or_commit(self.session)
+            await self.session.refresh(data)
+            return data
+        finally:
+            if self.cleanup_callback is not None:
+                await self.cleanup_callback()
 
     async def save_many(self, data: list[T]) -> list[T]:
-        self.session.add_all(data)
-        await self._flush_or_commit(self.session)
-        for batch_item in data:
-            await self.session.refresh(batch_item)
-        return data
+        try:
+            self.session.add_all(data)
+            await self._flush_or_commit(self.session)
+            for batch_item in data:
+                await self.session.refresh(batch_item)
+            return data
+        finally:
+            if self.cleanup_callback is not None:
+                await self.cleanup_callback()
 
 
 _T_co = TypeVar("_T_co", covariant=True)
 
 
 class _SessionMaker(Protocol[_T_co]):
-    @staticmethod
-    def __call__() -> _T_co: ...
+    def __call__(self, **kwargs: Any) -> _T_co: ...
 
 
 class SQLAlchemyFactory(BaseFactory[T], Generic[T]):
@@ -414,12 +439,36 @@ def get_model_fields(cls) -> list[FieldMeta]:
     def _get_sync_persistence(cls) -> SyncPersistenceProtocol[T]:
         if cls.__session__ is not None:
             session = cls.__session__() if callable(cls.__session__) else cls.__session__
-            return SQLASyncPersistence(session, persistence_method=cls.__persistence_method__)
+
+            if isinstance(cls.__session__, sessionmaker):
+                cleanup_callback = session.close
+            elif isinstance(cls.__session__, scoped_session):
+                cleanup_callback = cls.__session__.remove
+            else:
+                cleanup_callback = None
+
+            return SQLASyncPersistence(
+                session,
+                persistence_method=cls.__persistence_method__,
+                cleanup_callback=cleanup_callback,
+            )
         return super()._get_sync_persistence()
 
     @classmethod
     def _get_async_persistence(cls) -> AsyncPersistenceProtocol[T]:
         if cls.__async_session__ is not None:
             session = cls.__async_session__() if callable(cls.__async_session__) else cls.__async_session__
-            return SQLAASyncPersistence(session, persistence_method=cls.__persistence_method__)
+
+            if isinstance(cls.__async_session__, async_sessionmaker):
+                cleanup_callback = session.close
+            elif isinstance(cls.__async_session__, async_scoped_session):
+                cleanup_callback = cls.__async_session__.remove
+            else:
+                cleanup_callback = None
+
+            return SQLAASyncPersistence(
+                session,
+                persistence_method=cls.__persistence_method__,
+                cleanup_callback=cleanup_callback,
+            )
         return super()._get_async_persistence()

tests/sqlalchemy_factory/test_sqlalchemy_factory_common.py

@@ -1,3 +1,4 @@
+import asyncio
 import warnings
 from collections.abc import Collection
 from datetime import datetime
@@ -19,9 +20,9 @@
     types,
 )
 from sqlalchemy.engine import Engine
-from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_scoped_session
 from sqlalchemy.ext.hybrid import hybrid_property
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, scoped_session, sessionmaker
 
 try:
     from sqlalchemy.orm.collections import (
@@ -36,6 +37,11 @@
     column_keyed_dict = _collections.column_mapped_collection  # type: ignore[attr-defined]
     keyfunc_mapping = _collections.mapped_collection  # type: ignore[attr-defined]
 
+try:
+    from sqlalchemy.ext.asyncio import async_sessionmaker
+except ImportError:
+    from sqlalchemy.orm import sessionmaker as async_sessionmaker  # type: ignore[assignment]
+
 from sqlalchemy.orm.decl_api import DeclarativeMeta, registry
 
 from polyfactory.exceptions import ConfigurationException, ParameterException
@@ -144,6 +150,51 @@ class ShapeFactory(SQLAlchemyFactory[Shape]):
     assert instance.area == pow(instance.side, 2)
 
 
+def test_sessionmaker_created_sync_session_is_closed(engine: Engine) -> None:
+    class TrackingSession(Session):
+        created_sessions: list["TrackingSession"] = []
+
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            super().__init__(*args, **kwargs)
+            self.close_calls = 0
+            TrackingSession.created_sessions.append(self)
+
+        def close(self) -> None:
+            self.close_calls += 1
+            super().close()
+
+    class ShapeFactory(SQLAlchemyFactory[Shape]):
+        __model__ = Shape
+        __session__ = sessionmaker(bind=engine, class_=TrackingSession)
+
+    ShapeFactory.create_sync()
+    assert len(TrackingSession.created_sessions) == 1
+    assert TrackingSession.created_sessions[0].close_calls == 1
+
+
+def test_scoped_sessionmaker_created_sync_session_is_closed(engine: Engine) -> None:
+    class TrackingSession(Session):
+        created_sessions: list["TrackingSession"] = []
+
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            super().__init__(*args, **kwargs)
+            self.close_calls = 0
+            TrackingSession.created_sessions.append(self)
+
+        def close(self) -> None:
+            self.close_calls += 1
+            super().close()
+
+    class ShapeFactory(SQLAlchemyFactory[Shape]):
+        __model__ = Shape
+        __session__ = scoped_session(sessionmaker(bind=engine, class_=TrackingSession))
+
+    ShapeFactory.create_sync()
+    assert len(TrackingSession.created_sessions) == 1
+    assert TrackingSession.created_sessions[0].close_calls == 1
+    assert not ShapeFactory.__session__.registry.has()
+
+
 async def test_computed_column_async_persistence(engine: Engine, async_engine: AsyncEngine) -> None:
     class ShapeFactory(SQLAlchemyFactory[Shape]):
         __model__ = Shape
@@ -153,6 +204,57 @@ class ShapeFactory(SQLAlchemyFactory[Shape]):
     assert instance.area == pow(instance.side, 2)
 
 
+async def test_async_sessionmaker_created_session_is_closed(async_engine: AsyncEngine) -> None:
+    class TrackingAsyncSession(AsyncSession):
+        created_sessions: list["TrackingAsyncSession"] = []
+
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            super().__init__(*args, **kwargs)
+            self.close_calls = 0
+            TrackingAsyncSession.created_sessions.append(self)
+
+        async def close(self) -> None:
+            self.close_calls += 1
+            await super().close()
+
+    class ShapeFactory(SQLAlchemyFactory[Shape]):
+        __model__ = Shape
+        __async_session__ = async_sessionmaker(bind=async_engine, class_=TrackingAsyncSession)
+
+    await ShapeFactory.create_async()
+    assert len(TrackingAsyncSession.created_sessions) == 1
+    created_session = TrackingAsyncSession.created_sessions[0]
+    assert created_session.close_calls == 1
+
+
+async def test_async_scoped_sessionmaker_created_session_is_closed(async_engine: AsyncEngine) -> None:
+    class TrackingAsyncSession(AsyncSession):
+        created_sessions: list["TrackingAsyncSession"] = []
+
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            super().__init__(*args, **kwargs)
+            self.close_calls = 0
+            TrackingAsyncSession.created_sessions.append(self)
+
+        async def close(self) -> None:
+            self.close_calls += 1
+            await super().close()
+
+    scoped_factory = async_scoped_session(
+        async_sessionmaker(bind=async_engine, class_=TrackingAsyncSession),
+        scopefunc=asyncio.current_task,
+    )
+
+    class ShapeFactory(SQLAlchemyFactory[Shape]):
+        __model__ = Shape
+        __async_session__ = scoped_factory
+
+    await ShapeFactory.create_async()
+    assert len(TrackingAsyncSession.created_sessions) == 1
+    assert TrackingAsyncSession.created_sessions[0].close_calls == 1
+    assert not scoped_factory.registry.has()
+
+
 def test_computed_column_no_persistence() -> None:
     class ShapeFactory(SQLAlchemyFactory[Shape]):
         __model__ = Shape