mirror of
https://github.com/aiogram/aiogram.git
synced 2026-04-08 16:37:47 +00:00
Added MongoStorage for FSM (#1434)
* Mongo storage included to storages test * Added few additional checks in storages test * Added MongoStorage for FSM * Added changes description * Fixed error message syntax Co-authored-by: Alex Root Junior <jroot.junior@gmail.com> * Resolved mypy check error * IF/ELSE statement simplified * Fix ruff linter error: RET505 Unnecessary `elif` after `return` statement * Fix ruff linter error: E501 Line too long (100 > 99) * Added mongo storage testing in CI * Refactoring while review * Refactoring while review * Storing FSM state and data together in MongoDB-storage * Fix CI - MongoDB container action is only supported on Linux * Refactoring while review * Enable Macos in pypy-tests section of CI * Refactoring while review * Makefile updated * redis and mongo storages tests do not run in pypy-tests job of CI * Fix docstring of DefaultKeyBuilder --------- Co-authored-by: Alex Root Junior <jroot.junior@gmail.com>
This commit is contained in:
Rishat-F
GitHub
parent
25c76b7d74
commit
1ef7655fd7
15 changed files with 522 additions and 129 deletions
|
|
@ -1,7 +1,7 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, AsyncGenerator, Dict, Optional, Union
|
||||
from typing import Any, AsyncGenerator, Dict, Literal, Optional, Union
|
||||
|
||||
from aiogram.fsm.state import State
|
||||
|
||||
|
|
@ -20,6 +20,85 @@ class StorageKey:
|
|||
destiny: str = DEFAULT_DESTINY
|
||||
|
||||
|
||||
class KeyBuilder(ABC):
|
||||
"""Base class for key builder."""
|
||||
|
||||
@abstractmethod
|
||||
def build(
|
||||
self,
|
||||
key: StorageKey,
|
||||
part: Optional[Literal["data", "state", "lock"]] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Build key to be used in storage's db queries
|
||||
|
||||
:param key: contextual key
|
||||
:param part: part of the record
|
||||
:return: key to be used in storage's db queries
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class DefaultKeyBuilder(KeyBuilder):
|
||||
"""
|
||||
Simple key builder with default prefix.
|
||||
|
||||
Generates a colon-joined string with prefix, chat_id, user_id,
|
||||
optional bot_id, business_connection_id, destiny and field.
|
||||
|
||||
Format:
|
||||
:code:`<prefix>:<bot_id?>:<business_connection_id?>:<chat_id>:<user_id>:<destiny?>:<field?>`
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
prefix: str = "fsm",
|
||||
separator: str = ":",
|
||||
with_bot_id: bool = False,
|
||||
with_business_connection_id: bool = False,
|
||||
with_destiny: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
:param prefix: prefix for all records
|
||||
:param separator: separator
|
||||
:param with_bot_id: include Bot id in the key
|
||||
:param with_business_connection_id: include business connection id
|
||||
:param with_destiny: include destiny key
|
||||
"""
|
||||
self.prefix = prefix
|
||||
self.separator = separator
|
||||
self.with_bot_id = with_bot_id
|
||||
self.with_business_connection_id = with_business_connection_id
|
||||
self.with_destiny = with_destiny
|
||||
|
||||
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))
|
||||
if self.with_business_connection_id and key.business_connection_id:
|
||||
parts.append(str(key.business_connection_id))
|
||||
parts.append(str(key.chat_id))
|
||||
if key.thread_id:
|
||||
parts.append(str(key.thread_id))
|
||||
parts.append(str(key.user_id))
|
||||
if self.with_destiny:
|
||||
parts.append(key.destiny)
|
||||
elif key.destiny != DEFAULT_DESTINY:
|
||||
error_message = (
|
||||
"Default key builder is not configured to use key destiny other than the default."
|
||||
"\n\nProbably, you should set `with_destiny=True` in for DefaultKeyBuilder."
|
||||
)
|
||||
raise ValueError(error_message)
|
||||
if part:
|
||||
parts.append(part)
|
||||
return self.separator.join(parts)
|
||||
|
||||
|
||||
class BaseStorage(ABC):
|
||||
"""
|
||||
Base class for all FSM storages
|
||||
|
|
|
|||
130
aiogram/fsm/storage/mongo.py
Normal file
130
aiogram/fsm/storage/mongo.py
Normal file
|
|
@ -0,0 +1,130 @@
|
|||
from typing import Any, Dict, Optional, cast
|
||||
|
||||
from motor.motor_asyncio import AsyncIOMotorClient
|
||||
|
||||
from aiogram.fsm.state import State
|
||||
from aiogram.fsm.storage.base import (
|
||||
BaseStorage,
|
||||
DefaultKeyBuilder,
|
||||
KeyBuilder,
|
||||
StateType,
|
||||
StorageKey,
|
||||
)
|
||||
|
||||
|
||||
class MongoStorage(BaseStorage):
|
||||
"""
|
||||
MongoDB storage required :code:`motor` package installed (:code:`pip install motor`)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: AsyncIOMotorClient,
|
||||
key_builder: Optional[KeyBuilder] = None,
|
||||
db_name: str = "aiogram_fsm",
|
||||
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 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
|
||||
) -> "MongoStorage":
|
||||
"""
|
||||
Create an instance of :class:`MongoStorage` with specifying the connection string
|
||||
|
||||
:param url: for example :code:`mongodb://user:password@host:port`
|
||||
:param connection_kwargs: see :code:`motor` docs
|
||||
:param kwargs: arguments to be passed to :class:`MongoStorage`
|
||||
:return: an instance of :class:`MongoStorage`
|
||||
"""
|
||||
if connection_kwargs is None:
|
||||
connection_kwargs = {}
|
||||
client = AsyncIOMotorClient(url, **connection_kwargs)
|
||||
return cls(client=client, **kwargs)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Cleanup client resources and disconnect from MongoDB."""
|
||||
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 document.get("state")
|
||||
|
||||
async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None:
|
||||
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: Dict[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 update_result.get("data", {})
|
||||
|
|
@ -1,7 +1,6 @@
|
|||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, AsyncGenerator, Callable, Dict, Literal, Optional, cast
|
||||
from typing import Any, AsyncGenerator, Callable, Dict, Optional, cast
|
||||
|
||||
from redis.asyncio.client import Redis
|
||||
from redis.asyncio.connection import ConnectionPool
|
||||
|
|
@ -10,9 +9,10 @@ from redis.typing import ExpiryT
|
|||
|
||||
from aiogram.fsm.state import State
|
||||
from aiogram.fsm.storage.base import (
|
||||
DEFAULT_DESTINY,
|
||||
BaseEventIsolation,
|
||||
BaseStorage,
|
||||
DefaultKeyBuilder,
|
||||
KeyBuilder,
|
||||
StateType,
|
||||
StorageKey,
|
||||
)
|
||||
|
|
@ -22,79 +22,6 @@ _JsonLoads = Callable[..., Any]
|
|||
_JsonDumps = Callable[..., str]
|
||||
|
||||
|
||||
class KeyBuilder(ABC):
|
||||
"""
|
||||
Base class for Redis key builder
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def build(self, key: StorageKey, part: Literal["data", "state", "lock"]) -> str:
|
||||
"""
|
||||
This method should be implemented in subclasses
|
||||
|
||||
:param key: contextual key
|
||||
:param part: part of the record
|
||||
:return: key to be used in Redis queries
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class DefaultKeyBuilder(KeyBuilder):
|
||||
"""
|
||||
Simple Redis key builder with default prefix.
|
||||
|
||||
Generates a colon-joined string with prefix, chat_id, user_id,
|
||||
optional bot_id, business_connection_id and destiny.
|
||||
|
||||
Format:
|
||||
:code:`<prefix>:<bot_id?>:<business_connection_id?>:<chat_id>:<user_id>:<destiny?>:<field>`
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
prefix: str = "fsm",
|
||||
separator: str = ":",
|
||||
with_bot_id: bool = False,
|
||||
with_business_connection_id: bool = False,
|
||||
with_destiny: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
:param prefix: prefix for all records
|
||||
:param separator: separator
|
||||
:param with_bot_id: include Bot id in the key
|
||||
:param with_business_connection_id: include business connection id
|
||||
:param with_destiny: include a destiny key
|
||||
"""
|
||||
self.prefix = prefix
|
||||
self.separator = separator
|
||||
self.with_bot_id = with_bot_id
|
||||
self.with_business_connection_id = with_business_connection_id
|
||||
self.with_destiny = with_destiny
|
||||
|
||||
def build(self, key: StorageKey, part: Literal["data", "state", "lock"]) -> str:
|
||||
parts = [self.prefix]
|
||||
if self.with_bot_id:
|
||||
parts.append(str(key.bot_id))
|
||||
if self.with_business_connection_id and key.business_connection_id:
|
||||
parts.append(str(key.business_connection_id))
|
||||
parts.append(str(key.chat_id))
|
||||
if key.thread_id:
|
||||
parts.append(str(key.thread_id))
|
||||
parts.append(str(key.user_id))
|
||||
if self.with_destiny:
|
||||
parts.append(key.destiny)
|
||||
elif key.destiny != DEFAULT_DESTINY:
|
||||
raise ValueError(
|
||||
"Redis key builder is not configured to use key destiny other the default.\n"
|
||||
"\n"
|
||||
"Probably, you should set `with_destiny=True` in for DefaultKeyBuilder.\n"
|
||||
"E.g: `RedisStorage(redis, key_builder=DefaultKeyBuilder(with_destiny=True))`"
|
||||
)
|
||||
parts.append(part)
|
||||
return self.separator.join(parts)
|
||||
|
||||
|
||||
class RedisStorage(BaseStorage):
|
||||
"""
|
||||
Redis storage required :code:`redis` package installed (:code:`pip install redis`)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue