mirror of
https://github.com/aiogram/aiogram.git
synced 2026-04-08 16:37:47 +00:00
Storing FSM state and data together in MongoDB-storage
This commit is contained in:
Rishat Fayzullin
parent
7c18c87b95
commit
ce48b825e8
4 changed files with 199 additions and 44 deletions
|
|
@ -23,7 +23,11 @@ class KeyBuilder(ABC):
|
|||
"""Base class for key builder."""
|
||||
|
||||
@abstractmethod
|
||||
def build(self, key: StorageKey, part: Literal["data", "state", "lock"]) -> str:
|
||||
def build(
|
||||
self,
|
||||
key: StorageKey,
|
||||
part: Optional[Literal["data", "state", "lock"]] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Build key to be used in storage's db queries
|
||||
|
||||
|
|
@ -61,7 +65,11 @@ class DefaultKeyBuilder(KeyBuilder):
|
|||
self.with_bot_id = with_bot_id
|
||||
self.with_destiny = with_destiny
|
||||
|
||||
def build(self, key: StorageKey, part: Literal["data", "state", "lock"]) -> str:
|
||||
def build(
|
||||
self,
|
||||
key: StorageKey,
|
||||
part: Optional[Literal["data", "state", "lock"]] = None,
|
||||
) -> str:
|
||||
parts = [self.prefix]
|
||||
if self.with_bot_id:
|
||||
parts.append(str(key.bot_id))
|
||||
|
|
@ -77,7 +85,8 @@ class DefaultKeyBuilder(KeyBuilder):
|
|||
"\n\nProbably, you should set `with_destiny=True` in for DefaultKeyBuilder."
|
||||
)
|
||||
raise ValueError(error_message)
|
||||
parts.append(part)
|
||||
if part:
|
||||
parts.append(part)
|
||||
return self.separator.join(parts)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -22,24 +22,19 @@ class MongoStorage(BaseStorage):
|
|||
client: AsyncIOMotorClient,
|
||||
key_builder: Optional[KeyBuilder] = None,
|
||||
db_name: str = "aiogram_fsm",
|
||||
states_collection_name: str = "states",
|
||||
data_collection_name: str = "data",
|
||||
collection_name: str = "states_and_data",
|
||||
) -> None:
|
||||
"""
|
||||
:param client: Instance of AsyncIOMotorClient
|
||||
:param key_builder: builder that helps to convert contextual key to string
|
||||
:param db_name: name of the MongoDB database for FSM
|
||||
:param states_collection_name: name of the collection for storing FSM states.
|
||||
:param data_collection_name: name of the collection for storing additional data.
|
||||
: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._db_name = db_name
|
||||
self._states_collection: AsyncIOMotorCollection = self._client[db_name][
|
||||
states_collection_name
|
||||
]
|
||||
self._data_collection: AsyncIOMotorCollection = self._client[db_name][data_collection_name]
|
||||
self._collection: AsyncIOMotorCollection = self._client[db_name][collection_name]
|
||||
self._key_builder = key_builder
|
||||
|
||||
@classmethod
|
||||
|
|
@ -71,50 +66,65 @@ class MongoStorage(BaseStorage):
|
|||
return str(value)
|
||||
|
||||
async def set_state(self, key: StorageKey, state: StateType = None) -> None:
|
||||
document_id = self._key_builder.build(key, "state")
|
||||
document_id = self._key_builder.build(key)
|
||||
if state is None:
|
||||
await self._states_collection.delete_one({"_id": document_id})
|
||||
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._states_collection.update_one(
|
||||
{"_id": document_id},
|
||||
{"$set": {"state": self.resolve_state(state)}},
|
||||
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, "state")
|
||||
document = await self._states_collection.find_one({"_id": document_id})
|
||||
if document is None or document["state"] is None:
|
||||
document_id = self._key_builder.build(key)
|
||||
document = await self._collection.find_one({"_id": document_id})
|
||||
if document is None or document.get("state") is None:
|
||||
return None
|
||||
return str(document["state"])
|
||||
|
||||
async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None:
|
||||
document_id = self._key_builder.build(key, "data")
|
||||
document_id = self._key_builder.build(key)
|
||||
if not data:
|
||||
await self._data_collection.delete_one({"_id": document_id})
|
||||
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._data_collection.update_one(
|
||||
{"_id": document_id},
|
||||
{"$set": data},
|
||||
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, "data")
|
||||
document = await self._data_collection.find_one({"_id": document_id}, {"_id": 0})
|
||||
if not document:
|
||||
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)
|
||||
return cast(Dict[str, Any], document["data"])
|
||||
|
||||
async def update_data(self, key: StorageKey, data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
document_id = self._key_builder.build(key, "data")
|
||||
update_result = await self._data_collection.find_one_and_update(
|
||||
{"_id": document_id},
|
||||
{"$set": data},
|
||||
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._data_collection.delete_one({"_id": document_id})
|
||||
return cast(Dict[str, Any], update_result)
|
||||
await self._collection.delete_one({"_id": document_id})
|
||||
return update_result.get("data", {})
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
from typing import Literal, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from aiogram.fsm.storage.base import DEFAULT_DESTINY, StorageKey
|
||||
|
|
@ -13,22 +15,29 @@ FIELD = "data"
|
|||
|
||||
class TestDefaultKeyBuilder:
|
||||
@pytest.mark.parametrize(
|
||||
"with_bot_id,with_destiny,result",
|
||||
"with_bot_id,with_destiny,field,result",
|
||||
[
|
||||
[False, False, f"{PREFIX}:{CHAT_ID}:{USER_ID}:{FIELD}"],
|
||||
[True, False, f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{FIELD}"],
|
||||
[True, True, f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}"],
|
||||
[False, True, f"{PREFIX}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}"],
|
||||
[False, False, FIELD, f"{PREFIX}:{CHAT_ID}:{USER_ID}:{FIELD}"],
|
||||
[True, False, FIELD, f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{FIELD}"],
|
||||
[True, True, FIELD, f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}"],
|
||||
[False, True, FIELD, f"{PREFIX}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}"],
|
||||
[False, False, None, f"{PREFIX}:{CHAT_ID}:{USER_ID}"],
|
||||
],
|
||||
)
|
||||
async def test_generate_key(self, with_bot_id: bool, with_destiny: bool, result: str):
|
||||
async def test_generate_key(
|
||||
self,
|
||||
with_bot_id: bool,
|
||||
with_destiny: bool,
|
||||
field: Optional[Literal["data", "state", "lock"]],
|
||||
result: str,
|
||||
):
|
||||
key_builder = DefaultKeyBuilder(
|
||||
prefix=PREFIX,
|
||||
with_bot_id=with_bot_id,
|
||||
with_destiny=with_destiny,
|
||||
)
|
||||
key = StorageKey(chat_id=CHAT_ID, user_id=USER_ID, bot_id=BOT_ID, destiny=DEFAULT_DESTINY)
|
||||
assert key_builder.build(key, FIELD) == result
|
||||
assert key_builder.build(key, field) == result
|
||||
|
||||
async def test_destiny_check(self):
|
||||
key_builder = DefaultKeyBuilder(
|
||||
|
|
|
|||
|
|
@ -4,20 +4,147 @@ from aiogram.fsm.state import State
|
|||
from aiogram.fsm.storage.mongo import MongoStorage, StorageKey
|
||||
from tests.mocked_bot import MockedBot
|
||||
|
||||
PREFIX = "fsm"
|
||||
CHAT_ID = -42
|
||||
USER_ID = 42
|
||||
|
||||
|
||||
@pytest.fixture(name="storage_key")
|
||||
def create_storage_key(bot: MockedBot):
|
||||
return StorageKey(chat_id=-42, user_id=42, bot_id=bot.id)
|
||||
return StorageKey(chat_id=CHAT_ID, user_id=USER_ID, bot_id=bot.id)
|
||||
|
||||
|
||||
async def test_update_not_existing_data_with_empty_dictionary(
|
||||
mongo_storage: MongoStorage,
|
||||
storage_key: StorageKey,
|
||||
):
|
||||
assert await mongo_storage._data_collection.find_one({}) is None
|
||||
assert await mongo_storage._collection.find_one({}) is None
|
||||
assert await mongo_storage.get_data(key=storage_key) == {}
|
||||
assert await mongo_storage.update_data(key=storage_key, data={}) == {}
|
||||
assert await mongo_storage._data_collection.find_one({}) is None
|
||||
assert await mongo_storage._collection.find_one({}) is None
|
||||
|
||||
|
||||
async def test_document_life_cycle(
|
||||
mongo_storage: MongoStorage,
|
||||
storage_key: StorageKey,
|
||||
):
|
||||
assert await mongo_storage._collection.find_one({}) is None
|
||||
await mongo_storage.set_state(storage_key, "test")
|
||||
await mongo_storage.set_data(storage_key, {"key": "value"})
|
||||
assert await mongo_storage._collection.find_one({}) == {
|
||||
"_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}",
|
||||
"state": "test",
|
||||
"data": {"key": "value"},
|
||||
}
|
||||
await mongo_storage.set_state(storage_key, None)
|
||||
assert await mongo_storage._collection.find_one({}) == {
|
||||
"_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}",
|
||||
"data": {"key": "value"},
|
||||
}
|
||||
await mongo_storage.set_data(storage_key, {})
|
||||
assert await mongo_storage._collection.find_one({}) is None
|
||||
|
||||
|
||||
class TestStateAndDataDoNotAffectEachOther:
|
||||
async def test_state_and_data_do_not_affect_each_other_while_getting(
|
||||
self,
|
||||
mongo_storage: MongoStorage,
|
||||
storage_key: StorageKey,
|
||||
):
|
||||
assert await mongo_storage._collection.find_one({}) is None
|
||||
await mongo_storage.set_state(storage_key, "test")
|
||||
await mongo_storage.set_data(storage_key, {"key": "value"})
|
||||
assert await mongo_storage.get_state(storage_key) == "test"
|
||||
assert await mongo_storage.get_data(storage_key) == {"key": "value"}
|
||||
|
||||
async def test_data_do_not_affect_to_deleted_state_getting(
|
||||
self,
|
||||
mongo_storage: MongoStorage,
|
||||
storage_key: StorageKey,
|
||||
):
|
||||
await mongo_storage.set_state(storage_key, "test")
|
||||
await mongo_storage.set_data(storage_key, {"key": "value"})
|
||||
await mongo_storage.set_state(storage_key, None)
|
||||
assert await mongo_storage.get_state(storage_key) is None
|
||||
|
||||
async def test_state_do_not_affect_to_deleted_data_getting(
|
||||
self,
|
||||
mongo_storage: MongoStorage,
|
||||
storage_key: StorageKey,
|
||||
):
|
||||
await mongo_storage.set_state(storage_key, "test")
|
||||
await mongo_storage.set_data(storage_key, {"key": "value"})
|
||||
await mongo_storage.set_data(storage_key, {})
|
||||
assert await mongo_storage.get_data(storage_key) == {}
|
||||
|
||||
async def test_state_do_not_affect_to_updating_not_existing_data_with_empty_dictionary(
|
||||
self,
|
||||
mongo_storage: MongoStorage,
|
||||
storage_key: StorageKey,
|
||||
):
|
||||
await mongo_storage.set_state(storage_key, "test")
|
||||
assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == {
|
||||
"state": "test"
|
||||
}
|
||||
assert await mongo_storage.update_data(key=storage_key, data={}) == {}
|
||||
assert await mongo_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,
|
||||
mongo_storage: MongoStorage,
|
||||
storage_key: StorageKey,
|
||||
):
|
||||
await mongo_storage.set_state(storage_key, "test")
|
||||
assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == {
|
||||
"state": "test"
|
||||
}
|
||||
assert await mongo_storage.update_data(
|
||||
key=storage_key,
|
||||
data={"key": "value"},
|
||||
) == {"key": "value"}
|
||||
assert await mongo_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,
|
||||
mongo_storage: MongoStorage,
|
||||
storage_key: StorageKey,
|
||||
):
|
||||
await mongo_storage.set_state(storage_key, "test")
|
||||
await mongo_storage.set_data(storage_key, {"key": "value"})
|
||||
assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == {
|
||||
"state": "test",
|
||||
"data": {"key": "value"},
|
||||
}
|
||||
assert await mongo_storage.update_data(key=storage_key, data={}) == {"key": "value"}
|
||||
assert await mongo_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,
|
||||
mongo_storage: MongoStorage,
|
||||
storage_key: StorageKey,
|
||||
):
|
||||
await mongo_storage.set_state(storage_key, "test")
|
||||
await mongo_storage.set_data(storage_key, {"key": "value"})
|
||||
assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == {
|
||||
"state": "test",
|
||||
"data": {"key": "value"},
|
||||
}
|
||||
assert await mongo_storage.update_data(
|
||||
key=storage_key,
|
||||
data={"key": "VALUE", "key_2": "value_2"},
|
||||
) == {"key": "VALUE", "key_2": "value_2"}
|
||||
assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == {
|
||||
"state": "test",
|
||||
"data": {"key": "VALUE", "key_2": "value_2"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue