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

Сleanup storage (#587)

Browse source 
* remove key from storage, when not needed

* fix redis1 delete

* fix new line

* check reset

* fix redis store close check
This commit is contained in:
Andrey Tikhonov 2021-06-13 01:22:47 +03:00 • committed by GitHub
parent 7eb32785f6
commit 0b1c22b7b0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 103 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

9
aiogram/contrib/fsm_storage/memory.py
View file

@ -66,6 +66,7 @@ class MemoryStorage(BaseStorage):
data: typing.Dict = None):
chat, user = self.resolve_address(chat=chat, user=user)
self.data[chat][user]['data'] = copy.deepcopy(data)
self._cleanup(chat, user)
async def reset_state(self, *,
chat: typing.Union[str, int, None] = None,
@ -74,6 +75,7 @@ class MemoryStorage(BaseStorage):
await self.set_state(chat=chat, user=user, state=None)
if with_data:
await self.set_data(chat=chat, user=user, data={})
self._cleanup(chat, user)
def has_bucket(self):
return True
@ -91,6 +93,7 @@ class MemoryStorage(BaseStorage):
bucket: typing.Dict = None):
chat, user = self.resolve_address(chat=chat, user=user)
self.data[chat][user]['bucket'] = copy.deepcopy(bucket)
self._cleanup(chat, user)
async def update_bucket(self, *,
chat: typing.Union[str, int, None] = None,
@ -100,3 +103,9 @@ class MemoryStorage(BaseStorage):
bucket = {}
chat, user = self.resolve_address(chat=chat, user=user)
self.data[chat][user]['bucket'].update(bucket, **kwargs)
def _cleanup(self, chat, user):
if self.data[chat][user] == {'state': None, 'data': {}, 'bucket': {}}:
del self.data[chat][user]
if not self.data[chat]:
del self.data[chat]

20
aiogram/contrib/fsm_storage/redis.py
View file

@ -110,10 +110,12 @@ class RedisStorage(BaseStorage):
chat, user = self.check_address(chat=chat, user=user)
addr = f"fsm:{chat}:{user}"
record = {'state': state, 'data': data, 'bucket': bucket}
conn = await self.redis()
await conn.execute('SET', addr, json.dumps(record))
if state is None and data == bucket == {}:
await conn.execute('DEL', addr)
else:
record = {'state': state, 'data': data, 'bucket': bucket}
await conn.execute('SET', addr, json.dumps(record))
async def get_state(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None,
default: typing.Optional[str] = None) -> typing.Optional[str]:
@ -222,7 +224,7 @@ class RedisStorage2(BaseStorage):
await dp.storage.wait_closed()
"""
def __init__(self, host: str = 'localhost', port=6379, db=None, password=None,
def __init__(self, host: str = 'localhost', port=6379, db=None, password=None,
ssl=None, pool_size=10, loop=None, prefix='fsm',
state_ttl: int = 0,
data_ttl: int = 0,
@ -304,7 +306,10 @@ class RedisStorage2(BaseStorage):
chat, user = self.check_address(chat=chat, user=user)
key = self.generate_key(chat, user, STATE_DATA_KEY)
redis = await self.redis()
await redis.set(key, json.dumps(data), expire=self._data_ttl)
if data:
await redis.set(key, json.dumps(data), expire=self._data_ttl)
else:
await redis.delete(key)
async def update_data(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None,
data: typing.Dict = None, **kwargs):
@ -332,7 +337,10 @@ class RedisStorage2(BaseStorage):
chat, user = self.check_address(chat=chat, user=user)
key = self.generate_key(chat, user, STATE_BUCKET_KEY)
redis = await self.redis()
await redis.set(key, json.dumps(bucket), expire=self._bucket_ttl)
if bucket:
await redis.set(key, json.dumps(bucket), expire=self._bucket_ttl)
else:
await redis.delete(key)
async def update_bucket(self, *, chat: typing.Union[str, int, None] = None,
user: typing.Union[str, int, None] = None,