from operator import attrgetter

from asyncpg.connection import Connection

from smb.common.pgswim import SwimEngine

from .enums import CurrencyType, RateType

__all__ = ["DB"]


class DB(SwimEngine):
    @staticmethod
    async def _set_codecs(con: Connection):
        enum_encoder = attrgetter("name")
        await con.set_type_codec(
            "currency_type", encoder=enum_encoder, decoder=lambda x: CurrencyType[x]
        )
        await con.set_type_codec(
            "rate_type", encoder=enum_encoder, decoder=lambda x: RateType[x]
        )
