xiaoxia-saas/packages/adapters/sqlalchemy_impl/daily_usage_repository.py

from datetime import UTC, date, datetime

from sqlalchemy.orm import Session

from packages.adapters.sqlalchemy_impl.models import DailyUsageRecordModel
from packages.domain.daily_usage_record import DailyUsageRecord


class SQLAlchemyDailyUsageRepository:
    def __init__(self, session: Session):
        self.session = session

    def create(self, record: DailyUsageRecord) -> DailyUsageRecord:
        model = DailyUsageRecordModel(
            id=record.id,
            user_id=record.user_id,
            usage_date=record.usage_date,
            usage_type=record.usage_type,
            count=record.count,
            updated_at=record.updated_at,
        )
        self.session.add(model)
        self.session.commit()
        return record

    def get_by_user_and_date(
        self, user_id: str, usage_date: date, usage_type: str = "free_clip"
    ) -> DailyUsageRecord | None:
        model = (
            self.session.query(DailyUsageRecordModel)
            .filter(
                DailyUsageRecordModel.user_id == user_id,
                DailyUsageRecordModel.usage_date == usage_date,
                DailyUsageRecordModel.usage_type == usage_type,
            )
            .first()
        )
        if model is None:
            return None
        return self._to_domain(model)

    def update_count(self, record: DailyUsageRecord) -> DailyUsageRecord:
        model = self.session.query(DailyUsageRecordModel).filter(DailyUsageRecordModel.id == record.id).first()
        if model is None:
            return record
        model.count = record.count
        model.updated_at = datetime.now(UTC)
        self.session.add(model)
        self.session.commit()
        return record

    def upsert(self, user_id: str, usage_date: date, usage_type: str = "free_clip") -> DailyUsageRecord:
        """Increment usage count for the given user/date/type, creating if needed."""
        model = (
            self.session.query(DailyUsageRecordModel)
            .filter(
                DailyUsageRecordModel.user_id == user_id,
                DailyUsageRecordModel.usage_date == usage_date,
                DailyUsageRecordModel.usage_type == usage_type,
            )
            .first()
        )
        if model is None:
            record = DailyUsageRecord.create(user_id=user_id, usage_date=usage_date, usage_type=usage_type)
            record.count = 1
            model = DailyUsageRecordModel(
                id=record.id,
                user_id=record.user_id,
                usage_date=record.usage_date,
                usage_type=record.usage_type,
                count=1,
                updated_at=datetime.now(UTC),
            )
            self.session.add(model)
            self.session.commit()
            return record

        model.count += 1
        model.updated_at = datetime.now(UTC)
        self.session.add(model)
        self.session.commit()
        return self._to_domain(model)

    @staticmethod
    def _to_domain(model: DailyUsageRecordModel) -> DailyUsageRecord:
        return DailyUsageRecord(
            id=model.id,
            user_id=model.user_id,
            usage_date=model.usage_date,
            usage_type=model.usage_type,
            count=model.count,
            updated_at=model.updated_at,
        )