24 lines
645 B
Python
24 lines
645 B
Python
from aiogram.dispatcher.middlewares import LifetimeControllerMiddleware
|
|
from asyncpg.pool import Pool
|
|
|
|
from services.repository import Repo
|
|
|
|
|
|
class DbMiddleware(LifetimeControllerMiddleware):
|
|
skip_patterns = ["error", "update"]
|
|
|
|
def __init__(self, pool: Pool):
|
|
super().__init__()
|
|
self.pool = pool
|
|
|
|
async def pre_process(self, obj, data, *args):
|
|
conn = await self.pool.acquire()
|
|
data["conn"] = conn
|
|
data["repo"] = Repo(conn)
|
|
|
|
async def post_process(self, obj, data, *args):
|
|
del data["repo"]
|
|
conn = data.get("conn")
|
|
if conn:
|
|
await self.pool.release(conn)
|