python/aiogram
1
0
Fork 0
mirror of https://github.com/aiogram/aiogram.git synced 2025-12-15 11:25:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add max_conn param

Browse source
This commit is contained in:
Arslan Sakhapov 2018-02-22 21:20:01 +05:00
parent 99f5a89f70
commit 9ef7ef4643
1 changed files with 9 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

13
aiogram/contrib/fsm_storage/rethinkdb.py
View file

@ -37,7 +37,7 @@ class RethinkDBStorage(BaseStorage):
""" """
def __init__(self, host='localhost', port=28015, db='aiogram', table='aiogram', auth_key=None, def __init__(self, host='localhost', port=28015, db='aiogram', table='aiogram', auth_key=None,
user=None, password=None, timeout=20, ssl=None, loop=None): user=None, password=None, timeout=20, ssl=None, max_conn=10, loop=None):
self._host = host self._host = host
self._port = port self._port = port
self._db = db self._db = db
@ -48,7 +48,7 @@ class RethinkDBStorage(BaseStorage):
self._timeout = timeout self._timeout = timeout
self._ssl = ssl or {} self._ssl = ssl or {}
self._queue = asyncio.Queue() self._queue = asyncio.Queue(max_conn)
self._outstanding_connections = weakref.WeakSet() self._outstanding_connections = weakref.WeakSet()
self._loop = loop or asyncio.get_event_loop() self._loop = loop or asyncio.get_event_loop()
@ -66,9 +66,12 @@ class RethinkDBStorage(BaseStorage):
except r.ReqlError: except r.ReqlError:
raise ConnectionNotClosed('Exception was caught while closing connection') raise ConnectionNotClosed('Exception was caught while closing connection')
except asyncio.QueueEmpty: except asyncio.QueueEmpty:
if len(self._outstanding_connections) < self._queue.maxsize:
conn = await r.connect(host=self._host, port=self._port, db=self._db, conn = 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, auth_key=self._auth_key, user=self._user, password=self._password,
ssl=self._ssl) timeout=self._timeout, ssl=self._ssl)
else:
conn = await self._queue.get()
self._outstanding_connections.add(conn) self._outstanding_connections.add(conn)
return conn return conn
@ -185,7 +188,7 @@ class RethinkDBStorage(BaseStorage):
await r.table(self._table).insert({'id': chat, user: {'bucket': bucket}}).run(conn) await r.table(self._table).insert({'id': chat, user: {'bucket': bucket}}).run(conn)
await self.put_connection(conn) await self.put_connection(conn)
async def get_states_list(self) -> typing.List[typing.Tuple[int]]: async def get_states_list(self) -> typing.List[typing.Tuple[int, int]]:
""" """
Get list of all stored chat's and user's Get list of all stored chat's and user's