SH review

Author: NoahStapp

pymongo/asynchronous/client_session.py

@@ -182,25 +182,18 @@
 
 _IS_SYNC = False
 
-_SESSION: ContextVar[Optional[_AsyncBoundClientSession]] = ContextVar("SESSION", default=None)
+_SESSION: ContextVar[Optional[AsyncClientSession]] = ContextVar("SESSION", default=None)
 
 
-class _AsyncBoundClientSession:
-    def __init__(self, session: AsyncClientSession, client_id: int):
-        self.session = session
-        self.client_id = client_id
-
-
-class AsyncBoundSessionContext:
+class _AsyncBoundSessionContext:
     """Context manager returned by AsyncClientSession.bind() that manages bound state."""
 
     def __init__(self, session: AsyncClientSession) -> None:
         self._session = session
-        self._session_token: Optional[Token[_AsyncBoundClientSession]] = None
+        self._session_token: Optional[Token[AsyncClientSession]] = None
 
     async def __aenter__(self) -> AsyncClientSession:
-        bound_session = _AsyncBoundClientSession(self._session, id(self._session._client))
-        self._session_token = _SESSION.set(bound_session)  # type: ignore[assignment]
+        self._session_token = _SESSION.set(self._session)  # type: ignore[assignment]
         return self._session
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
@@ -574,12 +567,12 @@ def _check_ended(self) -> None:
         if self._server_session is None:
             raise InvalidOperation("Cannot use ended session")
 
-    def bind(self) -> AsyncBoundSessionContext:
+    def bind(self) -> _AsyncBoundSessionContext:
         """Bind this session so it is implicitly passed to all database operations within the returned context.
 
         .. versionadded:: 4.17
         """
-        return AsyncBoundSessionContext(self)
+        return _AsyncBoundSessionContext(self)
 
     async def __aenter__(self) -> AsyncClientSession:
         return self

pymongo/asynchronous/mongo_client.py

@@ -2308,8 +2308,8 @@ async def _process_response(
     def _get_bound_session(self) -> Optional[AsyncClientSession]:
         bound_session = _SESSION.get()
         if bound_session:
-            if bound_session.client_id == id(self):
-                return bound_session.session
+            if bound_session.client is self:
+                return bound_session
             else:
                 raise InvalidOperation(
                     "Only the client that created the bound session can perform operations within its context block. See <PLACEHOLDER> for more information."

pymongo/synchronous/client_session.py

@@ -181,25 +181,18 @@
 
 _IS_SYNC = True
 
-_SESSION: ContextVar[Optional[_BoundClientSession]] = ContextVar("SESSION", default=None)
-
-
-class _BoundClientSession:
-    def __init__(self, session: ClientSession, client_id: int):
-        self.session = session
-        self.client_id = client_id
+_SESSION: ContextVar[Optional[ClientSession]] = ContextVar("SESSION", default=None)
 
 
 class BoundSessionContext:
     """Context manager returned by ClientSession.bind() that manages bound state."""
 
     def __init__(self, session: ClientSession) -> None:
         self._session = session
-        self._session_token: Optional[Token[_BoundClientSession]] = None
+        self._session_token: Optional[Token[ClientSession]] = None
 
     def __enter__(self) -> ClientSession:
-        bound_session = _BoundClientSession(self._session, id(self._session._client))
-        self._session_token = _SESSION.set(bound_session)  # type: ignore[assignment]
+        self._session_token = _SESSION.set(self._session)  # type: ignore[assignment]
         return self._session
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:

pymongo/synchronous/mongo_client.py

@@ -2302,8 +2302,8 @@ def _process_response(self, reply: Mapping[str, Any], session: Optional[ClientSe
     def _get_bound_session(self) -> Optional[ClientSession]:
         bound_session = _SESSION.get()
         if bound_session:
-            if bound_session.client_id == id(self):
-                return bound_session.session
+            if bound_session.client is self:
+                return bound_session
             else:
                 raise InvalidOperation(
                     "Only the client that created the bound session can perform operations within its context block. See <PLACEHOLDER> for more information."

tools/synchro.py

@@ -37,8 +37,7 @@
     "AsyncRawBatchCursor": "RawBatchCursor",
     "AsyncRawBatchCommandCursor": "RawBatchCommandCursor",
     "AsyncClientSession": "ClientSession",
-    "_AsyncBoundClientSession": "_BoundClientSession",
-    "AsyncBoundSessionContext": "BoundSessionContext",
+    "_AsyncBoundSessionContext": "BoundSessionContext",
     "AsyncChangeStream": "ChangeStream",
     "AsyncCollectionChangeStream": "CollectionChangeStream",
     "AsyncDatabaseChangeStream": "DatabaseChangeStream",