From f321c69f469f3affb825f316f145c6e3660e6365 Mon Sep 17 00:00:00 2001 From: Arslan 'Ars2014' Sakhapov Date: Wed, 10 Jan 2018 14:33:27 +0500 Subject: [PATCH] Add RethinkDB-based storage for FSM --- aiogram/contrib/fsm_storage/redis.py | 6 +- aiogram/contrib/fsm_storage/rethinkdb.py | 142 +++++++++++++++++++++++ dev_requirements.txt | 1 + 3 files changed, 147 insertions(+), 2 deletions(-) create mode 100644 aiogram/contrib/fsm_storage/rethinkdb.py diff --git a/aiogram/contrib/fsm_storage/redis.py b/aiogram/contrib/fsm_storage/redis.py index ecf81afa..58d44127 100644 --- a/aiogram/contrib/fsm_storage/redis.py +++ b/aiogram/contrib/fsm_storage/redis.py @@ -377,8 +377,10 @@ async def migrate_redis1_to_redis2(storage1: RedisStorage, storage2: RedisStorag :param storage2: instance of RedisStorage2 :return: """ - assert isinstance(storage1, RedisStorage) - assert isinstance(storage2, RedisStorage2) + if not isinstance(storage1, RedisStorage): # better than assertion + raise TypeError(f'{type(storage1)} is not RedisStorage instance.') + if not isinstance(storage2, RedisStorage): + raise TypeError(f'{type(storage2)} is not RedisStorage instance.') log = logging.getLogger('aiogram.RedisStorage') diff --git a/aiogram/contrib/fsm_storage/rethinkdb.py b/aiogram/contrib/fsm_storage/rethinkdb.py new file mode 100644 index 00000000..f27fb4ae --- /dev/null +++ b/aiogram/contrib/fsm_storage/rethinkdb.py @@ -0,0 +1,142 @@ +# -*- coding:utf-8; -*- +__all__ = ['RethinkDBStorage'] + +import asyncio +import logging +import typing + +import rethinkdb as r +from rethinkdb.asyncio_net.net_asyncio import Connection + +from ...dispatcher import BaseStorage + + +r.set_loop_type('asyncio') + + +class ConnectionNotClosed(Exception): + """ + Indicates that DB connection wasn't closed. + """ + + +class RethinkDBStorage(BaseStorage): + """ + RethinkDB-based storage for FSM. + + Usage: + + ..code-block:: python3 + + storage = RethinkDBStorage(db='aiogram', table='aiogram', user='aiogram', password='aiogram_secret') + dispatcher = Dispatcher(bot, storage=storage) + + And need to close connection when shutdown + + ..code-clock:: python3 + + await storage.close() + + """ + def __init__(self, host='localhost', port=28015, db='aiogram', table='aiogram', auth_key=None, + user=None, password=None, timeout=20, ssl=None, loop=None): + self._host = host + self._port = port + self._db = db + self._table = table + self._auth_key = auth_key + self._user = user + self._password = password + self._timeout = timeout + self._ssl = ssl or {} + + self._connection: Connection = None + self._loop = loop or asyncio.get_event_loop() + self._lock = asyncio.Lock(loop=self._loop) + + async def connection(self): + """ + Get or create connection. + """ + if not self._connection: + async with self._lock: # thread-safe + self._connection = await r.connect(host=self._host, port=self._port, db=self._db, auth_key=self._auth_key, user=self._user, + password=self._password, timeout=self._timeout, ssl=self._ssl, io_loop=self._loop) + return self._connection + + async def close(self): + """ + Close connection. + """ + if self._connection and self._connection.is_open(): + await self._connection.close() + del self._connection + self._connection = None + + async def wait_closed(self): + """ + Checks if connection is closed. + """ + if self._connection: + raise ConnectionNotClosed + return True + + 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]: + chat, user = map(str, self.check_address(chat=chat, user=user)) + conn = await self.connection() + return await r.table(self._table).get(chat)[user]['state'].default(default or '').run(conn) + + async def get_data(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, + default: typing.Optional[str] = None) -> typing.Dict: + chat, user = map(str, self.check_address(chat=chat, user=user)) + conn = await self.connection() + return await r.table(self._table).get(chat)[user]['data'].default(default or {}).run(conn) + + 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): + chat, user = map(str, self.check_address(chat=chat, user=user)) + conn = await self.connection() + # https://stackoverflow.com/questions/24306933/how-to-make-a-rethinkdb-atomic-update-if-document-exists-insert-otherwise + await r.table(self._table).insert( + {'id': chat, user: {'state': state}}, + conflict='update').run(conn) + + async def set_data(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, data: typing.Dict = None): + chat, user = map(str, self.check_address(chat=chat, user=user)) + conn = await self.connection() + await r.table(self._table).insert( + {'id': chat, user: {'data': r.literal(data)}}, + conflict='update').run(conn) + + async def update_data(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, data: typing.Dict = None, + **kwargs): + chat, user = map(str, self.check_address(chat=chat, user=user)) + conn = await self.connection() + await r.table(self._table).insert( + {'id': chat, user: {'data': data}}, + conflict='update').run(conn) + + def has_bucket(self): + return True + + async def get_bucket(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, + default: typing.Optional[dict] = None) -> typing.Dict: + chat, user = map(str, self.check_address(chat=chat, user=user)) + conn = await self.connection() + return await r.table(self._table).get(chat)[user]['bucket'].default(default or {}).run(conn) + + async def set_bucket(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, bucket: typing.Dict = None): + chat, user = map(str, self.check_address(chat=chat, user=user)) + conn = await self.connection() + await r.table(self._table).insert( + {'id': chat, user: {'bucket': r.literal(bucket)}}, + conflict='update').run(conn) + + async def update_bucket(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, bucket: typing.Dict = None, + **kwargs): + chat, user = map(str, self.check_address(chat=chat, user=user)) + conn = await self.connection() + await r.table(self._table).insert( + {'id': chat, user: {'data': bucket}}, + conflict='update').run(conn) diff --git a/dev_requirements.txt b/dev_requirements.txt index 7fd66b33..92eaaa4b 100644 --- a/dev_requirements.txt +++ b/dev_requirements.txt @@ -5,3 +5,4 @@ pytest pytest-asyncio uvloop aioredis +rethinkdb