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"), ], )