Add context manager

Author: aclark4life

pymongo/asynchronous/client_session.py

@@ -139,7 +139,8 @@
 import time
 import uuid
 from collections.abc import Mapping as _Mapping
-from contextvars import ContextVar
+from contextlib import AbstractAsyncContextManager
+from contextvars import ContextVar, Token
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -1071,6 +1072,24 @@ def __copy__(self) -> NoReturn:
 _SESSION: ContextVar[Optional[AsyncClientSession]] = ContextVar("SESSION", default=None)
 
 
+class _BindSession(AbstractAsyncContextManager):
+    def __init__(self, session: AsyncClientSession) -> None:
+        self.session = session
+        self.token: Optional[Token[Optional[AsyncClientSession]]] = None
+
+    async def __aenter__(self) -> None:
+        self.token = _SESSION.set(self.session)
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> Optional[bool]:
+        if self.token is not None:
+            _SESSION.reset(self.token)
+
+
 class _EmptyServerSession:
     __slots__ = "dirty", "started_retryable_write"
 

pymongo/asynchronous/mongo_client.py

@@ -65,7 +65,7 @@
 from pymongo.asynchronous import client_session, database, uri_parser
 from pymongo.asynchronous.change_stream import AsyncChangeStream, AsyncClusterChangeStream
 from pymongo.asynchronous.client_bulk import _AsyncClientBulk
-from pymongo.asynchronous.client_session import _SESSION, _EmptyServerSession
+from pymongo.asynchronous.client_session import _BindSession, _EmptyServerSession
 from pymongo.asynchronous.command_cursor import AsyncCommandCursor
 from pymongo.asynchronous.settings import TopologySettings
 from pymongo.asynchronous.topology import Topology, _ErrorContext
@@ -1358,7 +1358,7 @@ def _start_session(self, implicit: bool, **kwargs: Any) -> AsyncClientSession:
         bind = opts._bind
         session = client_session.AsyncClientSession(self, server_session, opts, implicit)
         if bind:
-            _SESSION.set(session)
+            session = _BindSession(session)
         return session
 
     def start_session(

pymongo/synchronous/client_session.py

@@ -139,7 +139,8 @@
 import time
 import uuid
 from collections.abc import Mapping as _Mapping
-from contextvars import ContextVar
+from contextlib import AbstractContextManager
+from contextvars import ContextVar, Token
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -1066,6 +1067,24 @@ def __copy__(self) -> NoReturn:
 _SESSION: ContextVar[Optional[ClientSession]] = ContextVar("SESSION", default=None)
 
 
+class _BindSession(AbstractContextManager):
+    def __init__(self, session: ClientSession) -> None:
+        self.session = session
+        self.token: Optional[Token[Optional[ClientSession]]] = None
+
+    def __enter__(self) -> None:
+        self.token = _SESSION.set(self.session)
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> Optional[bool]:
+        if self.token is not None:
+            _SESSION.reset(self.token)
+
+
 class _EmptyServerSession:
     __slots__ = "dirty", "started_retryable_write"
 

pymongo/synchronous/mongo_client.py

@@ -107,7 +107,7 @@
 from pymongo.synchronous import client_session, database, uri_parser
 from pymongo.synchronous.change_stream import ChangeStream, ClusterChangeStream
 from pymongo.synchronous.client_bulk import _ClientBulk
-from pymongo.synchronous.client_session import _SESSION, _EmptyServerSession
+from pymongo.synchronous.client_session import _BindSession, _EmptyServerSession
 from pymongo.synchronous.command_cursor import CommandCursor
 from pymongo.synchronous.settings import TopologySettings
 from pymongo.synchronous.topology import Topology, _ErrorContext
@@ -1356,7 +1356,7 @@ def _start_session(self, implicit: bool, **kwargs: Any) -> ClientSession:
         bind = opts._bind
         session = client_session.ClientSession(self, server_session, opts, implicit)
         if bind:
-            _SESSION.set(session)
+            session = _BindSession(session)
         return session
 
     def start_session(