life-echo/api/app/features/auth/repo.py

from datetime import datetime

from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.db import utc_now
from app.features.auth.models import RefreshToken, SmsVerificationCode
from app.features.user.models import User


async def get_user_by_phone(phone: str, db: AsyncSession) -> User | None:
    stmt = select(User).where(User.phone == phone)
    result = await db.execute(stmt)
    return result.scalar_one_or_none()


async def get_user_by_id(user_id: str, db: AsyncSession) -> User | None:
    return await db.get(User, user_id)


async def get_user_by_email(email: str, db: AsyncSession) -> User | None:
    stmt = select(User).where(User.email == email)
    result = await db.execute(stmt)
    return result.scalar_one_or_none()


async def get_user_by_openid(openid: str, db: AsyncSession) -> User | None:
    stmt = select(User).where(User.openid == openid)
    result = await db.execute(stmt)
    return result.scalar_one_or_none()


async def get_refresh_token_by_token(
    token_str: str, db: AsyncSession
) -> RefreshToken | None:
    stmt = select(RefreshToken).where(RefreshToken.token == token_str)
    result = await db.execute(stmt)
    return result.scalar_one_or_none()


async def get_refresh_token_by_id(
    token_id: str, db: AsyncSession
) -> RefreshToken | None:
    return await db.get(RefreshToken, token_id)


async def link_refresh_rotation(
    old_token_id: str,
    new_token_id: str,
    rotated_at: datetime,
    db: AsyncSession,
) -> None:
    """Record lineage on the consumed refresh token row."""
    stmt = (
        update(RefreshToken)
        .where(RefreshToken.id == old_token_id)
        .values(replaced_by_token_id=new_token_id, rotated_at=rotated_at)
    )
    await db.execute(stmt)


async def get_active_tokens_for_user(
    user_id: str, db: AsyncSession
) -> list[RefreshToken]:
    stmt = select(RefreshToken).where(
        RefreshToken.user_id == user_id,
        RefreshToken.is_revoked == False,  # noqa: E712
    )
    result = await db.execute(stmt)
    return list(result.scalars().all())


async def create_user(user: User, db: AsyncSession) -> None:
    db.add(user)


async def create_refresh_token(token: RefreshToken, db: AsyncSession) -> None:
    db.add(token)


async def try_consume_refresh_token(
    token_str: str, db: AsyncSession
) -> RefreshToken | None:
    """Atomically revoke a valid refresh token; returns row or None."""
    now = utc_now()
    stmt = (
        update(RefreshToken)
        .where(
            RefreshToken.token == token_str,
            RefreshToken.is_revoked.is_(False),
            RefreshToken.expires_at > now,
        )
        .values(is_revoked=True)
        .returning(RefreshToken)
    )
    result = await db.execute(stmt)
    return result.scalar_one_or_none()


# ── SMS verification code ─────────────────────────────────────


async def create_verification_code(
    record: SmsVerificationCode, db: AsyncSession
) -> None:
    db.add(record)


async def get_recent_code_for_rate_limit(
    phone: str, db: AsyncSession
) -> SmsVerificationCode | None:
    """Latest verification code record for the phone (for rate limit check)."""
    stmt = (
        select(SmsVerificationCode)
        .where(
            SmsVerificationCode.phone == phone,
            SmsVerificationCode.is_expired.is_(False),
        )
        .order_by(SmsVerificationCode.created_at.desc())
        .limit(1)
    )
    result = await db.execute(stmt)
    return result.scalar_one_or_none()


async def get_latest_unused_code(
    phone: str, purpose: str, db: AsyncSession
) -> SmsVerificationCode | None:
    """Latest unused, non-expired verification code for phone+purpose."""
    stmt = (
        select(SmsVerificationCode)
        .where(
            SmsVerificationCode.phone == phone,
            SmsVerificationCode.purpose == purpose,
            SmsVerificationCode.is_used.is_(False),
            SmsVerificationCode.is_expired.is_(False),
        )
        .order_by(SmsVerificationCode.created_at.desc())
        .limit(1)
    )
    result = await db.execute(stmt)
    return result.scalar_one_or_none()


async def mark_verification_code_expired(code_id: str, db: AsyncSession) -> None:
    """Mark a verification code expired (e.g. after SMS provider failure)."""
    stmt = (
        update(SmsVerificationCode)
        .where(SmsVerificationCode.id == code_id)
        .values(is_expired=True)
    )
    await db.execute(stmt)


async def try_consume_verification_code(
    phone: str,
    code: str,
    purpose: str,
    db: AsyncSession,
) -> SmsVerificationCode | None:
    """Atomically mark matching unused code as used; returns row or None."""
    now = utc_now()
    stmt = (
        update(SmsVerificationCode)
        .where(
            SmsVerificationCode.phone == phone,
            SmsVerificationCode.purpose == purpose,
            SmsVerificationCode.code == code,
            SmsVerificationCode.is_used.is_(False),
            SmsVerificationCode.is_expired.is_(False),
            SmsVerificationCode.expires_at > now,
        )
        .values(is_used=True, verified_at=now)
        .returning(SmsVerificationCode)
    )
    result = await db.execute(stmt)
    return result.scalar_one_or_none()