Joey
142 lines
4.9 KiB
Python
142 lines
4.9 KiB
Python
"""套图生成 API：创建任务 / 查询状态 / 导出 ZIP。"""
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
import zipfile
|
||
from uuid import UUID
|
||
|
||
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
|
||
from fastapi.responses import StreamingResponse
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from config import get_settings
|
||
from db import get_db
|
||
from models import Product, ProductAsset, Suite, SuiteImage, STATUS_OK
|
||
from schemas import PLATFORM_SPECS, SUPPORTED_TYPES, SuiteCreateRequest, SuiteCreateResponse, SuiteImageOut, SuiteOut
|
||
from services import storage
|
||
from services.generator import run_suite
|
||
|
||
router = APIRouter(prefix="/api", tags=["suites"])
|
||
|
||
|
||
async def _suite_out(db: AsyncSession, suite: Suite) -> SuiteOut:
|
||
images = (await db.scalars(
|
||
select(SuiteImage).where(SuiteImage.suite_id == suite.id)
|
||
.order_by(SuiteImage.created_at)
|
||
)).all()
|
||
return SuiteOut(
|
||
id=str(suite.id),
|
||
product_id=str(suite.product_id),
|
||
status=suite.status,
|
||
style_set=suite.style_set,
|
||
platform=suite.platform,
|
||
lang=suite.lang,
|
||
ratio=suite.ratio,
|
||
types=list(suite.types or []),
|
||
provider=suite.provider,
|
||
images=[
|
||
SuiteImageOut(
|
||
type_id=i.type_id, name=i.name, url=i.stored_url or "",
|
||
status=i.status, error=i.error,
|
||
) for i in images
|
||
],
|
||
error=suite.error,
|
||
)
|
||
|
||
|
||
@router.post("/products/{product_id}/suites", response_model=SuiteCreateResponse)
|
||
async def create_suite(
|
||
product_id: str,
|
||
req: SuiteCreateRequest,
|
||
background: BackgroundTasks,
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
product = await db.get(Product, UUID(product_id))
|
||
if product is None:
|
||
raise HTTPException(status_code=404, detail="商品不存在")
|
||
|
||
# 主图组至少一张图（不要求转存完成：生图可直接用源站 URL 代理解析）
|
||
ok_assets = (await db.scalars(
|
||
select(ProductAsset.id).where(
|
||
ProductAsset.product_id == product.id,
|
||
ProductAsset.group_key == "main",
|
||
ProductAsset.type == "img",
|
||
)
|
||
)).all()
|
||
if not ok_assets:
|
||
raise HTTPException(status_code=400, detail="商品没有主图，无法生成")
|
||
|
||
bad = [t for t in req.types if t not in SUPPORTED_TYPES]
|
||
if bad:
|
||
raise HTTPException(status_code=400, detail=f"不支持的图类型: {bad}")
|
||
if req.platform not in PLATFORM_SPECS:
|
||
raise HTTPException(status_code=400, detail=f"不支持的目标平台: {req.platform}（ozon | wb | cn）")
|
||
spec = PLATFORM_SPECS[req.platform]
|
||
|
||
settings = get_settings()
|
||
suite = Suite(
|
||
product_id=product.id,
|
||
style_set=req.style_set,
|
||
platform=req.platform,
|
||
lang=spec["lang"],
|
||
ratio=spec["ratio"],
|
||
types=req.types,
|
||
provider=req.provider or settings.image_provider,
|
||
)
|
||
db.add(suite)
|
||
await db.commit()
|
||
await db.refresh(suite)
|
||
|
||
background.add_task(run_suite, str(suite.id))
|
||
return SuiteCreateResponse(suite_id=str(suite.id))
|
||
|
||
|
||
@router.get("/suites/{suite_id}", response_model=SuiteOut)
|
||
async def get_suite(suite_id: str, db: AsyncSession = Depends(get_db)):
|
||
suite = await db.get(Suite, UUID(suite_id))
|
||
if suite is None:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
return await _suite_out(db, suite)
|
||
|
||
|
||
@router.get("/products/{product_id}/suites")
|
||
async def list_suites(product_id: str, db: AsyncSession = Depends(get_db)):
|
||
suites = (await db.scalars(
|
||
select(Suite).where(Suite.product_id == UUID(product_id))
|
||
.order_by(Suite.created_at.desc())
|
||
)).all()
|
||
return [await _suite_out(db, s) for s in suites]
|
||
|
||
|
||
@router.get("/suites/{suite_id}/zip")
|
||
async def download_suite_zip(suite_id: str, db: AsyncSession = Depends(get_db)):
|
||
"""把任务内所有成功图打包成 ZIP（中文文件名）。"""
|
||
suite = await db.get(Suite, UUID(suite_id))
|
||
if suite is None:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
images = (await db.scalars(
|
||
select(SuiteImage).where(
|
||
SuiteImage.suite_id == suite.id, SuiteImage.status == STATUS_OK,
|
||
).order_by(SuiteImage.created_at)
|
||
)).all()
|
||
|
||
buf = io.BytesIO()
|
||
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
|
||
seen: set[str] = set()
|
||
for i, img in enumerate(images):
|
||
path = storage.local_path(img.stored_url or "")
|
||
if path is None:
|
||
continue
|
||
filename = img.name or img.type_id
|
||
if filename in seen: # 同类型多张时加序号防覆盖
|
||
filename = f"{filename}-{i + 1}"
|
||
seen.add(filename)
|
||
zf.write(path, f"{filename}{path.suffix or '.jpg'}")
|
||
buf.seek(0)
|
||
return StreamingResponse(
|
||
buf,
|
||
media_type="application/zip",
|
||
headers={"Content-Disposition": f'attachment; filename="suite-{suite_id}.zip"'},
|
||
)
|