PYTHON-4542 Improved sessions API

- Via context variable.

Author: aclark4life

pymongo/asynchronous/client_session.py

@@ -139,6 +139,7 @@
 import time
 import uuid
 from collections.abc import Mapping as _Mapping
+from contextvars import ContextVar
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -204,6 +205,7 @@ def __init__(
         causal_consistency: Optional[bool] = None,
         default_transaction_options: Optional[TransactionOptions] = None,
         snapshot: Optional[bool] = False,
+        bind: Optional[bool] = False,
     ) -> None:
         if snapshot:
             if causal_consistency:
@@ -222,6 +224,7 @@ def __init__(
                 )
         self._default_transaction_options = default_transaction_options
         self._snapshot = snapshot
+        self._bind = bind
 
     @property
     def causal_consistency(self) -> bool:
@@ -1065,6 +1068,9 @@ def __copy__(self) -> NoReturn:
         raise TypeError("A AsyncClientSession cannot be copied, create a new session instead")
 
 
+SESSION: ContextVar[Optional[AsyncClientSession]] = ContextVar("SESSION", default=None)
+
+
 class _EmptyServerSession:
     __slots__ = "dirty", "started_retryable_write"
 

pymongo/asynchronous/cursor.py

@@ -63,7 +63,7 @@
     from _typeshed import SupportsItems
 
     from bson.codec_options import CodecOptions
-    from pymongo.asynchronous.client_session import AsyncClientSession
+    from pymongo.asynchronous.client_session import SESSION, AsyncClientSession
     from pymongo.asynchronous.collection import AsyncCollection
     from pymongo.asynchronous.pool import AsyncConnection
     from pymongo.read_preferences import _ServerMode
@@ -136,9 +136,14 @@ def __init__(
         self._killed = False
         self._session: Optional[AsyncClientSession]
 
+        _SESSION = SESSION.get()
+
         if session:
             self._session = session
             self._explicit_session = True
+        elif _SESSION:
+            self._session = _SESSION
+            self._explicit_session = True
         else:
             self._session = None
             self._explicit_session = False

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 _EmptyServerSession
+from pymongo.asynchronous.client_session import SESSION, _EmptyServerSession
 from pymongo.asynchronous.command_cursor import AsyncCommandCursor
 from pymongo.asynchronous.settings import TopologySettings
 from pymongo.asynchronous.topology import Topology, _ErrorContext
@@ -1355,13 +1355,18 @@ def _close_cursor_soon(
     def _start_session(self, implicit: bool, **kwargs: Any) -> AsyncClientSession:
         server_session = _EmptyServerSession()
         opts = client_session.SessionOptions(**kwargs)
-        return client_session.AsyncClientSession(self, server_session, opts, implicit)
+        bind = opts._bind
+        session = client_session.AsyncClientSession(self, server_session, opts, implicit)
+        if bind:
+            SESSION.set(session)
+        return session
 
     def start_session(
         self,
         causal_consistency: Optional[bool] = None,
         default_transaction_options: Optional[client_session.TransactionOptions] = None,
         snapshot: Optional[bool] = False,
+        bind: Optional[bool] = False,
     ) -> client_session.AsyncClientSession:
         """Start a logical session.
 
@@ -1384,6 +1389,7 @@ def start_session(
             causal_consistency=causal_consistency,
             default_transaction_options=default_transaction_options,
             snapshot=snapshot,
+            bind=bind,
         )
 
     def _ensure_session(

pymongo/synchronous/client_session.py

@@ -139,6 +139,7 @@
 import time
 import uuid
 from collections.abc import Mapping as _Mapping
+from contextvars import ContextVar
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -203,6 +204,7 @@ def __init__(
         causal_consistency: Optional[bool] = None,
         default_transaction_options: Optional[TransactionOptions] = None,
         snapshot: Optional[bool] = False,
+        bind: Optional[bool] = False,
     ) -> None:
         if snapshot:
             if causal_consistency:
@@ -221,6 +223,7 @@ def __init__(
                 )
         self._default_transaction_options = default_transaction_options
         self._snapshot = snapshot
+        self._bind = bind
 
     @property
     def causal_consistency(self) -> bool:
@@ -1060,6 +1063,9 @@ def __copy__(self) -> NoReturn:
         raise TypeError("A ClientSession cannot be copied, create a new session instead")
 
 
+SESSION: ContextVar[Optional[ClientSession]] = ContextVar("SESSION", default=None)
+
+
 class _EmptyServerSession:
     __slots__ = "dirty", "started_retryable_write"
 

pymongo/synchronous/cursor.py

@@ -64,7 +64,7 @@
 
     from bson.codec_options import CodecOptions
     from pymongo.read_preferences import _ServerMode
-    from pymongo.synchronous.client_session import ClientSession
+    from pymongo.synchronous.client_session import SESSION, ClientSession
     from pymongo.synchronous.collection import Collection
     from pymongo.synchronous.pool import Connection
 
@@ -136,9 +136,14 @@ def __init__(
         self._killed = False
         self._session: Optional[ClientSession]
 
+        _SESSION = SESSION.get()
+
         if session:
             self._session = session
             self._explicit_session = True
+        elif _SESSION:
+            self._session = _SESSION
+            self._explicit_session = True
         else:
             self._session = None
             self._explicit_session = False

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 _EmptyServerSession
+from pymongo.synchronous.client_session import SESSION, _EmptyServerSession
 from pymongo.synchronous.command_cursor import CommandCursor
 from pymongo.synchronous.settings import TopologySettings
 from pymongo.synchronous.topology import Topology, _ErrorContext
@@ -1353,13 +1353,18 @@ def _close_cursor_soon(
     def _start_session(self, implicit: bool, **kwargs: Any) -> ClientSession:
         server_session = _EmptyServerSession()
         opts = client_session.SessionOptions(**kwargs)
-        return client_session.ClientSession(self, server_session, opts, implicit)
+        bind = opts._bind
+        session = client_session.ClientSession(self, server_session, opts, implicit)
+        if bind:
+            SESSION.set(session)
+        return session
 
     def start_session(
         self,
         causal_consistency: Optional[bool] = None,
         default_transaction_options: Optional[client_session.TransactionOptions] = None,
         snapshot: Optional[bool] = False,
+        bind: Optional[bool] = False,
     ) -> client_session.ClientSession:
         """Start a logical session.
 
@@ -1382,6 +1387,7 @@ def start_session(
             causal_consistency=causal_consistency,
             default_transaction_options=default_transaction_options,
             snapshot=snapshot,
+            bind=bind,
         )
 
     def _ensure_session(self, session: Optional[ClientSession] = None) -> Optional[ClientSession]:

test/asynchronous/test_session.py

@@ -380,6 +380,13 @@ async def test_cursor_clone(self):
             clone = cursor.clone()
             self.assertTrue(clone.session is s)
 
+        # Explicit session via context variable.
+        async with self.client.start_session(bind=True) as s:
+            cursor = coll.find()
+            self.assertTrue(cursor.session is s)
+            clone = cursor.clone()
+            self.assertTrue(clone.session is s)
+
         # No explicit session.
         cursor = coll.find(batch_size=2)
         await anext(cursor)

test/test_session.py

@@ -380,6 +380,13 @@ def test_cursor_clone(self):
             clone = cursor.clone()
             self.assertTrue(clone.session is s)
 
+        # Explicit session via context variable.
+        with self.client.start_session(bind=True) as s:
+            cursor = coll.find()
+            self.assertTrue(cursor.session is s)
+            clone = cursor.clone()
+            self.assertTrue(clone.session is s)
+
         # No explicit session.
         cursor = coll.find(batch_size=2)
         next(cursor)