-
Notifications
You must be signed in to change notification settings - Fork 1.1k
Python 4542 - Improved sessions API #2712
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
2c0e189
5c21939
969abb2
f241650
df3040d
d6b883b
13f1a15
5edd9a0
6eaf094
b0afe91
a1234ed
e00ac18
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -139,6 +139,7 @@ | |
| import time | ||
| import uuid | ||
| from collections.abc import Mapping as _Mapping | ||
| from contextvars import ContextVar, Token | ||
| from typing import ( | ||
| TYPE_CHECKING, | ||
| Any, | ||
|
|
@@ -181,6 +182,28 @@ | |
|
|
||
| _IS_SYNC = False | ||
|
|
||
| _SESSION: ContextVar[Optional[AsyncClientSession]] = ContextVar("SESSION", default=None) | ||
|
|
||
|
|
||
| class _AsyncBoundSessionContext: | ||
| """Context manager returned by AsyncClientSession.bind() that manages bound state.""" | ||
|
|
||
| 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] | ||
| return self._session | ||
|
|
||
| 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] | ||
NoahStapp marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| self._session_token = None | ||
| if self._end_session: | ||
| await self._session.end_session() | ||
|
|
||
|
|
||
| class SessionOptions: | ||
| """Options for a new :class:`AsyncClientSession`. | ||
|
|
@@ -547,6 +570,24 @@ def _check_ended(self) -> None: | |
| if self._server_session is None: | ||
| raise InvalidOperation("Cannot use ended session") | ||
|
|
||
| def bind(self, end_session: bool = True) -> _AsyncBoundSessionContext: | ||
| """Bind this session so it is implicitly passed to all database operations within the returned context. | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Could you add a short docs explain here? We still do that right?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. What do you mean?
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Like a |
||
|
|
||
| .. code-block:: python | ||
|
|
||
| async with client.start_session() as s: | ||
| async with s.bind(): | ||
| # session=s is passed implicitly | ||
| await client.db.collection.insert_one({"x": 1}) | ||
|
|
||
| :param end_session: Whether to end the session on exiting the returned context. Defaults to True. | ||
| If set to False, :meth:`~pymongo.asynchronous.client_session.AsyncClientSession.end_session()` must be called | ||
| once the session is no longer used. | ||
|
|
||
| .. versionadded:: 4.17 | ||
| """ | ||
| return _AsyncBoundSessionContext(self, end_session) | ||
|
|
||
| async def __aenter__(self) -> AsyncClientSession: | ||
| return self | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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 | ||
|
|
@@ -1408,7 +1408,8 @@ def start_session( | |
| def _ensure_session( | ||
| self, session: Optional[AsyncClientSession] = None | ||
| ) -> Optional[AsyncClientSession]: | ||
| """If provided session is None, lend a temporary session.""" | ||
| """If provided session and bound session are None, lend a temporary session.""" | ||
| session = session or self._get_bound_session() | ||
| if session: | ||
| return session | ||
|
|
||
|
|
@@ -2267,11 +2268,14 @@ async def _tmp_session( | |
| self, session: Optional[client_session.AsyncClientSession] | ||
| ) -> AsyncGenerator[Optional[client_session.AsyncClientSession], None]: | ||
| """If provided session is None, lend a temporary session.""" | ||
| if session is not None: | ||
| if not isinstance(session, client_session.AsyncClientSession): | ||
| raise ValueError( | ||
| f"'session' argument must be an AsyncClientSession or None, not {type(session)}" | ||
| ) | ||
| if session is not None and not isinstance(session, client_session.AsyncClientSession): | ||
| raise ValueError( | ||
| f"'session' argument must be an AsyncClientSession or None, not {type(session)}" | ||
| ) | ||
|
|
||
| # Check for a bound session. If one exists, treat it as an explicitly passed session. | ||
| session = session or self._get_bound_session() | ||
| if session: | ||
NoahStapp marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| # Don't call end_session. | ||
| yield session | ||
| return | ||
|
|
@@ -2301,6 +2305,18 @@ async def _process_response( | |
| if session is not None: | ||
| session._process_response(reply) | ||
|
|
||
| def _get_bound_session(self) -> Optional[AsyncClientSession]: | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This is encapsulated in a separate utility function because cursor operations call |
||
| bound_session = _SESSION.get() | ||
| if bound_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." | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Another
NoahStapp marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| ) | ||
| else: | ||
| return None | ||
|
|
||
| async def server_info( | ||
| self, session: Optional[client_session.AsyncClientSession] = None | ||
| ) -> dict[str, Any]: | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -108,7 +108,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 | ||
|
|
@@ -1406,7 +1406,8 @@ def start_session( | |
| ) | ||
|
|
||
| def _ensure_session(self, session: Optional[ClientSession] = None) -> Optional[ClientSession]: | ||
| """If provided session is None, lend a temporary session.""" | ||
| """If provided session and bound session are None, lend a temporary session.""" | ||
| session = session or self._get_bound_session() | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I don't understand what's going on here before the changes. How does: lend a temporary session?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We call |
||
| if session: | ||
| return session | ||
|
|
||
|
|
@@ -2263,11 +2264,14 @@ def _tmp_session( | |
| self, session: Optional[client_session.ClientSession] | ||
| ) -> Generator[Optional[client_session.ClientSession], None]: | ||
| """If provided session is None, lend a temporary session.""" | ||
| if session is not None: | ||
| if not isinstance(session, client_session.ClientSession): | ||
| raise ValueError( | ||
| f"'session' argument must be a ClientSession or None, not {type(session)}" | ||
| ) | ||
| if session is not None and not isinstance(session, client_session.ClientSession): | ||
| raise ValueError( | ||
| f"'session' argument must be a ClientSession or None, not {type(session)}" | ||
| ) | ||
|
|
||
| # Check for a bound session. If one exists, treat it as an explicitly passed session. | ||
| session = session or self._get_bound_session() | ||
| if session: | ||
NoahStapp marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| # Don't call end_session. | ||
| yield session | ||
| return | ||
|
|
@@ -2295,6 +2299,18 @@ def _process_response(self, reply: Mapping[str, Any], session: Optional[ClientSe | |
| if session is not None: | ||
| session._process_response(reply) | ||
|
|
||
| def _get_bound_session(self) -> Optional[ClientSession]: | ||
| bound_session = _SESSION.get() | ||
| if bound_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." | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Third
NoahStapp marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| ) | ||
| else: | ||
| return None | ||
|
|
||
| def server_info(self, session: Optional[client_session.ClientSession] = None) -> dict[str, Any]: | ||
| """Get information about the MongoDB server we're connected to. | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
<PLACEHOLDER>should be the MongoDB docs ?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah once we have examples and such added to a page I'll update these spots.