diff --git a/polyfactory/factories/sqlalchemy_factory.py b/polyfactory/factories/sqlalchemy_factory.py index f597f5a5..22bb06f9 100644 --- a/polyfactory/factories/sqlalchemy_factory.py +++ b/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,14 +84,22 @@ 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]): @@ -91,10 +107,12 @@ 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() diff --git a/tests/sqlalchemy_factory/test_sqlalchemy_factory_common.py b/tests/sqlalchemy_factory/test_sqlalchemy_factory_common.py index 36b39fc6..9b1d6641 100644 --- a/tests/sqlalchemy_factory/test_sqlalchemy_factory_common.py +++ b/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