webapp/app/services/short_url_service.py

# app/services/short_url_service.py

import string
import random
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from sqlalchemy.exc import IntegrityError

from app.core.biz_exception import NotFoundError, BizLogicError
from app.models.short_url import ShortUrl


class ShortUrlService:

    @staticmethod
    def generate_short_key(length: int = 8) -> str:
        """生成随机短 Key（字母 + 数字）"""
        chars = string.ascii_letters + string.digits
        return ''.join(random.choices(chars, k=length))

    @staticmethod
    async def create_short_url(
        db: AsyncSession,
        long_url: str
    ) -> ShortUrl:
        """创建短链接（异步 + 并发安全）"""

        # 1️⃣ 是否已有相同 long_url
        stmt = select(ShortUrl).where(ShortUrl.long_url == long_url)
        existing = await db.scalar(stmt)
        if existing:
            raise BizLogicError("Short url already exist")

        # 2️⃣ 生成 short_key（循环直到成功）
        while True:
            short_key = ShortUrlService.generate_short_key()

            db_obj = ShortUrl(
                short_key=short_key,
                long_url=long_url
            )
            db.add(db_obj)

            try:
                await db.commit()
                await db.refresh(db_obj)
                return db_obj

            except IntegrityError:
                # short_key 冲突，重试
                await db.rollback()
                continue

    @staticmethod
    async def get_long_url(
        db: AsyncSession,
        short_key: str
    ) -> str:
        """通过短 key 获取原始长链接（异步）"""

        stmt = select(ShortUrl).where(ShortUrl.short_key == short_key)
        record = await db.scalar(stmt)

        if not record:
            raise NotFoundError("Short url not found")

        return record.long_url