Implement Session Management and Transactional Decorator for CRUD Operations

- Introduced `SessionContextHolder` to manage SQLAlchemy sessions using context variables.
- Added `Transactional` decorator to ensure session commits and rollbacks during CRUD operations.
- Refactored `CrudRepository` methods to utilize the new session management, enhancing code clarity and reducing session handling boilerplate.
- Updated tests to reflect changes in the repository methods, ensuring proper functionality without session parameters.

Author: William-W-Chen

py_spring_model/__init__.py

@@ -1,4 +1,5 @@
 from py_spring_model.core.model import PySpringModel, Field
+from py_spring_model.core.session_context_holder import SessionContextHolder
 from py_spring_model.py_spring_model_provider import provide_py_spring_model
 from py_spring_model.repository.crud_repository import CrudRepository
 from py_spring_model.repository.repository_base import RepositoryBase

py_spring_model/core/session_context_holder.py

@@ -0,0 +1,50 @@
+from contextvars import ContextVar
+from functools import wraps
+from typing import ClassVar, Optional
+from sqlmodel import Session
+
+from py_spring_model.core.model import PySpringModel
+
+def Transactional(func):
+    """
+    A decorator that wraps a function and commits the session if the function is successful.
+    If the function raises an exception, the session is rolled back.
+    The session is then closed.
+    """
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        session = SessionContextHolder.get_or_create_session()
+        try:
+            result = func(*args, **kwargs)
+            session.commit()
+            return result
+        except Exception:
+            session.rollback()
+            raise
+        finally:
+            SessionContextHolder.clear_session()
+    return wrapper
+
+class SessionContextHolder:
+    """
+    A context holder for the session.
+    This is used to store the session in a context variable so that it can be accessed by the query service.
+    This is useful for the query service to be able to access the session without having to pass it in as an argument.
+    This is also useful for the query service to be able to access the session without having to pass it in as an argument.
+    """
+    _session: ClassVar[ContextVar[Optional[Session]]] = ContextVar("session", default=None)
+    @classmethod
+    def get_or_create_session(cls) -> Session:
+        optional_session = cls._session.get()
+        if optional_session is None:
+            session = PySpringModel.create_session()
+            cls._session.set(session)
+            return session
+        return optional_session
+    @classmethod
+    def clear_session(cls):
+        optional_session = cls._session.get()
+        if optional_session is None:
+            return
+        optional_session.close()
+        cls._session.set(None)

py_spring_model/repository/crud_repository.py

@@ -16,6 +16,7 @@
 from sqlmodel.sql.expression import Select, SelectOfScalar
 
 from py_spring_model.core.model import PySpringModel
+from py_spring_model.core.session_context_holder import SessionContextHolder, Transactional
 from py_spring_model.repository.repository_base import RepositoryBase
 
 T = TypeVar("T", bound=PySpringModel)
@@ -61,120 +62,124 @@ def _get_model_id_type_with_class(cls) -> tuple[Type[ID], Type[T]]:
     def _find_by_statement(
         self,
         statement: Union[Select, SelectOfScalar],
-        session: Optional[Session] = None,
-    ) -> tuple[Session, Optional[T]]:
-        session = session or self._create_session()
+    ) -> Optional[T]:
+        session = SessionContextHolder.get_or_create_session()
+
+        return session.exec(statement).first()
 
-        return session, session.exec(statement).first()
 
     def _find_by_query(
         self,
         query_by: dict[str, Any],
-        session: Optional[Session] = None,
-    ) -> tuple[Session, Optional[T]]:
-        session = session or self._create_session()
+    ) -> Optional[T]:
+        session = SessionContextHolder.get_or_create_session()
         statement = select(self.model_class).filter_by(**query_by)
-        return session, session.exec(statement).first()
+        return session.exec(statement).first()
 
     def _find_all_by_query(
         self,
         query_by: dict[str, Any],
-        session: Optional[Session] = None,
     ) -> tuple[Session, list[T]]:
-        session = session or self._create_session()
+        session = SessionContextHolder.get_or_create_session()
         statement = select(self.model_class).filter_by(**query_by)
         return session, list(session.exec(statement).fetchall())
 
     def _find_all_by_statement(
         self,
         statement: Union[Select, SelectOfScalar],
-        session: Optional[Session] = None,
-    ) -> tuple[Session, list[T]]:
-        session = session or self._create_session()
-        return session, list(session.exec(statement).fetchall())
+    ) -> list[T]:
+        session = SessionContextHolder.get_or_create_session()
+        return list(session.exec(statement).fetchall())
 
     def find_by_id(self, id: ID) -> Optional[T]:
-        with self.create_managed_session() as session:
-            statement = select(self.model_class).where(self.model_class.id == id)  # type: ignore
-            optional_entity = session.exec(statement).first()
-            if optional_entity is None:
-                return
-
-            return optional_entity.clone()  # type: ignore
+        session = SessionContextHolder.get_or_create_session()
+        statement = select(self.model_class).where(self.model_class.id == id)  # type: ignore
+        optional_entity = session.exec(statement).first()
+        if optional_entity is None:
+            return
 
+        return optional_entity
     def find_all_by_ids(self, ids: list[ID]) -> list[T]:
-        with self.create_managed_session() as session:
-            statement = select(self.model_class).where(self.model_class.id.in_(ids))  # type: ignore
-            return [entity.clone() for entity in session.exec(statement).all()]  # type: ignore
+        session = SessionContextHolder.get_or_create_session()
+        statement = select(self.model_class).where(self.model_class.id.in_(ids))  # type: ignore
+        return [entity for entity in session.exec(statement).all()]  # type: ignore
 
     def find_all(self) -> list[T]:
-        with self.create_managed_session() as session:
-            statement = select(self.model_class)  # type: ignore
-            return [entity.clone() for entity in session.exec(statement).all()]  # type: ignore
+        session = SessionContextHolder.get_or_create_session()
+        statement = select(self.model_class)  # type: ignore
+        return [entity for entity in session.exec(statement).all()]  # type: ignore
 
+    @Transactional
     def save(self, entity: T) -> T:
-        with self.create_managed_session() as session:
-            session.add(entity)
-        return entity.clone()  # type: ignore
+        session = SessionContextHolder.get_or_create_session()
+        session.add(entity)
+        return entity
 
+    @Transactional
     def save_all(
         self,
         entities: Iterable[T],
     ) -> bool:
-        with self.create_managed_session() as session:
-            session.add_all(entities)
+        session = SessionContextHolder.get_or_create_session()
+        session.add_all(entities)
         return True
 
+    @Transactional
     def delete(self, entity: T) -> bool:
-        with self.create_managed_session() as session:
-            _, optional_intance = self._find_by_query(entity.model_dump(), session)
-            if optional_intance is None:
-                return False
-            session.delete(optional_intance)
+        session = SessionContextHolder.get_or_create_session()
+        optional_intance = self._find_by_query(entity.model_dump())
+        if optional_intance is None:
+            return False
+        session.delete(optional_intance)
         return True
 
+    @Transactional
     def delete_all(self, entities: Iterable[T]) -> bool:
-        with self.create_managed_session() as session:
-            ids = [entity.id for entity in entities] # type: ignore
-            
-            statement = select(self.model_class).where(self.model_class.id.in_(ids))  # type: ignore
-            _, deleted_entities = self._find_all_by_statement(statement, session)
-            if deleted_entities is None:
-                return False
-            
-            for entity in deleted_entities:
-                session.delete(entity)
+        session = SessionContextHolder.get_or_create_session()
+        ids = [entity.id for entity in entities] # type: ignore
+        
+        statement = select(self.model_class).where(self.model_class.id.in_(ids))  # type: ignore
+        deleted_entities = self._find_all_by_statement(statement)
+        if deleted_entities is None:
+            return False
+        
+        for entity in deleted_entities:
+            session.delete(entity)
 
         return True
 
+
+    @Transactional
     def delete_by_id(self, _id: ID) -> bool:
-        with self.create_managed_session() as session:
-            _, entity = self._find_by_query({"id": _id}, session)
-            if entity is None:
-                return False
-            session.delete(entity)
+        session = SessionContextHolder.get_or_create_session()
+        entity = self._find_by_query({"id": _id})
+        if entity is None:
+            return False
+        session.delete(entity)
         return True
 
+    @Transactional
     def delete_all_by_ids(self, ids: list[ID]) -> bool:
-        with self.create_managed_session() as session:
-            statement = select(self.model_class).where(self.model_class.id.in_(ids))  # type: ignore
-            _, deleted_entities = self._find_all_by_statement(statement, session)
-            if deleted_entities is None:
-                return False
-            for entity in deleted_entities:
-                session.delete(entity)
-            return True
+        session = SessionContextHolder.get_or_create_session()
+        statement = select(self.model_class).where(self.model_class.id.in_(ids))  # type: ignore
+        deleted_entities = self._find_all_by_statement(statement)
+        if deleted_entities is None:
+            return False
+        for entity in deleted_entities:
+            session.delete(entity)
+        return True
 
+    @Transactional
     def upsert(self, entity: T, query_by: dict[str, Any]) -> T:
-        with self.create_managed_session() as session:
-            statement = select(self.model_class).filter_by(**query_by)  # type: ignore
-            _, existing_entity = self._find_by_statement(statement, session)
-            if existing_entity is not None:
-                # If the entity exists, update its attributes
-                for key, value in entity.model_dump().items():
-                    setattr(existing_entity, key, value)
-                session.add(existing_entity)
-            else:
-                # If the entity does not exist, insert it
-                session.add(entity)
-            return entity
+        session = SessionContextHolder.get_or_create_session()
+        statement = select(self.model_class).filter_by(**query_by)  # type: ignore
+        existing_entity = self._find_by_statement(statement)
+        if existing_entity is not None:
+            # If the entity exists, update its attributes
+            for key, value in entity.model_dump().items():
+                setattr(existing_entity, key, value)
+            session.add(existing_entity)
+        else:
+            # If the entity does not exist, insert it
+            session.add(entity)
+        return entity

tests/test_crud_repository.py

@@ -57,12 +57,12 @@ def test_find_all(self, user_repository: UserRepository):
 
     def test_find_by_query(self, user_repository: UserRepository):
         self.create_test_user(user_repository)
-        _, user = user_repository._find_by_query({"name": "John Doe"})
+        user = user_repository._find_by_query({"name": "John Doe"})
         assert user is not None
         assert user.id == 1
         assert user.name == "John Doe"
 
-        _, email_user = user_repository._find_by_query({"email": "john@example.com"})
+        email_user = user_repository._find_by_query({"email": "john@example.com"})
         assert email_user is not None
         assert email_user.id == 1
         assert email_user.email == "john@example.com"