Blame for database.py

# TANZG 量价分析引擎 - 数据库模块
import json
import math
import os
import re
import sqlite3
from contextlib import contextmanager

# 默认仍使用本地 SQLite；配置 DATABASE_URL 或 PG* 环境变量后自动切换 PostgreSQL。
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DB_PATH = os.getenv("SQLITE_PATH") or os.getenv("DB_PATH") or os.path.join(BASE_DIR, "history.db")
DATABASE_URL = os.getenv("DATABASE_URL") or os.getenv("POSTGRES_URL") or os.getenv("POSTGRES_DSN")
USE_POSTGRES = bool(
    DATABASE_URL
    or os.getenv("PGHOST")
    or os.getenv("POSTGRES_HOST")
)


def _to_native(obj):
    """递归将 numpy / pandas 类型转换为 Python 原生类型，兼容 PostgreSQL 写入。"""
    try:
        import numpy as np
        if isinstance(obj, np.floating):
            val = float(obj)
            return None if (math.isnan(val) or math.isinf(val)) else val
        if isinstance(obj, np.integer):
            return int(obj)
        if isinstance(obj, np.bool_):
            return bool(obj)
        if isinstance(obj, np.ndarray):
            return _to_native(obj.tolist())
    except ImportError:
        pass

    try:
        import pandas as pd
        if obj is pd.NA:
            return None
        if hasattr(pd, "isna") and not isinstance(obj, (dict, list, tuple, str, bytes)):
            try:
                if pd.isna(obj):
                    return None
            except Exception:
                pass
    except Exception:
        pass

    if isinstance(obj, dict):
        return {k: _to_native(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_to_native(x) for x in obj]
    if isinstance(obj, float) and (math.isnan(obj) or math.isinf(obj)):
        return None
    return obj


def _pg_connect():
    import psycopg2
    if DATABASE_URL:
        return psycopg2.connect(DATABASE_URL)
    return psycopg2.connect(
        host=os.getenv("PGHOST") or os.getenv("POSTGRES_HOST", "localhost"),
        port=int(os.getenv("PGPORT") or os.getenv("POSTGRES_PORT", "5432")),
        dbname=os.getenv("PGDATABASE") or os.getenv("POSTGRES_DB", "tanzg"),
        user=os.getenv("PGUSER") or os.getenv("POSTGRES_USER", "postgres"),
        password=os.getenv("PGPASSWORD") or os.getenv("POSTGRES_PASSWORD", ""),
    )


@contextmanager
def _connect():
    if USE_POSTGRES:
        conn = _pg_connect()
    else:
        db_dir = os.path.dirname(DB_PATH)
        if db_dir:
            os.makedirs(db_dir, exist_ok=True)
        conn = sqlite3.connect(DB_PATH)
        conn.row_factory = sqlite3.Row
    try:
        yield conn
        conn.commit()
    finally:
        conn.close()


def _convert_sql(sql):
    """将项目里 SQLite 风格 SQL 转为 PostgreSQL 兼容 SQL。"""
    if not USE_POSTGRES:
        return sql
    converted = sql
    converted = converted.replace("INTEGER PRIMARY KEY AUTOINCREMENT", "SERIAL PRIMARY KEY")
    converted = converted.replace("?", "%s")
    return converted


def _execute(conn, sql, params=None):
    cur = conn.cursor()
    cur.execute(_convert_sql(sql), tuple(_to_native(params or ())))
    return cur


def _fetchall_dict(cur):
    rows = cur.fetchall()
    if not rows:
        return []
    if not USE_POSTGRES:
        return [dict(r) for r in rows]
    cols = [d[0] for d in cur.description]
    return [dict(zip(cols, r)) for r in rows]


def _fetchone_dict(cur):
    row = cur.fetchone()
    if not row:
        return None
    if not USE_POSTGRES:
        return dict(row)
    cols = [d[0] for d in cur.description]
    return dict(zip(cols, row))


def _insert_returning_id(conn, sql, params=None):
    if USE_POSTGRES:
        sql = sql.rstrip().rstrip(";") + " RETURNING id"
        cur = _execute(conn, sql, params)
        return cur.fetchone()[0]
    cur = _execute(conn, sql, params)
    return cur.lastrowid


def init_db():
    with _connect() as conn:
        _execute(conn, """
            CREATE TABLE IF NOT EXISTS analysis_history (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                ts_code TEXT NOT NULL,
                name TEXT,
                created_at TEXT,
                latest_close REAL,
                pct_change REAL,
                ma5 REAL,
                ma10 REAL,
                ma20 REAL,
                boll_upper REAL,
                boll_mid REAL,
                boll_lower REAL,
                boll_pct REAL,
                vol_ratio REAL,
                high_volume_days TEXT,
                trigger_rules TEXT,
                signal TEXT,
                signal_emoji TEXT,
                summary TEXT,
                kline_10d TEXT,
                raw_data TEXT
            )
        """)


def init_watchlist_db():
    with _connect() as conn:
        _execute(conn, """
            CREATE TABLE IF NOT EXISTS watchlist (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                ts_code TEXT NOT NULL UNIQUE,
                name TEXT,
                position INTEGER DEFAULT 0,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP
            )
        """)


def init_recommendations_db():
    with _connect() as conn:
        _execute(conn, """
            CREATE TABLE IF NOT EXISTS recommendations (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                created_at TEXT,
                sentiment TEXT,
                sentiment_label TEXT,
                hot_sectors TEXT,
                stocks TEXT,
                raw_report TEXT
            )
        """)


def get_watchlist():
    with _connect() as conn:
        cur = _execute(conn, "SELECT id, ts_code, name, position FROM watchlist ORDER BY position ASC, id ASC")
        return _fetchall_dict(cur)


def add_to_watchlist(ts_code, name):
    with _connect() as conn:
        cur = _execute(conn, "SELECT COALESCE(MAX(position), -1) + 1 FROM watchlist")
        next_pos = cur.fetchone()[0]
        _execute(conn, "INSERT INTO watchlist (ts_code, name, position) VALUES (?, ?, ?)", (ts_code, name, next_pos))
        return True


def remove_from_watchlist(ts_code):
    with _connect() as conn:
        _execute(conn, "DELETE FROM watchlist WHERE ts_code = ?", (ts_code,))
        return True


def update_watchlist_order(stocks):
    with _connect() as conn:
        _execute(conn, "DELETE FROM watchlist")
        for idx, stock in enumerate(stocks):
            _execute(conn, "INSERT INTO watchlist (ts_code, name, position) VALUES (?, ?, ?)", (stock['code'], stock.get('name', ''), idx))
        return True


def move_watchlist_item(ts_code, direction):
    with _connect() as conn:
        cur = _execute(conn, "SELECT id, ts_code, position FROM watchlist ORDER BY position ASC")
        items = _fetchall_dict(cur)
        if not items:
            return False, "关注列表为空"
        target_idx = next((i for i, item in enumerate(items) if item['ts_code'] == ts_code), -1)
        if target_idx == -1:
            return False, "未找到该股票"
        if direction == 'up':
            if target_idx == 0:
                return False, "已在最上方"
            swap_idx = target_idx - 1
        else:
            if target_idx == len(items) - 1:
                return False, "已在最下方"
            swap_idx = target_idx + 1
        items[target_idx], items[swap_idx] = items[swap_idx], items[target_idx]
        for idx, item in enumerate(items):
            _execute(conn, "UPDATE watchlist SET position = ? WHERE id = ?", (idx, item['id']))
        return True, "移动成功"


def watchlist_exists(ts_code):
    with _connect() as conn:
        cur = _execute(conn, "SELECT COUNT(*) FROM watchlist WHERE ts_code = ?", (ts_code,))
        return cur.fetchone()[0] > 0


def save_report(report: dict) -> int:
    report = _to_native(report)
    with _connect() as conn:
        return _insert_returning_id(conn, """
            INSERT INTO analysis_history
            (ts_code, name, created_at, latest_close, pct_change,
             ma5, ma10, ma20, boll_upper, boll_mid, boll_lower, boll_pct,
             vol_ratio, high_volume_days, trigger_rules, signal, signal_emoji,
             summary, kline_10d, raw_data)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """, (
            report.get("ts_code"),
            report.get("name"),
            report.get("created_at"),
            report.get("latest_close"),
            report.get("pct_change"),
            report.get("ma5"),
            report.get("ma10"),
            report.get("ma20"),
            report.get("boll_upper"),
            report.get("boll_mid"),
            report.get("boll_lower"),
            report.get("boll_pct"),
            report.get("vol_ratio"),
            json.dumps(_to_native(report.get("high_volume_days", [])), ensure_ascii=False),
            json.dumps(_to_native(report.get("trigger_rules", {})), ensure_ascii=False),
            report.get("signal"),
            report.get("signal_emoji"),
            report.get("summary"),
            json.dumps(_to_native(report.get("kline_10d", [])), ensure_ascii=False),
            json.dumps(_to_native(report.get("raw_data", [])), ensure_ascii=False),
        ))


def get_history(limit=50, ts_code=None, start_date=None, end_date=None):
    with _connect() as conn:
        where = []
        params = []
        if ts_code:
            where.append("ts_code = ?")
            params.append(ts_code)
        if start_date:
            where.append("created_at >= ?")
            params.append(start_date)
        if end_date:
            where.append("created_at <= ?")
            params.append(end_date)
        sql = "SELECT * FROM analysis_history"
        if where:
            sql += " WHERE " + " AND ".join(where)
        sql += " ORDER BY created_at DESC LIMIT ?"
        params.append(int(limit))
        cur = _execute(conn, sql, params)
        return _fetchall_dict(cur)


def get_history_by_id(report_id: int):
    with _connect() as conn:
        cur = _execute(conn, "SELECT * FROM analysis_history WHERE id = ?", (report_id,))
        return _fetchone_dict(cur)


def delete_history(report_id: int):
    with _connect() as conn:
        _execute(conn, "DELETE FROM analysis_history WHERE id = ?", (report_id,))


def save_recommendation(rec: dict) -> int:
    rec = _to_native(rec)
    with _connect() as conn:
        return _insert_returning_id(conn, """
            INSERT INTO recommendations
            (created_at, sentiment, sentiment_label, hot_sectors, stocks, raw_report)
            VALUES (?, ?, ?, ?, ?, ?)
        """, (
            rec.get("created_at"),
            rec.get("sentiment", "中性"),
            rec.get("sentiment_label", "中性"),
            json.dumps(_to_native(rec.get("hot_sectors", [])), ensure_ascii=False),
            json.dumps(_to_native(rec.get("stocks", [])), ensure_ascii=False),
            rec.get("raw_report", ""),
        ))


def get_recommendation_history(limit=20):
    with _connect() as conn:
        cur = _execute(conn, "SELECT * FROM recommendations ORDER BY created_at DESC LIMIT ?", (int(limit),))
        return _fetchall_dict(cur)


def get_recommendation_by_id(rid: int):
    with _connect() as conn:
        cur = _execute(conn, "SELECT * FROM recommendations WHERE id = ?", (rid,))
        return _fetchone_dict(cur)


init_db()
init_recommendations_db()
init_watchlist_db()