python/aiogram
1
0
Fork 0
mirror of https://github.com/aiogram/aiogram.git synced 2026-04-08 16:37:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

storage cleanup (#1144)

Browse source 
* storage cleanup

* storage cleanup
This commit is contained in:
RootShinobi 2023-04-08 18:01:11 +03:00 • committed by GitHub
parent d8a977f357
commit dbaf6fabcb
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 54 additions and 75 deletions
Download patch file Download diff file Expand all files Collapse all files

14
aiogram/fsm/context.py
View file

@ -1,33 +1,31 @@
from typing import Any, Dict, Optional
from aiogram import Bot
from aiogram.fsm.storage.base import BaseStorage, StateType, StorageKey
class FSMContext:
def __init__(self, bot: Bot, storage: BaseStorage, key: StorageKey) -> None:
self.bot = bot
def __init__(self, storage: BaseStorage, key: StorageKey) -> None:
self.storage = storage
self.key = key
async def set_state(self, state: StateType = None) -> None:
await self.storage.set_state(bot=self.bot, key=self.key, state=state)
await self.storage.set_state(key=self.key, state=state)
async def get_state(self) -> Optional[str]:
return await self.storage.get_state(bot=self.bot, key=self.key)
return await self.storage.get_state(key=self.key)
async def set_data(self, data: Dict[str, Any]) -> None:
await self.storage.set_data(bot=self.bot, key=self.key, data=data)
await self.storage.set_data(key=self.key, data=data)
async def get_data(self) -> Dict[str, Any]:
return await self.storage.get_data(bot=self.bot, key=self.key)
return await self.storage.get_data(key=self.key)
async def update_data(
self, data: Optional[Dict[str, Any]] = None, **kwargs: Any
) -> Dict[str, Any]:
if data:
kwargs.update(data)
return await self.storage.update_data(bot=self.bot, key=self.key, data=kwargs)
return await self.storage.update_data(key=self.key, data=kwargs)
async def clear(self) -> None:
await self.set_state(state=None)

3
aiogram/fsm/middleware.py
View file

@ -35,7 +35,7 @@ class FSMContextMiddleware(BaseMiddleware):
data["fsm_storage"] = self.storage
if context:
data.update({"state": context, "raw_state": await context.get_state()})
async with self.events_isolation.lock(bot=bot, key=context.key):
async with self.events_isolation.lock(key=context.key):
return await handler(event, data)
return await handler(event, data)
@ -76,7 +76,6 @@ class FSMContextMiddleware(BaseMiddleware):
destiny: str = DEFAULT_DESTINY,
) -> FSMContext:
return FSMContext(
bot=bot,
storage=self.storage,
key=StorageKey(
user_id=user_id,

23
aiogram/fsm/storage/base.py
View file

@ -3,7 +3,6 @@ from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Any, AsyncGenerator, Dict, Optional, Union
from aiogram import Bot
from aiogram.fsm.state import State
StateType = Optional[Union[str, State]]
@ -25,61 +24,56 @@ class BaseStorage(ABC):
"""
@abstractmethod
async def set_state(self, bot: Bot, key: StorageKey, state: StateType = None) -> None:
async def set_state(self, key: StorageKey, state: StateType = None) -> None:
"""
Set state for specified key
:param bot: instance of the current bot
:param key: storage key
:param state: new state
"""
pass
@abstractmethod
async def get_state(self, bot: Bot, key: StorageKey) -> Optional[str]:
async def get_state(self, key: StorageKey) -> Optional[str]:
"""
Get key state
:param bot: instance of the current bot
:param key: storage key
:return: current state
"""
pass
@abstractmethod
async def set_data(self, bot: Bot, key: StorageKey, data: Dict[str, Any]) -> None:
async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None:
"""
Write data (replace)
:param bot: instance of the current bot
:param key: storage key
:param data: new data
"""
pass
@abstractmethod
async def get_data(self, bot: Bot, key: StorageKey) -> Dict[str, Any]:
async def get_data(self, key: StorageKey) -> Dict[str, Any]:
"""
Get current data for key
:param bot: instance of the current bot
:param key: storage key
:return: current data
"""
pass
async def update_data(self, bot: Bot, key: StorageKey, data: Dict[str, Any]) -> Dict[str, Any]:
async def update_data(self, key: StorageKey, data: Dict[str, Any]) -> Dict[str, Any]:
"""
Update date in the storage for key (like dict.update)
:param bot: instance of the current bot
:param key: storage key
:param data: partial data
:return: new data
"""
current_data = await self.get_data(bot=bot, key=key)
current_data = await self.get_data(key=key)
current_data.update(data)
await self.set_data(bot=bot, key=key, data=current_data)
await self.set_data(key=key, data=current_data)
return current_data.copy()
@abstractmethod
@ -93,12 +87,11 @@ class BaseStorage(ABC):
class BaseEventIsolation(ABC):
@abstractmethod
@asynccontextmanager
async def lock(self, bot: Bot, key: StorageKey) -> AsyncGenerator[None, None]:
async def lock(self, key: StorageKey) -> AsyncGenerator[None, None]:
"""
Isolate events with lock.
Will be used as context manager
:param bot: instance of the current bot
:param key: storage key
:return: An async generator
"""

13
aiogram/fsm/storage/memory.py
View file

@ -4,7 +4,6 @@ from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import Any, AsyncGenerator, DefaultDict, Dict, Hashable, Optional
from aiogram import Bot
from aiogram.fsm.state import State
from aiogram.fsm.storage.base import (
BaseEventIsolation,
@ -38,22 +37,22 @@ class MemoryStorage(BaseStorage):
async def close(self) -> None:
pass
async def set_state(self, bot: Bot, key: StorageKey, state: StateType = None) -> None:
async def set_state(self, key: StorageKey, state: StateType = None) -> None:
self.storage[key].state = state.state if isinstance(state, State) else state
async def get_state(self, bot: Bot, key: StorageKey) -> Optional[str]:
async def get_state(self, key: StorageKey) -> Optional[str]:
return self.storage[key].state
async def set_data(self, bot: Bot, key: StorageKey, data: Dict[str, Any]) -> None:
async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None:
self.storage[key].data = data.copy()
async def get_data(self, bot: Bot, key: StorageKey) -> Dict[str, Any]:
async def get_data(self, key: StorageKey) -> Dict[str, Any]:
return self.storage[key].data.copy()
class DisabledEventIsolation(BaseEventIsolation):
@asynccontextmanager
async def lock(self, bot: Bot, key: StorageKey) -> AsyncGenerator[None, None]:
async def lock(self, key: StorageKey) -> AsyncGenerator[None, None]:
yield
async def close(self) -> None:
@ -66,7 +65,7 @@ class SimpleEventIsolation(BaseEventIsolation):
self._locks: DefaultDict[Hashable, Lock] = defaultdict(Lock)
@asynccontextmanager
async def lock(self, bot: Bot, key: StorageKey) -> AsyncGenerator[None, None]:
async def lock(self, key: StorageKey) -> AsyncGenerator[None, None]:
lock = self._locks[key]
async with lock:
yield

20
aiogram/fsm/storage/redis.py
View file

@ -1,13 +1,13 @@
import json
from abc import ABC, abstractmethod
from contextlib import asynccontextmanager
from typing import Any, AsyncGenerator, Dict, Literal, Optional, cast
from typing import Any, AsyncGenerator, Callable, Dict, Literal, Optional, cast
from redis.asyncio.client import Redis
from redis.asyncio.connection import ConnectionPool
from redis.asyncio.lock import Lock
from redis.typing import ExpiryT
from aiogram import Bot
from aiogram.fsm.state import State
from aiogram.fsm.storage.base import (
DEFAULT_DESTINY,
@ -18,6 +18,8 @@ from aiogram.fsm.storage.base import (
)
DEFAULT_REDIS_LOCK_KWARGS = {"timeout": 60}
_JsonLoads = Callable[..., Any]
_JsonDumps = Callable[..., str]
class KeyBuilder(ABC):
@ -93,13 +95,14 @@ class RedisStorage(BaseStorage):
key_builder: Optional[KeyBuilder] = None,
state_ttl: Optional[ExpiryT] = None,
data_ttl: Optional[ExpiryT] = None,
json_loads: _JsonLoads = json.loads,
json_dumps: _JsonDumps = json.dumps,
) -> None:
"""
:param redis: Instance of Redis connection
:param key_builder: builder that helps to convert contextual key to string
:param state_ttl: TTL for state records
:param data_ttl: TTL for data records
:param lock_kwargs: Custom arguments for Redis lock
"""
if key_builder is None:
key_builder = DefaultKeyBuilder()
@ -107,6 +110,8 @@ class RedisStorage(BaseStorage):
self.key_builder = key_builder
self.state_ttl = state_ttl
self.data_ttl = data_ttl
self.json_loads = json_loads
self.json_dumps = json_dumps
@classmethod
def from_url(
@ -134,7 +139,6 @@ class RedisStorage(BaseStorage):
async def set_state(
self,
bot: Bot,
key: StorageKey,
state: StateType = None,
) -> None:
@ -150,7 +154,6 @@ class RedisStorage(BaseStorage):
async def get_state(
self,
bot: Bot,
key: StorageKey,
) -> Optional[str]:
redis_key = self.key_builder.build(key, "state")
@ -161,7 +164,6 @@ class RedisStorage(BaseStorage):
async def set_data(
self,
bot: Bot,
key: StorageKey,
data: Dict[str, Any],
) -> None:
@ -171,13 +173,12 @@ class RedisStorage(BaseStorage):
return
await self.redis.set(
redis_key,
bot.session.json_dumps(data),
self.json_dumps(data),
ex=self.data_ttl,
)
async def get_data(
self,
bot: Bot,
key: StorageKey,
) -> Dict[str, Any]:
redis_key = self.key_builder.build(key, "data")
@ -186,7 +187,7 @@ class RedisStorage(BaseStorage):
return {}
if isinstance(value, bytes):
value = value.decode("utf-8")
return cast(Dict[str, Any], bot.session.json_loads(value))
return cast(Dict[str, Any], self.json_loads(value))
class RedisEventIsolation(BaseEventIsolation):
@ -220,7 +221,6 @@ class RedisEventIsolation(BaseEventIsolation):
@asynccontextmanager
async def lock(
self,
bot: Bot,
key: StorageKey,
) -> AsyncGenerator[None, None]:
redis_key = self.key_builder.build(key, "lock")