mirror of
https://github.com/nexus-stc/hyperboria
synced 2024-06-13 08:30:11 +02:00
dd23846059
- [nexus] Switch bot - [bot] Added extra receivers functionality GitOrigin-RevId: 68fc32d3e79ff411758f54f435fe8680fc42dead
52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
from typing import Optional
|
|
|
|
from aiokit import AioThing
|
|
from psycopg.rows import tuple_row
|
|
from psycopg_pool import AsyncConnectionPool
|
|
|
|
|
|
class AioPostgresPoolHolder(AioThing):
|
|
def __init__(self, conninfo, timeout=30, min_size=1, max_size=4):
|
|
super().__init__()
|
|
self.pool = None
|
|
self.fn = lambda: AsyncConnectionPool(
|
|
conninfo=conninfo,
|
|
timeout=timeout,
|
|
min_size=min_size,
|
|
max_size=max_size,
|
|
)
|
|
|
|
async def start(self):
|
|
if not self.pool:
|
|
self.pool = self.fn()
|
|
|
|
async def stop(self):
|
|
if self.pool:
|
|
await self.pool.close()
|
|
self.pool = None
|
|
|
|
async def iterate(
|
|
self,
|
|
stmt: str,
|
|
values=None,
|
|
row_factory=tuple_row,
|
|
cursor_name: Optional[str] = None,
|
|
itersize: Optional[int] = None,
|
|
):
|
|
if not self.pool:
|
|
raise RuntimeError('AioPostgresPoolHolder has not been started')
|
|
async with self.pool.connection() as conn:
|
|
async with conn.cursor(name=cursor_name, row_factory=row_factory) as cur:
|
|
if itersize is not None:
|
|
cur.itersize = itersize
|
|
await cur.execute(stmt, values)
|
|
async for row in cur:
|
|
yield row
|
|
|
|
async def execute(self, stmt: str, values=None, cursor_name: Optional[str] = None, row_factory=tuple_row):
|
|
if not self.pool:
|
|
raise RuntimeError('AioPostgresPoolHolder has not been started')
|
|
async with self.pool.connection() as conn:
|
|
async with conn.cursor(name=cursor_name, row_factory=row_factory) as cur:
|
|
await cur.execute(stmt, values)
|