PYTHON-4700 - Convert CSFLE tests to async

.evergreen/run-tests.sh

@@ -257,9 +257,9 @@ if [ -z "$GREEN_FRAMEWORK" ]; then
     # Use --capture=tee-sys so pytest prints test output inline:
     # https://docs.pytest.org/en/stable/how-to/capture-stdout-stderr.html
     if [ -z "$TEST_SUITES" ]; then
-      python -m pytest -v --capture=tee-sys --durations=5 --maxfail=10 $TEST_ARGS
+      python -m pytest -v --capture=tee-sys --durations=5 $TEST_ARGS
     else
-      python -m pytest -v --capture=tee-sys --durations=5 --maxfail=10 -m $TEST_SUITES $TEST_ARGS
+      python -m pytest -v --capture=tee-sys --durations=5 -m $TEST_SUITES $TEST_ARGS
     fi
 else
     python green_framework_test.py $GREEN_FRAMEWORK -v $TEST_ARGS

pymongo/asynchronous/collection.py

@@ -816,6 +816,7 @@ async def _insert_command(
         ) -> None:
             if bypass_doc_val:
                 command["bypassDocumentValidation"] = True
+            print("Running _insert_command")
 
             result = await conn.command(
                 self._database.name,
@@ -886,6 +887,7 @@ async def insert_one(
         common.validate_is_document_type("document", document)
         if not (isinstance(document, RawBSONDocument) or "_id" in document):
             document["_id"] = ObjectId()  # type: ignore[index]
+        print("Running insert_one")
 
         write_concern = self._write_concern_for(session)
         return InsertOneResult(

pymongo/asynchronous/database.py

@@ -1082,6 +1082,7 @@ async def _list_collections(
             AsyncCollection[MutableMapping[str, Any]],
             self.get_collection("$cmd", read_preference=read_preference),
         )
+        print("Calling _list_collections")
         cmd = {"listCollections": 1, "cursor": {}}
         cmd.update(kwargs)
         async with self._client._tmp_session(session, close=False) as tmp_session:
@@ -1125,13 +1126,14 @@ async def _list_collections_helper(
 
         .. versionadded:: 3.6
         """
+        print("Calling _list_collections_helper")
         if filter is not None:
             kwargs["filter"] = filter
         read_pref = (session and session._txn_read_preference()) or ReadPreference.PRIMARY
         if comment is not None:
             kwargs["comment"] = comment
 
-        async def _cmd(
+        async def _cmd_list_collections(
             session: Optional[AsyncClientSession],
             _server: Server,
             conn: AsyncConnection,
@@ -1141,8 +1143,9 @@ async def _cmd(
                 conn, session, read_preference=read_preference, **kwargs
             )
 
+        print("Calling _list_collections_helper _cmd")
         return await self._client._retryable_read(
-            _cmd, read_pref, session, operation=_Op.LIST_COLLECTIONS
+            _cmd_list_collections, read_pref, session, operation=_Op.LIST_COLLECTIONS
         )
 
     async def list_collections(
@@ -1171,6 +1174,7 @@ async def list_collections(
 
         .. versionadded:: 3.6
         """
+        print("Calling list_collections")
         return await self._list_collections_helper(session, filter, comment, **kwargs)
 
     async def _list_collection_names(

pymongo/asynchronous/encryption.py

@@ -18,6 +18,7 @@
 import contextlib
 import enum
 import socket
+import traceback
 import uuid
 import weakref
 from copy import deepcopy
@@ -142,6 +143,37 @@ def __init__(
         self.opts = opts
         self._spawned = False
 
+    async def _async_kms_request(
+        self,
+        kms_context: MongoCryptKmsContext,
+        host: str,
+        port: Optional[int],
+        opts: PoolOptions,
+        message: bytes,
+    ) -> None:
+        from pymongo.network_layer import async_receive_data_socket  # type: ignore[attr-defined]
+
+        try:
+            conn = await _configured_socket((host, port), opts)
+            try:
+                await async_sendall(conn, message)
+                while kms_context.bytes_needed > 0:
+                    # CSOT: update timeout.
+                    conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
+                    data = await async_receive_data_socket(conn, kms_context.bytes_needed)
+                    kms_context.feed(data)
+            except OSError as err:
+                raise OSError("KMS connection closed") from err
+            except BLOCKING_IO_ERRORS:
+                raise socket.timeout("timed out") from None
+            finally:
+                conn.close()
+        except (PyMongoError, MongoCryptError):
+            raise  # Propagate pymongo errors directly.
+        except Exception as error:
+            # Wrap I/O errors in PyMongo exceptions.
+            _raise_connection_failure((host, port), error)
+
     async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
         """Complete a KMS request.
 
@@ -174,20 +206,23 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
         )
         host, port = parse_host(endpoint, _HTTPS_PORT)
         try:
-            conn = await _configured_socket((host, port), opts)
-            try:
-                await async_sendall(conn, message)
-                while kms_context.bytes_needed > 0:
-                    # CSOT: update timeout.
-                    conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
-                    data = conn.recv(kms_context.bytes_needed)
-                    if not data:
-                        raise OSError("KMS connection closed")
-                    kms_context.feed(data)
-            except BLOCKING_IO_ERRORS:
-                raise socket.timeout("timed out") from None
-            finally:
-                conn.close()
+            if _IS_SYNC:
+                conn = await _configured_socket((host, port), opts)
+                try:
+                    await async_sendall(conn, message)
+                    while kms_context.bytes_needed > 0:
+                        # CSOT: update timeout.
+                        conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
+                        data = conn.recv(kms_context.bytes_needed)
+                        if not data:
+                            raise OSError("KMS connection closed")
+                        kms_context.feed(data)
+                except BLOCKING_IO_ERRORS:
+                    raise socket.timeout("timed out") from None
+                finally:
+                    conn.close()
+            else:
+                await self._async_kms_request(kms_context, host, port, opts, message)
         except (PyMongoError, MongoCryptError):
             raise  # Propagate pymongo errors directly.
         except Exception as error:

pymongo/asynchronous/mongo_client.py

@@ -34,6 +34,8 @@
 
 import contextlib
 import os
+import threading
+import traceback
 import warnings
 import weakref
 from collections import defaultdict
@@ -2565,7 +2567,12 @@ async def run(self) -> T:
         while True:
             self._check_last_error(check_csot=True)
             try:
-                return await self._read() if self._is_read else await self._write()
+                if self._operation in (_Op.LIST_COLLECTIONS, _Op.INSERT):
+                    print(f"{threading.current_thread().name} -- Retryable read for {self._operation}")
+                res = await self._read() if self._is_read else await self._write()
+                if self._operation in (_Op.LIST_COLLECTIONS, _Op.INSERT):
+                    print(f"Result for {self._operation}: {res}")
+                return res
             except ServerSelectionTimeoutError:
                 # The application may think the write was never attempted
                 # if we raise ServerSelectionTimeoutError on the retry
@@ -2574,8 +2581,12 @@ async def run(self) -> T:
                 # A ServerSelectionTimeoutError error indicates that there may
                 # be a persistent outage. Attempting to retry in this case will
                 # most likely be a waste of time.
+                if self._operation in (_Op.LIST_COLLECTIONS, _Op.INSERT):
+                    print(f"Raising error for {self._operation} at {traceback.print_stack(limit=1)}")
                 raise
             except PyMongoError as exc:
+                if self._operation in (_Op.LIST_COLLECTIONS, _Op.INSERT):
+                    print(f"Error for {self._operation}: {exc}, {traceback.print_exception(exc)}")
                 # Execute specialized catch on read
                 if self._is_read:
                     if isinstance(exc, (ConnectionFailure, OperationFailure)):
@@ -2585,15 +2596,21 @@ async def run(self) -> T:
                             isinstance(exc, OperationFailure)
                             and exc_code not in helpers_shared._RETRYABLE_ERROR_CODES
                         ):
+                            if self._operation in (_Op.LIST_COLLECTIONS, _Op.INSERT):
+                                print(f"Raising error for {self._operation} at {traceback.print_stack(limit=1)}")
                             raise
                         self._retrying = True
                         self._last_error = exc
                     else:
+                        if self._operation in (_Op.LIST_COLLECTIONS, _Op.INSERT):
+                            print(f"Raising error for {self._operation} at {traceback.print_stack(limit=1)}")
                         raise
 
                 # Specialized catch on write operation
                 if not self._is_read:
                     if not self._retryable:
+                        if self._operation in (_Op.LIST_COLLECTIONS, _Op.INSERT):
+                            print(f"Raising error for {self._operation} at {traceback.print_stack(limit=1)}")
                         raise
                     if isinstance(exc, ClientBulkWriteException) and exc.error:
                         retryable_write_error_exc = isinstance(
@@ -2606,8 +2623,12 @@ async def run(self) -> T:
                         await self._session._unpin()
                     if not retryable_write_error_exc or self._is_not_eligible_for_retry():
                         if exc.has_error_label("NoWritesPerformed") and self._last_error:
+                            if self._operation in (_Op.LIST_COLLECTIONS, _Op.INSERT):
+                                print(f"Raising error for {self._operation} at {traceback.print_stack(limit=1)}")
                             raise self._last_error from exc
                         else:
+                            if self._operation in (_Op.LIST_COLLECTIONS, _Op.INSERT):
+                                print(f"Raising error for {self._operation} at {traceback.print_stack(limit=1)}")
                             raise
                     if self._bulk:
                         self._bulk.retrying = True
@@ -2670,10 +2691,12 @@ async def _write(self) -> T:
 
         :return: Output for func()'s call
         """
+        print(f"Calling _write for {self._func.__name__}")
         try:
             max_wire_version = 0
             is_mongos = False
             self._server = await self._get_server()
+            print(f"Got server for _write for {self._func.__name__}")
             async with self._client._checkout(self._server, self._session) as conn:
                 max_wire_version = conn.max_wire_version
                 sessions_supported = (
@@ -2689,6 +2712,7 @@ async def _write(self) -> T:
                     self._retryable = False
                 return await self._func(self._session, conn, self._retryable)  # type: ignore
         except PyMongoError as exc:
+            print(f"Got error for _write for {self._func.__name__}")
             if not self._retryable:
                 raise
             # Add the RetryableWriteError label, if applicable.
@@ -2700,12 +2724,15 @@ async def _read(self) -> T:
 
         :return: Output for func()'s call
         """
+        print(f"Calling _read for {self._func.__name__}")
         self._server = await self._get_server()
         assert self._read_pref is not None, "Read Preference required on read calls"
+        print(f"Got server for _read for {self._func.__name__}")
         async with self._client._conn_from_server(self._read_pref, self._server, self._session) as (
             conn,
             read_pref,
         ):
+            print(f"Got conn for _read for {self._func.__name__}")
             if self._retrying and not self._retryable:
                 self._check_last_error()
             return await self._func(self._session, self._server, conn, read_pref)  # type: ignore

pymongo/asynchronous/network.py

@@ -17,6 +17,7 @@
 
 import datetime
 import logging
+import threading
 import time
 from typing import (
     TYPE_CHECKING,
@@ -117,6 +118,8 @@ async def command(
     name = next(iter(spec))
     ns = dbname + ".$cmd"
     speculative_hello = False
+    if name == "insertOne":
+        print(f"Running insertOne with {spec}")
 
     # Publish the original command document, perhaps with lsid and $clusterTime.
     orig = spec
@@ -132,6 +135,8 @@ async def command(
         spec["collation"] = collation
 
     publish = listeners is not None and listeners.enabled_for_commands
+    if name == "listCollections":
+        print(f"{threading.current_thread().name} -- Publishing is {publish} listeners: {listeners} {listeners.enabled_for_commands if listeners is not None else 'NONE'}!")
     start = datetime.datetime.now()
     if publish:
         speculative_hello = _is_speculative_authenticate(name, spec)

pymongo/asynchronous/pool.py

@@ -838,10 +838,12 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket
             sock.settimeout(timeout)
             sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, True)
             _set_keepalive_times(sock)
+            print(f"Socket: {sock}, {sock.timeout}")
             sock.connect(sa)
             return sock
         except OSError as e:
             err = e
+            print(f"Socket err: {e}")
             sock.close()
 
     if err is not None:

pymongo/network_layer.py

@@ -130,7 +130,7 @@ def _is_ready(fut: Future) -> None:
                         loop.remove_writer(fd)
 
     async def _async_receive_ssl(
-        conn: _sslConn, length: int, loop: AbstractEventLoop
+        conn: _sslConn, length: int, loop: AbstractEventLoop, once: Optional[bool] = False
     ) -> memoryview:
         mv = memoryview(bytearray(length))
         total_read = 0
@@ -145,6 +145,8 @@ def _is_ready(fut: Future) -> None:
                 read = conn.recv_into(mv[total_read:])
                 if read == 0:
                     raise OSError("connection closed")
+                if once:
+                    return mv[:read]
                 total_read += read
             except BLOCKING_IO_ERRORS as exc:
                 fd = conn.fileno()
@@ -275,6 +277,28 @@ async def async_receive_data(
         sock.settimeout(sock_timeout)
 
 
+async def async_receive_data_socket(
+    sock: Union[socket.socket, _sslConn], length: int
+) -> memoryview:
+    sock_timeout = sock.gettimeout()
+    timeout = sock_timeout
+
+    sock.settimeout(0.0)
+    loop = asyncio.get_event_loop()
+    try:
+        if _HAVE_SSL and isinstance(sock, (SSLSocket, _sslConn)):
+            return await asyncio.wait_for(
+                _async_receive_ssl(sock, length, loop, once=True),  # type: ignore[arg-type]
+                timeout=timeout,
+            )
+        else:
+            return await asyncio.wait_for(_async_receive(sock, length, loop), timeout=timeout)  # type: ignore[arg-type]
+    except asyncio.TimeoutError as err:
+        raise socket.timeout("timed out") from err
+    finally:
+        sock.settimeout(sock_timeout)
+
+
 async def _async_receive(conn: socket.socket, length: int, loop: AbstractEventLoop) -> memoryview:
     mv = memoryview(bytearray(length))
     bytes_read = 0