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 5d0f6a9dfb.

* added optional dependency to pymongo to pyproject.toml

* Revert "migrated mongo storage from using deprecated motor to PyMongo"

This reverts commit 1c0207e1d1.

* 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 86170e1cb9.

* 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

@@ -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.

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`)
     """
 

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", {}))

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

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

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()

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

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