diff --git a/pymongo/_csot.py b/pymongo/_csot.py index 06c6b68ac9..fc91f57436 100644 --- a/pymongo/_csot.py +++ b/pymongo/_csot.py @@ -16,13 +16,24 @@ from __future__ import annotations +import contextlib import functools import inspect import time from collections import deque from contextlib import AbstractContextManager from contextvars import ContextVar, Token -from typing import TYPE_CHECKING, Any, Callable, Deque, MutableMapping, Optional, TypeVar, cast +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Deque, + Generator, + MutableMapping, + Optional, + TypeVar, + cast, +) if TYPE_CHECKING: from pymongo.write_concern import WriteConcern @@ -54,6 +65,17 @@ def remaining() -> Optional[float]: return DEADLINE.get() - time.monotonic() +@contextlib.contextmanager +def reset() -> Generator: + timeout = get_timeout() + if timeout is None: + deadline_token = DEADLINE.set(DEADLINE.get()) + else: + deadline_token = DEADLINE.set(DEADLINE.get() + timeout) + yield + DEADLINE.reset(deadline_token) + + def clamp_remaining(max_timeout: float) -> float: """Return the remaining timeout clamped to a max value.""" timeout = remaining() diff --git a/pymongo/asynchronous/client_session.py b/pymongo/asynchronous/client_session.py index d80495d804..e3a7d92fb1 100644 --- a/pymongo/asynchronous/client_session.py +++ b/pymongo/asynchronous/client_session.py @@ -473,7 +473,11 @@ def _max_time_expired_error(exc: PyMongoError) -> bool: def _within_time_limit(start_time: float) -> bool: """Are we within the with_transaction retry limit?""" - return time.monotonic() - start_time < _WITH_TRANSACTION_RETRY_TIME_LIMIT + timeout = _csot.get_timeout() + if timeout: + return time.monotonic() - start_time < timeout + else: + return time.monotonic() - start_time < _WITH_TRANSACTION_RETRY_TIME_LIMIT _T = TypeVar("_T") @@ -512,6 +516,7 @@ def __init__( # Is this an implicitly created session? self._implicit = implicit self._transaction = _Transaction(None, client) + self._timeout = client.options.timeout async def end_session(self) -> None: """Finish this session. If a transaction has started, abort it. @@ -597,6 +602,7 @@ def _inherit_option(self, name: str, val: _T) -> _T: return parent_val return getattr(self.client, name) + @_csot.apply async def with_transaction( self, callback: Callable[[AsyncClientSession], Coroutine[Any, Any, _T]], @@ -697,7 +703,8 @@ async def callback(session, custom_arg, custom_kwarg=None): ret = await callback(self) except Exception as exc: if self.in_transaction: - await self.abort_transaction() + with _csot.reset(): + await self.abort_transaction() if ( isinstance(exc, PyMongoError) and exc.has_error_label("TransientTransactionError") @@ -816,6 +823,7 @@ async def commit_transaction(self) -> None: finally: self._transaction.state = _TxnState.COMMITTED + @_csot.apply async def abort_transaction(self) -> None: """Abort a multi-statement transaction. diff --git a/pymongo/asynchronous/topology.py b/pymongo/asynchronous/topology.py index 82af4257ba..a51dd8a98f 100644 --- a/pymongo/asynchronous/topology.py +++ b/pymongo/asynchronous/topology.py @@ -249,7 +249,8 @@ def get_server_selection_timeout(self) -> float: timeout = _csot.remaining() if timeout is None: return self._settings.server_selection_timeout - return timeout + else: + return min(timeout, self._settings.server_selection_timeout) async def select_servers( self, diff --git a/pymongo/synchronous/client_session.py b/pymongo/synchronous/client_session.py index f1d680fc0a..ade46a9799 100644 --- a/pymongo/synchronous/client_session.py +++ b/pymongo/synchronous/client_session.py @@ -472,7 +472,11 @@ def _max_time_expired_error(exc: PyMongoError) -> bool: def _within_time_limit(start_time: float) -> bool: """Are we within the with_transaction retry limit?""" - return time.monotonic() - start_time < _WITH_TRANSACTION_RETRY_TIME_LIMIT + timeout = _csot.get_timeout() + if timeout: + return time.monotonic() - start_time < timeout + else: + return time.monotonic() - start_time < _WITH_TRANSACTION_RETRY_TIME_LIMIT _T = TypeVar("_T") @@ -511,6 +515,7 @@ def __init__( # Is this an implicitly created session? self._implicit = implicit self._transaction = _Transaction(None, client) + self._timeout = client.options.timeout def end_session(self) -> None: """Finish this session. If a transaction has started, abort it. @@ -596,6 +601,7 @@ def _inherit_option(self, name: str, val: _T) -> _T: return parent_val return getattr(self.client, name) + @_csot.apply def with_transaction( self, callback: Callable[[ClientSession], _T], @@ -694,7 +700,8 @@ def callback(session, custom_arg, custom_kwarg=None): ret = callback(self) except Exception as exc: if self.in_transaction: - self.abort_transaction() + with _csot.reset(): + self.abort_transaction() if ( isinstance(exc, PyMongoError) and exc.has_error_label("TransientTransactionError") @@ -813,6 +820,7 @@ def commit_transaction(self) -> None: finally: self._transaction.state = _TxnState.COMMITTED + @_csot.apply def abort_transaction(self) -> None: """Abort a multi-statement transaction. diff --git a/pymongo/synchronous/topology.py b/pymongo/synchronous/topology.py index a350c1702e..9e5284c9d0 100644 --- a/pymongo/synchronous/topology.py +++ b/pymongo/synchronous/topology.py @@ -249,7 +249,8 @@ def get_server_selection_timeout(self) -> float: timeout = _csot.remaining() if timeout is None: return self._settings.server_selection_timeout - return timeout + else: + return min(timeout, self._settings.server_selection_timeout) def select_servers( self, diff --git a/test/asynchronous/test_client.py b/test/asynchronous/test_client.py index 590154b857..623dcd086e 100644 --- a/test/asynchronous/test_client.py +++ b/test/asynchronous/test_client.py @@ -64,6 +64,7 @@ from test.utils import ( NTHREADS, CMAPListener, + EventListener, FunctionCallRecorder, async_get_pool, async_wait_until, @@ -114,7 +115,13 @@ ServerSelectionTimeoutError, WriteConcernError, ) -from pymongo.monitoring import ServerHeartbeatListener, ServerHeartbeatStartedEvent +from pymongo.monitoring import ( + ConnectionClosedEvent, + ConnectionCreatedEvent, + ConnectionReadyEvent, + ServerHeartbeatListener, + ServerHeartbeatStartedEvent, +) from pymongo.pool_options import _MAX_METADATA_SIZE, _METADATA, ENV_VAR_K8S, PoolOptions from pymongo.read_preferences import ReadPreference from pymongo.server_description import ServerDescription @@ -2585,5 +2592,44 @@ async def test_direct_client_maintains_pool_to_arbiter(self): self.assertEqual(listener.event_count(monitoring.PoolReadyEvent), 1) +# https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#4-background-connection-pooling +class TestClientCSOTProse(AsyncIntegrationTest): + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#timeoutms-is-refreshed-for-each-handshake-command + @async_client_context.require_auth + @async_client_context.require_version_min(4, 4, -1) + @async_client_context.require_failCommand_appName + async def test_02_timeoutMS_refreshed_for_each_handshake_command(self): + listener = CMAPListener() + + async with self.fail_point( + { + "mode": {"times": 1}, + "data": { + "failCommands": ["hello", "isMaster", "saslContinue"], + "blockConnection": True, + "blockTimeMS": 15, + "appName": "refreshTimeoutBackgroundPoolTest", + }, + } + ): + _ = await self.async_single_client( + minPoolSize=1, + timeoutMS=20, + appname="refreshTimeoutBackgroundPoolTest", + event_listeners=[listener], + ) + + async def predicate(): + return ( + listener.event_count(ConnectionCreatedEvent) == 1 + and listener.event_count(ConnectionReadyEvent) == 1 + ) + + await async_wait_until( + predicate, + "didn't ever see a ConnectionCreatedEvent and a ConnectionReadyEvent", + ) + + if __name__ == "__main__": unittest.main() diff --git a/test/asynchronous/test_collection.py b/test/asynchronous/test_collection.py index db52bad4ac..3b8da88602 100644 --- a/test/asynchronous/test_collection.py +++ b/test/asynchronous/test_collection.py @@ -64,6 +64,7 @@ InvalidDocument, InvalidName, InvalidOperation, + NetworkTimeout, OperationFailure, WriteConcernError, ) @@ -2277,6 +2278,32 @@ async def afind(*args, **kwargs): for helper, args in helpers: await helper(*args, let={}) # type: ignore + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#1-multi-batch-inserts + @async_client_context.require_standalone + @async_client_context.require_version_min(4, 4, -1) + @async_client_context.require_failCommand_fail_point + async def test_01_multi_batch_inserts(self): + client = await self.async_single_client(read_preference=ReadPreference.PRIMARY_PREFERRED) + await client.db.coll.drop() + + async with self.fail_point( + { + "mode": {"times": 2}, + "data": {"failCommands": ["insert"], "blockConnection": True, "blockTimeMS": 1010}, + } + ): + listener = OvertCommandListener() + client2 = await self.async_single_client( + timeoutMS=2000, + read_preference=ReadPreference.PRIMARY_PREFERRED, + event_listeners=[listener], + ) + docs = [{"a": "b" * 1000000} for _ in range(50)] + with self.assertRaises(NetworkTimeout): + await client2.db.coll.insert_many(docs) + + self.assertEqual(2, len(listener.started_events)) + if __name__ == "__main__": unittest.main() diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index e42c85aa7a..eedab46e78 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -86,6 +86,7 @@ EncryptedCollectionError, EncryptionError, InvalidOperation, + NetworkTimeout, OperationFailure, ServerSelectionTimeoutError, WriteError, @@ -3133,5 +3134,113 @@ async def test_explicit_session_errors_when_unsupported(self): await self.mongocryptd_client.db.test.insert_one({"x": 1}, session=s) +# https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#3-clientencryption +class TestCSOTProse(AsyncEncryptionIntegrationTest): + mongocryptd_client: AsyncMongoClient + MONGOCRYPTD_PORT = 27020 + LOCAL_MASTERKEY = Binary( + base64.b64decode( + b"Mng0NCt4ZHVUYUJCa1kxNkVyNUR1QURhZ2h2UzR2d2RrZzh0cFBwM3R6NmdWMDFBMUN3YkQ5aXRRMkhGRGdQV09wOGVNYUMxT2k3NjZKelhaQmRCZGJkTXVyZG9uSjFk" + ), + UUID_SUBTYPE, + ) + + async def asyncSetUp(self) -> None: + self.listener = OvertCommandListener() + self.client = await self.async_single_client( + read_preference=ReadPreference.PRIMARY_PREFERRED, event_listeners=[self.listener] + ) + await self.client.keyvault.datakeys.drop() + self.key_vault_client = await self.async_rs_or_single_client( + timeoutMS=50, event_listeners=[self.listener] + ) + self.client_encryption = self.create_client_encryption( + key_vault_namespace="keyvault.datakeys", + kms_providers={"local": {"key": self.LOCAL_MASTERKEY}}, + key_vault_client=self.key_vault_client, + codec_options=OPTS, + ) + + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#createdatakey + @async_client_context.require_failCommand_fail_point + @async_client_context.require_version_min(4, 4, -1) + async def test_01_create_data_key(self): + async with self.fail_point( + { + "mode": {"times": 1}, + "data": {"failCommands": ["insert"], "blockConnection": True, "blockTimeMS": 100}, + } + ): + self.listener.reset() + with self.assertRaisesRegex(EncryptionError, "timed out"): + await self.client_encryption.create_data_key("local") + + events = self.listener.started_events + self.assertEqual(1, len(events)) + self.assertEqual("insert", events[0].command_name) + + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#encrypt + @async_client_context.require_failCommand_fail_point + @async_client_context.require_version_min(4, 4, -1) + async def test_02_encrypt(self): + data_key_id = await self.client_encryption.create_data_key("local") + self.assertEqual(4, data_key_id.subtype) + async with self.fail_point( + { + "mode": {"times": 1}, + "data": {"failCommands": ["find"], "blockConnection": True, "blockTimeMS": 100}, + } + ): + self.listener.reset() + with self.assertRaisesRegex(EncryptionError, "timed out"): + await self.client_encryption.encrypt( + "hello", + Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, + key_id=data_key_id, + ) + + events = self.listener.started_events + self.assertEqual(1, len(events)) + self.assertEqual("find", events[0].command_name) + + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#decrypt + @async_client_context.require_failCommand_fail_point + @async_client_context.require_version_min(4, 4, -1) + async def test_03_decrypt(self): + data_key_id = await self.client_encryption.create_data_key("local") + self.assertEqual(4, data_key_id.subtype) + + encrypted = await self.client_encryption.encrypt( + "hello", Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, key_id=data_key_id + ) + self.assertEqual(6, encrypted.subtype) + + await self.key_vault_client.close() + self.key_vault_client = await self.async_rs_or_single_client( + timeoutMS=50, event_listeners=[self.listener] + ) + await self.client_encryption.close() + self.client_encryption = self.create_client_encryption( + key_vault_namespace="keyvault.datakeys", + kms_providers={"local": {"key": self.LOCAL_MASTERKEY}}, + key_vault_client=self.key_vault_client, + codec_options=OPTS, + ) + + async with self.fail_point( + { + "mode": {"times": 1}, + "data": {"failCommands": ["find"], "blockConnection": True, "blockTimeMS": 100}, + } + ): + self.listener.reset() + with self.assertRaisesRegex(EncryptionError, "timed out"): + await self.client_encryption.decrypt(encrypted) + + events = self.listener.started_events + self.assertEqual(1, len(events)) + self.assertEqual("find", events[0].command_name) + + if __name__ == "__main__": unittest.main() diff --git a/test/asynchronous/test_session.py b/test/asynchronous/test_session.py index b432621798..ba991a6ee8 100644 --- a/test/asynchronous/test_session.py +++ b/test/asynchronous/test_session.py @@ -48,7 +48,7 @@ from pymongo.asynchronous.cursor import AsyncCursor from pymongo.asynchronous.helpers import anext from pymongo.common import _MAX_END_SESSIONS -from pymongo.errors import ConfigurationError, InvalidOperation, OperationFailure +from pymongo.errors import ConfigurationError, InvalidOperation, NetworkTimeout, OperationFailure from pymongo.operations import IndexModel, InsertOne, UpdateOne from pymongo.read_concern import ReadConcern diff --git a/test/asynchronous/test_transactions.py b/test/asynchronous/test_transactions.py index b5d0686417..f2bd94e25c 100644 --- a/test/asynchronous/test_transactions.py +++ b/test/asynchronous/test_transactions.py @@ -43,6 +43,7 @@ ConfigurationError, ConnectionFailure, InvalidOperation, + NetworkTimeout, OperationFailure, ) from pymongo.operations import IndexModel, InsertOne @@ -386,6 +387,45 @@ async def find_raw_batches(*args, **kwargs): if isinstance(res, (AsyncCommandCursor, AsyncCursor)): await res.to_list() + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#10-convenient-transactions + @async_client_context.require_transactions + @async_client_context.require_version_min(4, 4, -1) + @async_client_context.require_failCommand_fail_point + async def test_10_convenient_transactions_csot(self): + await self.client.db.coll.drop() + + listener = OvertCommandListener() + + async with self.fail_point( + { + "mode": {"times": 2}, + "data": { + "failCommands": ["insert", "abortTransaction"], + "blockConnection": True, + "blockTimeMS": 200, + }, + } + ): + client = await self.async_rs_or_single_client( + timeoutMS=150, + event_listeners=[listener], + ) + session = client.start_session() + + async def callback(s): + await client.db.coll.insert_one({"_id": 1}, session=s) + + with self.assertRaises(NetworkTimeout): + await session.with_transaction(callback) + + started = listener.started_command_names() + failed = listener.failed_command_names() + + self.assertIn("insert", started) + self.assertIn("abortTransaction", started) + self.assertIn("insert", failed) + self.assertIn("abortTransaction", failed) + class PatchSessionTimeout: """Patches the client_session's with_transaction timeout for testing.""" diff --git a/test/test_client.py b/test/test_client.py index 5bbb5bd751..63fd843bed 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -63,6 +63,7 @@ from test.utils import ( NTHREADS, CMAPListener, + EventListener, FunctionCallRecorder, assertRaisesExactly, delay, @@ -102,7 +103,13 @@ ServerSelectionTimeoutError, WriteConcernError, ) -from pymongo.monitoring import ServerHeartbeatListener, ServerHeartbeatStartedEvent +from pymongo.monitoring import ( + ConnectionClosedEvent, + ConnectionCreatedEvent, + ConnectionReadyEvent, + ServerHeartbeatListener, + ServerHeartbeatStartedEvent, +) from pymongo.pool_options import _MAX_METADATA_SIZE, _METADATA, ENV_VAR_K8S, PoolOptions from pymongo.read_preferences import ReadPreference from pymongo.server_description import ServerDescription @@ -2541,5 +2548,44 @@ def test_direct_client_maintains_pool_to_arbiter(self): self.assertEqual(listener.event_count(monitoring.PoolReadyEvent), 1) +# https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#4-background-connection-pooling +class TestClientCSOTProse(IntegrationTest): + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#timeoutms-is-refreshed-for-each-handshake-command + @client_context.require_auth + @client_context.require_version_min(4, 4, -1) + @client_context.require_failCommand_appName + def test_02_timeoutMS_refreshed_for_each_handshake_command(self): + listener = CMAPListener() + + with self.fail_point( + { + "mode": {"times": 1}, + "data": { + "failCommands": ["hello", "isMaster", "saslContinue"], + "blockConnection": True, + "blockTimeMS": 15, + "appName": "refreshTimeoutBackgroundPoolTest", + }, + } + ): + _ = self.single_client( + minPoolSize=1, + timeoutMS=20, + appname="refreshTimeoutBackgroundPoolTest", + event_listeners=[listener], + ) + + def predicate(): + return ( + listener.event_count(ConnectionCreatedEvent) == 1 + and listener.event_count(ConnectionReadyEvent) == 1 + ) + + wait_until( + predicate, + "didn't ever see a ConnectionCreatedEvent and a ConnectionReadyEvent", + ) + + if __name__ == "__main__": unittest.main() diff --git a/test/test_collection.py b/test/test_collection.py index 84a900d45b..87ed4d459a 100644 --- a/test/test_collection.py +++ b/test/test_collection.py @@ -59,6 +59,7 @@ InvalidDocument, InvalidName, InvalidOperation, + NetworkTimeout, OperationFailure, WriteConcernError, ) @@ -2254,6 +2255,32 @@ def afind(*args, **kwargs): for helper, args in helpers: helper(*args, let={}) # type: ignore + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#1-multi-batch-inserts + @client_context.require_standalone + @client_context.require_version_min(4, 4, -1) + @client_context.require_failCommand_fail_point + def test_01_multi_batch_inserts(self): + client = self.single_client(read_preference=ReadPreference.PRIMARY_PREFERRED) + client.db.coll.drop() + + with self.fail_point( + { + "mode": {"times": 2}, + "data": {"failCommands": ["insert"], "blockConnection": True, "blockTimeMS": 1010}, + } + ): + listener = OvertCommandListener() + client2 = self.single_client( + timeoutMS=2000, + read_preference=ReadPreference.PRIMARY_PREFERRED, + event_listeners=[listener], + ) + docs = [{"a": "b" * 1000000} for _ in range(50)] + with self.assertRaises(NetworkTimeout): + client2.db.coll.insert_many(docs) + + self.assertEqual(2, len(listener.started_events)) + if __name__ == "__main__": unittest.main() diff --git a/test/test_csot.py b/test/test_csot.py index 64210b4d64..dc7c999e08 100644 --- a/test/test_csot.py +++ b/test/test_csot.py @@ -12,11 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Test the CSOT unified spec tests.""" +"""Test the CSOT unified spec and prose tests.""" from __future__ import annotations import os import sys +from test.utils import OvertCommandListener + +from pymongo.read_concern import ReadConcern sys.path[0:0] = [""] @@ -24,8 +27,8 @@ from test.unified_format import generate_test_classes import pymongo -from pymongo import _csot -from pymongo.errors import PyMongoError +from pymongo import ReadPreference, WriteConcern, _csot +from pymongo.errors import NetworkTimeout, PyMongoError # Location of JSON test specifications. TEST_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "csot") diff --git a/test/test_encryption.py b/test/test_encryption.py index 0806f91a06..ee0bc434ae 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -83,6 +83,7 @@ EncryptedCollectionError, EncryptionError, InvalidOperation, + NetworkTimeout, OperationFailure, ServerSelectionTimeoutError, WriteError, @@ -3115,5 +3116,113 @@ def test_explicit_session_errors_when_unsupported(self): self.mongocryptd_client.db.test.insert_one({"x": 1}, session=s) +# https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#3-clientencryption +class TestCSOTProse(EncryptionIntegrationTest): + mongocryptd_client: MongoClient + MONGOCRYPTD_PORT = 27020 + LOCAL_MASTERKEY = Binary( + base64.b64decode( + b"Mng0NCt4ZHVUYUJCa1kxNkVyNUR1QURhZ2h2UzR2d2RrZzh0cFBwM3R6NmdWMDFBMUN3YkQ5aXRRMkhGRGdQV09wOGVNYUMxT2k3NjZKelhaQmRCZGJkTXVyZG9uSjFk" + ), + UUID_SUBTYPE, + ) + + def setUp(self) -> None: + self.listener = OvertCommandListener() + self.client = self.single_client( + read_preference=ReadPreference.PRIMARY_PREFERRED, event_listeners=[self.listener] + ) + self.client.keyvault.datakeys.drop() + self.key_vault_client = self.rs_or_single_client( + timeoutMS=50, event_listeners=[self.listener] + ) + self.client_encryption = self.create_client_encryption( + key_vault_namespace="keyvault.datakeys", + kms_providers={"local": {"key": self.LOCAL_MASTERKEY}}, + key_vault_client=self.key_vault_client, + codec_options=OPTS, + ) + + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#createdatakey + @client_context.require_failCommand_fail_point + @client_context.require_version_min(4, 4, -1) + def test_01_create_data_key(self): + with self.fail_point( + { + "mode": {"times": 1}, + "data": {"failCommands": ["insert"], "blockConnection": True, "blockTimeMS": 100}, + } + ): + self.listener.reset() + with self.assertRaisesRegex(EncryptionError, "timed out"): + self.client_encryption.create_data_key("local") + + events = self.listener.started_events + self.assertEqual(1, len(events)) + self.assertEqual("insert", events[0].command_name) + + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#encrypt + @client_context.require_failCommand_fail_point + @client_context.require_version_min(4, 4, -1) + def test_02_encrypt(self): + data_key_id = self.client_encryption.create_data_key("local") + self.assertEqual(4, data_key_id.subtype) + with self.fail_point( + { + "mode": {"times": 1}, + "data": {"failCommands": ["find"], "blockConnection": True, "blockTimeMS": 100}, + } + ): + self.listener.reset() + with self.assertRaisesRegex(EncryptionError, "timed out"): + self.client_encryption.encrypt( + "hello", + Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, + key_id=data_key_id, + ) + + events = self.listener.started_events + self.assertEqual(1, len(events)) + self.assertEqual("find", events[0].command_name) + + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#decrypt + @client_context.require_failCommand_fail_point + @client_context.require_version_min(4, 4, -1) + def test_03_decrypt(self): + data_key_id = self.client_encryption.create_data_key("local") + self.assertEqual(4, data_key_id.subtype) + + encrypted = self.client_encryption.encrypt( + "hello", Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, key_id=data_key_id + ) + self.assertEqual(6, encrypted.subtype) + + self.key_vault_client.close() + self.key_vault_client = self.rs_or_single_client( + timeoutMS=50, event_listeners=[self.listener] + ) + self.client_encryption.close() + self.client_encryption = self.create_client_encryption( + key_vault_namespace="keyvault.datakeys", + kms_providers={"local": {"key": self.LOCAL_MASTERKEY}}, + key_vault_client=self.key_vault_client, + codec_options=OPTS, + ) + + with self.fail_point( + { + "mode": {"times": 1}, + "data": {"failCommands": ["find"], "blockConnection": True, "blockTimeMS": 100}, + } + ): + self.listener.reset() + with self.assertRaisesRegex(EncryptionError, "timed out"): + self.client_encryption.decrypt(encrypted) + + events = self.listener.started_events + self.assertEqual(1, len(events)) + self.assertEqual("find", events[0].command_name) + + if __name__ == "__main__": unittest.main() diff --git a/test/test_gridfs_bucket.py b/test/test_gridfs_bucket.py index 28adb7051a..4322b6b74a 100644 --- a/test/test_gridfs_bucket.py +++ b/test/test_gridfs_bucket.py @@ -37,6 +37,7 @@ from gridfs.errors import CorruptGridFile, NoFile from pymongo.errors import ( ConfigurationError, + NetworkTimeout, NotPrimaryError, ServerSelectionTimeoutError, WriteConcernError, @@ -525,5 +526,84 @@ def test_gridfs_secondary_lazy(self): self.assertRaises(NotPrimaryError, gfs.upload_from_stream, "test_filename", b"data") +# https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#6-gridfs---upload +class TestGridFsCSOT(IntegrationTest): + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#uploads-via-openuploadstream-can-be-timed-out + @client_context.require_failCommand_fail_point + @client_context.require_version_min(4, 4, -1) + def test_06_01_uploads_via_open_upload_stream_can_be_timed_out(self): + self.client.db.fs.files.drop() + self.client.db.fs.chunks.drop() + + with self.fail_point( + { + "mode": {"times": 1}, + "data": {"failCommands": ["insert"], "blockConnection": True, "blockTimeMS": 200}, + } + ): + client = self.single_client(timeoutMS=150) + bucket = gridfs.GridFSBucket(client.db) + upload_stream = bucket.open_upload_stream("filename") + upload_stream.write(b"0x12") + with self.assertRaises(NetworkTimeout): + upload_stream.close() + + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#aborting-an-upload-stream-can-be-timed-out + @client_context.require_failCommand_fail_point + @client_context.require_version_min(4, 4, -1) + def test_06_02_aborting_an_upload_stream_can_be_timed_out(self): + self.client.db.fs.files.drop() + self.client.db.fs.chunks.drop() + + with self.fail_point( + { + "mode": {"times": 1}, + "data": {"failCommands": ["delete"], "blockConnection": True, "blockTimeMS": 200}, + } + ): + client = self.single_client(timeoutMS=150) + bucket = gridfs.GridFSBucket(client.db, chunk_size_bytes=2) + upload_stream = bucket.open_upload_stream("filename") + upload_stream.write(b"0x010x020x030x04") + with self.assertRaises(NetworkTimeout): + upload_stream.abort() + + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#7-gridfs---download + @client_context.require_failCommand_fail_point + @client_context.require_version_min(4, 4, -1) + def test_07_gridfs_download_csot(self): + self.client.db.fs.files.drop() + self.client.db.fs.chunks.drop() + + id = ObjectId("000000000000000000000005") + + self.client.db.fs.files.insert_one( + { + "_id": id, + "length": 10, + "chunkSize": 4, + "uploadDate": {"$date": "1970-01-01T00:00:00.000Z"}, + "md5": "57d83cd477bfb1ccd975ab33d827a92b", + "filename": "length-10", + "contentType": "application/octet-stream", + "aliases": [], + "metadata": {}, + } + ) + + client = self.single_client(timeoutMS=150) + bucket = gridfs.GridFSBucket(client.db) + download_stream = bucket.open_download_stream(id) + + with self.fail_point( + { + "mode": {"times": 1}, + "data": {"failCommands": ["find"], "blockConnection": True, "blockTimeMS": 200}, + } + ): + with self.assertRaises(NetworkTimeout): + download_stream.read() + + if __name__ == "__main__": unittest.main() diff --git a/test/test_server_selection.py b/test/test_server_selection.py index 984b967f50..ab6f7ae14e 100644 --- a/test/test_server_selection.py +++ b/test/test_server_selection.py @@ -17,9 +17,10 @@ import os import sys +import time from pymongo import MongoClient, ReadPreference -from pymongo.errors import ServerSelectionTimeoutError +from pymongo.errors import NetworkTimeout, ServerSelectionTimeoutError from pymongo.hello import HelloCompat from pymongo.operations import _Op from pymongo.server_selectors import writable_server_selector @@ -200,5 +201,103 @@ def test_server_selector_bypassed(self): self.assertEqual(selector.call_count, 0) +# https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#8-server-selection +class TestServerSelectionCSOT(IntegrationTest): + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#serverselectiontimeoutms-honored-if-timeoutms-is-not-set + @client_context.require_version_min(4, 4, -1) + def test_08_01_server_selection_timeoutMS_honored(self): + client = self.single_client("mongodb://invalid/?serverSelectionTimeoutMS=10") + with self.assertRaises(ServerSelectionTimeoutError): + start = time.time_ns() * 1000 + client.admin.command({"ping": 1}) + + end = time.time_ns() * 1000 + + self.assertLessEqual(start - end, 15) + + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#timeoutms-honored-for-server-selection-if-its-lower-than-serverselectiontimeoutms + @client_context.require_version_min(4, 4, -1) + def test_08_02_timeoutMS_honored_for_server_selection_if_lower(self): + client = self.single_client("mongodb://invalid/?timeoutMS=10&serverSelectionTimeoutMS=20") + with self.assertRaises(ServerSelectionTimeoutError): + start = time.time_ns() * 1_000_000 + client.admin.command({"ping": 1}) + end = time.time_ns() * 1_000_000 + + self.assertLessEqual(start - end, 15) + + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#serverselectiontimeoutms-honored-for-server-selection-if-its-lower-than-timeoutms + @client_context.require_version_min(4, 4, -1) + def test_08_03_serverselectiontimeoutms_honored_for_server_selection_if_lower(self): + client = self.single_client("mongodb://invalid/?timeoutMS=20&serverSelectionTimeoutMS=10") + with self.assertRaises(ServerSelectionTimeoutError): + start = time.time_ns() * 1_000_000 + client.admin.command({"ping": 1}) + + end = time.time_ns() * 1_000_000 + + self.assertLessEqual(start - end, 15) + + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#serverselectiontimeoutms-honored-for-server-selection-if-timeoutms0 + @client_context.require_version_min(4, 4, -1) + def test_08_04_serverselectiontimeoutms_honored_for_server_selection_if_zero_timeoutms(self): + client = self.single_client("mongodb://invalid/?timeoutMS=0&serverSelectionTimeoutMS=10") + with self.assertRaises(ServerSelectionTimeoutError): + start = time.time_ns() * 1_000_000 + client.admin.command({"ping": 1}) + + end = time.time_ns() * 1_000_000 + + self.assertLessEqual(start - end, 15) + + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#timeoutms-honored-for-connection-handshake-commands-if-its-lower-than-serverselectiontimeoutms + @client_context.require_auth + @client_context.require_version_min(4, 4, -1) + @client_context.require_failCommand_fail_point + def test_08_05_timeoutms_honored_for_handshake_if_lower(self): + with self.fail_point( + { + "mode": {"times": 1}, + "data": { + "failCommands": ["saslContinue"], + "blockConnection": True, + "blockTimeMS": 15, + }, + } + ): + client = self.single_client(timeoutMS=10, serverSelectionTimeoutMS=20) + with self.assertRaises(NetworkTimeout): + start = time.time_ns() * 1_000_000 + client.db.coll.insert_one({"x": 1}) + + end = time.time_ns() * 1_000_000 + + self.assertLessEqual(start - end, 15) + + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#serverselectiontimeoutms-honored-for-connection-handshake-commands-if-its-lower-than-timeoutms + @client_context.require_auth + @client_context.require_version_min(4, 4, -1) + @client_context.require_failCommand_fail_point + def test_08_06_serverSelectionTimeoutMS_honored_for_handshake_if_lower(self): + with self.fail_point( + { + "mode": {"times": 1}, + "data": { + "failCommands": ["saslContinue"], + "blockConnection": True, + "blockTimeMS": 15, + }, + } + ): + client = self.single_client(timeoutMS=20, serverSelectionTimeoutMS=10) + with self.assertRaises(NetworkTimeout): + start = time.time_ns() * 1_000_000 + client.db.coll.insert_one({"x": 1}) + + end = time.time_ns() * 1_000_000 + + self.assertLessEqual(start - end, 15) + + if __name__ == "__main__": unittest.main() diff --git a/test/test_session.py b/test/test_session.py index d0bbb075a8..5349e9604a 100644 --- a/test/test_session.py +++ b/test/test_session.py @@ -44,7 +44,7 @@ from gridfs.synchronous.grid_file import GridFS, GridFSBucket from pymongo import ASCENDING, MongoClient, monitoring from pymongo.common import _MAX_END_SESSIONS -from pymongo.errors import ConfigurationError, InvalidOperation, OperationFailure +from pymongo.errors import ConfigurationError, InvalidOperation, NetworkTimeout, OperationFailure from pymongo.operations import IndexModel, InsertOne, UpdateOne from pymongo.read_concern import ReadConcern from pymongo.synchronous.command_cursor import CommandCursor diff --git a/test/test_transactions.py b/test/test_transactions.py index 3cecbe9d38..a7cab34d78 100644 --- a/test/test_transactions.py +++ b/test/test_transactions.py @@ -38,6 +38,7 @@ ConfigurationError, ConnectionFailure, InvalidOperation, + NetworkTimeout, OperationFailure, ) from pymongo.operations import IndexModel, InsertOne @@ -378,6 +379,45 @@ def find_raw_batches(*args, **kwargs): if isinstance(res, (CommandCursor, Cursor)): res.to_list() + # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#10-convenient-transactions + @client_context.require_transactions + @client_context.require_version_min(4, 4, -1) + @client_context.require_failCommand_fail_point + def test_10_convenient_transactions_csot(self): + self.client.db.coll.drop() + + listener = OvertCommandListener() + + with self.fail_point( + { + "mode": {"times": 2}, + "data": { + "failCommands": ["insert", "abortTransaction"], + "blockConnection": True, + "blockTimeMS": 200, + }, + } + ): + client = self.rs_or_single_client( + timeoutMS=150, + event_listeners=[listener], + ) + session = client.start_session() + + def callback(s): + client.db.coll.insert_one({"_id": 1}, session=s) + + with self.assertRaises(NetworkTimeout): + session.with_transaction(callback) + + started = listener.started_command_names() + failed = listener.failed_command_names() + + self.assertIn("insert", started) + self.assertIn("abortTransaction", started) + self.assertIn("insert", failed) + self.assertIn("abortTransaction", failed) + class PatchSessionTimeout: """Patches the client_session's with_transaction timeout for testing.""" diff --git a/test/utils.py b/test/utils.py index 3eac4fa509..177a2940fb 100644 --- a/test/utils.py +++ b/test/utils.py @@ -178,6 +178,10 @@ def started_command_names(self) -> List[str]: """Return list of command names started.""" return [event.command_name for event in self.started_events] + def failed_command_names(self) -> List[str]: + """Return list of command names failed.""" + return [event.command_name for event in self.failed_events] + def reset(self) -> None: """Reset the state of this listener.""" self.results.clear()