Files
2026-07-04 00:17:11 +08:00

131 lines
3.8 KiB
Python

from datetime import date
from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session
from ..database import get_db
from ..models import Collection, CollectionStock
from ..schemas import (
CollectionCreate,
CollectionUpdate,
CollectionOut,
CollectionListItem,
CollectionStockOut,
StockInput,
ShareResponse,
)
import secrets
import string
router = APIRouter()
def _generate_short_code(length=8):
chars = string.ascii_letters + string.digits
return "".join(secrets.choice(chars) for _ in range(length))
@router.post("", response_model=CollectionOut)
def create_collection(data: CollectionCreate, db: Session = Depends(get_db)):
col = Collection(name=data.name)
db.add(col)
db.flush()
for s in data.stocks:
db.add(CollectionStock(
collection_id=col.id,
stock_code=s.code,
stock_name=s.name,
added_date=date.fromisoformat(s.added_date),
))
db.commit()
db.refresh(col)
return col
@router.get("", response_model=list[CollectionListItem])
def list_collections(db: Session = Depends(get_db)):
cols = db.query(Collection).order_by(Collection.created_at.desc()).all()
return [
CollectionListItem(
id=c.id,
name=c.name,
stock_count=len(c.stocks),
short_code=c.short_code,
created_at=c.created_at,
)
for c in cols
]
@router.get("/{col_id}", response_model=CollectionOut)
def get_collection(col_id: int, db: Session = Depends(get_db)):
col = db.query(Collection).filter(Collection.id == col_id).first()
if not col:
raise HTTPException(404, "集合不存在")
return col
@router.put("/{col_id}", response_model=CollectionOut)
def update_collection(col_id: int, data: CollectionUpdate, db: Session = Depends(get_db)):
col = db.query(Collection).filter(Collection.id == col_id).first()
if not col:
raise HTTPException(404, "集合不存在")
if data.name is not None:
col.name = data.name
db.commit()
db.refresh(col)
return col
@router.delete("/{col_id}")
def delete_collection(col_id: int, db: Session = Depends(get_db)):
col = db.query(Collection).filter(Collection.id == col_id).first()
if not col:
raise HTTPException(404, "集合不存在")
db.delete(col)
db.commit()
return {"ok": True}
@router.post("/{col_id}/stocks", response_model=CollectionStockOut)
def add_stock(col_id: int, data: StockInput, db: Session = Depends(get_db)):
col = db.query(Collection).filter(Collection.id == col_id).first()
if not col:
raise HTTPException(404, "集合不存在")
stock = CollectionStock(
collection_id=col_id,
stock_code=data.code,
stock_name=data.name,
added_date=date.fromisoformat(data.added_date),
)
db.add(stock)
db.commit()
db.refresh(stock)
return stock
@router.delete("/{col_id}/stocks/{stock_code}")
def remove_stock(col_id: int, stock_code: str, db: Session = Depends(get_db)):
stock = db.query(CollectionStock).filter(
CollectionStock.collection_id == col_id,
CollectionStock.stock_code == stock_code,
).first()
if not stock:
raise HTTPException(404, "该股票不在集合中")
db.delete(stock)
db.commit()
return {"ok": True}
@router.post("/{col_id}/share", response_model=ShareResponse)
def share_collection(col_id: int, db: Session = Depends(get_db)):
col = db.query(Collection).filter(Collection.id == col_id).first()
if not col:
raise HTTPException(404, "集合不存在")
if not col.short_code:
col.short_code = _generate_short_code()
db.commit()
db.refresh(col)
return ShareResponse(
short_code=col.short_code,
url=f"/s/{col.short_code}",
)