sanguo_vnpy_v2/sanguo_backtest/result_store.py

"""Backtest result storage using SQLite + parquet files."""
import sqlite3
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
import pandas as pd


@dataclass
class BacktestResult:
    """Backtest result data structure."""
    task_id: str
    type: str
    status: str
    strategy: str
    symbol: str
    params: dict
    start: str
    end: str
    statistics: dict
    equity_curve: Optional[pd.DataFrame] = None
    trades: Optional[pd.DataFrame] = None
    error_msg: Optional[str] = None
    id: Optional[int] = None


# SQLite schema for backtest stats
_SCHEMA = """CREATE TABLE IF NOT EXISTS backtest_stats (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    task_id TEXT,
    type TEXT,
    status TEXT,
    strategy TEXT,
    symbol TEXT,
    params TEXT,
    start TEXT,
    end TEXT,
    statistics TEXT,
    equity_path TEXT,
    trades_path TEXT,
    error_msg TEXT,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);"""


def _connect(db_path: str) -> sqlite3.Connection:
    """Create database connection and initialize schema."""
    conn = sqlite3.connect(db_path)
    conn.executescript(_SCHEMA)
    return conn


def save_result(result: BacktestResult, db_path: str, file_dir: Optional[str] = None) -> int:
    """
    Save backtest result to database and optionally save DataFrame fields to parquet files.

    Args:
        result: BacktestResult object to save
        db_path: Path to SQLite database file
        file_dir: Optional directory to save parquet files (equity_curve and trades)

    Returns:
        int: The ID of the inserted record
    """
    conn = _connect(db_path)
    try:
        equity_path = None
        trades_path = None

        # Save DataFrames to parquet if file_dir is provided
        if file_dir:
            fdir = Path(file_dir)
            fdir.mkdir(parents=True, exist_ok=True)

            if result.equity_curve is not None and not result.equity_curve.empty:
                equity_path = str(fdir / f"{result.task_id}_equity.json")
                result.equity_curve.to_json(equity_path, orient="records", date_format="iso", force_ascii=False)

            if result.trades is not None and not result.trades.empty:
                trades_path = str(fdir / f"{result.task_id}_trades.json")
                result.trades.to_json(trades_path, orient="records", date_format="iso", force_ascii=False)

        # Insert record into database
        cur = conn.execute(
            """INSERT INTO backtest_stats
               (task_id, type, status, strategy, symbol, params, start, end,
                statistics, equity_path, trades_path, error_msg)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                result.task_id,
                result.type,
                result.status,
                result.strategy,
                result.symbol,
                json.dumps(result.params),
                result.start,
                result.end,
                json.dumps(result.statistics),
                equity_path,
                trades_path,
                result.error_msg
            )
        )
        conn.commit()
        result.id = cur.lastrowid
        return cur.lastrowid
    finally:
        conn.close()


def load_result(rid: int, db_path: str) -> BacktestResult:
    """
    Load backtest result by ID from database.

    Args:
        rid: Result ID to load
        db_path: Path to SQLite database file

    Returns:
        BacktestResult: Loaded result object

    Raises:
        KeyError: If result ID not found
    """
    conn = _connect(db_path)
    try:
        row = conn.execute("SELECT * FROM backtest_stats WHERE id=?", (rid,)).fetchone()
        if not row:
            raise KeyError(f"result {rid} not found")

        # Get column names from table description
        cols = [d[0] for d in conn.execute("SELECT * FROM backtest_stats LIMIT 0").description]
        d = dict(zip(cols, row))

        # Load JSON files if paths exist (equity_curve/trades persisted as JSON)
        equity = pd.read_json(d["equity_path"], orient="records") if d.get("equity_path") else None
        trades = pd.read_json(d["trades_path"], orient="records") if d.get("trades_path") else None

        return BacktestResult(
            task_id=d["task_id"],
            type=d["type"],
            status=d["status"],
            strategy=d["strategy"],
            symbol=d["symbol"],
            params=json.loads(d["params"]),
            start=d["start"],
            end=d["end"],
            statistics=json.loads(d["statistics"]) if d["statistics"] else {},
            equity_curve=equity,
            trades=trades,
            error_msg=d.get("error_msg")
        )
    finally:
        conn.close()


def load_result_by_task_id(task_id: str, db_path: str) -> BacktestResult | None:
    """Load the most recent result for a task_id (historical lookup after restart)."""
    conn = _connect(db_path)
    try:
        row = conn.execute(
            "SELECT id FROM backtest_stats WHERE task_id=? ORDER BY id DESC LIMIT 1",
            (task_id,),
        ).fetchone()
        if not row:
            return None
        return load_result(row[0], db_path)
    finally:
        conn.close()


def list_results(type_filter: Optional[str] = None, db_path: str = "") -> list[BacktestResult]:
    """
    List all backtest results, optionally filtered by type.

    Args:
        type_filter: Optional filter for result type (e.g., 'cta', 'factor')
        db_path: Path to SQLite database file

    Returns:
        list[BacktestResult]: List of loaded result objects
    """
    conn = _connect(db_path)
    try:
        query = "SELECT id FROM backtest_stats"
        args = ()

        if type_filter:
            query = query + " WHERE type=?"
            args = (type_filter,)

        return [load_result(row[0], db_path) for row in conn.execute(query, args).fetchall()]
    finally:
        conn.close()