Add end_session parameter to ClientSession.bind()

Author: NoahStapp

pymongo/asynchronous/client_session.py

@@ -188,9 +188,10 @@
 class _AsyncBoundSessionContext:
     """Context manager returned by AsyncClientSession.bind() that manages bound state."""
 
-    def __init__(self, session: AsyncClientSession) -> None:
+    def __init__(self, session: AsyncClientSession, end_session: bool) -> None:
         self._session = session
         self._session_token: Optional[Token[AsyncClientSession]] = None
+        self._end_session = end_session
 
     async def __aenter__(self) -> AsyncClientSession:
         self._session_token = _SESSION.set(self._session)  # type: ignore[assignment]
@@ -200,6 +201,8 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         if self._session_token:
             _SESSION.reset(self._session_token)  # type: ignore[arg-type]
             self._session_token = None
+        if self._end_session:
+            await self._session.end_session()
 
 
 class SessionOptions:
@@ -567,12 +570,14 @@ def _check_ended(self) -> None:
         if self._server_session is None:
             raise InvalidOperation("Cannot use ended session")
 
-    def bind(self) -> _AsyncBoundSessionContext:
+    def bind(self, end_session: bool = False) -> _AsyncBoundSessionContext:
         """Bind this session so it is implicitly passed to all database operations within the returned context.
 
+        :param end_session: Whether to end the session on exiting the returned context. Defaults to False.
+
         .. versionadded:: 4.17
         """
-        return _AsyncBoundSessionContext(self)
+        return _AsyncBoundSessionContext(self, end_session)
 
     async def __aenter__(self) -> AsyncClientSession:
         return self

pymongo/synchronous/client_session.py

@@ -187,9 +187,10 @@
 class BoundSessionContext:
     """Context manager returned by ClientSession.bind() that manages bound state."""
 
-    def __init__(self, session: ClientSession) -> None:
+    def __init__(self, session: ClientSession, end_session: bool) -> None:
         self._session = session
         self._session_token: Optional[Token[ClientSession]] = None
+        self._end_session = end_session
 
     def __enter__(self) -> ClientSession:
         self._session_token = _SESSION.set(self._session)  # type: ignore[assignment]
@@ -199,6 +200,8 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         if self._session_token:
             _SESSION.reset(self._session_token)  # type: ignore[arg-type]
             self._session_token = None
+        if self._end_session:
+            self._session.end_session()
 
 
 class SessionOptions:
@@ -566,12 +569,14 @@ def _check_ended(self) -> None:
         if self._server_session is None:
             raise InvalidOperation("Cannot use ended session")
 
-    def bind(self) -> BoundSessionContext:
+    def bind(self, end_session: bool = False) -> BoundSessionContext:
         """Bind this session so it is implicitly passed to all database operations within the returned context.
 
+        :param end_session: Whether to end the session on exiting the returned context. Defaults to False.
+
         .. versionadded:: 4.17
         """
-        return BoundSessionContext(self)
+        return BoundSessionContext(self, end_session)
 
     def __enter__(self) -> ClientSession:
         return self

test/asynchronous/test_session.py

@@ -922,6 +922,22 @@ async def test_nested_session_binding(self):
             await session1.end_session()
             await session2.end_session()
 
+    async def test_session_binding_end_session(self):
+        coll = self.client.pymongo_test.test
+        await coll.insert_one({"x": 1})
+
+        async with self.client.start_session().bind(end_session=True) as s1:
+            await coll.find_one()
+
+        self.assertTrue(s1.has_ended)
+
+        async with self.client.start_session().bind() as s2:
+            await coll.find_one()
+
+        self.assertFalse(s2.has_ended)
+
+        await s2.end_session()
+
 
 class TestCausalConsistency(AsyncUnitTest):
     listener: SessionTestListener

test/test_session.py

@@ -922,6 +922,22 @@ def test_nested_session_binding(self):
             session1.end_session()
             session2.end_session()
 
+    def test_session_binding_end_session(self):
+        coll = self.client.pymongo_test.test
+        coll.insert_one({"x": 1})
+
+        with self.client.start_session().bind(end_session=True) as s1:
+            coll.find_one()
+
+        self.assertTrue(s1.has_ended)
+
+        with self.client.start_session().bind() as s2:
+            coll.find_one()
+
+        self.assertFalse(s2.has_ended)
+
+        s2.end_session()
+
 
 class TestCausalConsistency(UnitTest):
     listener: SessionTestListener