from aiogram.dispatcher.middlewares import LifetimeControllerMiddleware
from asyncpg.pool import Pool

from services.repositories import Repos


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"] = Repos(conn)

    async def post_process(self, obj, data, *args):
        del data["repo"]
        conn = data.get("conn")
        if conn:
            await self.pool.release(conn)