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

from asyncpg import Connection

from maps_adv.adv_store.v2.lib.db.engine import DB
from smb.common.pgswim.lib.engine import PoolType


class BaseDataManager(ABC):
    __slots__ = ("_db",)
    _db: DB

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

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


__all__ = ["BaseDataManager"]
