From a15188f11d91186ba17bc6417a10558b692f5846 Mon Sep 17 00:00:00 2001 From: Noah Stapp Date: Tue, 1 Sep 2026 11:28:33 -0400 Subject: [PATCH 1/2] PYTHON-6019 Consolidate mongo_client module-level helpers into shared file --- pymongo/asynchronous/mongo_client.py | 89 ++------------------- pymongo/mongo_client_shared.py | 111 +++++++++++++++++++++++++++ pymongo/synchronous/mongo_client.py | 89 ++------------------- 3 files changed, 123 insertions(+), 166 deletions(-) create mode 100644 pymongo/mongo_client_shared.py diff --git a/pymongo/asynchronous/mongo_client.py b/pymongo/asynchronous/mongo_client.py index faa4c376f9..364e74eebe 100644 --- a/pymongo/asynchronous/mongo_client.py +++ b/pymongo/asynchronous/mongo_client.py @@ -69,28 +69,27 @@ from pymongo.driver_info import DriverInfo from pymongo.errors import ( AutoReconnect, - BulkWriteError, ClientBulkWriteException, ConfigurationError, ConnectionFailure, InvalidOperation, - NotPrimaryError, OperationFailure, PyMongoError, ServerSelectionTimeoutError, - WaitQueueTimeoutError, ) from pymongo.lock import ( _HAS_REGISTER_AT_FORK, _async_create_lock, - _release_locks, ) from pymongo.logger import ( - _CLIENT_LOGGER, _log_client_error, - _log_or_warn, ) from pymongo.message import _CursorAddress, _GetMore, _Query +from pymongo.mongo_client_shared import ( + _add_retryable_write_error, + _after_fork_child, + _detect_external_db, +) from pymongo.monitoring import ConnectionClosedReason, _EventListeners from pymongo.operations import ( DeleteMany, @@ -2573,45 +2572,6 @@ async def bulk_write( return await blk.execute(session, _Op.BULK_WRITE) -def _retryable_error_doc(exc: PyMongoError) -> Optional[Mapping[str, Any]]: - """Return the server response from PyMongo exception or None.""" - if isinstance(exc, (BulkWriteError, ClientBulkWriteException)): - # Check the last writeConcernError to determine if this - # BulkWriteError is retryable. - wces = exc.details["writeConcernErrors"] - return wces[-1] if wces else None - if isinstance(exc, (NotPrimaryError, OperationFailure)): - return cast(Mapping[str, Any], exc.details) - return None - - -def _add_retryable_write_error(exc: PyMongoError) -> None: - doc = _retryable_error_doc(exc) - if doc: - code = doc.get("code", 0) - # retryWrites on MMAPv1 should raise an actionable error. - if code == 20 and str(exc).startswith("Transaction numbers"): - errmsg = ( - "This MongoDB deployment does not support " - "retryable writes. Please add retryWrites=false " - "to your connection string." - ) - raise OperationFailure(errmsg, code, exc.details) # type: ignore[attr-defined] - for label in doc.get("errorLabels", []): - exc._add_error_label(label) - - # AsyncConnection errors are always retryable except NotPrimaryError and WaitQueueTimeoutError which is - # handled above. - if isinstance(exc, ClientBulkWriteException): - exc_to_check = exc.error - else: - exc_to_check = exc - if isinstance(exc_to_check, ConnectionFailure) and not isinstance( - exc_to_check, (NotPrimaryError, WaitQueueTimeoutError) - ): - exc_to_check._add_error_label("RetryableWriteError") - - class _ClientCheckout: """Context manager for checking out a connection from the pool. @@ -3143,45 +3103,8 @@ async def _read(self) -> T: return await self._func(self._session, self._server, conn, read_pref) # type: ignore -def _after_fork_child() -> None: - """Releases the locks in child process and resets the - topologies in all MongoClients. - """ - # Reinitialize locks - _release_locks() - - # Perform cleanup in clients (i.e. get rid of topology) - for _, client in AsyncMongoClient._clients.items(): - client._after_fork() - - -def _detect_external_db(entity: str) -> bool: - """Detects external database hosts and logs an informational message at the INFO level.""" - entity = entity.lower() - cosmos_db_hosts = [".cosmos.azure.com"] - document_db_hosts = [".docdb.amazonaws.com", ".docdb-elastic.amazonaws.com"] - - for host in cosmos_db_hosts: - if entity.endswith(host): - _log_or_warn( - _CLIENT_LOGGER, - "You appear to be connected to a CosmosDB cluster. For more information regarding feature " - "compatibility and support please visit https://www.mongodb.com/supportability/cosmosdb", - ) - return True - for host in document_db_hosts: - if entity.endswith(host): - _log_or_warn( - _CLIENT_LOGGER, - "You appear to be connected to a DocumentDB cluster. For more information regarding feature " - "compatibility and support please visit https://www.mongodb.com/supportability/documentdb", - ) - return True - return False - - if _HAS_REGISTER_AT_FORK: # This will run in the same thread as the fork was called. # If we fork in a critical region on the same thread, it should break. # This is fine since we would never call fork directly from a critical region. - os.register_at_fork(after_in_child=_after_fork_child) + os.register_at_fork(after_in_child=lambda: _after_fork_child(AsyncMongoClient._clients)) diff --git a/pymongo/mongo_client_shared.py b/pymongo/mongo_client_shared.py new file mode 100644 index 0000000000..e9e9b89357 --- /dev/null +++ b/pymongo/mongo_client_shared.py @@ -0,0 +1,111 @@ +# Copyright 2009-present MongoDB, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you +# may not use this file except in compliance with the License. You +# may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or +# implied. See the License for the specific language governing +# permissions and limitations under the License. + +"""Internal helpers for MongoClient, shared between the asynchronous and synchronous APIs.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, Optional, cast + +from pymongo.errors import ( + BulkWriteError, + ClientBulkWriteException, + ConnectionFailure, + NotPrimaryError, + OperationFailure, + PyMongoError, + WaitQueueTimeoutError, +) +from pymongo.lock import _release_locks +from pymongo.logger import ( + _CLIENT_LOGGER, + _log_or_warn, +) + + +def _retryable_error_doc(exc: PyMongoError) -> Optional[Mapping[str, Any]]: + """Return the server response from PyMongo exception or None.""" + if isinstance(exc, (BulkWriteError, ClientBulkWriteException)): + # Check the last writeConcernError to determine if this + # BulkWriteError is retryable. + wces = exc.details["writeConcernErrors"] + return wces[-1] if wces else None + if isinstance(exc, (NotPrimaryError, OperationFailure)): + return cast(Mapping[str, Any], exc.details) + return None + + +def _add_retryable_write_error(exc: PyMongoError) -> None: + doc = _retryable_error_doc(exc) + if doc: + code = doc.get("code", 0) + # retryWrites on MMAPv1 should raise an actionable error. + if code == 20 and str(exc).startswith("Transaction numbers"): + errmsg = ( + "This MongoDB deployment does not support " + "retryable writes. Please add retryWrites=false " + "to your connection string." + ) + raise OperationFailure(errmsg, code, exc.details) # type: ignore[attr-defined] + for label in doc.get("errorLabels", []): + exc._add_error_label(label) + + # Connection errors are always retryable except NotPrimaryError and WaitQueueTimeoutError which is + # handled above. + if isinstance(exc, ClientBulkWriteException): + exc_to_check = exc.error + else: + exc_to_check = exc + if isinstance(exc_to_check, ConnectionFailure) and not isinstance( + exc_to_check, (NotPrimaryError, WaitQueueTimeoutError) + ): + exc_to_check._add_error_label("RetryableWriteError") + + +def _after_fork_child(clients: Mapping[Any, Any]) -> None: + """Releases the locks in child process and resets the + topologies in all MongoClients. + """ + # Reinitialize locks + _release_locks() + + # Perform cleanup in clients (i.e. get rid of topology) + for _, client in clients.items(): + client._after_fork() + + +def _detect_external_db(entity: str) -> bool: + """Detects external database hosts and logs an informational message at the INFO level.""" + entity = entity.lower() + cosmos_db_hosts = [".cosmos.azure.com"] + document_db_hosts = [".docdb.amazonaws.com", ".docdb-elastic.amazonaws.com"] + + for host in cosmos_db_hosts: + if entity.endswith(host): + _log_or_warn( + _CLIENT_LOGGER, + "You appear to be connected to a CosmosDB cluster. For more information regarding feature " + "compatibility and support please visit https://www.mongodb.com/supportability/cosmosdb", + ) + return True + for host in document_db_hosts: + if entity.endswith(host): + _log_or_warn( + _CLIENT_LOGGER, + "You appear to be connected to a DocumentDB cluster. For more information regarding feature " + "compatibility and support please visit https://www.mongodb.com/supportability/documentdb", + ) + return True + return False diff --git a/pymongo/synchronous/mongo_client.py b/pymongo/synchronous/mongo_client.py index 69e079b1fa..73b3f22fc8 100644 --- a/pymongo/synchronous/mongo_client.py +++ b/pymongo/synchronous/mongo_client.py @@ -62,28 +62,27 @@ from pymongo.driver_info import DriverInfo from pymongo.errors import ( AutoReconnect, - BulkWriteError, ClientBulkWriteException, ConfigurationError, ConnectionFailure, InvalidOperation, - NotPrimaryError, OperationFailure, PyMongoError, ServerSelectionTimeoutError, - WaitQueueTimeoutError, ) from pymongo.lock import ( _HAS_REGISTER_AT_FORK, _create_lock, - _release_locks, ) from pymongo.logger import ( - _CLIENT_LOGGER, _log_client_error, - _log_or_warn, ) from pymongo.message import _CursorAddress, _GetMore, _Query +from pymongo.mongo_client_shared import ( + _add_retryable_write_error, + _after_fork_child, + _detect_external_db, +) from pymongo.monitoring import ConnectionClosedReason, _EventListeners from pymongo.operations import ( DeleteMany, @@ -2562,45 +2561,6 @@ def bulk_write( return blk.execute(session, _Op.BULK_WRITE) -def _retryable_error_doc(exc: PyMongoError) -> Optional[Mapping[str, Any]]: - """Return the server response from PyMongo exception or None.""" - if isinstance(exc, (BulkWriteError, ClientBulkWriteException)): - # Check the last writeConcernError to determine if this - # BulkWriteError is retryable. - wces = exc.details["writeConcernErrors"] - return wces[-1] if wces else None - if isinstance(exc, (NotPrimaryError, OperationFailure)): - return cast(Mapping[str, Any], exc.details) - return None - - -def _add_retryable_write_error(exc: PyMongoError) -> None: - doc = _retryable_error_doc(exc) - if doc: - code = doc.get("code", 0) - # retryWrites on MMAPv1 should raise an actionable error. - if code == 20 and str(exc).startswith("Transaction numbers"): - errmsg = ( - "This MongoDB deployment does not support " - "retryable writes. Please add retryWrites=false " - "to your connection string." - ) - raise OperationFailure(errmsg, code, exc.details) # type: ignore[attr-defined] - for label in doc.get("errorLabels", []): - exc._add_error_label(label) - - # Connection errors are always retryable except NotPrimaryError and WaitQueueTimeoutError which is - # handled above. - if isinstance(exc, ClientBulkWriteException): - exc_to_check = exc.error - else: - exc_to_check = exc - if isinstance(exc_to_check, ConnectionFailure) and not isinstance( - exc_to_check, (NotPrimaryError, WaitQueueTimeoutError) - ): - exc_to_check._add_error_label("RetryableWriteError") - - class _ClientCheckout: """Context manager for checking out a connection from the pool. @@ -3132,45 +3092,8 @@ def _read(self) -> T: return self._func(self._session, self._server, conn, read_pref) # type: ignore -def _after_fork_child() -> None: - """Releases the locks in child process and resets the - topologies in all MongoClients. - """ - # Reinitialize locks - _release_locks() - - # Perform cleanup in clients (i.e. get rid of topology) - for _, client in MongoClient._clients.items(): - client._after_fork() - - -def _detect_external_db(entity: str) -> bool: - """Detects external database hosts and logs an informational message at the INFO level.""" - entity = entity.lower() - cosmos_db_hosts = [".cosmos.azure.com"] - document_db_hosts = [".docdb.amazonaws.com", ".docdb-elastic.amazonaws.com"] - - for host in cosmos_db_hosts: - if entity.endswith(host): - _log_or_warn( - _CLIENT_LOGGER, - "You appear to be connected to a CosmosDB cluster. For more information regarding feature " - "compatibility and support please visit https://www.mongodb.com/supportability/cosmosdb", - ) - return True - for host in document_db_hosts: - if entity.endswith(host): - _log_or_warn( - _CLIENT_LOGGER, - "You appear to be connected to a DocumentDB cluster. For more information regarding feature " - "compatibility and support please visit https://www.mongodb.com/supportability/documentdb", - ) - return True - return False - - if _HAS_REGISTER_AT_FORK: # This will run in the same thread as the fork was called. # If we fork in a critical region on the same thread, it should break. # This is fine since we would never call fork directly from a critical region. - os.register_at_fork(after_in_child=_after_fork_child) + os.register_at_fork(after_in_child=lambda: _after_fork_child(MongoClient._clients)) From a286c91fab539d9ae56fa05a3c8519080335903d Mon Sep 17 00:00:00 2001 From: Noah Stapp Date: Tue, 1 Sep 2026 12:23:19 -0400 Subject: [PATCH 2/2] IH review --- pymongo/mongo_client_shared.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymongo/mongo_client_shared.py b/pymongo/mongo_client_shared.py index e9e9b89357..29a3db6e7e 100644 --- a/pymongo/mongo_client_shared.py +++ b/pymongo/mongo_client_shared.py @@ -76,7 +76,7 @@ def _add_retryable_write_error(exc: PyMongoError) -> None: def _after_fork_child(clients: Mapping[Any, Any]) -> None: """Releases the locks in child process and resets the - topologies in all MongoClients. + topologies in the passed clients. """ # Reinitialize locks _release_locks()