quant-web/strategy_db.py

"""SQLite persistence for user-saved strategies (builtin + custom Python)."""

from __future__ import annotations

import json
import os
import sqlite3
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any


@dataclass(frozen=True)
class SavedStrategy:
    id: int
    username: str
    name: str
    ticker: str
    strategy_key: str
    params: dict[str, Any]
    source_code: str | None
    created_at: str


def _db_path() -> str:
    return os.environ.get("STRATEGY_DB_PATH", "/data/strategies/strategies.db")


def _migrate(conn: sqlite3.Connection) -> None:
    cols = {row[1] for row in conn.execute("PRAGMA table_info(strategies)")}
    if "strategy_key" not in cols:
        conn.execute(
            "ALTER TABLE strategies ADD COLUMN strategy_key TEXT NOT NULL DEFAULT 'ma_crossover'"
        )
    if "source_code" not in cols:
        conn.execute("ALTER TABLE strategies ADD COLUMN source_code TEXT")


def init_db() -> None:
    path = _db_path()
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with _connect() as conn:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS strategies (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT NOT NULL,
                name TEXT NOT NULL,
                ticker TEXT NOT NULL,
                strategy_key TEXT NOT NULL DEFAULT 'ma_crossover',
                params_json TEXT NOT NULL,
                source_code TEXT,
                created_at TEXT NOT NULL,
                UNIQUE(username, name)
            )
            """
        )
        _migrate(conn)
        conn.commit()


@contextmanager
def _connect():
    conn = sqlite3.connect(_db_path())
    conn.row_factory = sqlite3.Row
    try:
        yield conn
    finally:
        conn.close()


def _row_to_strategy(row: sqlite3.Row) -> SavedStrategy:
    return SavedStrategy(
        id=row["id"],
        username=row["username"],
        name=row["name"],
        ticker=row["ticker"],
        strategy_key=row["strategy_key"] if "strategy_key" in row.keys() else "ma_crossover",
        params=json.loads(row["params_json"]),
        source_code=row["source_code"] if "source_code" in row.keys() else None,
        created_at=row["created_at"],
    )


def save_strategy(
    username: str,
    name: str,
    ticker: str,
    strategy_key: str,
    params: dict[str, Any],
    source_code: str | None = None,
) -> None:
    created_at = datetime.now(timezone.utc).isoformat()
    with _connect() as conn:
        conn.execute(
            """
            INSERT INTO strategies (
                username, name, ticker, strategy_key, params_json, source_code, created_at
            )
            VALUES (?, ?, ?, ?, ?, ?, ?)
            ON CONFLICT(username, name) DO UPDATE SET
                ticker = excluded.ticker,
                strategy_key = excluded.strategy_key,
                params_json = excluded.params_json,
                source_code = excluded.source_code,
                created_at = excluded.created_at
            """,
            (
                username,
                name.strip(),
                ticker.upper(),
                strategy_key,
                json.dumps(params),
                source_code,
                created_at,
            ),
        )
        conn.commit()


def list_strategies(username: str) -> list[SavedStrategy]:
    with _connect() as conn:
        rows = conn.execute(
            """
            SELECT id, username, name, ticker, strategy_key, params_json, source_code, created_at
            FROM strategies
            WHERE username = ?
            ORDER BY created_at DESC
            """,
            (username,),
        ).fetchall()
    return [_row_to_strategy(row) for row in rows]


def load_strategy(username: str, name: str) -> SavedStrategy | None:
    with _connect() as conn:
        row = conn.execute(
            """
            SELECT id, username, name, ticker, strategy_key, params_json, source_code, created_at
            FROM strategies
            WHERE username = ? AND name = ?
            """,
            (username, name),
        ).fetchone()
    return _row_to_strategy(row) if row else None


def delete_strategy(username: str, name: str) -> None:
    with _connect() as conn:
        conn.execute(
            "DELETE FROM strategies WHERE username = ? AND name = ?",
            (username, name),
        )
        conn.commit()