feat: Implement explicit session depth tracking in SessionContextHolder

- Added TransactionalDepth enum to define transaction levels.
- Enhanced SessionContextHolder with methods to manage session depth.
- Updated @transactional decorator to utilize session depth for commit/rollback operations.
- Introduced tests for session depth tracking and behavior in nested transactions.

Author: William-W-Chen

py_spring_model/core/session_context_holder.py

@@ -1,11 +1,16 @@
 from contextvars import ContextVar
+from enum import IntEnum
 from functools import wraps
 from typing import Any, Callable, ClassVar, Optional
 
 from py_spring_model.core.py_spring_session import PySpringSession
 
 from py_spring_model.core.model import PySpringModel
 
+class TransactionalDepth(IntEnum):
+    OUTERMOST = 1
+    ON_EXIT = 0
+
 def Transactional(func: Callable[..., Any]) -> Callable[..., Any]:
     """
     Decorator for managing database transactions in a nested-safe manner.
@@ -45,30 +50,34 @@ def update_account():
     """
     @wraps(func)
     def wrapper(*args, **kwargs):
-        is_outermost_transaction = not SessionContextHolder.has_session()
+        # Increment session depth and get session
+        session_depth = SessionContextHolder.enter_session()
         session = SessionContextHolder.get_or_create_session()
         try:
             result = func(*args, **kwargs)
-            if is_outermost_transaction:
+            # Only commit at the outermost level (session_depth == 1)
+            if session_depth == TransactionalDepth.OUTERMOST.value:
                 session.commit()
             return result
         except Exception as error:
-            if is_outermost_transaction:
+            # Only rollback at the outermost level (session_depth == 1)
+            if session_depth == TransactionalDepth.OUTERMOST.value:
                 session.rollback()
             raise error
         finally:
-            if is_outermost_transaction:
-                SessionContextHolder.clear_session()
+            # Decrement depth and clean up session if needed
+            SessionContextHolder.exit_session()
     return wrapper
 
 class SessionContextHolder:
     """
-    A context holder for the session.
+    A context holder for the session with explicit depth tracking.
     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.
+    The depth counter ensures that only the outermost transaction manages commit/rollback operations.
     """
     _session: ClassVar[ContextVar[Optional[PySpringSession]]] = ContextVar("session", default=None)
+    _session_depth: ClassVar[ContextVar[int]] = ContextVar("session_depth", default=0)
+    
     @classmethod
     def get_or_create_session(cls) -> PySpringSession:
         optional_session = cls._session.get()
@@ -82,9 +91,44 @@ def get_or_create_session(cls) -> PySpringSession:
     def has_session(cls) -> bool:
         return cls._session.get() is not None
 
+    @classmethod
+    def get_session_depth(cls) -> int:
+        """Get the current session depth."""
+        return cls._session_depth.get()
+    
+    @classmethod
+    def enter_session(cls) -> int:
+        """
+        Enter a new session context and increment the depth counter.
+        Returns the new depth level.
+        """
+        current_depth = cls._session_depth.get()
+        new_depth = current_depth + 1
+        cls._session_depth.set(new_depth)
+        return new_depth
+    
+    @classmethod
+    def exit_session(cls) -> int:
+        """
+        Exit the current session context and decrement the depth counter.
+        If depth reaches 0, clear the session.
+        Returns the new depth level.
+        """
+        current_depth = cls._session_depth.get()
+        new_depth = max(0, current_depth - 1)  # Prevent negative depth
+        cls._session_depth.set(new_depth)
+        
+        # Clear session only when depth reaches 0 (outermost level)
+        if new_depth == 0:
+            cls.clear_session()
+        
+        return new_depth
+    
     @classmethod
     def clear_session(cls):
+        """Clear the session and reset depth to 0."""
         session = cls._session.get()
         if session is not None:
             session.close()
-        cls._session.set(None)
+        cls._session.set(None)
+        cls._session_depth.set(TransactionalDepth.ON_EXIT.value)

tests/test_session_depth.py

@@ -0,0 +1,112 @@
+import pytest
+from py_spring_model.core.session_context_holder import SessionContextHolder, Transactional
+from py_spring_model.core.model import PySpringModel
+
+
+class TestSessionDepth:
+    """Test the explicit session depth tracking functionality"""
+
+    def setup_method(self):
+        """Clean up any existing sessions before each test"""
+        SessionContextHolder.clear_session()
+
+    def teardown_method(self):
+        """Clean up after each test"""
+        SessionContextHolder.clear_session()
+
+    def test_session_depth_starts_at_zero(self):
+        """Test that session depth starts at 0"""
+        assert SessionContextHolder.get_session_depth() == 0
+
+    def test_session_depth_increments_and_decrements(self):
+        """Test that session depth properly increments and decrements"""
+        # Initially 0
+        assert SessionContextHolder.get_session_depth() == 0
+        
+        # Enter first level
+        depth1 = SessionContextHolder.enter_session()
+        assert depth1 == 1
+        assert SessionContextHolder.get_session_depth() == 1
+        
+        # Enter second level
+        depth2 = SessionContextHolder.enter_session()
+        assert depth2 == 2
+        assert SessionContextHolder.get_session_depth() == 2
+        
+        # Exit second level
+        depth_after_exit1 = SessionContextHolder.exit_session()
+        assert depth_after_exit1 == 1
+        assert SessionContextHolder.get_session_depth() == 1
+        
+        # Exit first level
+        depth_after_exit2 = SessionContextHolder.exit_session()
+        assert depth_after_exit2 == 0
+        assert SessionContextHolder.get_session_depth() == 0
+
+    def test_session_cleared_only_at_outermost_level(self):
+        """Test that session is only cleared when depth reaches 0"""
+        # Enter first level and create session
+        SessionContextHolder.enter_session()
+        session = SessionContextHolder.get_or_create_session()
+        assert SessionContextHolder.has_session()
+        
+        # Enter second level - session should still exist
+        SessionContextHolder.enter_session()
+        assert SessionContextHolder.has_session()
+        assert SessionContextHolder.get_session_depth() == 2
+        
+        # Exit second level - session should still exist
+        SessionContextHolder.exit_session()
+        assert SessionContextHolder.has_session()
+        assert SessionContextHolder.get_session_depth() == 1
+        
+        # Exit first level - session should be cleared
+        SessionContextHolder.exit_session()
+        assert not SessionContextHolder.has_session()
+        assert SessionContextHolder.get_session_depth() == 0
+
+    def test_clear_session_resets_depth(self):
+        """Test that clear_session() resets the depth to 0"""
+        SessionContextHolder.enter_session()
+        SessionContextHolder.enter_session()
+        assert SessionContextHolder.get_session_depth() == 2
+        
+        SessionContextHolder.clear_session()
+        assert SessionContextHolder.get_session_depth() == 0
+
+    @pytest.mark.parametrize("nesting_levels", [1, 2, 3, 5])
+    def test_transactional_depth_tracking(self, nesting_levels):
+        """Test that @Transactional properly tracks depth at various nesting levels"""
+        depth_records = []
+        
+        def create_nested_function(level: int):
+            @Transactional
+            def nested_func():
+                current_depth = SessionContextHolder.get_session_depth()
+                depth_records.append(current_depth)
+                if level > 1:
+                    create_nested_function(level - 1)()
+            return nested_func
+        
+        # Create and execute nested function
+        outermost_func = create_nested_function(nesting_levels)
+        outermost_func()
+        
+        # Verify depth progression
+        assert len(depth_records) == nesting_levels
+        for i, recorded_depth in enumerate(depth_records):
+            expected_depth = i + 1
+            assert recorded_depth == expected_depth, f"At level {i+1}, expected depth {expected_depth}, got {recorded_depth}"
+        
+        # Verify session is cleaned up
+        assert SessionContextHolder.get_session_depth() == 0
+        assert not SessionContextHolder.has_session()
+
+    def test_depth_prevents_negative_values(self):
+        """Test that depth counter prevents going below 0"""
+        assert SessionContextHolder.get_session_depth() == 0
+        
+        # Try to exit when already at 0
+        new_depth = SessionContextHolder.exit_session()
+        assert new_depth == 0
+        assert SessionContextHolder.get_session_depth() == 0