From 99fa2460dad7d2a743b057ad2d8c925d74a10f56 Mon Sep 17 00:00:00 2001 From: kievzenit Date: Sun, 17 Aug 2025 19:16:47 +0300 Subject: [PATCH] Migrate motor to pymongo (#1705) * migrated mongo storage from using deprecated motor to PyMongo * added storages to __init__.py file to improve DX * changelog file created * Revert "added storages to __init__.py file to improve DX" This reverts commit 5d0f6a9dfbb3ad4ea21abe508ec20b649c6ff069. * added optional dependency to pymongo to pyproject.toml * Revert "migrated mongo storage from using deprecated motor to PyMongo" This reverts commit 1c0207e1d1dd7dc31732d3325b31d170e7624e19. * added deprecation warning to mongo storage * created pymongo storage * added entry for PyMongoStorage to documentation in fsm.storages * updated changelog to have information about how to migrate from MongoStorage to PyMongoStorage * added test for pymongo storage (copied from mongo storage test) * fixed formatting using black and isort * fixed bug in close method of PyMongoStorage (client close method was not awaited) * added test for PyMongoStorage that checks if storage could be properly closed * pymongo package changed to be lower case in PyMongoStorage * added fixture registration for pymongo storage * test for pymongo is now using proper test fixtures * removed redundant call to get_data, because we have checked this condition in the previous line * added more tests to pymongo test, to check for all possible cases of using update_data method * fixed PyMongoStorage update_data method implementation * added pymongo tests to test_storages * fixed pymongo tests, update_data method should not delete document when {} was passed * Revert "fixed PyMongoStorage update_data method implementation" This reverts commit 86170e1cb91cce4967118f16e38eb263197caef0. * fixed linting issues in PyMongoStorage * changed allowed versions of pymongo, to be compatible with motor * pinned the upper version of pymongo to <4.11 --- CHANGES/1705.misc.rst | 1 + aiogram/fsm/storage/mongo.py | 6 + aiogram/fsm/storage/pymongo.py | 136 +++++++++++ .../finite_state_machine/storages.rst | 4 + pyproject.toml | 1 + tests/conftest.py | 31 +++ tests/test_fsm/storage/test_pymongo.py | 218 ++++++++++++++++++ tests/test_fsm/storage/test_storages.py | 1 + 8 files changed, 398 insertions(+) create mode 100644 CHANGES/1705.misc.rst create mode 100644 aiogram/fsm/storage/pymongo.py create mode 100644 tests/test_fsm/storage/test_pymongo.py diff --git a/CHANGES/1705.misc.rst b/CHANGES/1705.misc.rst new file mode 100644 index 00000000..0173044e --- /dev/null +++ b/CHANGES/1705.misc.rst @@ -0,0 +1 @@ +Migrated `MongoStorage` from relying on deprecated `motor` package to using new async `PyMongo`. To use mongo storage with new async `PyMongo`, you need to install the `PyMongo` package instead of `motor` and just substitute deprecated `MongoStorage` with `PyMongoStorage` class, no other action needed. diff --git a/aiogram/fsm/storage/mongo.py b/aiogram/fsm/storage/mongo.py index b44aeeb9..ab1bf9fe 100644 --- a/aiogram/fsm/storage/mongo.py +++ b/aiogram/fsm/storage/mongo.py @@ -15,6 +15,12 @@ from aiogram.fsm.storage.base import ( class MongoStorage(BaseStorage): """ + + .. warning:: + DEPRECATED: Use :class:`PyMongoStorage` instead. + This class will be removed in future versions. + + MongoDB storage required :code:`motor` package installed (:code:`pip install motor`) """ diff --git a/aiogram/fsm/storage/pymongo.py b/aiogram/fsm/storage/pymongo.py new file mode 100644 index 00000000..89db25ce --- /dev/null +++ b/aiogram/fsm/storage/pymongo.py @@ -0,0 +1,136 @@ +from typing import Any, Dict, Mapping, Optional, cast + +from pymongo import AsyncMongoClient + +from aiogram.exceptions import DataNotDictLikeError +from aiogram.fsm.state import State +from aiogram.fsm.storage.base import ( + BaseStorage, + DefaultKeyBuilder, + KeyBuilder, + StateType, + StorageKey, +) + + +class PyMongoStorage(BaseStorage): + """ + MongoDB storage required :code:`pymongo` package installed (:code:`pip install pymongo`). + """ + + def __init__( + self, + client: AsyncMongoClient[Any], + key_builder: Optional[KeyBuilder] = None, + db_name: str = "aiogram_fsm", + collection_name: str = "states_and_data", + ) -> None: + """ + :param client: Instance of AsyncMongoClient + :param key_builder: builder that helps to convert contextual key to string + :param db_name: name of the MongoDB database for FSM + :param collection_name: name of the collection for storing FSM states and data + """ + if key_builder is None: + key_builder = DefaultKeyBuilder() + self._client = client + self._database = self._client[db_name] + self._collection = self._database[collection_name] + self._key_builder = key_builder + + @classmethod + def from_url( + cls, url: str, connection_kwargs: Optional[Dict[str, Any]] = None, **kwargs: Any + ) -> "PyMongoStorage": + """ + Create an instance of :class:`PyMongoStorage` with specifying the connection string + + :param url: for example :code:`mongodb://user:password@host:port` + :param connection_kwargs: see :code:`pymongo` docs + :param kwargs: arguments to be passed to :class:`PyMongoStorage` + :return: an instance of :class:`PyMongoStorage` + """ + if connection_kwargs is None: + connection_kwargs = {} + client: AsyncMongoClient[Any] = AsyncMongoClient(url, **connection_kwargs) + return cls(client=client, **kwargs) + + async def close(self) -> None: + """Cleanup client resources and disconnect from MongoDB.""" + return await self._client.close() + + def resolve_state(self, value: StateType) -> Optional[str]: + if value is None: + return None + if isinstance(value, State): + return value.state + return str(value) + + async def set_state(self, key: StorageKey, state: StateType = None) -> None: + document_id = self._key_builder.build(key) + if state is None: + updated = await self._collection.find_one_and_update( + filter={"_id": document_id}, + update={"$unset": {"state": 1}}, + projection={"_id": 0}, + return_document=True, + ) + if updated == {}: + await self._collection.delete_one({"_id": document_id}) + else: + await self._collection.update_one( + filter={"_id": document_id}, + update={"$set": {"state": self.resolve_state(state)}}, + upsert=True, + ) + + async def get_state(self, key: StorageKey) -> Optional[str]: + document_id = self._key_builder.build(key) + document = await self._collection.find_one({"_id": document_id}) + if document is None: + return None + return cast(Optional[str], document.get("state")) + + async def set_data(self, key: StorageKey, data: Mapping[str, Any]) -> None: + if not isinstance(data, dict): + raise DataNotDictLikeError( + f"Data must be a dict or dict-like object, got {type(data).__name__}" + ) + + document_id = self._key_builder.build(key) + if not data: + updated = await self._collection.find_one_and_update( + filter={"_id": document_id}, + update={"$unset": {"data": 1}}, + projection={"_id": 0}, + return_document=True, + ) + if updated == {}: + await self._collection.delete_one({"_id": document_id}) + else: + await self._collection.update_one( + filter={"_id": document_id}, + update={"$set": {"data": data}}, + upsert=True, + ) + + async def get_data(self, key: StorageKey) -> Dict[str, Any]: + document_id = self._key_builder.build(key) + document = await self._collection.find_one({"_id": document_id}) + if document is None or not document.get("data"): + return {} + return cast(Dict[str, Any], document["data"]) + + async def update_data(self, key: StorageKey, data: Mapping[str, Any]) -> Dict[str, Any]: + document_id = self._key_builder.build(key) + update_with = {f"data.{key}": value for key, value in data.items()} + update_result = await self._collection.find_one_and_update( + filter={"_id": document_id}, + update={"$set": update_with}, + upsert=True, + return_document=True, + projection={"_id": 0}, + ) + if not update_result: + await self._collection.delete_one({"_id": document_id}) + return cast(Dict[str, Any], update_result.get("data", {})) diff --git a/docs/dispatcher/finite_state_machine/storages.rst b/docs/dispatcher/finite_state_machine/storages.rst index 3795d2f1..ef39c4c3 100644 --- a/docs/dispatcher/finite_state_machine/storages.rst +++ b/docs/dispatcher/finite_state_machine/storages.rst @@ -22,6 +22,10 @@ RedisStorage MongoStorage ------------ +.. autoclass:: aiogram.fsm.storage.pymongo.PyMongoStorage + :members: __init__, from_url + :member-order: bysource + .. autoclass:: aiogram.fsm.storage.mongo.MongoStorage :members: __init__, from_url :member-order: bysource diff --git a/pyproject.toml b/pyproject.toml index 646b1d98..19db6402 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,7 @@ redis = [ ] mongo = [ "motor>=3.3.2,<3.7.0", + "pymongo>4.5,<4.11", ] proxy = [ "aiohttp-socks~=0.8.3", diff --git a/tests/conftest.py b/tests/conftest.py index fe5c3d22..6a0c37f4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -17,6 +17,7 @@ from aiogram.fsm.storage.memory import ( SimpleEventIsolation, ) from aiogram.fsm.storage.mongo import MongoStorage +from aiogram.fsm.storage.pymongo import PyMongoStorage from aiogram.fsm.storage.redis import RedisStorage from tests.mocked_bot import MockedBot @@ -102,6 +103,36 @@ async def mongo_storage(mongo_server): await storage.close() +@pytest.fixture() +def pymongo_server(request): + mongo_uri = request.config.getoption("--mongo") + if mongo_uri is None: + pytest.skip(SKIP_MESSAGE_PATTERN.format(db="mongo")) + else: + return mongo_uri + + +@pytest.fixture() +async def pymongo_storage(pymongo_server): + try: + parse_mongo_url(pymongo_server) + except InvalidURI as e: + raise UsageError(INVALID_URI_PATTERN.format(db="mongo", uri=pymongo_server, err=e)) + storage = PyMongoStorage.from_url( + url=pymongo_server, + connection_kwargs={"serverSelectionTimeoutMS": 2000}, + ) + try: + await storage._client.server_info() + except PyMongoError as e: + pytest.fail(str(e)) + else: + yield storage + await storage._client.drop_database(storage._database) + finally: + await storage.close() + + @pytest.fixture() async def memory_storage(): storage = MemoryStorage() diff --git a/tests/test_fsm/storage/test_pymongo.py b/tests/test_fsm/storage/test_pymongo.py new file mode 100644 index 00000000..574e7357 --- /dev/null +++ b/tests/test_fsm/storage/test_pymongo.py @@ -0,0 +1,218 @@ +import pytest +from pymongo.errors import PyMongoError + +from aiogram.fsm.state import State +from aiogram.fsm.storage.pymongo import PyMongoStorage, StorageKey +from tests.conftest import CHAT_ID, USER_ID + +PREFIX = "fsm" + + +async def test_get_storage_passing_only_url(pymongo_server): + storage = PyMongoStorage.from_url(url=pymongo_server) + try: + await storage._client.server_info() + except PyMongoError as e: + pytest.fail(str(e)) + + +async def test_pymongo_storage_close_does_not_throw(pymongo_server): + storage = PyMongoStorage.from_url(url=pymongo_server) + try: + assert await storage.close() is None + except Exception as e: + pytest.fail(f"close() raised an exception: {e}") + + +async def test_update_not_existing_data_with_empty_dictionary( + pymongo_storage: PyMongoStorage, + storage_key: StorageKey, +): + assert await pymongo_storage._collection.find_one({}) is None + assert await pymongo_storage.update_data(key=storage_key, data={}) == {} + assert await pymongo_storage._collection.find_one({}) is None + + +async def test_update_not_existing_data_with_non_empty_dictionary( + pymongo_storage: PyMongoStorage, + storage_key: StorageKey, +): + assert await pymongo_storage._collection.find_one({}) is None + assert await pymongo_storage.update_data(key=storage_key, data={"key": "value"}) == { + "key": "value" + } + assert await pymongo_storage._collection.find_one({}) == { + "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}", + "data": {"key": "value"}, + } + await pymongo_storage._collection.delete_one({}) + + +async def test_update_existing_data_with_empty_dictionary( + pymongo_storage: PyMongoStorage, + storage_key: StorageKey, +): + assert await pymongo_storage._collection.find_one({}) is None + await pymongo_storage.set_data(key=storage_key, data={"key": "value"}) + assert await pymongo_storage.update_data(key=storage_key, data={}) == {"key": "value"} + assert await pymongo_storage._collection.find_one({}) == { + "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}", + "data": {"key": "value"}, + } + await pymongo_storage._collection.delete_one({}) + + +async def test_update_existing_data_with_non_empty_dictionary( + pymongo_storage: PyMongoStorage, + storage_key: StorageKey, +): + assert await pymongo_storage._collection.find_one({}) is None + await pymongo_storage.set_data(key=storage_key, data={"key": "value"}) + assert await pymongo_storage.update_data(key=storage_key, data={"key": "new_value"}) == { + "key": "new_value" + } + assert await pymongo_storage._collection.find_one({}) == { + "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}", + "data": {"key": "new_value"}, + } + await pymongo_storage._collection.delete_one({}) + + +async def test_document_life_cycle( + pymongo_storage: PyMongoStorage, + storage_key: StorageKey, +): + assert await pymongo_storage._collection.find_one({}) is None + await pymongo_storage.set_state(storage_key, "test") + await pymongo_storage.set_data(storage_key, {"key": "value"}) + assert await pymongo_storage._collection.find_one({}) == { + "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}", + "state": "test", + "data": {"key": "value"}, + } + await pymongo_storage.set_state(storage_key, None) + assert await pymongo_storage._collection.find_one({}) == { + "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}", + "data": {"key": "value"}, + } + await pymongo_storage.set_data(storage_key, {}) + assert await pymongo_storage._collection.find_one({}) is None + + +class TestStateAndDataDoNotAffectEachOther: + async def test_state_and_data_do_not_affect_each_other_while_getting( + self, + pymongo_storage: PyMongoStorage, + storage_key: StorageKey, + ): + assert await pymongo_storage._collection.find_one({}) is None + await pymongo_storage.set_state(storage_key, "test") + await pymongo_storage.set_data(storage_key, {"key": "value"}) + assert await pymongo_storage.get_state(storage_key) == "test" + assert await pymongo_storage.get_data(storage_key) == {"key": "value"} + + async def test_data_do_not_affect_to_deleted_state_getting( + self, + pymongo_storage: PyMongoStorage, + storage_key: StorageKey, + ): + await pymongo_storage.set_state(storage_key, "test") + await pymongo_storage.set_data(storage_key, {"key": "value"}) + await pymongo_storage.set_state(storage_key, None) + assert await pymongo_storage.get_state(storage_key) is None + + async def test_state_do_not_affect_to_deleted_data_getting( + self, + pymongo_storage: PyMongoStorage, + storage_key: StorageKey, + ): + await pymongo_storage.set_state(storage_key, "test") + await pymongo_storage.set_data(storage_key, {"key": "value"}) + await pymongo_storage.set_data(storage_key, {}) + assert await pymongo_storage.get_data(storage_key) == {} + + async def test_state_do_not_affect_to_updating_not_existing_data_with_empty_dictionary( + self, + pymongo_storage: PyMongoStorage, + storage_key: StorageKey, + ): + await pymongo_storage.set_state(storage_key, "test") + assert await pymongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test" + } + assert await pymongo_storage.update_data(key=storage_key, data={}) == {} + assert await pymongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test" + } + + async def test_state_do_not_affect_to_updating_not_existing_data_with_non_empty_dictionary( + self, + pymongo_storage: PyMongoStorage, + storage_key: StorageKey, + ): + await pymongo_storage.set_state(storage_key, "test") + assert await pymongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test" + } + assert await pymongo_storage.update_data( + key=storage_key, + data={"key": "value"}, + ) == {"key": "value"} + assert await pymongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test", + "data": {"key": "value"}, + } + + async def test_state_do_not_affect_to_updating_existing_data_with_empty_dictionary( + self, + pymongo_storage: PyMongoStorage, + storage_key: StorageKey, + ): + await pymongo_storage.set_state(storage_key, "test") + await pymongo_storage.set_data(storage_key, {"key": "value"}) + assert await pymongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test", + "data": {"key": "value"}, + } + assert await pymongo_storage.update_data(key=storage_key, data={}) == {"key": "value"} + assert await pymongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test", + "data": {"key": "value"}, + } + + async def test_state_do_not_affect_to_updating_existing_data_with_non_empty_dictionary( + self, + pymongo_storage: PyMongoStorage, + storage_key: StorageKey, + ): + await pymongo_storage.set_state(storage_key, "test") + await pymongo_storage.set_data(storage_key, {"key": "value"}) + assert await pymongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test", + "data": {"key": "value"}, + } + assert await pymongo_storage.update_data( + key=storage_key, + data={"key": "VALUE", "key_2": "value_2"}, + ) == {"key": "VALUE", "key_2": "value_2"} + assert await pymongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test", + "data": {"key": "VALUE", "key_2": "value_2"}, + } + + +@pytest.mark.parametrize( + "value,result", + [ + [None, None], + ["", ""], + ["text", "text"], + [State(), None], + [State(state="*"), "*"], + [State("text"), "@:text"], + [State("test", group_name="Test"), "Test:test"], + [[1, 2, 3], "[1, 2, 3]"], + ], +) +def test_resolve_state(value, result, pymongo_storage: PyMongoStorage): + assert pymongo_storage.resolve_state(value) == result diff --git a/tests/test_fsm/storage/test_storages.py b/tests/test_fsm/storage/test_storages.py index 44424ed2..884f6874 100644 --- a/tests/test_fsm/storage/test_storages.py +++ b/tests/test_fsm/storage/test_storages.py @@ -11,6 +11,7 @@ from aiogram.fsm.storage.base import BaseStorage, StorageKey [ pytest.lazy_fixture("redis_storage"), pytest.lazy_fixture("mongo_storage"), + pytest.lazy_fixture("pymongo_storage"), pytest.lazy_fixture("memory_storage"), ], )