前端提交

tools/seed_sim_account_demo.py

@@ -3,7 +3,7 @@
 按 D4 决策：1 个客户 + 10 万初始资金 + 2 只基金的初始持仓。
 
 执行：
-    python -m tools.seed_sim_account_demo [--customer-id N]
+    python -m tools.seed_sim_account_demo [--customer-id N] [--quotes-only]
 
 ## 已知表结构问题（2026-09-12 实测）
 
@@ -26,7 +26,7 @@ import sys
 from datetime import UTC, datetime, timedelta
 from decimal import Decimal
 
-from sqlalchemy import func, insert, select
+from sqlalchemy import func, insert, select, update
 from sqlalchemy.orm import Session
 
 from app.core.config import get_settings
@@ -119,6 +119,7 @@ async def _upsert_product(session: Session, spec: dict) -> int:
 
 async def _upsert_market_price(session: Session, product_id: int, spec: dict) -> None:
     today = datetime.now(UTC).date()
+    now = datetime.now(UTC).replace(tzinfo=None)
     existing = (
         await session.execute(
             select(FundMarketPrice.id).where(
@@ -128,8 +129,22 @@ async def _upsert_market_price(session: Session, product_id: int, spec: dict) ->
         )
     ).scalar_one_or_none()
     if existing is not None:
+        await session.execute(
+            update(FundMarketPrice)
+            .where(FundMarketPrice.id == existing)
+            .values(
+                open_price=spec["close_price"],
+                high_price=spec["close_price"] + Decimal("0.05"),
+                low_price=spec["close_price"] - Decimal("0.05"),
+                close_price=spec["close_price"],
+                volume=Decimal("1000000"),
+                turnover_amount=spec["close_price"] * Decimal("1000000"),
+                total_fund_shares=spec["total_fund_shares"],
+                source="eastmoney_demo_seed",
+                source_updated_at=now,
+            )
+        )
         return
-    now = datetime.now(UTC).replace(tzinfo=None)
     next_id = await _next_id(session, FundMarketPrice)
     stmt = insert(FundMarketPrice).values(
         id=next_id,
@@ -149,19 +164,21 @@ async def _upsert_market_price(session: Session, product_id: int, spec: dict) ->
     await session.execute(stmt)
 
 
-async def _upsert_account(session: Session, customer_id: int) -> FundSimAccount:
+async def _upsert_account(
+    session: Session, customer_id: int
+) -> tuple[FundSimAccount, bool]:
     existing = (
         await session.execute(
             select(FundSimAccount.id).where(FundSimAccount.customer_id == customer_id)
         )
     ).scalar_one_or_none()
     if existing is not None:
-        # 返回完整对象
-        return (
+        account = (
             await session.execute(
                 select(FundSimAccount).where(FundSimAccount.customer_id == customer_id)
             )
         ).scalar_one()
+        return account, False
     now = datetime.now(UTC).replace(tzinfo=None)
     next_id = await _next_id(session, FundSimAccount)
     stmt = insert(FundSimAccount).values(
@@ -179,16 +196,17 @@ async def _upsert_account(session: Session, customer_id: int) -> FundSimAccount:
         updated_at=now,
     )
     await session.execute(stmt)
-    return (
+    account = (
         await session.execute(
             select(FundSimAccount).where(FundSimAccount.customer_id == customer_id)
         )
     ).scalar_one()
+    return account, True
 
 
 async def _upsert_holding(
     session: Session, customer_id: int, product_id: int, spec: dict
-) -> None:
+) -> Decimal:
     existing = (
         await session.execute(
             select(FundHolding.id).where(
@@ -198,7 +216,7 @@ async def _upsert_holding(
         )
     ).scalar_one_or_none()
     if existing is not None:
-        return
+        return Decimal("0")
     qty = spec["initial_quantity"]
     cost = (qty * spec["close_price"]).quantize(Decimal("0.01"))
     now = datetime.now(UTC).replace(tzinfo=None)
@@ -227,9 +245,10 @@ async def _upsert_holding(
         updated_at=now,
     )
     await session.execute(stmt)
+    return cost
 
 
-async def run(customer_id: int) -> None:
+async def run(customer_id: int, *, quotes_only: bool = False) -> None:
     s = get_settings()
     print(f"数据库：{s.mysql_dsn.split('@')[-1]}")
     print(f"目标客户 ID = {customer_id}")
@@ -240,38 +259,44 @@ async def run(customer_id: int) -> None:
                 pid = await _upsert_product(session, spec)
                 await _upsert_market_price(session, pid, spec)
                 product_ids.append(pid)
-            print(f"  ✓ 演示产品 {len(product_ids)} 个 + 当日行情")
-            account = await _upsert_account(session, customer_id)
+            print(f"  [完成] 演示产品 {len(product_ids)} 个 + 当日行情")
+            if quotes_only:
+                return
+            account, account_created = await _upsert_account(session, customer_id)
+            added_cost = Decimal("0")
             for pid, spec in zip(product_ids, DEMO_PRODUCTS, strict=True):
-                await _upsert_holding(session, customer_id, pid, spec)
-            total_cost = sum(
-                (spec["initial_quantity"] * spec["close_price"]).quantize(Decimal("0.01"))
-                for spec in DEMO_PRODUCTS
-            )
-            from sqlalchemy import update as sa_update
-            account.cash_balance = (INITIAL_BALANCE - total_cost).quantize(Decimal("0.01"))
-            account.available_cash = account.cash_balance
-            account.updated_at = datetime.now(UTC).replace(tzinfo=None)
-            await session.execute(
-                sa_update(FundSimAccount)
-                .where(FundSimAccount.id == account.id)
-                .values(
-                    cash_balance=account.cash_balance,
-                    available_cash=account.available_cash,
-                    updated_at=account.updated_at,
+                added_cost += await _upsert_holding(session, customer_id, pid, spec)
+            if added_cost:
+                starting_cash = INITIAL_BALANCE if account_created else account.cash_balance
+                account.cash_balance = (starting_cash - added_cost).quantize(Decimal("0.01"))
+                account.available_cash = account.cash_balance
+                account.updated_at = datetime.now(UTC).replace(tzinfo=None)
+                await session.execute(
+                    update(FundSimAccount)
+                    .where(FundSimAccount.id == account.id)
+                    .values(
+                        cash_balance=account.cash_balance,
+                        available_cash=account.available_cash,
+                        updated_at=account.updated_at,
+                    )
                 )
-            )
-        print(f"  ✓ 虚拟账户 {account.account_no} 初始余额 ¥{INITIAL_BALANCE}")
-        print(f"  ✓ 持仓已建立，账户剩余 ¥{account.cash_balance}（已扣持仓成本 ¥{total_cost}）")
+        print(f"  [完成] 虚拟账户 {account.account_no} 初始余额 CNY {INITIAL_BALANCE}")
+        print(f"  [完成] 持仓已建立，本次新增持仓成本 CNY {added_cost}")
+        print(f"  [完成] 当前账户现金 CNY {account.cash_balance}")
 
 
 def main() -> int:
     parser = argparse.ArgumentParser(description="场内模拟交易演示种子")
     parser.add_argument("--customer-id", type=int, default=CUSTOMER_ID)
+    parser.add_argument(
+        "--quotes-only",
+        action="store_true",
+        help="仅刷新演示产品与当日行情，不修改账户、持仓或余额",
+    )
     args = parser.parse_args()
-    asyncio.run(run(args.customer_id))
+    asyncio.run(run(args.customer_id, quotes_only=args.quotes_only))
     return 0
 
 
 if __name__ == "__main__":
-    sys.exit(main())
+    sys.exit(main())