Merge branch 'master' of github.com:mongodb/mongo-python-driver into add-justfile-jan

Author: blink1073

bson/_cbsonmodule.c

@@ -1644,6 +1644,56 @@ static int write_raw_doc(buffer_t buffer, PyObject* raw, PyObject* _raw_str) {
     return bytes_written;
 }
 
+
+/* Update Invalid Document error message to include doc.
+ */
+void handle_invalid_doc_error(PyObject* dict) {
+    PyObject *etype = NULL, *evalue = NULL, *etrace = NULL;
+    PyObject *msg = NULL, *dict_str = NULL, *new_msg = NULL;
+    PyErr_Fetch(&etype, &evalue, &etrace);
+    PyObject *InvalidDocument = _error("InvalidDocument");
+    if (InvalidDocument == NULL) {
+        goto cleanup;
+    }
+
+    if (evalue && PyErr_GivenExceptionMatches(etype, InvalidDocument)) {
+        PyObject *msg = PyObject_Str(evalue);
+        if (msg) {
+            // Prepend doc to the existing message
+            PyObject *dict_str = PyObject_Str(dict);
+            if (dict_str == NULL) {
+                goto cleanup;
+            }
+            const char * dict_str_utf8 = PyUnicode_AsUTF8(dict_str);
+            if (dict_str_utf8 == NULL) {
+                goto cleanup;
+            }
+            const char * msg_utf8 = PyUnicode_AsUTF8(msg);
+            if (msg_utf8 == NULL) {
+                goto cleanup;
+            }
+            PyObject *new_msg = PyUnicode_FromFormat("Invalid document %s | %s", dict_str_utf8, msg_utf8);
+            Py_DECREF(evalue);
+            Py_DECREF(etype);
+            etype = InvalidDocument;
+            InvalidDocument = NULL;
+            if (new_msg) {
+                evalue = new_msg;
+            } else {
+                evalue = msg;
+            }
+        }
+        PyErr_NormalizeException(&etype, &evalue, &etrace);
+    }
+cleanup:
+    PyErr_Restore(etype, evalue, etrace);
+    Py_XDECREF(msg);
+    Py_XDECREF(InvalidDocument);
+    Py_XDECREF(dict_str);
+    Py_XDECREF(new_msg);
+}
+
+
 /* returns the number of bytes written or 0 on failure */
 int write_dict(PyObject* self, buffer_t buffer,
                PyObject* dict, unsigned char check_keys,
@@ -1743,40 +1793,8 @@ int write_dict(PyObject* self, buffer_t buffer,
         while (PyDict_Next(dict, &pos, &key, &value)) {
             if (!decode_and_write_pair(self, buffer, key, value,
                                     check_keys, options, top_level)) {
-                if (PyErr_Occurred()) {
-                    PyObject *etype = NULL, *evalue = NULL, *etrace = NULL;
-                    PyErr_Fetch(&etype, &evalue, &etrace);
-                    PyObject *InvalidDocument = _error("InvalidDocument");
-
-                    if (top_level && InvalidDocument && PyErr_GivenExceptionMatches(etype, InvalidDocument)) {
-
-                        Py_DECREF(etype);
-                        etype = InvalidDocument;
-
-                        if (evalue) {
-                            PyObject *msg = PyObject_Str(evalue);
-                            Py_DECREF(evalue);
-
-                            if (msg) {
-                                // Prepend doc to the existing message
-                                PyObject *dict_str = PyObject_Str(dict);
-                                PyObject *new_msg = PyUnicode_FromFormat("Invalid document %s | %s", PyUnicode_AsUTF8(dict_str), PyUnicode_AsUTF8(msg));
-                                Py_DECREF(dict_str);
-
-                                if (new_msg) {
-                                    evalue = new_msg;
-                                }
-                                else {
-                                evalue = msg;
-                                }
-                            }
-                        }
-                        PyErr_NormalizeException(&etype, &evalue, &etrace);
-                    }
-                    else {
-                        Py_DECREF(InvalidDocument);
-                    }
-                    PyErr_Restore(etype, evalue, etrace);
+                if (PyErr_Occurred() && top_level) {
+                    handle_invalid_doc_error(dict);
                 }
                 return 0;
             }
@@ -1796,6 +1814,9 @@ int write_dict(PyObject* self, buffer_t buffer,
             }
             if (!decode_and_write_pair(self, buffer, key, value,
                                     check_keys, options, top_level)) {
+                if (PyErr_Occurred() && top_level) {
+                    handle_invalid_doc_error(dict);
+                }
                 Py_DECREF(key);
                 Py_DECREF(value);
                 Py_DECREF(iter);

pymongo/asynchronous/auth.py

@@ -38,6 +38,7 @@
     _authenticate_oidc,
     _get_authenticator,
 )
+from pymongo.asynchronous.helpers import _getaddrinfo
 from pymongo.auth_shared import (
     MongoCredential,
     _authenticate_scram_start,
@@ -177,15 +178,22 @@ def _auth_key(nonce: str, username: str, password: str) -> str:
     return md5hash.hexdigest()
 
 
-def _canonicalize_hostname(hostname: str, option: str | bool) -> str:
+async def _canonicalize_hostname(hostname: str, option: str | bool) -> str:
     """Canonicalize hostname following MIT-krb5 behavior."""
     # https://github.com/krb5/krb5/blob/d406afa363554097ac48646a29249c04f498c88e/src/util/k5test.py#L505-L520
     if option in [False, "none"]:
         return hostname
 
-    af, socktype, proto, canonname, sockaddr = socket.getaddrinfo(
-        hostname, None, 0, 0, socket.IPPROTO_TCP, socket.AI_CANONNAME
-    )[0]
+    af, socktype, proto, canonname, sockaddr = (
+        await _getaddrinfo(
+            hostname,
+            None,
+            family=0,
+            type=0,
+            proto=socket.IPPROTO_TCP,
+            flags=socket.AI_CANONNAME,
+        )
+    )[0]  # type: ignore[index]
 
     # For forward just to resolve the cname as dns.lookup() will not return it.
     if option == "forward":
@@ -213,7 +221,7 @@ async def _authenticate_gssapi(credentials: MongoCredential, conn: AsyncConnecti
         # Starting here and continuing through the while loop below - establish
         # the security context. See RFC 4752, Section 3.1, first paragraph.
         host = props.service_host or conn.address[0]
-        host = _canonicalize_hostname(host, props.canonicalize_host_name)
+        host = await _canonicalize_hostname(host, props.canonicalize_host_name)
         service = props.service_name + "@" + host
         if props.service_realm is not None:
             service = service + "@" + props.service_realm

pymongo/asynchronous/encryption.py

@@ -219,7 +219,14 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                 # Wrap I/O errors in PyMongo exceptions.
                 if isinstance(exc, BLOCKING_IO_ERRORS):
                     exc = socket.timeout("timed out")
-                _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
+                # Async raises an OSError instead of returning empty bytes.
+                if isinstance(exc, OSError):
+                    msg_prefix = "KMS connection closed"
+                else:
+                    msg_prefix = None
+                _raise_connection_failure(
+                    address, exc, msg_prefix=msg_prefix, timeout_details=_get_timeout_details(opts)
+                )
             finally:
                 conn.close()
         except MongoCryptError:

pymongo/asynchronous/helpers.py

@@ -15,7 +15,9 @@
 """Miscellaneous pieces that need to be synchronized."""
 from __future__ import annotations
 
+import asyncio
 import builtins
+import socket
 import sys
 from typing import (
     Any,
@@ -68,6 +70,24 @@ async def inner(*args: Any, **kwargs: Any) -> Any:
     return cast(F, inner)
 
 
+async def _getaddrinfo(
+    host: Any, port: Any, **kwargs: Any
+) -> list[
+    tuple[
+        socket.AddressFamily,
+        socket.SocketKind,
+        int,
+        str,
+        tuple[str, int] | tuple[str, int, int, int],
+    ]
+]:
+    if not _IS_SYNC:
+        loop = asyncio.get_running_loop()
+        return await loop.getaddrinfo(host, port, **kwargs)  # type: ignore[return-value]
+    else:
+        return socket.getaddrinfo(host, port, **kwargs)
+
+
 if sys.version_info >= (3, 10):
     anext = builtins.anext
     aiter = builtins.aiter

pymongo/asynchronous/pool.py

@@ -40,7 +40,7 @@
 from bson import DEFAULT_CODEC_OPTIONS
 from pymongo import _csot, helpers_shared
 from pymongo.asynchronous.client_session import _validate_session_write_concern
-from pymongo.asynchronous.helpers import _handle_reauth
+from pymongo.asynchronous.helpers import _getaddrinfo, _handle_reauth
 from pymongo.asynchronous.network import command, receive_message
 from pymongo.common import (
     MAX_BSON_SIZE,
@@ -783,7 +783,7 @@ def __repr__(self) -> str:
         )
 
 
-def _create_connection(address: _Address, options: PoolOptions) -> socket.socket:
+async def _create_connection(address: _Address, options: PoolOptions) -> socket.socket:
     """Given (host, port) and PoolOptions, connect and return a socket object.
 
     Can raise socket.error.
@@ -814,7 +814,7 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket
         family = socket.AF_UNSPEC
 
     err = None
-    for res in socket.getaddrinfo(host, port, family, socket.SOCK_STREAM):
+    for res in await _getaddrinfo(host, port, family=family, type=socket.SOCK_STREAM):  # type: ignore[attr-defined]
         af, socktype, proto, dummy, sa = res
         # SOCK_CLOEXEC was new in CPython 3.2, and only available on a limited
         # number of platforms (newer Linux and *BSD). Starting with CPython 3.4
@@ -863,7 +863,7 @@ async def _configured_socket(
 
     Sets socket's SSL and timeout options.
     """
-    sock = _create_connection(address, options)
+    sock = await _create_connection(address, options)
     ssl_context = options._ssl_context
 
     if ssl_context is None:

pymongo/message.py

@@ -24,7 +24,6 @@
 import datetime
 import random
 import struct
-from collections import ChainMap
 from io import BytesIO as _BytesIO
 from typing import (
     TYPE_CHECKING,
@@ -1116,18 +1115,8 @@ def _check_doc_size_limits(
         # key and the index of its namespace within ns_info as its value.
         op_doc[op_type] = ns_info[namespace]  # type: ignore[index]
 
-        # Since the data document itself is nested within the insert document
-        # it won't be automatically re-ordered by the BSON conversion.
-        # We use ChainMap here to make the _id field the first field instead.
-        doc_to_encode = op_doc
-        if real_op_type == "insert":
-            doc = op_doc["document"]
-            if not isinstance(doc, RawBSONDocument):
-                doc_to_encode = op_doc.copy()  # type: ignore[attr-defined] # Shallow copy
-                doc_to_encode["document"] = ChainMap(doc, {"_id": doc["_id"]})  # type: ignore[index]
-
         # Encode current operation doc and, if newly added, namespace doc.
-        op_doc_encoded = _dict_to_bson(doc_to_encode, False, opts)
+        op_doc_encoded = _dict_to_bson(op_doc, False, opts)
         op_length = len(op_doc_encoded)
         if ns_doc:
             ns_doc_encoded = _dict_to_bson(ns_doc, False, opts)

pymongo/network_layer.py

@@ -72,7 +72,7 @@
 async def async_sendall(sock: Union[socket.socket, _sslConn], buf: bytes) -> None:
     timeout = sock.gettimeout()
     sock.settimeout(0.0)
-    loop = asyncio.get_event_loop()
+    loop = asyncio.get_running_loop()
     try:
         if _HAVE_SSL and isinstance(sock, (SSLSocket, _sslConn)):
             await asyncio.wait_for(_async_sendall_ssl(sock, buf, loop), timeout=timeout)
@@ -259,7 +259,7 @@ async def async_receive_data(
         timeout = sock_timeout
 
     sock.settimeout(0.0)
-    loop = asyncio.get_event_loop()
+    loop = asyncio.get_running_loop()
     cancellation_task = create_task(_poll_cancellation(conn))
     try:
         if _HAVE_SSL and isinstance(sock, (SSLSocket, _sslConn)):
@@ -290,7 +290,7 @@ async def async_receive_data_socket(
     timeout = sock_timeout
 
     sock.settimeout(0.0)
-    loop = asyncio.get_event_loop()
+    loop = asyncio.get_running_loop()
     try:
         if _HAVE_SSL and isinstance(sock, (SSLSocket, _sslConn)):
             return await asyncio.wait_for(

pymongo/pyopenssl_context.py

@@ -273,7 +273,7 @@ def __set_check_ocsp_endpoint(self, value: bool) -> None:
 
     check_ocsp_endpoint = property(__get_check_ocsp_endpoint, __set_check_ocsp_endpoint)
 
-    def __get_options(self) -> None:
+    def __get_options(self) -> int:
         # Calling set_options adds the option to the existing bitmask and
         # returns the new bitmask.
         # https://www.pyopenssl.org/en/stable/api/ssl.html#OpenSSL.SSL.Context.set_options

pymongo/synchronous/auth.py

@@ -45,6 +45,7 @@
     _authenticate_oidc,
     _get_authenticator,
 )
+from pymongo.synchronous.helpers import _getaddrinfo
 
 if TYPE_CHECKING:
     from pymongo.hello import Hello
@@ -180,9 +181,16 @@ def _canonicalize_hostname(hostname: str, option: str | bool) -> str:
     if option in [False, "none"]:
         return hostname
 
-    af, socktype, proto, canonname, sockaddr = socket.getaddrinfo(
-        hostname, None, 0, 0, socket.IPPROTO_TCP, socket.AI_CANONNAME
-    )[0]
+    af, socktype, proto, canonname, sockaddr = (
+        _getaddrinfo(
+            hostname,
+            None,
+            family=0,
+            type=0,
+            proto=socket.IPPROTO_TCP,
+            flags=socket.AI_CANONNAME,
+        )
+    )[0]  # type: ignore[index]
 
     # For forward just to resolve the cname as dns.lookup() will not return it.
     if option == "forward":

pymongo/synchronous/encryption.py

@@ -219,7 +219,14 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                 # Wrap I/O errors in PyMongo exceptions.
                 if isinstance(exc, BLOCKING_IO_ERRORS):
                     exc = socket.timeout("timed out")
-                _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
+                # Async raises an OSError instead of returning empty bytes.
+                if isinstance(exc, OSError):
+                    msg_prefix = "KMS connection closed"
+                else:
+                    msg_prefix = None
+                _raise_connection_failure(
+                    address, exc, msg_prefix=msg_prefix, timeout_details=_get_timeout_details(opts)
+                )
             finally:
                 conn.close()
         except MongoCryptError: