From c6accd1a53d43e4cd63c105a370a4830bd168b4a Mon Sep 17 00:00:00 2001 From: Gabben Date: Thu, 27 Jun 2019 19:34:38 +0500 Subject: [PATCH 1/4] Add "expire" argument Add "expire" argument to all set_ and update_ methods in RedisStorage2 --- aiogram/contrib/fsm_storage/redis.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/aiogram/contrib/fsm_storage/redis.py b/aiogram/contrib/fsm_storage/redis.py index c3a91f00..6b96869b 100644 --- a/aiogram/contrib/fsm_storage/redis.py +++ b/aiogram/contrib/fsm_storage/redis.py @@ -287,29 +287,29 @@ class RedisStorage2(BaseStorage): return default or {} async def set_state(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, - state: typing.Optional[typing.AnyStr] = None): + state: typing.Optional[typing.AnyStr] = None, expire: int = 0): chat, user = self.check_address(chat=chat, user=user) key = self.generate_key(chat, user, STATE_KEY) redis = await self.redis() if state is None: await redis.delete(key) else: - await redis.set(key, state) + await redis.set(key, state, expire=expire) async def set_data(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, - data: typing.Dict = None): + data: typing.Dict = None, expire: int = 0): 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)) + await redis.set(key, json.dumps(data), expire=expire) async def update_data(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, - data: typing.Dict = None, **kwargs): + data: typing.Dict = None, expire: int = 0, **kwargs): if data is None: data = {} temp_data = await self.get_data(chat=chat, user=user, default={}) temp_data.update(data, **kwargs) - await self.set_data(chat=chat, user=user, data=temp_data) + await self.set_data(chat=chat, user=user, data=temp_data, expire=expire) def has_bucket(self): return True @@ -325,20 +325,20 @@ class RedisStorage2(BaseStorage): return default or {} async def set_bucket(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, - bucket: typing.Dict = None): + bucket: typing.Dict = None, expire: int = 0): 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)) + await redis.set(key, json.dumps(bucket), expire=expire) async def update_bucket(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, - bucket: typing.Dict = None, **kwargs): + bucket: typing.Dict = None, expire: int = 0, **kwargs): if bucket is None: bucket = {} temp_bucket = await self.get_bucket(chat=chat, user=user) temp_bucket.update(bucket, **kwargs) - await self.set_bucket(chat=chat, user=user, data=temp_bucket) + await self.set_bucket(chat=chat, user=user, data=temp_bucket, expire=expire) async def reset_all(self, full=True): """ From e2842944fa1e051648341b07cb920489f4f55e29 Mon Sep 17 00:00:00 2001 From: Gabben Date: Thu, 27 Jun 2019 20:39:26 +0500 Subject: [PATCH 2/4] Update redis.py Update docstrings --- aiogram/contrib/fsm_storage/redis.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/aiogram/contrib/fsm_storage/redis.py b/aiogram/contrib/fsm_storage/redis.py index 6b96869b..e75e7f40 100644 --- a/aiogram/contrib/fsm_storage/redis.py +++ b/aiogram/contrib/fsm_storage/redis.py @@ -62,8 +62,6 @@ class RedisStorage(BaseStorage): async def redis(self) -> aioredis.RedisConnection: """ Get Redis connection - - This property is awaitable. """ # Use thread-safe asyncio Lock because this method without that is not safe async with self._connection_lock: @@ -241,8 +239,6 @@ class RedisStorage2(BaseStorage): async def redis(self) -> aioredis.Redis: """ Get Redis connection - - This property is awaitable. """ # Use thread-safe asyncio Lock because this method without that is not safe async with self._connection_lock: From 8a819eb7ed0758f7342e234c1b10f6c9282e72f1 Mon Sep 17 00:00:00 2001 From: Gabben Date: Thu, 27 Jun 2019 21:00:25 +0500 Subject: [PATCH 3/4] Fix update_bucket --- aiogram/contrib/fsm_storage/redis.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiogram/contrib/fsm_storage/redis.py b/aiogram/contrib/fsm_storage/redis.py index e75e7f40..0e89eaea 100644 --- a/aiogram/contrib/fsm_storage/redis.py +++ b/aiogram/contrib/fsm_storage/redis.py @@ -334,7 +334,7 @@ class RedisStorage2(BaseStorage): bucket = {} temp_bucket = await self.get_bucket(chat=chat, user=user) temp_bucket.update(bucket, **kwargs) - await self.set_bucket(chat=chat, user=user, data=temp_bucket, expire=expire) + await self.set_bucket(chat=chat, user=user, bucket=temp_bucket, expire=expire) async def reset_all(self, full=True): """ From 607e3ea13514942bc10192a81604e00404f209c8 Mon Sep 17 00:00:00 2001 From: Gabben Date: Fri, 28 Jun 2019 15:26:07 +0500 Subject: [PATCH 4/4] Update redis.py --- aiogram/contrib/fsm_storage/redis.py | 34 ++++++++++++++++------------ 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/aiogram/contrib/fsm_storage/redis.py b/aiogram/contrib/fsm_storage/redis.py index 0e89eaea..106a7b97 100644 --- a/aiogram/contrib/fsm_storage/redis.py +++ b/aiogram/contrib/fsm_storage/redis.py @@ -35,7 +35,6 @@ class RedisStorage(BaseStorage): await dp.storage.wait_closed() """ - def __init__(self, host='localhost', port=6379, db=None, password=None, ssl=None, loop=None, **kwargs): self._host = host self._port = port @@ -220,9 +219,12 @@ class RedisStorage2(BaseStorage): await dp.storage.wait_closed() """ - - def __init__(self, host='localhost', port=6379, db=None, password=None, ssl=None, - pool_size=10, loop=None, prefix='fsm', **kwargs): + 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, + bucket_ttl: int = 0, + **kwargs): self._host = host self._port = port self._db = db @@ -233,6 +235,10 @@ class RedisStorage2(BaseStorage): self._kwargs = kwargs self._prefix = (prefix,) + self._state_ttl = state_ttl + self._data_ttl = data_ttl + self._bucket_ttl = bucket_ttl + self._redis: aioredis.RedisConnection = None self._connection_lock = asyncio.Lock(loop=self._loop) @@ -283,29 +289,29 @@ class RedisStorage2(BaseStorage): return default or {} async def set_state(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, - state: typing.Optional[typing.AnyStr] = None, expire: int = 0): + state: typing.Optional[typing.AnyStr] = None): chat, user = self.check_address(chat=chat, user=user) key = self.generate_key(chat, user, STATE_KEY) redis = await self.redis() if state is None: await redis.delete(key) else: - await redis.set(key, state, expire=expire) + await redis.set(key, state, expire=self._state_ttl) async def set_data(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, - data: typing.Dict = None, expire: int = 0): + data: typing.Dict = None): 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=expire) + await redis.set(key, json.dumps(data), expire=self._data_ttl) async def update_data(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, - data: typing.Dict = None, expire: int = 0, **kwargs): + data: typing.Dict = None, **kwargs): if data is None: data = {} temp_data = await self.get_data(chat=chat, user=user, default={}) temp_data.update(data, **kwargs) - await self.set_data(chat=chat, user=user, data=temp_data, expire=expire) + await self.set_data(chat=chat, user=user, data=temp_data) def has_bucket(self): return True @@ -321,20 +327,20 @@ class RedisStorage2(BaseStorage): return default or {} async def set_bucket(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, - bucket: typing.Dict = None, expire: int = 0): + bucket: typing.Dict = None): 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=expire) + await redis.set(key, json.dumps(bucket), expire=self._bucket_ttl) async def update_bucket(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, - bucket: typing.Dict = None, expire: int = 0, **kwargs): + bucket: typing.Dict = None, **kwargs): if bucket is None: bucket = {} temp_bucket = await self.get_bucket(chat=chat, user=user) temp_bucket.update(bucket, **kwargs) - await self.set_bucket(chat=chat, user=user, bucket=temp_bucket, expire=expire) + await self.set_bucket(chat=chat, user=user, bucket=temp_bucket) async def reset_all(self, full=True): """