Flag implicit sessions used by retryable operations that also use CommandCursors

Author: NoahStapp

pymongo/asynchronous/client_session.py

@@ -514,6 +514,7 @@ def __init__(
         self._implicit = implicit
         self._transaction = _Transaction(None, client)
         self._attached_to_cursor = False
+        self._leave_alive = False
 
     async def end_session(self) -> None:
         """Finish this session. If a transaction has started, abort it.
@@ -536,7 +537,7 @@ async def _end_session(self, lock: bool) -> None:
 
     def _end_implicit_session(self) -> None:
         # Implicit sessions can't be part of transactions or pinned connections
-        if self._server_session is not None:
+        if not self._leave_alive and self._server_session is not None:
             # print(f"Ending session {self}, implicit: {self._implicit}, attached: {self._attached_to_cursor}")
             self._client._return_server_session(self._server_session)
             self._server_session = None
@@ -606,6 +607,18 @@ def attached_to_cursor(self) -> bool:
     def attached_to_cursor(self, value: bool) -> None:
         self._attached_to_cursor = value
 
+    @property
+    def leave_alive(self) -> bool:
+        """Whether to leave this session alive when it is
+        no longer in use.
+        Typically used for implicit sessions that are used for multiple operations within a single larger operation.
+        """
+        return self._leave_alive
+
+    @leave_alive.setter
+    def leave_alive(self, value: bool) -> None:
+        self._leave_alive = value
+
     def _inherit_option(self, name: str, val: _T) -> _T:
         """Return the inherited TransactionOption value."""
         if val:

pymongo/asynchronous/command_cursor.py

@@ -65,6 +65,7 @@ def __init__(
         max_await_time_ms: Optional[int] = None,
         session: Optional[AsyncClientSession] = None,
         comment: Any = None,
+        leave_session_alive: bool = False,
     ) -> None:
         """Create a new command cursor."""
         self._sock_mgr: Any = None
@@ -83,6 +84,7 @@ def __init__(
             self._session.attached_to_cursor = True
         self._killed = self._id == 0
         self._comment = comment
+        self._leave_session_alive = leave_session_alive
         if self._killed:
             self._end_session()
 
@@ -240,8 +242,9 @@ async def _die_lock(self) -> None:
         self._sock_mgr = None
 
     def _end_session(self) -> None:
-        if self._session:
-            if self._session.implicit:
+        if self._session and self._session.implicit:
+            self._session.attached_to_cursor = False
+            if not self._leave_session_alive:
                 # print(f"Ending session {self}, session: {self._session}")
                 self._session._end_implicit_session()
                 self._session = None

pymongo/asynchronous/database.py

@@ -611,6 +611,7 @@ async def create_collection(
             common.validate_is_mapping("clusteredIndex", clustered_index)
 
         async with self._client._tmp_session(session) as s:
+            s.leave_alive = True
             # Skip this check in a transaction where listCollections is not
             # supported.
             if (
@@ -705,7 +706,6 @@ async def aggregate(
                 AsyncCommandCursor,
                 pipeline,
                 kwargs,
-                session is not None,
                 user_fields={"cursor": {"firstBatch": 1}},
             )
             return await self.client._retryable_read(

pymongo/synchronous/client_session.py

@@ -513,6 +513,7 @@ def __init__(
         self._implicit = implicit
         self._transaction = _Transaction(None, client)
         self._attached_to_cursor = False
+        self._leave_alive = False
 
     def end_session(self) -> None:
         """Finish this session. If a transaction has started, abort it.
@@ -535,7 +536,7 @@ def _end_session(self, lock: bool) -> None:
 
     def _end_implicit_session(self) -> None:
         # Implicit sessions can't be part of transactions or pinned connections
-        if self._server_session is not None:
+        if not self._leave_alive and self._server_session is not None:
             # print(f"Ending session {self}, implicit: {self._implicit}, attached: {self._attached_to_cursor}")
             self._client._return_server_session(self._server_session)
             self._server_session = None
@@ -605,6 +606,18 @@ def attached_to_cursor(self) -> bool:
     def attached_to_cursor(self, value: bool) -> None:
         self._attached_to_cursor = value
 
+    @property
+    def leave_alive(self) -> bool:
+        """Whether to leave this session alive when it is
+        no longer in use.
+        Typically used for implicit sessions that are used for multiple operations within a single larger operation.
+        """
+        return self._leave_alive
+
+    @leave_alive.setter
+    def leave_alive(self, value: bool) -> None:
+        self._leave_alive = value
+
     def _inherit_option(self, name: str, val: _T) -> _T:
         """Return the inherited TransactionOption value."""
         if val:

pymongo/synchronous/command_cursor.py

@@ -65,6 +65,7 @@ def __init__(
         max_await_time_ms: Optional[int] = None,
         session: Optional[ClientSession] = None,
         comment: Any = None,
+        leave_session_alive: bool = False,
     ) -> None:
         """Create a new command cursor."""
         self._sock_mgr: Any = None
@@ -83,6 +84,7 @@ def __init__(
             self._session.attached_to_cursor = True
         self._killed = self._id == 0
         self._comment = comment
+        self._leave_session_alive = leave_session_alive
         if self._killed:
             self._end_session()
 
@@ -240,8 +242,9 @@ def _die_lock(self) -> None:
         self._sock_mgr = None
 
     def _end_session(self) -> None:
-        if self._session:
-            if self._session.implicit:
+        if self._session and self._session.implicit:
+            self._session.attached_to_cursor = False
+            if not self._leave_session_alive:
                 # print(f"Ending session {self}, session: {self._session}")
                 self._session._end_implicit_session()
                 self._session = None

pymongo/synchronous/database.py

@@ -611,6 +611,7 @@ def create_collection(
             common.validate_is_mapping("clusteredIndex", clustered_index)
 
         with self._client._tmp_session(session) as s:
+            s.leave_alive = True
             # Skip this check in a transaction where listCollections is not
             # supported.
             if (
@@ -705,7 +706,6 @@ def aggregate(
                 CommandCursor,
                 pipeline,
                 kwargs,
-                session is not None,
                 user_fields={"cursor": {"firstBatch": 1}},
             )
             return self.client._retryable_read(