PYTHON-6019 Consolidate mongo_client module-level helpers into shared…

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))

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 the passed clients.
+    """
+    # 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

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))