Merge branch 'master' of github.com:mongodb/mongo-python-driver

Author: blink1073

.evergreen/run-tests.sh

@@ -30,6 +30,7 @@ set -o xtrace
 
 AUTH=${AUTH:-noauth}
 SSL=${SSL:-nossl}
+TEST_SUITES="test/ test/asynchronous/"
 TEST_ARGS="${*:1}"
 
 export PIP_QUIET=1  # Quiet by default
@@ -95,7 +96,7 @@ if [ -n "$TEST_LOADBALANCER" ]; then
     export LOAD_BALANCER=1
     export SINGLE_MONGOS_LB_URI="${SINGLE_MONGOS_LB_URI:-mongodb://127.0.0.1:8000/?loadBalanced=true}"
     export MULTI_MONGOS_LB_URI="${MULTI_MONGOS_LB_URI:-mongodb://127.0.0.1:8001/?loadBalanced=true}"
-    export TEST_ARGS="test/test_load_balancer.py"
+    export TEST_SUITES="test/test_load_balancer.py"
 fi
 
 if [ "$SSL" != "nossl" ]; then
@@ -171,9 +172,7 @@ if [ -n "$TEST_ENCRYPTION" ]; then
         export PATH=$CRYPT_SHARED_DIR:$PATH
     fi
     # Only run the encryption tests.
-    if [ -z "$TEST_ARGS" ]; then
-        TEST_ARGS="test/test_encryption.py"
-    fi
+    TEST_SUITES="test/test_encryption.py"
 fi
 
 if [ -n "$TEST_FLE_AZURE_AUTO" ] || [ -n "$TEST_FLE_GCP_AUTO" ]; then
@@ -187,9 +186,7 @@ if [ -n "$TEST_FLE_AZURE_AUTO" ] || [ -n "$TEST_FLE_GCP_AUTO" ]; then
       exit 1
     fi
 
-    if [ -z "$TEST_ARGS" ]; then
-        TEST_ARGS="test/test_on_demand_csfle.py"
-    fi
+    TEST_SUITES="test/test_on_demand_csfle.py"
 fi
 
 if [ -n "$TEST_INDEX_MANAGEMENT" ]; then
@@ -198,36 +195,36 @@ if [ -n "$TEST_INDEX_MANAGEMENT" ]; then
     set +x
     export DB_PASSWORD="${DRIVERS_ATLAS_LAMBDA_PASSWORD}"
     set -x
-    TEST_ARGS="test/test_index_management.py"
+    TEST_SUITES="test/test_index_management.py"
 fi
 
 if [ -n "$TEST_DATA_LAKE" ] && [ -z "$TEST_ARGS" ]; then
-    TEST_ARGS="test/test_data_lake.py"
+    TEST_SUITES="test/test_data_lake.py"
 fi
 
 if [ -n "$TEST_ATLAS" ]; then
-    TEST_ARGS="test/atlas/test_connection.py"
+    TEST_SUITES="test/atlas/test_connection.py"
 fi
 
 if [ -n "$TEST_OCSP" ]; then
     python -m pip install ".[ocsp]"
-    TEST_ARGS="test/ocsp/test_ocsp.py"
+    TEST_SUITES="test/ocsp/test_ocsp.py"
 fi
 
 if [ -n "$TEST_AUTH_AWS" ]; then
     python -m pip install ".[aws]"
-    TEST_ARGS="test/auth_aws/test_auth_aws.py"
+    TEST_SUITES="test/auth_aws/test_auth_aws.py"
 fi
 
 if [ -n "$TEST_AUTH_OIDC" ]; then
     python -m pip install ".[aws]"
-    TEST_ARGS="test/auth_oidc/test_auth_oidc.py $TEST_ARGS"
+    TEST_SUITES="test/auth_oidc/test_auth_oidc.py $TEST_ARGS"
 fi
 
 if [ -n "$PERF_TEST" ]; then
     python -m pip install simplejson
     start_time=$(date +%s)
-    TEST_ARGS="test/performance/perf_test.py"
+    TEST_SUITES="test/performance/perf_test.py"
 fi
 
 echo "Running $AUTH tests over $SSL with python $(which python)"
@@ -257,8 +254,7 @@ PIP_QUIET=0 python -m pip list
 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
-    python -m pytest -v --capture=tee-sys --durations=5 --maxfail=10 $TEST_ARGS
-    python -m pytest -v --capture=tee-sys --durations=5 --maxfail=10 test/asynchronous/ $TEST_ARGS
+    python -m pytest -v --capture=tee-sys --durations=5 --maxfail=10 $TEST_SUITES $TEST_ARGS
 else
     python green_framework_test.py $GREEN_FRAMEWORK -v $TEST_ARGS
 fi

.github/workflows/test-python.yml

@@ -205,12 +205,11 @@ jobs:
             python-version: '3.8'
         - name: Start MongoDB
           uses: supercharge/[email protected]
-        - name: Run Test
+        - name: Run connect test from sdist
           shell: bash
           run: |
             cd sdist/test
             ls
             which python
             pip install -e ".[test]"
-            pytest -v
-            pytest -v test/asynchronous/
+            PYMONGO_MUST_CONNECT=1 pytest -v test/test_client_context.py

gridfs/asynchronous/grid_file.py

@@ -154,7 +154,7 @@ async def put(self, data: Any, **kwargs: Any) -> Any:
         """
         async with AsyncGridIn(self._collection, **kwargs) as grid_file:
             await grid_file.write(data)
-            return await grid_file._id
+            return grid_file._id
 
     async def get(self, file_id: Any, session: Optional[AsyncClientSession] = None) -> AsyncGridOut:
         """Get a file from GridFS by ``"_id"``.
@@ -220,7 +220,7 @@ async def get_version(
             query["filename"] = filename
 
         _disallow_transactions(session)
-        cursor = await self._files.find(query, session=session)
+        cursor = self._files.find(query, session=session)
         if version is None:
             version = -1
         if version < 0:
@@ -923,7 +923,7 @@ async def open_download_stream_by_name(
         validate_string("filename", filename)
         query = {"filename": filename}
         _disallow_transactions(session)
-        cursor = await self._files.find(query, session=session)
+        cursor = self._files.find(query, session=session)
         if revision < 0:
             skip = abs(revision) - 1
             cursor.limit(-1).skip(skip).sort("uploadDate", DESCENDING)
@@ -1760,12 +1760,12 @@ def expected_chunk_length(self, chunk_n: int) -> int:
     def __aiter__(self) -> _AsyncGridOutChunkIterator:
         return self
 
-    async def _create_cursor(self) -> None:
+    def _create_cursor(self) -> None:
         filter = {"files_id": self._id}
         if self._next_chunk > 0:
             filter["n"] = {"$gte": self._next_chunk}
         _disallow_transactions(self._session)
-        self._cursor = await self._chunks.find(filter, sort=[("n", 1)], session=self._session)
+        self._cursor = self._chunks.find(filter, sort=[("n", 1)], session=self._session)
 
     async def _next_with_retry(self) -> Mapping[str, Any]:
         """Return the next chunk and retry once on CursorNotFound.
@@ -1775,13 +1775,13 @@ async def _next_with_retry(self) -> Mapping[str, Any]:
         server's default cursor timeout).
         """
         if self._cursor is None:
-            await self._create_cursor()
+            self._create_cursor()
             assert self._cursor is not None
         try:
             return await self._cursor.next()
         except CursorNotFound:
             await self._cursor.close()
-            await self._create_cursor()
+            self._create_cursor()
             return await self._cursor.next()
 
     async def next(self) -> Mapping[str, Any]:

pymongo/asynchronous/client_session.py

@@ -144,6 +144,7 @@
     Any,
     AsyncContextManager,
     Callable,
+    Coroutine,
     Mapping,
     MutableMapping,
     NoReturn,
@@ -598,7 +599,7 @@ def _inherit_option(self, name: str, val: _T) -> _T:
 
     async def with_transaction(
         self,
-        callback: Callable[[AsyncClientSession], _T],
+        callback: Callable[[AsyncClientSession], Coroutine[Any, Any, _T]],
         read_concern: Optional[ReadConcern] = None,
         write_concern: Optional[WriteConcern] = None,
         read_preference: Optional[_ServerMode] = None,
@@ -693,7 +694,7 @@ async def callback(session, custom_arg, custom_kwarg=None):
                 read_concern, write_concern, read_preference, max_commit_time_ms
             )
             try:
-                ret = callback(self)
+                ret = await callback(self)
             except Exception as exc:
                 if self.in_transaction:
                     await self.abort_transaction()

pymongo/asynchronous/collection.py

@@ -1700,18 +1700,18 @@ async def find_one(
         """
         if filter is not None and not isinstance(filter, abc.Mapping):
             filter = {"_id": filter}
-        cursor = await self.find(filter, *args, **kwargs)
+        cursor = self.find(filter, *args, **kwargs)
         async for result in cursor.limit(-1):
             return result
         return None
 
-    async def find(self, *args: Any, **kwargs: Any) -> AsyncCursor[_DocumentType]:
+    def find(self, *args: Any, **kwargs: Any) -> AsyncCursor[_DocumentType]:
         """Query the database.
 
         The `filter` argument is a query document that all results
         must match. For example:
 
-        >>> await db.test.find({"hello": "world"})
+        >>> db.test.find({"hello": "world"})
 
         only matches documents that have a key "hello" with value
         "world".  Matches can have other keys *in addition* to
@@ -1891,9 +1891,7 @@ async def find(self, *args: Any, **kwargs: Any) -> AsyncCursor[_DocumentType]:
 
         .. seealso:: The MongoDB documentation on `find <https://dochub.mongodb.org/core/find>`_.
         """
-        cursor = AsyncCursor(self, *args, **kwargs)
-        await cursor._supports_exhaust()
-        return cursor
+        return AsyncCursor(self, *args, **kwargs)
 
     async def find_raw_batches(
         self, *args: Any, **kwargs: Any

pymongo/asynchronous/command_cursor.py

@@ -424,4 +424,4 @@ def _unpack_response(  # type: ignore[override]
         return raw_response  # type: ignore[return-value]
 
     def __getitem__(self, index: int) -> NoReturn:
-        raise InvalidOperation("Cannot call __getitem__ on RawBatchCursor")
+        raise InvalidOperation("Cannot call __getitem__ on AsyncRawBatchCommandCursor")

pymongo/asynchronous/cursor.py

@@ -237,6 +237,14 @@ def __init__(
         self._dbname = collection.database.name
         self._collname = collection.name
 
+        # Checking exhaust cursor support requires network IO
+        if _IS_SYNC:
+            self._exhaust_checked = True
+            self._supports_exhaust()  # type: ignore[unused-coroutine]
+        else:
+            self._exhaust = cursor_type == CursorType.EXHAUST
+            self._exhaust_checked = False
+
     async def _supports_exhaust(self) -> None:
         # Exhaust cursor support
         if self._cursor_type == CursorType.EXHAUST:
@@ -1242,6 +1250,9 @@ async def rewind(self) -> AsyncCursor[_DocumentType]:
 
     async def next(self) -> _DocumentType:
         """Advance the cursor."""
+        if not self._exhaust_checked:
+            self._exhaust_checked = True
+            await self._supports_exhaust()
         if self._empty:
             raise StopAsyncIteration
         if len(self._data) or await self._refresh():
@@ -1308,4 +1319,4 @@ async def explain(self) -> _DocumentType:
         return await clone.explain()
 
     def __getitem__(self, index: Any) -> NoReturn:
-        raise InvalidOperation("Cannot call __getitem__ on RawBatchCursor")
+        raise InvalidOperation("Cannot call __getitem__ on AsyncRawBatchCursor")

pymongo/asynchronous/encryption.py

@@ -260,7 +260,7 @@ async def fetch_keys(self, filter: bytes) -> AsyncGenerator[bytes, None]:
         :return: A generator which yields the requested keys from the key vault.
         """
         assert self.key_vault_coll is not None
-        async with await self.key_vault_coll.find(RawBSONDocument(filter)) as cursor:
+        async with self.key_vault_coll.find(RawBSONDocument(filter)) as cursor:
             async for key in cursor:
                 yield key.raw
 
@@ -975,7 +975,7 @@ async def get_key(self, id: Binary) -> Optional[RawBSONDocument]:
         assert self._key_vault_coll is not None
         return await self._key_vault_coll.find_one({"_id": id})
 
-    async def get_keys(self) -> AsyncCursor[RawBSONDocument]:
+    def get_keys(self) -> AsyncCursor[RawBSONDocument]:
         """Get all of the data keys.
 
         :return: An instance of :class:`~pymongo.cursor.Cursor` over the data key
@@ -985,7 +985,7 @@ async def get_keys(self) -> AsyncCursor[RawBSONDocument]:
         """
         self._check_closed()
         assert self._key_vault_coll is not None
-        return await self._key_vault_coll.find({})
+        return self._key_vault_coll.find({})
 
     async def delete_key(self, id: Binary) -> DeleteResult:
         """Delete a key document in the key vault collection that has the given ``key_id``.

pymongo/asynchronous/pool.py

@@ -14,8 +14,10 @@
 
 from __future__ import annotations
 
+import asyncio
 import collections
 import contextlib
+import functools
 import logging
 import os
 import socket
@@ -876,12 +878,23 @@ async def _configured_socket(
             if _IS_SYNC:
                 ssl_sock = ssl_context.wrap_socket(sock, server_hostname=host)
             else:
-                ssl_sock = await ssl_context.a_wrap_socket(sock, server_hostname=host)  # type: ignore[assignment, misc]
+                if hasattr(ssl_context, "a_wrap_socket"):
+                    ssl_sock = await ssl_context.a_wrap_socket(sock, server_hostname=host)  # type: ignore[assignment, misc]
+                else:
+                    loop = asyncio.get_running_loop()
+                    ssl_sock = await loop.run_in_executor(
+                        None,
+                        functools.partial(ssl_context.wrap_socket, sock, server_hostname=host),  # type: ignore[assignment, misc]
+                    )
         else:
             if _IS_SYNC:
                 ssl_sock = ssl_context.wrap_socket(sock)
             else:
-                ssl_sock = await ssl_context.a_wrap_socket(sock)  # type: ignore[assignment, misc]
+                if hasattr(ssl_context, "a_wrap_socket"):
+                    ssl_sock = await ssl_context.a_wrap_socket(sock)  # type: ignore[assignment, misc]
+                else:
+                    loop = asyncio.get_running_loop()
+                    ssl_sock = await loop.run_in_executor(None, ssl_context.wrap_socket, sock)  # type: ignore[assignment, misc]
     except _CertificateError:
         sock.close()
         # Raise _CertificateError directly like we do after match_hostname

pymongo/synchronous/collection.py

@@ -1710,7 +1710,7 @@ def find(self, *args: Any, **kwargs: Any) -> Cursor[_DocumentType]:
         The `filter` argument is a query document that all results
         must match. For example:
 
-        >>> await db.test.find({"hello": "world"})
+        >>> db.test.find({"hello": "world"})
 
         only matches documents that have a key "hello" with value
         "world".  Matches can have other keys *in addition* to
@@ -1890,9 +1890,7 @@ def find(self, *args: Any, **kwargs: Any) -> Cursor[_DocumentType]:
 
         .. seealso:: The MongoDB documentation on `find <https://dochub.mongodb.org/core/find>`_.
         """
-        cursor = Cursor(self, *args, **kwargs)
-        cursor._supports_exhaust()
-        return cursor
+        return Cursor(self, *args, **kwargs)
 
     def find_raw_batches(self, *args: Any, **kwargs: Any) -> RawBatchCursor[_DocumentType]:
         """Query the database and retrieve batches of raw BSON.