buff3.0/backend-client/app/utils/nonce_store.py

"""
DB-first nonce dedup helper with Redis fallback.
Table expected: api_nonces(nonce PK, endpoint, session_id, created_at, expires_at)
"""
from typing import Optional
from datetime import datetime, timedelta
from sqlalchemy import text
from sqlalchemy.exc import IntegrityError
from app.core.redis_cache import redis_cache
from app.core.logger import logger
from app.utils.timezone import beijing_now_naive, beijing_now

async def check_and_set_nonce(db, nonce: str, endpoint: str, session_id: Optional[str], ttl_seconds: int) -> bool:
    """
    Try DB insert first (primary). If insert succeeds => first use => return True.
    If insert fails due to IntegrityError => duplicate => return False.
    If DB insert raises other errors, attempt Redis fallback: SET NX with ttl.
    If Redis available, use it (True if set, False if existed). If both fail, raise.
    """
    try:
        expires_at = beijing_now_naive() + timedelta(seconds=ttl_seconds)
        sql = text(
            "INSERT INTO api_nonces (nonce, endpoint, session_id, created_at, expires_at) "
            "VALUES (:nonce, :endpoint, :session_id, :created_at, :expires_at)"
        )
        db.execute(sql, {
            "nonce": nonce,
            "endpoint": endpoint,
            "session_id": session_id,
            "created_at": beijing_now_naive(),
            "expires_at": expires_at
        })
        db.commit()
        return True
    except IntegrityError:
        # duplicate
        try:
            db.rollback()
        except Exception:
            pass
        return False
    except Exception as e:
        logger.exception("DB 插入 nonce 失败，尝试 Redis 回退")
        try:
            # Redis fallback (async)
            key = f"nonce:{nonce}"
            ttl = ttl_seconds
            ok = await redis_cache.set(key, {"ts": int(beijing_now().timestamp())}, ttl)
            # redis_cache.set returns truthy on success; if key already exists, exists() would be True
            if not ok:
                # check exists -> duplicate
                exists = await redis_cache.exists(key)
                if exists:
                    return False
            return True
        except Exception:
            logger.exception("Redis 回退也失败，nonce 去重无法完成")
            # Conservative: treat as failure to set -> raise to let caller decide to reject
            raise