"""
Начну с того, зачем вообще придуман данный декоратор.

Как ты прописываешь обращения к СУБД в методах?
Вероятнее всего, ты возьмешь соединение (при учатии пула соединений), а затем откроешь транзакцию
(ведь в асинхронном приложении возможен race condition, да и в целом тебе могут понадобиться
свойства транзакций (требования ACID), например атомарность (блок дейтвий будет выполнен весь, либо не выполнен вовсе))
И вот, когда есть соединение, и начата транзакция, ты пишешь SQL-запрос.

В коде это выглядело бы как-то так:
"""



class DBMS:

    def __init__(...):
        self.pool = ...

    async def fetch_data(self, data_id):
        async with self.pool.acquire() as connect:
            async with connect.transaction():
                return await connect.fetch(""" SELECT * FROM table_name WHERE data_id = $1; """, data_id)


"""
В целом, всё выглядит логично.
А теперь добавим еще методы...
"""


    async def insert_data(self, data):
        async with self.pool.acquire() as connect:
            async with connect.transaction():
                await connect.execute(""" INSERT INTO table_name VALUES ($1); """, data)

    async def update_data(self, data_id, data):
        async with self.pool.acquire() as connect:
            async with connect.transaction():
                await connect.execute(""" SELECT * FROM table_name WHERE data_id = $1 FOR UPDATE; """, data_id)
                await connect.execute(""" UPDATE table_name SET data_column = $1 WHERE data_id = $2; """, data, data_id)


"""
Тут уже видна избыточность. Мы делаем "подготовительную" часть в каждом методе.
А теперь вспомним для чего были придуманы декораторы. Они как раз берут на себя действия,
которые выполняются до и после какого-либо нашего кода.
Выносим повторяющуюся часть в отдельное место:
"""


def acquire_connection(method: Callable) -> Callable:
    async def wrapper(self, *args, **kwargs) -> Any:  # не забываем про аргумент self, ведь мы оборачиваем методы
        async with self._pool.acquire() as connect:  # берем соединение (именно берем, а не открываем!)
            async with connect.transaction():  # начинаем транзакцию
                # выполняем код метода, передавая в него готовый _connect
                # _connect специально подчеркнут слева, чтобы ты не забыл, что его не нужно "прокидывать" извне
                return await method(self, _connect=connect, *args, **kwargs)
    return wrapper



"""
Теперь посмотрим на тот же класс с декоратором:
"""


from asyncpg.connection import Connection  # тайпинг проставляется для удобства в методах


class DBMS:

    def __init__(...):
        self.pool = ...

    @acquire_connection
    async def fetch_data(self, data_id, _connect: Connection):
        # _connect "подставляется" сам, потому что вся операция выполняется в декораторе
        # необходимо его только прописывать, как ты это делаешь, например, в хэндлерах бота, куда "попадает" апдейт
        return await _connect.fetch(""" SELECT * FROM table_name WHERE data_id = $1; """, data_id)

    @acquire_connection
    async def insert_data(self, data, _connect: Connection):
        await _connect.execute(""" INSERT INTO table_name VALUES ($1); """, data)

    @acquire_connection
    async def update_data(self, data_id, data, _connect: Connection):
        await _connect.execute(""" SELECT * FROM table_name WHERE data_id = $1 FOR UPDATE; """, data_id)
        await _connect.execute(""" UPDATE table_name SET data_column = $1 WHERE data_id = $2; """, data, data_id)



"""
Настоятельно рекомендую передавать именованные аргументы, или хотя-бы совать _connect в конец,
чтобы какой-нибудь user_id не встал вдруг на место _connect (как это было у тебя, ибо _connect ты поместил в начале,
передав на его место другую вещь, сам того незаметив)
"""

dbms = DBMS(...)

data = await dbms.fetch_data(data_id=12345)
await dbms.insert_data(data=...)
await dbms.update_data(data_id=12345, data=...)