WIP

Author: NoahStapp

pymongo/asynchronous/aggregation.py

@@ -50,7 +50,6 @@ def __init__(
         cursor_class: type[AsyncCommandCursor[Any]],
         pipeline: _Pipeline,
         options: MutableMapping[str, Any],
-        explicit_session: bool,
         let: Optional[Mapping[str, Any]] = None,
         user_fields: Optional[MutableMapping[str, Any]] = None,
         result_processor: Optional[Callable[[Mapping[str, Any], AsyncConnection], None]] = None,
@@ -92,7 +91,6 @@ def __init__(
             self._options["cursor"]["batchSize"] = self._batch_size
 
         self._cursor_class = cursor_class
-        self._explicit_session = explicit_session
         self._user_fields = user_fields
         self._result_processor = result_processor
 
@@ -197,7 +195,6 @@ async def get_cursor(
             batch_size=self._batch_size or 0,
             max_await_time_ms=self._max_await_time_ms,
             session=session,
-            explicit_session=self._explicit_session,
             comment=self._options.get("comment"),
         )
         await cmd_cursor._maybe_pin_connection(conn)

pymongo/asynchronous/change_stream.py

@@ -236,7 +236,7 @@ def _process_result(self, result: Mapping[str, Any], conn: AsyncConnection) -> N
                     )
 
     async def _run_aggregation_cmd(
-        self, session: Optional[AsyncClientSession], explicit_session: bool
+        self, session: Optional[AsyncClientSession]
     ) -> AsyncCommandCursor:  # type: ignore[type-arg]
         """Run the full aggregation pipeline for this AsyncChangeStream and return
         the corresponding AsyncCommandCursor.
@@ -246,7 +246,6 @@ async def _run_aggregation_cmd(
             AsyncCommandCursor,
             self._aggregation_pipeline(),
             self._command_options(),
-            explicit_session,
             result_processor=self._process_result,
             comment=self._comment,
         )
@@ -258,10 +257,8 @@ async def _run_aggregation_cmd(
         )
 
     async def _create_cursor(self) -> AsyncCommandCursor:  # type: ignore[type-arg]
-        async with self._client._tmp_session(self._session, close=False) as s:
-            return await self._run_aggregation_cmd(
-                session=s, explicit_session=self._session is not None
-            )
+        async with self._client._tmp_session(self._session) as s:
+            return await self._run_aggregation_cmd(session=s)
 
     async def _resume(self) -> None:
         """Reestablish this change stream after a resumable error."""

pymongo/asynchronous/client_bulk.py

@@ -449,7 +449,6 @@ async def _process_results_cursor(
                 result["cursor"],
                 conn.address,
                 session=session,
-                explicit_session=session is not None,
                 comment=self.comment,
             )
             await cmd_cursor._maybe_pin_connection(conn)

pymongo/asynchronous/client_session.py

@@ -137,6 +137,7 @@
 
 import collections
 import time
+import traceback
 import uuid
 from collections.abc import Mapping as _Mapping
 from typing import (
@@ -513,6 +514,7 @@ def __init__(
         # Is this an implicitly created session?
         self._implicit = implicit
         self._transaction = _Transaction(None, client)
+        self._attached_to_cursor = False
 
     async def end_session(self) -> None:
         """Finish this session. If a transaction has started, abort it.
@@ -536,11 +538,14 @@ 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:
+            # 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
 
     def _check_ended(self) -> None:
         if self._server_session is None:
+            # print(f"Session {self} is already ended")
+            # print(f"Session {self} is already ended, implicit: {self._implicit}, attached: {self._attached_to_cursor}: {''.join(traceback.format_stack())}")
             raise InvalidOperation("Cannot use ended session")
 
     async def __aenter__(self) -> AsyncClientSession:
@@ -588,6 +593,20 @@ def operation_time(self) -> Optional[Timestamp]:
         """
         return self._operation_time
 
+    @property
+    def implicit(self) -> bool:
+        """Whether this session was implicitly created by the driver."""
+        return self._implicit
+
+    @property
+    def attached_to_cursor(self) -> bool:
+        """Whether this session is owned by a cursor."""
+        return self._attached_to_cursor
+
+    @attached_to_cursor.setter
+    def attached_to_cursor(self, value: bool) -> None:
+        self._attached_to_cursor = value
+
     def _inherit_option(self, name: str, val: _T) -> _T:
         """Return the inherited TransactionOption value."""
         if val:

pymongo/asynchronous/collection.py

@@ -2549,7 +2549,6 @@ async def _list_indexes(
             self.with_options(codec_options=codec_options, read_preference=ReadPreference.PRIMARY),
         )
         read_pref = (session and session._txn_read_preference()) or ReadPreference.PRIMARY
-        explicit_session = session is not None
 
         async def _cmd(
             session: Optional[AsyncClientSession],
@@ -2576,13 +2575,12 @@ async def _cmd(
                 cursor,
                 conn.address,
                 session=session,
-                explicit_session=explicit_session,
                 comment=cmd.get("comment"),
             )
             await cmd_cursor._maybe_pin_connection(conn)
             return cmd_cursor
 
-        async with self._database.client._tmp_session(session, False) as s:
+        async with self._database.client._tmp_session(session) as s:
             return await self._database.client._retryable_read(
                 _cmd, read_pref, s, operation=_Op.LIST_INDEXES
             )
@@ -2678,7 +2676,6 @@ async def list_search_indexes(
             AsyncCommandCursor,
             pipeline,
             kwargs,
-            explicit_session=session is not None,
             comment=comment,
             user_fields={"cursor": {"firstBatch": 1}},
         )
@@ -2900,7 +2897,6 @@ async def _aggregate(
         pipeline: _Pipeline,
         cursor_class: Type[AsyncCommandCursor],  # type: ignore[type-arg]
         session: Optional[AsyncClientSession],
-        explicit_session: bool,
         let: Optional[Mapping[str, Any]] = None,
         comment: Optional[Any] = None,
         **kwargs: Any,
@@ -2912,7 +2908,6 @@ async def _aggregate(
             cursor_class,
             pipeline,
             kwargs,
-            explicit_session,
             let,
             user_fields={"cursor": {"firstBatch": 1}},
         )
@@ -3018,13 +3013,12 @@ async def aggregate(
         .. _aggregate command:
             https://mongodb.com/docs/manual/reference/command/aggregate
         """
-        async with self._database.client._tmp_session(session, close=False) as s:
+        async with self._database.client._tmp_session(session) as s:
             return await self._aggregate(
                 _CollectionAggregationCommand,
                 pipeline,
                 AsyncCommandCursor,
                 session=s,
-                explicit_session=session is not None,
                 let=let,
                 comment=comment,
                 **kwargs,
@@ -3065,15 +3059,14 @@ async def aggregate_raw_batches(
             raise InvalidOperation("aggregate_raw_batches does not support auto encryption")
         if comment is not None:
             kwargs["comment"] = comment
-        async with self._database.client._tmp_session(session, close=False) as s:
+        async with self._database.client._tmp_session(session) as s:
             return cast(
                 AsyncRawBatchCursor[_DocumentType],
                 await self._aggregate(
                     _CollectionRawAggregationCommand,
                     pipeline,
                     AsyncRawBatchCommandCursor,
                     session=s,
-                    explicit_session=session is not None,
                     **kwargs,
                 ),
             )

pymongo/asynchronous/command_cursor.py

@@ -64,7 +64,6 @@ def __init__(
         batch_size: int = 0,
         max_await_time_ms: Optional[int] = None,
         session: Optional[AsyncClientSession] = None,
-        explicit_session: bool = False,
         comment: Any = None,
     ) -> None:
         """Create a new command cursor."""
@@ -80,7 +79,8 @@ def __init__(
         self._max_await_time_ms = max_await_time_ms
         self._timeout = self._collection.database.client.options.timeout
         self._session = session
-        self._explicit_session = explicit_session
+        if self._session is not None:
+            self._session.attached_to_cursor = True
         self._killed = self._id == 0
         self._comment = comment
         if self._killed:
@@ -197,7 +197,7 @@ def session(self) -> Optional[AsyncClientSession]:
 
         .. versionadded:: 3.6
         """
-        if self._explicit_session:
+        if not self._session.implicit:
             return self._session
         return None
 
@@ -218,10 +218,11 @@ def _die_no_lock(self) -> None:
         """Closes this cursor without acquiring a lock."""
         cursor_id, address = self._prepare_to_die()
         self._collection.database.client._cleanup_cursor_no_lock(
-            cursor_id, address, self._sock_mgr, self._session, self._explicit_session
+            cursor_id, address, self._sock_mgr, self._session
         )
-        if not self._explicit_session:
-            self._session = None
+        if self._session:
+            if self._session.implicit:
+                self._session = None
         self._sock_mgr = None
 
     async def _die_lock(self) -> None:
@@ -232,16 +233,18 @@ async def _die_lock(self) -> None:
             address,
             self._sock_mgr,
             self._session,
-            self._explicit_session,
         )
-        if not self._explicit_session:
-            self._session = None
+        if self._session:
+            if self._session.implicit:
+                self._session = None
         self._sock_mgr = None
 
     def _end_session(self) -> None:
-        if self._session and not self._explicit_session:
-            self._session._end_implicit_session()
-            self._session = None
+        if self._session:
+            if self._session.implicit:
+                # print(f"Ending session {self}, session: {self._session}")
+                self._session._end_implicit_session()
+                self._session = None
 
     async def close(self) -> None:
         """Explicitly close / kill this cursor."""
@@ -430,7 +433,6 @@ def __init__(
         batch_size: int = 0,
         max_await_time_ms: Optional[int] = None,
         session: Optional[AsyncClientSession] = None,
-        explicit_session: bool = False,
         comment: Any = None,
     ) -> None:
         """Create a new cursor / iterator over raw batches of BSON data.
@@ -449,7 +451,6 @@ def __init__(
             batch_size,
             max_await_time_ms,
             session,
-            explicit_session,
             comment,
         )
 

pymongo/asynchronous/cursor.py

@@ -138,10 +138,9 @@ def __init__(
 
         if session:
             self._session = session
-            self._explicit_session = True
+            self._session.attached_to_cursor = True
         else:
             self._session = None
-            self._explicit_session = False
 
         spec: Mapping[str, Any] = filter or {}
         validate_is_mapping("filter", spec)
@@ -150,7 +149,7 @@ def __init__(
         if not isinstance(limit, int):
             raise TypeError(f"limit must be an instance of int, not {type(limit)}")
         validate_boolean("no_cursor_timeout", no_cursor_timeout)
-        if no_cursor_timeout and not self._explicit_session:
+        if no_cursor_timeout and self._session and self._session.implicit:
             warnings.warn(
                 "use an explicit session with no_cursor_timeout=True "
                 "otherwise the cursor may still timeout after "
@@ -283,7 +282,7 @@ def clone(self) -> AsyncCursor[_DocumentType]:
     def _clone(self, deepcopy: bool = True, base: Optional[AsyncCursor] = None) -> AsyncCursor:  # type: ignore[type-arg]
         """Internal clone helper."""
         if not base:
-            if self._explicit_session:
+            if self._session and not self._session.implicit:
                 base = self._clone_base(self._session)
             else:
                 base = self._clone_base(None)
@@ -945,7 +944,7 @@ def session(self) -> Optional[AsyncClientSession]:
 
         .. versionadded:: 3.6
         """
-        if self._explicit_session:
+        if self._session and not self._session.implicit:
             return self._session
         return None
 
@@ -1034,10 +1033,11 @@ def _die_no_lock(self) -> None:
 
         cursor_id, address = self._prepare_to_die(already_killed)
         self._collection.database.client._cleanup_cursor_no_lock(
-            cursor_id, address, self._sock_mgr, self._session, self._explicit_session
+            cursor_id, address, self._sock_mgr, self._session
         )
-        if not self._explicit_session:
-            self._session = None
+        if self._session:
+            if self._session.implicit:
+                self._session = None
         self._sock_mgr = None
 
     async def _die_lock(self) -> None:
@@ -1054,10 +1054,10 @@ async def _die_lock(self) -> None:
             address,
             self._sock_mgr,
             self._session,
-            self._explicit_session,
         )
-        if not self._explicit_session:
-            self._session = None
+        if self._session:
+            if self._session.implicit:
+                self._session = None
         self._sock_mgr = None
 
     async def close(self) -> None:

pymongo/asynchronous/database.py

@@ -699,7 +699,7 @@ async def aggregate(
         .. _aggregate command:
             https://mongodb.com/docs/manual/reference/command/aggregate
         """
-        async with self.client._tmp_session(session, close=False) as s:
+        async with self.client._tmp_session(session) as s:
             cmd = _DatabaseAggregationCommand(
                 self,
                 AsyncCommandCursor,
@@ -1011,7 +1011,7 @@ async def cursor_command(
         else:
             command_name = next(iter(command))
 
-        async with self._client._tmp_session(session, close=False) as tmp_session:
+        async with self._client._tmp_session(session) as tmp_session:
             opts = codec_options or DEFAULT_CODEC_OPTIONS
 
             if read_preference is None:
@@ -1043,7 +1043,6 @@ async def cursor_command(
                         conn.address,
                         max_await_time_ms=max_await_time_ms,
                         session=tmp_session,
-                        explicit_session=session is not None,
                         comment=comment,
                     )
                     await cmd_cursor._maybe_pin_connection(conn)
@@ -1089,7 +1088,7 @@ async def _list_collections(
         )
         cmd = {"listCollections": 1, "cursor": {}}
         cmd.update(kwargs)
-        async with self._client._tmp_session(session, close=False) as tmp_session:
+        async with self._client._tmp_session(session) as tmp_session:
             cursor = (
                 await self._command(conn, cmd, read_preference=read_preference, session=tmp_session)
             )["cursor"]
@@ -1098,7 +1097,6 @@ async def _list_collections(
                 cursor,
                 conn.address,
                 session=tmp_session,
-                explicit_session=session is not None,
                 comment=cmd.get("comment"),
             )
         await cmd_cursor._maybe_pin_connection(conn)