harness_zbt/backend/app/infrastructure/mysql/attribution_repositories.py

from sqlalchemy import select

from app.domains.attribution.models import (
    LeadAttribution,
    OrderAttribution,
    PromotionCode,
    Salesperson,
)
from app.domains.attribution.repository import AttributionRepository
from app.infrastructure.mysql.core_models import (
    LeadAttributionRecord,
    OrderAttributionRecord,
    PromotionCodeRecord,
    SalespersonRecord,
)
from app.infrastructure.mysql.repositories import SessionFactory


class SqlAlchemyAttributionRepository(AttributionRepository):
    def __init__(self, session_factory: SessionFactory) -> None:
        self._session_factory = session_factory

    def save_salesperson(self, salesperson: Salesperson) -> None:
        with self._session_factory() as session:
            session.merge(
                SalespersonRecord(
                    id=salesperson.id,
                    admin_user_id=salesperson.admin_user_id,
                    salesperson_code=salesperson.code,
                    name=salesperson.name,
                    status=salesperson.status,
                    created_at=salesperson.created_at,
                )
            )
            session.commit()

    def list_salespersons(self) -> list[Salesperson]:
        with self._session_factory() as session:
            records = session.scalars(
                select(SalespersonRecord).order_by(SalespersonRecord.salesperson_code)
            ).all()
            return [
                Salesperson(
                    id=item.id,
                    admin_user_id=item.admin_user_id,
                    code=item.salesperson_code,
                    name=item.name,
                    status=item.status,
                    created_at=item.created_at,
                )
                for item in records
            ]

    def save_promotion_code(self, promotion_code: PromotionCode) -> None:
        with self._session_factory() as session:
            session.merge(
                PromotionCodeRecord(
                    id=promotion_code.id,
                    code=promotion_code.code,
                    salesperson_id=promotion_code.salesperson_id,
                    channel=promotion_code.channel,
                    status=promotion_code.status,
                    created_at=promotion_code.created_at,
                )
            )
            session.commit()

    def get_promotion_code(self, code: str) -> PromotionCode | None:
        with self._session_factory() as session:
            record = session.scalar(
                select(PromotionCodeRecord).where(PromotionCodeRecord.code == code.upper())
            )
            return self._code(record) if record else None

    def list_promotion_codes(self) -> list[PromotionCode]:
        with self._session_factory() as session:
            records = session.scalars(
                select(PromotionCodeRecord).order_by(PromotionCodeRecord.code)
            ).all()
            return [self._code(item) for item in records]

    def save_lead(self, lead: LeadAttribution) -> None:
        with self._session_factory() as session:
            session.merge(
                LeadAttributionRecord(
                    id=lead.id,
                    h5_user_id=lead.h5_user_id,
                    salesperson_id=lead.salesperson_id,
                    promotion_code_id=lead.promotion_code_id,
                    source_code=lead.source_code,
                    visit_count=lead.visit_count,
                    first_touch_at=lead.first_touch_at,
                    last_touch_at=lead.last_touch_at,
                )
            )
            session.commit()

    def get_lead_by_user(self, h5_user_id: str) -> LeadAttribution | None:
        with self._session_factory() as session:
            record = session.scalar(
                select(LeadAttributionRecord).where(LeadAttributionRecord.h5_user_id == h5_user_id)
            )
            return self._lead(record) if record else None

    def list_leads(self) -> list[LeadAttribution]:
        with self._session_factory() as session:
            records = session.scalars(select(LeadAttributionRecord)).all()
            return [self._lead(item) for item in records]

    def save_order_attribution(self, attribution: OrderAttribution) -> None:
        with self._session_factory() as session:
            session.merge(
                OrderAttributionRecord(
                    id=attribution.id,
                    order_id=attribution.order_id,
                    h5_user_id=attribution.h5_user_id,
                    salesperson_id=attribution.salesperson_id,
                    promotion_code_id=attribution.promotion_code_id,
                    source_code=attribution.source_code,
                    amount_cents=attribution.amount_cents,
                    attributed_at=attribution.attributed_at,
                )
            )
            session.commit()

    def get_order_attribution(self, order_id: str) -> OrderAttribution | None:
        with self._session_factory() as session:
            record = session.scalar(
                select(OrderAttributionRecord).where(OrderAttributionRecord.order_id == order_id)
            )
            return self._order(record) if record else None

    def list_order_attributions(self) -> list[OrderAttribution]:
        with self._session_factory() as session:
            records = session.scalars(select(OrderAttributionRecord)).all()
            return [self._order(item) for item in records]

    @staticmethod
    def _code(record: PromotionCodeRecord) -> PromotionCode:
        return PromotionCode(
            id=record.id,
            code=record.code,
            salesperson_id=record.salesperson_id,
            channel=record.channel,
            status=record.status,
            created_at=record.created_at,
        )

    @staticmethod
    def _lead(record: LeadAttributionRecord) -> LeadAttribution:
        return LeadAttribution(
            id=record.id,
            h5_user_id=record.h5_user_id,
            salesperson_id=record.salesperson_id,
            promotion_code_id=record.promotion_code_id,
            source_code=record.source_code,
            visit_count=record.visit_count,
            first_touch_at=record.first_touch_at,
            last_touch_at=record.last_touch_at,
        )

    @staticmethod
    def _order(record: OrderAttributionRecord) -> OrderAttribution:
        return OrderAttribution(
            id=record.id,
            order_id=record.order_id,
            h5_user_id=record.h5_user_id,
            salesperson_id=record.salesperson_id,
            promotion_code_id=record.promotion_code_id,
            source_code=record.source_code,
            amount_cents=record.amount_cents,
            attributed_at=record.attributed_at,
        )