diff --git a/sanguo_trader/persistence.py b/sanguo_trader/persistence.py new file mode 100644 index 0000000..73ca300 --- /dev/null +++ b/sanguo_trader/persistence.py @@ -0,0 +1,173 @@ +"""模拟盘 SQLite 持久化(4 表 + checkpoint + WAL,spec §8 / §9.1)。 + +表:paper_accounts / paper_trades / paper_positions / paper_daily_balance +WAL 模式支持 worker 进程写 + 主进程读(spec §9.1 共享 DB 进度)。 +""" +import json +import sqlite3 +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +SCHEMA = """ +CREATE TABLE IF NOT EXISTS paper_accounts ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + task_id TEXT, owner_id TEXT DEFAULT 'admin', name TEXT, + mode TEXT, interval TEXT, + symbols TEXT, strategies TEXT, + initial_capital REAL, rate REAL, slippage REAL, size REAL, pricetick REAL, + stamp_duty_rate REAL, transfer_fee_rate REAL, min_commission REAL, + status TEXT, start_date TEXT, end_date TEXT, + last_run_date TEXT, next_run_at TEXT, scheduler_job_id TEXT, + checkpoint_date TEXT, error_msg TEXT, + created_at TEXT, updated_at TEXT +); +CREATE TABLE IF NOT EXISTS paper_trades ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + account_id INTEGER, strategy_id TEXT, datetime TEXT, symbol TEXT, + direction TEXT, offset TEXT, match_session TEXT, + price REAL, volume INTEGER, + commission REAL, stamp_duty REAL, transfer_fee REAL, + rejected INTEGER, reject_reason TEXT, bar_date TEXT, + blocked_by_strategy_id TEXT +); +CREATE TABLE IF NOT EXISTS paper_positions ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + account_id INTEGER, scope TEXT, symbol TEXT, date TEXT, + volume INTEGER, frozen INTEGER, avg_price REAL, market_value REAL, + updated_at TEXT +); +CREATE TABLE IF NOT EXISTS paper_daily_balance ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + account_id INTEGER, date TEXT, + cash REAL, market_value REAL, total_equity REAL, + per_strategy_pnl TEXT, is_checkpoint INTEGER +); +""" + + +def _now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def init_db(db_path: str) -> None: + Path(db_path).parent.mkdir(parents=True, exist_ok=True) + with sqlite3.connect(db_path) as conn: + conn.executescript(SCHEMA) + conn.execute("PRAGMA journal_mode=WAL") + conn.commit() + + +def save_account(db_path: str, account: dict[str, Any]) -> int: + with sqlite3.connect(db_path) as conn: + cur = conn.execute( + """INSERT INTO paper_accounts + (task_id, owner_id, name, mode, interval, symbols, strategies, + initial_capital, rate, slippage, size, pricetick, + stamp_duty_rate, transfer_fee_rate, min_commission, + status, start_date, end_date, created_at, updated_at) + VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)""", + ( + account.get("task_id"), account.get("owner_id", "admin"), + account.get("name"), account.get("mode"), account.get("interval"), + json.dumps(account.get("symbols", [])), + json.dumps(account.get("strategies", [])), + account.get("initial_capital", 0), account.get("rate", 0.0003), + account.get("slippage", 0), account.get("size", 1), + account.get("pricetick", 0.01), + account.get("stamp_duty_rate", 0.0005), + account.get("transfer_fee_rate", 0.00001), + account.get("min_commission", 5.0), + account.get("status", "pending"), account.get("start_date"), + account.get("end_date"), _now(), _now(), + ), + ) + conn.commit() + return cur.lastrowid + + +def update_account_status(db_path: str, account_id: int, status: str, + error_msg: str = "") -> None: + with sqlite3.connect(db_path) as conn: + conn.execute( + "UPDATE paper_accounts SET status=?, error_msg=?, updated_at=? WHERE id=?", + (status, error_msg, _now(), account_id), + ) + conn.commit() + + +def save_trade(db_path: str, account_id: int, trade: dict[str, Any], + rejected: bool = False, reject_reason: str = "") -> int: + with sqlite3.connect(db_path) as conn: + cur = conn.execute( + """INSERT INTO paper_trades + (account_id, strategy_id, datetime, symbol, direction, offset, + match_session, price, volume, commission, stamp_duty, transfer_fee, + rejected, reject_reason, bar_date, blocked_by_strategy_id) + VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)""", + ( + account_id, trade.get("strategy_id", ""), trade.get("datetime", ""), + trade.get("symbol", ""), trade.get("direction", ""), + trade.get("offset", ""), trade.get("match_session", ""), + trade.get("price", 0), trade.get("volume", 0), + trade.get("commission", 0), trade.get("stamp_duty", 0), + trade.get("transfer_fee", 0), + 1 if rejected else 0, reject_reason, trade.get("bar_date", ""), + trade.get("blocked_by_strategy_id"), + ), + ) + conn.commit() + return cur.lastrowid + + +def save_daily_balance(db_path: str, account_id: int, date: str, cash: float, + market_value: float, total_equity: float, + per_strategy_pnl: dict | None = None, + is_checkpoint: bool = False) -> None: + with sqlite3.connect(db_path) as conn: + conn.execute( + """INSERT INTO paper_daily_balance + (account_id, date, cash, market_value, total_equity, + per_strategy_pnl, is_checkpoint) + VALUES (?,?,?,?,?,?,?)""", + (account_id, date, cash, market_value, total_equity, + json.dumps(per_strategy_pnl or {}), 1 if is_checkpoint else 0), + ) + conn.commit() + + +def update_checkpoint(db_path: str, account_id: int, date: str) -> None: + with sqlite3.connect(db_path) as conn: + conn.execute( + "UPDATE paper_accounts SET checkpoint_date=?, updated_at=? WHERE id=?", + (date, _now(), account_id), + ) + conn.commit() + + +def load_checkpoint(db_path: str, account_id: int) -> str | None: + with sqlite3.connect(db_path) as conn: + cur = conn.execute( + "SELECT checkpoint_date FROM paper_accounts WHERE id=?", (account_id,) + ) + row = cur.fetchone() + return row[0] if row else None + + +def list_trades(db_path: str, account_id: int) -> list[dict]: + with sqlite3.connect(db_path) as conn: + conn.row_factory = sqlite3.Row + cur = conn.execute( + "SELECT * FROM paper_trades WHERE account_id=? ORDER BY id", (account_id,) + ) + return [dict(r) for r in cur.fetchall()] + + +def list_daily_balance(db_path: str, account_id: int) -> list[dict]: + with sqlite3.connect(db_path) as conn: + conn.row_factory = sqlite3.Row + cur = conn.execute( + "SELECT * FROM paper_daily_balance WHERE account_id=? ORDER BY date", + (account_id,), + ) + return [dict(r) for r in cur.fetchall()] diff --git a/tests/trader/test_persistence.py b/tests/trader/test_persistence.py new file mode 100644 index 0000000..22b8888 --- /dev/null +++ b/tests/trader/test_persistence.py @@ -0,0 +1,61 @@ +"""persistence 4 表 round-trip 测试(spec §8)。""" +from sanguo_trader.persistence import ( + init_db, save_account, save_trade, save_daily_balance, + update_account_status, update_checkpoint, load_checkpoint, + list_trades, list_daily_balance, +) + + +def test_init_and_save_account(tmp_path): + db = str(tmp_path / "p.db") + init_db(db) + aid = save_account(db, {"name": "test", "mode": "replay", "initial_capital": 1_000_000}) + assert aid > 0 + + +def test_save_trade_and_list(tmp_path): + db = str(tmp_path / "p.db") + init_db(db) + aid = save_account(db, {"name": "t"}) + save_trade(db, aid, {"strategy_id": "s1", "symbol": "600000", + "price": 10.0, "volume": 100, "commission": 5.0}) + trades = list_trades(db, aid) + assert len(trades) == 1 + assert trades[0]["price"] == 10.0 + assert trades[0]["rejected"] == 0 + + +def test_save_reject_trade(tmp_path): + db = str(tmp_path / "p.db") + init_db(db) + aid = save_account(db, {"name": "t"}) + save_trade(db, aid, {"symbol": "300750"}, rejected=True, + reject_reason="limit_up_locked") + trades = list_trades(db, aid) + assert trades[0]["rejected"] == 1 + assert trades[0]["reject_reason"] == "limit_up_locked" + + +def test_daily_balance_and_checkpoint(tmp_path): + db = str(tmp_path / "p.db") + init_db(db) + aid = save_account(db, {"name": "t"}) + save_daily_balance(db, aid, "2024-01-01", 999_000, 1000, 1_000_000, + per_strategy_pnl={"s1": 100}, is_checkpoint=True) + save_daily_balance(db, aid, "2024-01-02", 998_000, 1100, 999_000, + is_checkpoint=False) + update_checkpoint(db, aid, "2024-01-01") + assert load_checkpoint(db, aid) == "2024-01-01" + balances = list_daily_balance(db, aid) + assert len(balances) == 2 + assert balances[0]["is_checkpoint"] == 1 + + +def test_update_account_status(tmp_path): + db = str(tmp_path / "p.db") + init_db(db) + aid = save_account(db, {"name": "t", "status": "pending"}) + update_account_status(db, aid, "done") + with __import__("sqlite3").connect(db) as conn: + row = conn.execute("SELECT status FROM paper_accounts WHERE id=?", (aid,)).fetchone() + assert row[0] == "done"