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)
|