from abc import ABC
from contextlib import asynccontextmanager
from typing import Optional

from asyncpg import Connection

from maps_adv.billing_proxy.lib.db.engine import DB, PoolType


class BaseDataManager(ABC):
    _db: DB

    def __init__(self, db: Optional[DB] = None):
        self._db = db

    @asynccontextmanager
    async def connection(
        self, con: Optional[Connection] = None, type: PoolType = PoolType.master
    ):
        if not con:
            con = await self._db.acquire(type)
            try:
                yield con
            finally:
                await self._db.release(con)
        else:
            yield con
