from operator import attrgetter

from asyncpg import Connection
from smb.common.pgswim import SwimEngine

from maps_adv.geosmb.promoter.server.lib.enums import (
    EventDataSource,
    EventType,
    LeadDataSource,
    Source,
)


class DB(SwimEngine):
    @staticmethod
    async def _set_codecs(con: Connection):
        await con.set_type_codec(
            "event_data_source",
            encoder=attrgetter("name"),
            decoder=lambda x: EventDataSource[x],
        )
        await con.set_type_codec(
            "lead_data_source",
            encoder=attrgetter("name"),
            decoder=lambda x: LeadDataSource[x],
        )
        await con.set_type_codec(
            "event_type", encoder=attrgetter("name"), decoder=lambda x: EventType[x]
        )
        await con.set_type_codec(
            "source",
            encoder=attrgetter("name"),
            decoder=lambda x: Source[x],
        )
