Set token

Author: aclark4life

pymongo/asynchronous/client_session.py

@@ -548,10 +548,12 @@ def _check_ended(self) -> None:
             raise InvalidOperation("Cannot use ended session")
 
     async def __aenter__(self) -> AsyncClientSession:
-        _SESSION.set(self)
+        self._token = _SESSION.set(self)
         return self
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        if self._token:
+            _SESSION.reset(self._token)
         await self._end_session(lock=True)
 
     @property

pymongo/synchronous/client_session.py

@@ -547,10 +547,12 @@ def _check_ended(self) -> None:
             raise InvalidOperation("Cannot use ended session")
 
     def __enter__(self) -> ClientSession:
-        _SESSION.set(self)
+        self._token = _SESSION.set(self)
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        if self._token:
+            _SESSION.reset(self._token)
         self._end_session(lock=True)
 
     @property