From 834f311966264c1fb1c3067f2c22866b8ff191b5 Mon Sep 17 00:00:00 2001 From: Iris Date: Thu, 1 Oct 2026 19:08:31 -0700 Subject: [PATCH 1/2] PYTHON-6025 Consolidate non-IO code in AsyncClientSession/ClientSession and _Transaction into shared base classes --- pymongo/asynchronous/client_session.py | 356 ++--------------------- pymongo/client_session_shared.py | 378 ++++++++++++++++++++++++- pymongo/message.py | 6 +- pymongo/synchronous/client_session.py | 353 +---------------------- pymongo/typings.py | 3 + 5 files changed, 415 insertions(+), 681 deletions(-) diff --git a/pymongo/asynchronous/client_session.py b/pymongo/asynchronous/client_session.py index acadb33ce3..61606ead75 100644 --- a/pymongo/asynchronous/client_session.py +++ b/pymongo/asynchronous/client_session.py @@ -138,56 +138,49 @@ import asyncio import random import time -from collections.abc import Awaitable, Mapping, MutableMapping -from collections.abc import Mapping as _Mapping +from collections.abc import Awaitable from contextlib import AbstractAsyncContextManager from contextvars import ContextVar, Token from typing import ( TYPE_CHECKING, Any, Callable, - NoReturn, Optional, TypeVar, ) -from bson.int64 import Int64 -from bson.timestamp import Timestamp from pymongo import _csot from pymongo.asynchronous.cursor_base import _ConnectionManager from pymongo.client_session_shared import ( _BACKOFF_INITIAL, _BACKOFF_MAX, _UNKNOWN_COMMIT_ERROR_CODES, - SessionOptions, + SessionOptions, # noqa: F401 (re-exported for pymongo.client_session) TransactionOptions, - _EmptyServerSession, + _AgnosticClientSessionBase, _make_timeout_error, _max_time_expired_error, _reraise_with_unknown_commit, + _TransactionBase, _TxnState, _within_time_limit, ) from pymongo.errors import ( - ConfigurationError, ConnectionFailure, InvalidOperation, OperationFailure, PyMongoError, WTimeoutError, ) -from pymongo.operations import _WRITES_WITH_CLUSTER_TIME from pymongo.read_concern import ReadConcern -from pymongo.read_preferences import ReadPreference, _ServerMode -from pymongo.server_type import SERVER_TYPE +from pymongo.read_preferences import _ServerMode from pymongo.write_concern import WriteConcern if TYPE_CHECKING: from types import TracebackType + from pymongo.asynchronous.mongo_client import AsyncMongoClient from pymongo.asynchronous.pool import AsyncConnection - from pymongo.asynchronous.server import Server - from pymongo.typings import ClusterTime, _Address _IS_SYNC = False @@ -236,45 +229,11 @@ async def __aexit__( await self.__session.abort_transaction() -class _Transaction: +class _Transaction(_TransactionBase["AsyncConnection"]): """Internal class to hold transaction information in a AsyncClientSession.""" - def __init__(self, opts: Optional[TransactionOptions], client: AsyncMongoClient[Any]): - self.opts = opts - self.state = _TxnState.NONE - self.sharded = False - self.pinned_address: Optional[_Address] = None - self.conn_mgr: Optional[_ConnectionManager] = None - self.recovery_token = None - self.attempt = 0 - self.client = client - self.has_completed_command = False - - def active(self) -> bool: - return self.state in (_TxnState.STARTING, _TxnState.IN_PROGRESS) - - def starting(self) -> bool: - return self.state == _TxnState.STARTING - - def set_starting(self) -> None: - self.state = _TxnState.STARTING - - def set_in_progress(self) -> None: - if self.state == _TxnState.STARTING: - self.state = _TxnState.IN_PROGRESS - - @property - def pinned_conn(self) -> Optional[AsyncConnection]: - if self.active() and self.conn_mgr: - return self.conn_mgr.conn - return None - - def pin(self, server: Server, conn: AsyncConnection) -> None: - self.sharded = True - self.pinned_address = server.description.address - if server.description.server_type == SERVER_TYPE.LoadBalancer: - conn.pin_txn() - self.conn_mgr = _ConnectionManager(conn, False) + _conn_mgr_cls = _ConnectionManager + conn_mgr: Optional[_ConnectionManager] async def unpin(self) -> None: self.pinned_address = None @@ -290,21 +249,15 @@ async def reset(self) -> None: self.attempt = 0 self.has_completed_command = False - def __del__(self) -> None: - if self.conn_mgr: - # Reuse the cursor closing machinery to return the socket to the - # pool soon. - self.client._close_cursor_soon(0, None, self.conn_mgr) - self.conn_mgr = None - _T = TypeVar("_T") -if TYPE_CHECKING: - from pymongo.asynchronous.mongo_client import AsyncMongoClient - -class AsyncClientSession: +class AsyncClientSession( + _AgnosticClientSessionBase[ + "AsyncMongoClient[Any]", "AsyncConnection", _AsyncBoundSessionContext + ] +): """A session for ordering sequential operations. :class:`AsyncClientSession` instances are **not thread-safe or fork-safe**. @@ -317,27 +270,10 @@ class AsyncClientSession: :meth:`~pymongo.asynchronous.mongo_client.AsyncMongoClient.start_session`. """ - def __init__( - self, - client: AsyncMongoClient[Any], - server_session: Any, - options: SessionOptions, - implicit: bool, - ) -> None: - # An AsyncMongoClient, a _ServerSession, a SessionOptions, and a set. - self._client: AsyncMongoClient[Any] = client - self._server_session = server_session - self._options = options - self._cluster_time: Optional[Mapping[str, Any]] = None - self._operation_time: Optional[Timestamp] = None - self._snapshot_time = None - # Is this an implicitly created session? - self._implicit = implicit - self._transaction = _Transaction(None, client) - # Is this session attached to a cursor? - self._attached_to_cursor = False - # Should we leave the session alive when the cursor is closed? - self._leave_alive = False + _transaction_cls = _Transaction + _bound_session_context_cls = _AsyncBoundSessionContext + _transaction: _Transaction + _client: AsyncMongoClient[Any] async def end_session(self) -> None: """Finish this session. If a transaction has started, abort it. @@ -358,89 +294,12 @@ async def _end_session(self, lock: bool) -> None: self._client._return_server_session(self._server_session) self._server_session = None - def _end_implicit_session(self) -> None: - # Implicit sessions can't be part of transactions or pinned connections - if not self._leave_alive and self._server_session is not None: - self._client._return_server_session(self._server_session) - self._server_session = None - - 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. - - .. 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 async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: await self._end_session(lock=True) - @property - def client(self) -> AsyncMongoClient[Any]: - """The :class:`~pymongo.asynchronous.mongo_client.AsyncMongoClient` this session was - created from. - """ - return self._client - - @property - def options(self) -> SessionOptions: - """The :class:`SessionOptions` this session was created with.""" - return self._options - - @property - def session_id(self) -> Mapping[str, Any]: - """A BSON document, the opaque server session identifier.""" - self._check_ended() - self._materialize(self._client.topology_description.logical_session_timeout_minutes) - return self._server_session.session_id - - @property - def _transaction_id(self) -> Int64: - """The current transaction id for the underlying server session.""" - self._materialize(self._client.topology_description.logical_session_timeout_minutes) - return self._server_session.transaction_id - - @property - def cluster_time(self) -> Optional[ClusterTime]: - """The cluster time returned by the last operation executed - in this session. - """ - return self._cluster_time - - @property - def operation_time(self) -> Optional[Timestamp]: - """The operation time returned by the last operation executed - in this session. - """ - return self._operation_time - - def _inherit_option(self, name: str, val: _T) -> _T: - """Return the inherited TransactionOption value.""" - if val: - return val - txn_opts = self.options.default_transaction_options - parent_val = txn_opts and getattr(txn_opts, name) - if parent_val: - return parent_val - return getattr(self.client, name) - async def with_transaction( self, callback: Callable[[AsyncClientSession], Awaitable[_T]], @@ -752,185 +611,6 @@ async def _finish_transaction(self, conn: AsyncConnection, command_name: str) -> conn, cmd, session=self, write_concern=wc, parse_write_concern_error=True ) - def _advance_cluster_time(self, cluster_time: Optional[Mapping[str, Any]]) -> None: - """Internal cluster time helper.""" - if self._cluster_time is None: - self._cluster_time = cluster_time - elif cluster_time is not None: - if cluster_time["clusterTime"] > self._cluster_time["clusterTime"]: - self._cluster_time = cluster_time - - def advance_cluster_time(self, cluster_time: Mapping[str, Any]) -> None: - """Update the cluster time for this session. - - :param cluster_time: The - :data:`~pymongo.asynchronous.client_session.AsyncClientSession.cluster_time` from - another `AsyncClientSession` instance. - """ - if not isinstance(cluster_time, _Mapping): - raise TypeError( - f"cluster_time must be a subclass of collections.Mapping, not {type(cluster_time)}" - ) - if not isinstance(cluster_time.get("clusterTime"), Timestamp): - raise ValueError("Invalid cluster_time") - self._advance_cluster_time(cluster_time) - - def _advance_operation_time(self, operation_time: Optional[Timestamp]) -> None: - """Internal operation time helper.""" - if self._operation_time is None: - self._operation_time = operation_time - elif operation_time is not None: - if operation_time > self._operation_time: - self._operation_time = operation_time - - def advance_operation_time(self, operation_time: Timestamp) -> None: - """Update the operation time for this session. - - :param operation_time: The - :data:`~pymongo.asynchronous.client_session.AsyncClientSession.operation_time` from - another `AsyncClientSession` instance. - """ - if not isinstance(operation_time, Timestamp): - raise TypeError( - f"operation_time must be an instance of bson.timestamp.Timestamp, not {type(operation_time)}" - ) - self._advance_operation_time(operation_time) - - def _process_response(self, reply: Mapping[str, Any]) -> None: - """Process a response to a command that was run with this session.""" - self._advance_cluster_time(reply.get("$clusterTime")) - self._advance_operation_time(reply.get("operationTime")) - if self._options.snapshot and self._snapshot_time is None: - if "cursor" in reply: - ct = reply["cursor"].get("atClusterTime") - else: - ct = reply.get("atClusterTime") - self._snapshot_time = ct - if self.in_transaction and self._transaction.sharded: - recovery_token = reply.get("recoveryToken") - if recovery_token: - self._transaction.recovery_token = recovery_token - - @property - def has_ended(self) -> bool: - """True if this session is finished.""" - return self._server_session is None - - @property - def in_transaction(self) -> bool: - """True if this session has an active multi-statement transaction. - - .. versionadded:: 3.10 - """ - return self._transaction.active() - - @property - def _starting_transaction(self) -> bool: - """True if this session is starting a multi-statement transaction.""" - return self._transaction.starting() - - @property - def _pinned_address(self) -> Optional[_Address]: - """The mongos address this transaction was created on.""" - if self._transaction.active(): - return self._transaction.pinned_address - return None - - @property - def _pinned_connection(self) -> Optional[AsyncConnection]: - """The connection this transaction was started on.""" - return self._transaction.pinned_conn - - def _pin(self, server: Server, conn: AsyncConnection) -> None: - """Pin this session to the given Server or to the given connection.""" - self._transaction.pin(server, conn) - async def _unpin(self) -> None: """Unpin this session from any pinned Server.""" await self._transaction.unpin() - - def _txn_read_preference(self) -> Optional[_ServerMode]: - """Return read preference of this transaction or None.""" - if self.in_transaction: - assert self._transaction.opts - return self._transaction.opts.read_preference - return None - - def _materialize(self, logical_session_timeout_minutes: Optional[int] = None) -> None: - if isinstance(self._server_session, _EmptyServerSession): - old = self._server_session - self._server_session = self._client._topology.get_server_session( - logical_session_timeout_minutes - ) - if old.started_retryable_write: - self._server_session.inc_transaction_id() - - def _apply_to( - self, - command: MutableMapping[str, Any], - is_retryable: bool, - read_preference: _ServerMode, - conn: AsyncConnection, - ) -> None: - # getMores must be sent with a session if the cursor was opened with one - operation = next(iter(command)) - if not conn.supports_sessions and ( - isinstance(self._server_session, _EmptyServerSession) or operation != "getMore" - ): - if not self._implicit: - raise ConfigurationError("Sessions are not supported by this MongoDB deployment") - return - self._check_ended() - self._materialize(conn.logical_session_timeout_minutes) - # Add afterClusterTime on snapshot reads or writes in causally-consistent sessions - if self.options.snapshot or ( - self.options.causal_consistency - and not self.in_transaction - and operation in _WRITES_WITH_CLUSTER_TIME - ): - self._update_read_concern(command, conn) - - self._server_session.last_use = time.monotonic() - command["lsid"] = self._server_session.session_id - - if is_retryable: - command["txnNumber"] = self._server_session.transaction_id - return - - if self.in_transaction: - if read_preference != ReadPreference.PRIMARY: - raise InvalidOperation( - f"read preference in a transaction must be primary, not: {read_preference!r}" - ) - - if self._transaction.state == _TxnState.STARTING: - # First command begins a new transaction. - command["startTransaction"] = True - - assert self._transaction.opts - if self._transaction.opts.read_concern: - rc = self._transaction.opts.read_concern.document - if rc: - command["readConcern"] = rc - self._update_read_concern(command, conn) - - command["txnNumber"] = self._server_session.transaction_id - command["autocommit"] = False - - def _start_retryable_write(self) -> None: - self._check_ended() - self._server_session.inc_transaction_id() - - def _update_read_concern(self, cmd: MutableMapping[str, Any], conn: AsyncConnection) -> None: - if self.options.causal_consistency and self.operation_time is not None: - cmd.setdefault("readConcern", {})["afterClusterTime"] = self.operation_time - if self.options.snapshot: - if conn.max_wire_version < 13: - raise ConfigurationError("Snapshot reads require MongoDB 5.0 or later") - rc = cmd.setdefault("readConcern", {}) - rc["level"] = "snapshot" - if self._snapshot_time is not None: - rc["atClusterTime"] = self._snapshot_time - - def __copy__(self) -> NoReturn: - raise TypeError("A AsyncClientSession cannot be copied, create a new session instead") diff --git a/pymongo/client_session_shared.py b/pymongo/client_session_shared.py index 467d069355..d2d086905a 100644 --- a/pymongo/client_session_shared.py +++ b/pymongo/client_session_shared.py @@ -19,33 +19,53 @@ import collections import time import uuid +from collections.abc import Mapping, MutableMapping from typing import ( TYPE_CHECKING, Any, + Callable, + Generic, NoReturn, Optional, + Protocol, TypeVar, + cast, ) from bson.binary import Binary from bson.int64 import Int64 +from bson.timestamp import Timestamp from pymongo import _csot from pymongo.errors import ( ConfigurationError, ExecutionTimeout, + InvalidOperation, NetworkTimeout, OperationFailure, PyMongoError, ) from pymongo.helpers_shared import _RETRYABLE_ERROR_CODES +from pymongo.operations import _WRITES_WITH_CLUSTER_TIME from pymongo.read_concern import ReadConcern -from pymongo.read_preferences import _ServerMode +from pymongo.read_preferences import ReadPreference, _ServerMode +from pymongo.server_type import SERVER_TYPE from pymongo.write_concern import WriteConcern if TYPE_CHECKING: - from pymongo.typings import _AgnosticClientSession + from pymongo.typings import ( + ClusterTime, + _Address, + _AgnosticClientSession, + _AgnosticConnection, + _AgnosticMongoClient, + _AgnosticServer, + ) _ClientSessionT = TypeVar("_ClientSessionT", bound="_AgnosticClientSession") +_T = TypeVar("_T") +_ClientT = TypeVar("_ClientT", bound="_AgnosticMongoClient") +_ConnectionT = TypeVar("_ConnectionT", bound="_AgnosticConnection") +_BoundSessionContextT = TypeVar("_BoundSessionContextT") class SessionOptions: @@ -393,3 +413,357 @@ def _clear_stale(self, session_timeout_minutes: Optional[int]) -> None: self.append(s) # The remaining sessions also haven't timed out. break + + +class _AgnosticClientSessionBase(Generic[_ClientT, _ConnectionT, _BoundSessionContextT]): + """A session for ordering sequential operations. + + Session instances are **not thread-safe or fork-safe**. They can only be used + by one thread or process at a time. A single session cannot be used to run + multiple operations concurrently. + + Should not be initialized directly by application developers - to create a + session, call the client's ``start_session`` method. + """ + + # Set by the async/sync subclasses to the IO-specific classes. + _transaction_cls: type[_TransactionBase[Any]] + _bound_session_context_cls: Callable[[Any, bool], _BoundSessionContextT] + + def __init__( + self, + client: _ClientT, + server_session: Any, + options: SessionOptions, + implicit: bool, + ) -> None: + # An _AgnosticMongoClient, a _ServerSession, a SessionOptions, and a set. + self._client: _ClientT = client + self._server_session = server_session + self._options = options + self._cluster_time: Optional[Mapping[str, Any]] = None + self._operation_time: Optional[Timestamp] = None + self._snapshot_time = None + # Is this an implicitly created session? + self._implicit = implicit + self._transaction: _TransactionBase[_ConnectionT] = self._transaction_cls(None, client) + # Is this session attached to a cursor? + self._attached_to_cursor = False + # Should we leave the session alive when the cursor is closed? + self._leave_alive = False + + def _end_implicit_session(self) -> None: + # Implicit sessions can't be part of transactions or pinned connections + if not self._leave_alive and self._server_session is not None: + self._client._return_server_session(self._server_session) + self._server_session = None + + def _check_ended(self) -> None: + if self._server_session is None: + raise InvalidOperation("Cannot use ended session") + + def bind(self, end_session: bool = True) -> _BoundSessionContextT: + """Bind this session so it is implicitly passed to all database operations within the returned context. + + .. code-block:: python + + # Synchronous + with client.start_session() as s: + with s.bind(): + # session=s is passed implicitly + client.db.collection.insert_one({"x": 1}) + + # Asynchronous + 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, ``end_session()`` must be called once the session is no longer used. + + .. versionadded:: 4.17 + """ + return self._bound_session_context_cls(self, end_session) + + @property + def client(self) -> _ClientT: + """The client this session was created from.""" + return self._client + + @property + def options(self) -> SessionOptions: + """The :class:`SessionOptions` this session was created with.""" + return self._options + + @property + def session_id(self) -> Mapping[str, Any]: + """A BSON document, the opaque server session identifier.""" + self._check_ended() + self._materialize(self._client.topology_description.logical_session_timeout_minutes) + return self._server_session.session_id + + @property + def _transaction_id(self) -> Int64: + """The current transaction id for the underlying server session.""" + self._materialize(self._client.topology_description.logical_session_timeout_minutes) + return self._server_session.transaction_id + + @property + def cluster_time(self) -> Optional[ClusterTime]: + """The cluster time returned by the last operation executed + in this session. + """ + return self._cluster_time + + @property + def operation_time(self) -> Optional[Timestamp]: + """The operation time returned by the last operation executed + in this session. + """ + return self._operation_time + + def _inherit_option(self, name: str, val: _T) -> _T: + """Return the inherited TransactionOption value.""" + if val: + return val + txn_opts = self.options.default_transaction_options + parent_val = txn_opts and getattr(txn_opts, name) + if parent_val: + return parent_val + return getattr(self.client, name) + + def _advance_cluster_time(self, cluster_time: Optional[Mapping[str, Any]]) -> None: + """Internal cluster time helper.""" + if self._cluster_time is None: + self._cluster_time = cluster_time + elif cluster_time is not None: + if cluster_time["clusterTime"] > self._cluster_time["clusterTime"]: + self._cluster_time = cluster_time + + def advance_cluster_time(self, cluster_time: Mapping[str, Any]) -> None: + """Update the cluster time for this session. + + :param cluster_time: The ``cluster_time`` from another session. + """ + if not isinstance(cluster_time, Mapping): + raise TypeError( + f"cluster_time must be a subclass of collections.Mapping, not {type(cluster_time)}" + ) + if not isinstance(cluster_time.get("clusterTime"), Timestamp): + raise ValueError("Invalid cluster_time") + self._advance_cluster_time(cluster_time) + + def _advance_operation_time(self, operation_time: Optional[Timestamp]) -> None: + """Internal operation time helper.""" + if self._operation_time is None: + self._operation_time = operation_time + elif operation_time is not None: + if operation_time > self._operation_time: + self._operation_time = operation_time + + def advance_operation_time(self, operation_time: Timestamp) -> None: + """Update the operation time for this session. + + :param operation_time: The ``operation_time`` from another session. + """ + if not isinstance(operation_time, Timestamp): + raise TypeError( + f"operation_time must be an instance of bson.timestamp.Timestamp, not {type(operation_time)}" + ) + self._advance_operation_time(operation_time) + + def _process_response(self, reply: Mapping[str, Any]) -> None: + """Process a response to a command that was run with this session.""" + self._advance_cluster_time(reply.get("$clusterTime")) + self._advance_operation_time(reply.get("operationTime")) + if self._options.snapshot and self._snapshot_time is None: + if "cursor" in reply: + ct = reply["cursor"].get("atClusterTime") + else: + ct = reply.get("atClusterTime") + self._snapshot_time = ct + if self.in_transaction and self._transaction.sharded: + recovery_token = reply.get("recoveryToken") + if recovery_token: + self._transaction.recovery_token = recovery_token + + @property + def has_ended(self) -> bool: + """True if this session is finished.""" + return self._server_session is None + + @property + def in_transaction(self) -> bool: + """True if this session has an active multi-statement transaction. + + .. versionadded:: 3.10 + """ + return self._transaction.active() + + @property + def _starting_transaction(self) -> bool: + """True if this session is starting a multi-statement transaction.""" + return self._transaction.starting() + + @property + def _pinned_address(self) -> Optional[_Address]: + """The mongos address this transaction was created on.""" + if self._transaction.active(): + return self._transaction.pinned_address + return None + + @property + def _pinned_connection(self) -> Optional[_ConnectionT]: + """The connection this transaction was started on.""" + return self._transaction.pinned_conn + + def _pin(self, server: _AgnosticServer, conn: _ConnectionT) -> None: + """Pin this session to the given Server or to the given connection.""" + self._transaction.pin(server, conn) + + def _txn_read_preference(self) -> Optional[_ServerMode]: + """Return read preference of this transaction or None.""" + if self.in_transaction: + assert self._transaction.opts + return self._transaction.opts.read_preference + return None + + def _materialize(self, logical_session_timeout_minutes: Optional[int] = None) -> None: + if isinstance(self._server_session, _EmptyServerSession): + old = self._server_session + self._server_session = self._client._topology.get_server_session( + logical_session_timeout_minutes + ) + if old.started_retryable_write: + self._server_session.inc_transaction_id() + + def _apply_to( + self, + command: MutableMapping[str, Any], + is_retryable: bool, + read_preference: _ServerMode, + conn: _AgnosticConnection, + ) -> None: + # getMores must be sent with a session if the cursor was opened with one + operation = next(iter(command)) + if not conn.supports_sessions and ( + isinstance(self._server_session, _EmptyServerSession) or operation != "getMore" + ): + if not self._implicit: + raise ConfigurationError("Sessions are not supported by this MongoDB deployment") + return + self._check_ended() + self._materialize(conn.logical_session_timeout_minutes) + # Add afterClusterTime on snapshot reads or writes in causally-consistent sessions + if self.options.snapshot or ( + self.options.causal_consistency + and not self.in_transaction + and operation in _WRITES_WITH_CLUSTER_TIME + ): + self._update_read_concern(command, conn) + + self._server_session.last_use = time.monotonic() + command["lsid"] = self._server_session.session_id + + if is_retryable: + command["txnNumber"] = self._server_session.transaction_id + return + + if self.in_transaction: + if read_preference != ReadPreference.PRIMARY: + raise InvalidOperation( + f"read preference in a transaction must be primary, not: {read_preference!r}" + ) + + if self._transaction.state == _TxnState.STARTING: + # First command begins a new transaction. + command["startTransaction"] = True + + assert self._transaction.opts + if self._transaction.opts.read_concern: + rc = self._transaction.opts.read_concern.document + if rc: + command["readConcern"] = rc + self._update_read_concern(command, conn) + + command["txnNumber"] = self._server_session.transaction_id + command["autocommit"] = False + + def _start_retryable_write(self) -> None: + self._check_ended() + self._server_session.inc_transaction_id() + + def _update_read_concern( + self, cmd: MutableMapping[str, Any], conn: _AgnosticConnection + ) -> None: + if self.options.causal_consistency and self.operation_time is not None: + cmd.setdefault("readConcern", {})["afterClusterTime"] = self.operation_time + if self.options.snapshot: + if conn.max_wire_version < 13: + raise ConfigurationError("Snapshot reads require MongoDB 5.0 or later") + rc = cmd.setdefault("readConcern", {}) + rc["level"] = "snapshot" + if self._snapshot_time is not None: + rc["atClusterTime"] = self._snapshot_time + + def __copy__(self) -> NoReturn: + raise TypeError(f"A {type(self).__name__} cannot be copied, create a new session instead") + + +class _ConnectionManagerProtocol(Protocol[_ConnectionT]): + """Protocol for the connection manager used to pin a connection to a transaction.""" + + conn: Optional[_ConnectionT] + + +class _TransactionBase(Generic[_ConnectionT]): + """Internal class to hold transaction information in a ClientSession.""" + + # Set by the async/sync subclasses to the correct _ConnectionManager class. + _conn_mgr_cls: Callable[[_ConnectionT, bool], _ConnectionManagerProtocol[_ConnectionT]] + + def __init__(self, opts: Optional[TransactionOptions], client: _AgnosticMongoClient): + self.opts = opts + self.state = _TxnState.NONE + self.sharded = False + self.pinned_address: Optional[_Address] = None + self.conn_mgr: Optional[_ConnectionManagerProtocol[_ConnectionT]] = None + self.recovery_token = None + self.attempt = 0 + self.client = client + self.has_completed_command = False + + def active(self) -> bool: + return self.state in (_TxnState.STARTING, _TxnState.IN_PROGRESS) + + def starting(self) -> bool: + return self.state == _TxnState.STARTING + + def set_starting(self) -> None: + self.state = _TxnState.STARTING + + def set_in_progress(self) -> None: + if self.state == _TxnState.STARTING: + self.state = _TxnState.IN_PROGRESS + + @property + def pinned_conn(self) -> Optional[_ConnectionT]: + if self.active() and self.conn_mgr: + return self.conn_mgr.conn + return None + + def pin(self, server: _AgnosticServer, conn: _ConnectionT) -> None: + self.sharded = True + self.pinned_address = server.description.address + if server.description.server_type == SERVER_TYPE.LoadBalancer: + conn.pin_txn() + self.conn_mgr = self._conn_mgr_cls(conn, False) + + def __del__(self) -> None: + if self.conn_mgr: + # Reuse the cursor closing machinery to return the socket to the + # pool soon. _close_cursor_soon expects the IO-specific + # _ConnectionManager, which conn_mgr is at runtime. + self.client._close_cursor_soon(0, None, cast(Any, self.conn_mgr)) + self.conn_mgr = None diff --git a/pymongo/message.py b/pymongo/message.py index df7d940532..b4b3e485d1 100644 --- a/pymongo/message.py +++ b/pymongo/message.py @@ -1292,10 +1292,10 @@ def as_command( cmd = {"explain": cmd} conn.add_server_api(cmd) if self.session: - self.session._apply_to(cmd, False, self.read_preference, conn) # type: ignore[arg-type] + self.session._apply_to(cmd, False, self.read_preference, conn) # Explain does not support readConcern. if not explain and not self.session.in_transaction: - self.session._update_read_concern(cmd, conn) # type: ignore[arg-type] + self.session._update_read_concern(cmd, conn) conn.send_cluster_time(cmd, self.session, self.client) # type: ignore[arg-type] # Support CSOT if apply_timeout: @@ -1404,7 +1404,7 @@ def as_command( self.comment, ) if self.session: - self.session._apply_to(cmd, False, self.read_preference, conn) # type: ignore[arg-type] + self.session._apply_to(cmd, False, self.read_preference, conn) conn.add_server_api(cmd) conn.send_cluster_time(cmd, self.session, self.client) # type: ignore[arg-type] # Support CSOT diff --git a/pymongo/synchronous/client_session.py b/pymongo/synchronous/client_session.py index 30cb17e14f..b536e053c8 100644 --- a/pymongo/synchronous/client_session.py +++ b/pymongo/synchronous/client_session.py @@ -137,56 +137,48 @@ import random import time -from collections.abc import Mapping, MutableMapping -from collections.abc import Mapping as _Mapping from contextlib import AbstractContextManager from contextvars import ContextVar, Token from typing import ( TYPE_CHECKING, Any, Callable, - NoReturn, Optional, TypeVar, ) -from bson.int64 import Int64 -from bson.timestamp import Timestamp from pymongo import _csot from pymongo.client_session_shared import ( _BACKOFF_INITIAL, _BACKOFF_MAX, _UNKNOWN_COMMIT_ERROR_CODES, - SessionOptions, + SessionOptions, # noqa: F401 (re-exported for pymongo.client_session) TransactionOptions, - _EmptyServerSession, + _AgnosticClientSessionBase, _make_timeout_error, _max_time_expired_error, _reraise_with_unknown_commit, + _TransactionBase, _TxnState, _within_time_limit, ) from pymongo.errors import ( - ConfigurationError, ConnectionFailure, InvalidOperation, OperationFailure, PyMongoError, WTimeoutError, ) -from pymongo.operations import _WRITES_WITH_CLUSTER_TIME from pymongo.read_concern import ReadConcern -from pymongo.read_preferences import ReadPreference, _ServerMode -from pymongo.server_type import SERVER_TYPE +from pymongo.read_preferences import _ServerMode from pymongo.synchronous.cursor_base import _ConnectionManager from pymongo.write_concern import WriteConcern if TYPE_CHECKING: from types import TracebackType + from pymongo.synchronous.mongo_client import MongoClient from pymongo.synchronous.pool import Connection - from pymongo.synchronous.server import Server - from pymongo.typings import ClusterTime, _Address _IS_SYNC = True @@ -235,45 +227,11 @@ def __exit__( self.__session.abort_transaction() -class _Transaction: +class _Transaction(_TransactionBase["Connection"]): """Internal class to hold transaction information in a ClientSession.""" - def __init__(self, opts: Optional[TransactionOptions], client: MongoClient[Any]): - self.opts = opts - self.state = _TxnState.NONE - self.sharded = False - self.pinned_address: Optional[_Address] = None - self.conn_mgr: Optional[_ConnectionManager] = None - self.recovery_token = None - self.attempt = 0 - self.client = client - self.has_completed_command = False - - def active(self) -> bool: - return self.state in (_TxnState.STARTING, _TxnState.IN_PROGRESS) - - def starting(self) -> bool: - return self.state == _TxnState.STARTING - - def set_starting(self) -> None: - self.state = _TxnState.STARTING - - def set_in_progress(self) -> None: - if self.state == _TxnState.STARTING: - self.state = _TxnState.IN_PROGRESS - - @property - def pinned_conn(self) -> Optional[Connection]: - if self.active() and self.conn_mgr: - return self.conn_mgr.conn - return None - - def pin(self, server: Server, conn: Connection) -> None: - self.sharded = True - self.pinned_address = server.description.address - if server.description.server_type == SERVER_TYPE.LoadBalancer: - conn.pin_txn() - self.conn_mgr = _ConnectionManager(conn, False) + _conn_mgr_cls = _ConnectionManager + conn_mgr: Optional[_ConnectionManager] def unpin(self) -> None: self.pinned_address = None @@ -289,21 +247,13 @@ def reset(self) -> None: self.attempt = 0 self.has_completed_command = False - def __del__(self) -> None: - if self.conn_mgr: - # Reuse the cursor closing machinery to return the socket to the - # pool soon. - self.client._close_cursor_soon(0, None, self.conn_mgr) - self.conn_mgr = None - _T = TypeVar("_T") -if TYPE_CHECKING: - from pymongo.synchronous.mongo_client import MongoClient - -class ClientSession: +class ClientSession( + _AgnosticClientSessionBase["MongoClient[Any]", "Connection", _BoundSessionContext] +): """A session for ordering sequential operations. :class:`ClientSession` instances are **not thread-safe or fork-safe**. @@ -316,27 +266,10 @@ class ClientSession: :meth:`~pymongo.mongo_client.MongoClient.start_session`. """ - def __init__( - self, - client: MongoClient[Any], - server_session: Any, - options: SessionOptions, - implicit: bool, - ) -> None: - # A MongoClient, a _ServerSession, a SessionOptions, and a set. - self._client: MongoClient[Any] = client - self._server_session = server_session - self._options = options - self._cluster_time: Optional[Mapping[str, Any]] = None - self._operation_time: Optional[Timestamp] = None - self._snapshot_time = None - # Is this an implicitly created session? - self._implicit = implicit - self._transaction = _Transaction(None, client) - # Is this session attached to a cursor? - self._attached_to_cursor = False - # Should we leave the session alive when the cursor is closed? - self._leave_alive = False + _transaction_cls = _Transaction + _bound_session_context_cls = _BoundSessionContext + _transaction: _Transaction + _client: MongoClient[Any] def end_session(self) -> None: """Finish this session. If a transaction has started, abort it. @@ -357,89 +290,12 @@ def _end_session(self, lock: bool) -> None: self._client._return_server_session(self._server_session) self._server_session = None - def _end_implicit_session(self) -> None: - # Implicit sessions can't be part of transactions or pinned connections - if not self._leave_alive and self._server_session is not None: - self._client._return_server_session(self._server_session) - self._server_session = None - - def _check_ended(self) -> None: - if self._server_session is None: - raise InvalidOperation("Cannot use ended session") - - def bind(self, end_session: bool = True) -> _BoundSessionContext: - """Bind this session so it is implicitly passed to all database operations within the returned context. - - .. code-block:: python - - with client.start_session() as s: - with s.bind(): - # session=s is passed implicitly - 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.client_session.ClientSession.end_session()` must be called - once the session is no longer used. - - .. versionadded:: 4.17 - """ - return _BoundSessionContext(self, end_session) - def __enter__(self) -> ClientSession: return self def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: self._end_session(lock=True) - @property - def client(self) -> MongoClient[Any]: - """The :class:`~pymongo.mongo_client.MongoClient` this session was - created from. - """ - return self._client - - @property - def options(self) -> SessionOptions: - """The :class:`SessionOptions` this session was created with.""" - return self._options - - @property - def session_id(self) -> Mapping[str, Any]: - """A BSON document, the opaque server session identifier.""" - self._check_ended() - self._materialize(self._client.topology_description.logical_session_timeout_minutes) - return self._server_session.session_id - - @property - def _transaction_id(self) -> Int64: - """The current transaction id for the underlying server session.""" - self._materialize(self._client.topology_description.logical_session_timeout_minutes) - return self._server_session.transaction_id - - @property - def cluster_time(self) -> Optional[ClusterTime]: - """The cluster time returned by the last operation executed - in this session. - """ - return self._cluster_time - - @property - def operation_time(self) -> Optional[Timestamp]: - """The operation time returned by the last operation executed - in this session. - """ - return self._operation_time - - def _inherit_option(self, name: str, val: _T) -> _T: - """Return the inherited TransactionOption value.""" - if val: - return val - txn_opts = self.options.default_transaction_options - parent_val = txn_opts and getattr(txn_opts, name) - if parent_val: - return parent_val - return getattr(self.client, name) - def with_transaction( self, callback: Callable[[ClientSession], _T], @@ -749,185 +605,6 @@ def _finish_transaction(self, conn: Connection, command_name: str) -> dict[str, conn, cmd, session=self, write_concern=wc, parse_write_concern_error=True ) - def _advance_cluster_time(self, cluster_time: Optional[Mapping[str, Any]]) -> None: - """Internal cluster time helper.""" - if self._cluster_time is None: - self._cluster_time = cluster_time - elif cluster_time is not None: - if cluster_time["clusterTime"] > self._cluster_time["clusterTime"]: - self._cluster_time = cluster_time - - def advance_cluster_time(self, cluster_time: Mapping[str, Any]) -> None: - """Update the cluster time for this session. - - :param cluster_time: The - :data:`~pymongo.client_session.ClientSession.cluster_time` from - another `ClientSession` instance. - """ - if not isinstance(cluster_time, _Mapping): - raise TypeError( - f"cluster_time must be a subclass of collections.Mapping, not {type(cluster_time)}" - ) - if not isinstance(cluster_time.get("clusterTime"), Timestamp): - raise ValueError("Invalid cluster_time") - self._advance_cluster_time(cluster_time) - - def _advance_operation_time(self, operation_time: Optional[Timestamp]) -> None: - """Internal operation time helper.""" - if self._operation_time is None: - self._operation_time = operation_time - elif operation_time is not None: - if operation_time > self._operation_time: - self._operation_time = operation_time - - def advance_operation_time(self, operation_time: Timestamp) -> None: - """Update the operation time for this session. - - :param operation_time: The - :data:`~pymongo.client_session.ClientSession.operation_time` from - another `ClientSession` instance. - """ - if not isinstance(operation_time, Timestamp): - raise TypeError( - f"operation_time must be an instance of bson.timestamp.Timestamp, not {type(operation_time)}" - ) - self._advance_operation_time(operation_time) - - def _process_response(self, reply: Mapping[str, Any]) -> None: - """Process a response to a command that was run with this session.""" - self._advance_cluster_time(reply.get("$clusterTime")) - self._advance_operation_time(reply.get("operationTime")) - if self._options.snapshot and self._snapshot_time is None: - if "cursor" in reply: - ct = reply["cursor"].get("atClusterTime") - else: - ct = reply.get("atClusterTime") - self._snapshot_time = ct - if self.in_transaction and self._transaction.sharded: - recovery_token = reply.get("recoveryToken") - if recovery_token: - self._transaction.recovery_token = recovery_token - - @property - def has_ended(self) -> bool: - """True if this session is finished.""" - return self._server_session is None - - @property - def in_transaction(self) -> bool: - """True if this session has an active multi-statement transaction. - - .. versionadded:: 3.10 - """ - return self._transaction.active() - - @property - def _starting_transaction(self) -> bool: - """True if this session is starting a multi-statement transaction.""" - return self._transaction.starting() - - @property - def _pinned_address(self) -> Optional[_Address]: - """The mongos address this transaction was created on.""" - if self._transaction.active(): - return self._transaction.pinned_address - return None - - @property - def _pinned_connection(self) -> Optional[Connection]: - """The connection this transaction was started on.""" - return self._transaction.pinned_conn - - def _pin(self, server: Server, conn: Connection) -> None: - """Pin this session to the given Server or to the given connection.""" - self._transaction.pin(server, conn) - def _unpin(self) -> None: """Unpin this session from any pinned Server.""" self._transaction.unpin() - - def _txn_read_preference(self) -> Optional[_ServerMode]: - """Return read preference of this transaction or None.""" - if self.in_transaction: - assert self._transaction.opts - return self._transaction.opts.read_preference - return None - - def _materialize(self, logical_session_timeout_minutes: Optional[int] = None) -> None: - if isinstance(self._server_session, _EmptyServerSession): - old = self._server_session - self._server_session = self._client._topology.get_server_session( - logical_session_timeout_minutes - ) - if old.started_retryable_write: - self._server_session.inc_transaction_id() - - def _apply_to( - self, - command: MutableMapping[str, Any], - is_retryable: bool, - read_preference: _ServerMode, - conn: Connection, - ) -> None: - # getMores must be sent with a session if the cursor was opened with one - operation = next(iter(command)) - if not conn.supports_sessions and ( - isinstance(self._server_session, _EmptyServerSession) or operation != "getMore" - ): - if not self._implicit: - raise ConfigurationError("Sessions are not supported by this MongoDB deployment") - return - self._check_ended() - self._materialize(conn.logical_session_timeout_minutes) - # Add afterClusterTime on snapshot reads or writes in causally-consistent sessions - if self.options.snapshot or ( - self.options.causal_consistency - and not self.in_transaction - and operation in _WRITES_WITH_CLUSTER_TIME - ): - self._update_read_concern(command, conn) - - self._server_session.last_use = time.monotonic() - command["lsid"] = self._server_session.session_id - - if is_retryable: - command["txnNumber"] = self._server_session.transaction_id - return - - if self.in_transaction: - if read_preference != ReadPreference.PRIMARY: - raise InvalidOperation( - f"read preference in a transaction must be primary, not: {read_preference!r}" - ) - - if self._transaction.state == _TxnState.STARTING: - # First command begins a new transaction. - command["startTransaction"] = True - - assert self._transaction.opts - if self._transaction.opts.read_concern: - rc = self._transaction.opts.read_concern.document - if rc: - command["readConcern"] = rc - self._update_read_concern(command, conn) - - command["txnNumber"] = self._server_session.transaction_id - command["autocommit"] = False - - def _start_retryable_write(self) -> None: - self._check_ended() - self._server_session.inc_transaction_id() - - def _update_read_concern(self, cmd: MutableMapping[str, Any], conn: Connection) -> None: - if self.options.causal_consistency and self.operation_time is not None: - cmd.setdefault("readConcern", {})["afterClusterTime"] = self.operation_time - if self.options.snapshot: - if conn.max_wire_version < 13: - raise ConfigurationError("Snapshot reads require MongoDB 5.0 or later") - rc = cmd.setdefault("readConcern", {}) - rc["level"] = "snapshot" - if self._snapshot_time is not None: - rc["atClusterTime"] = self._snapshot_time - - def __copy__(self) -> NoReturn: - raise TypeError("A ClientSession cannot be copied, create a new session instead") diff --git a/pymongo/typings.py b/pymongo/typings.py index 0281c57c59..3fcf516fbd 100644 --- a/pymongo/typings.py +++ b/pymongo/typings.py @@ -34,6 +34,7 @@ from pymongo.asynchronous.collection import AsyncCollection from pymongo.asynchronous.mongo_client import AsyncMongoClient from pymongo.asynchronous.pool import AsyncConnection + from pymongo.asynchronous.server import Server as AsyncServer from pymongo.collation import Collation from pymongo.synchronous.bulk import _Bulk from pymongo.synchronous.client_bulk import _ClientBulk @@ -41,6 +42,7 @@ from pymongo.synchronous.collection import Collection from pymongo.synchronous.mongo_client import MongoClient from pymongo.synchronous.pool import Connection + from pymongo.synchronous.server import Server # Common Shared Types. @@ -55,6 +57,7 @@ _AgnosticMongoClient = Union["AsyncMongoClient", "MongoClient"] # type: ignore[type-arg] _AgnosticCollection = Union["AsyncCollection[_DocumentType]", "Collection[_DocumentType]"] _AgnosticConnection = Union["AsyncConnection", "Connection"] +_AgnosticServer = Union["AsyncServer", "Server"] _AgnosticClientSession = Union["AsyncClientSession", "ClientSession"] _AgnosticBulk = Union["_AsyncBulk", "_Bulk"] _AgnosticClientBulk = Union["_AsyncClientBulk", "_ClientBulk"] From 66b0d6b346066fa99e0e18970ed36e014034a64a Mon Sep 17 00:00:00 2001 From: Iris Date: Wed, 7 Oct 2026 09:45:21 -0700 Subject: [PATCH 2/2] fix docs --- .../pymongo/asynchronous/client_session.rst | 1 + doc/api/pymongo/client_session.rst | 1 + pymongo/client_session_shared.py | 19 +++++++++++++++---- 3 files changed, 17 insertions(+), 4 deletions(-) diff --git a/doc/api/pymongo/asynchronous/client_session.rst b/doc/api/pymongo/asynchronous/client_session.rst index d8403325d7..0c784c45d0 100644 --- a/doc/api/pymongo/asynchronous/client_session.rst +++ b/doc/api/pymongo/asynchronous/client_session.rst @@ -4,3 +4,4 @@ .. automodule:: pymongo.asynchronous.client_session :members: + :inherited-members: diff --git a/doc/api/pymongo/client_session.rst b/doc/api/pymongo/client_session.rst index d9527b002b..3c9e067b1d 100644 --- a/doc/api/pymongo/client_session.rst +++ b/doc/api/pymongo/client_session.rst @@ -3,3 +3,4 @@ .. automodule:: pymongo.client_session :members: + :inherited-members: diff --git a/pymongo/client_session_shared.py b/pymongo/client_session_shared.py index d2d086905a..a5b057298d 100644 --- a/pymongo/client_session_shared.py +++ b/pymongo/client_session_shared.py @@ -480,7 +480,9 @@ def bind(self, end_session: bool = True) -> _BoundSessionContextT: 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, ``end_session()`` must be called once the session is no longer used. + If set to False, :meth:`~pymongo.client_session.ClientSession.end_session()` or + :meth:`~pymongo.asynchronous.client_session.AsyncClientSession.end_session()` must be + called once the session is no longer used. .. versionadded:: 4.17 """ @@ -488,7 +490,10 @@ def bind(self, end_session: bool = True) -> _BoundSessionContextT: @property def client(self) -> _ClientT: - """The client this session was created from.""" + """The :class:`~pymongo.mongo_client.MongoClient` or + :class:`~pymongo.asynchronous.mongo_client.AsyncMongoClient` this session was + created from. + """ return self._client @property @@ -544,7 +549,10 @@ def _advance_cluster_time(self, cluster_time: Optional[Mapping[str, Any]]) -> No def advance_cluster_time(self, cluster_time: Mapping[str, Any]) -> None: """Update the cluster time for this session. - :param cluster_time: The ``cluster_time`` from another session. + :param cluster_time: The + :data:`~pymongo.client_session.ClientSession.cluster_time` or + :data:`~pymongo.asynchronous.client_session.AsyncClientSession.cluster_time` from + another session. """ if not isinstance(cluster_time, Mapping): raise TypeError( @@ -565,7 +573,10 @@ def _advance_operation_time(self, operation_time: Optional[Timestamp]) -> None: def advance_operation_time(self, operation_time: Timestamp) -> None: """Update the operation time for this session. - :param operation_time: The ``operation_time`` from another session. + :param operation_time: The + :data:`~pymongo.client_session.ClientSession.operation_time` or + :data:`~pymongo.asynchronous.client_session.AsyncClientSession.operation_time` from + another session. """ if not isinstance(operation_time, Timestamp): raise TypeError(