stock-advisor/database.py

# -*- coding: utf-8 -*-
"""
SQLite 数据访问层：线程安全连接 + 轻量查询封装
表结构：stocks / stock_daily / news / institutions / inst_ratings / fund_holdings
       / watchlist(自选) / analysis_cache(深度报告缓存)
"""
import sqlite3
import threading
from contextlib import contextmanager

from config import DB_PATH

_local = threading.local()

SCHEMA = """
CREATE TABLE IF NOT EXISTS stocks (
  code TEXT PRIMARY KEY,
  name TEXT NOT NULL,
  industry TEXT DEFAULT '',
  board TEXT DEFAULT '主板',
  total_shares REAL DEFAULT 0,     -- 总股本（亿股）
  float_shares REAL DEFAULT 0,     -- 流通股本（亿股）
  list_date TEXT DEFAULT '',
  pe REAL DEFAULT 0,               -- 市盈率（动）
  pb REAL DEFAULT 0,               -- 市净率
  market_cap REAL DEFAULT 0,       -- 总市值（亿元）
  description TEXT DEFAULT ''      -- 公司简介
);

CREATE TABLE IF NOT EXISTS stock_daily (
  code TEXT NOT NULL,
  date TEXT NOT NULL,
  open REAL DEFAULT 0,
  high REAL DEFAULT 0,
  low REAL DEFAULT 0,
  close REAL DEFAULT 0,
  volume REAL DEFAULT 0,           -- 成交量（万股）
  amount REAL DEFAULT 0,           -- 成交额（万元）
  change_pct REAL DEFAULT 0,       -- 涨跌幅 %
  PRIMARY KEY (code, date)
);
CREATE INDEX IF NOT EXISTS idx_daily_date ON stock_daily(date);

CREATE TABLE IF NOT EXISTS news (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  title TEXT NOT NULL,
  content TEXT DEFAULT '',
  source TEXT DEFAULT '',
  category TEXT DEFAULT '市场',
  publish_date TEXT DEFAULT '',
  related_stocks TEXT DEFAULT '',  -- 关联股票代码，逗号分隔
  sentiment REAL DEFAULT 0,        -- 情感 -1~1
  is_positive INTEGER DEFAULT 1,
  created_at TEXT DEFAULT (datetime('now','localtime'))
);
CREATE INDEX IF NOT EXISTS idx_news_date ON news(publish_date);

CREATE TABLE IF NOT EXISTS institutions (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  name TEXT UNIQUE,
  type TEXT DEFAULT '',            -- 公募基金/券商研究/保险资管/外资机构/私募游资
  description TEXT DEFAULT ''
);

CREATE TABLE IF NOT EXISTS inst_ratings (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  inst_id INTEGER,
  inst_name TEXT,
  stock_code TEXT,
  rating TEXT DEFAULT '中性',      -- 买入/增持/中性/减持
  target_price REAL DEFAULT 0,
  rating_date TEXT DEFAULT '',
  prev_rating TEXT DEFAULT ''
);
CREATE INDEX IF NOT EXISTS idx_ratings_code ON inst_ratings(stock_code);
CREATE INDEX IF NOT EXISTS idx_ratings_date ON inst_ratings(rating_date);

CREATE TABLE IF NOT EXISTS fund_holdings (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  inst_id INTEGER,
  inst_name TEXT,
  stock_code TEXT,
  quarter TEXT,                    -- 2026Q2
  hold_shares REAL DEFAULT 0,      -- 持仓（万股）
  hold_value REAL DEFAULT 0,       -- 持仓市值（万元）
  change_shares REAL DEFAULT 0,    -- 环比增减（万股）
  change_pct REAL DEFAULT 0,       -- 环比增减 %
  UNIQUE(inst_id, stock_code, quarter)
);
CREATE INDEX IF NOT EXISTS idx_holdings_code ON fund_holdings(stock_code);

CREATE TABLE IF NOT EXISTS watchlist (
  code TEXT PRIMARY KEY,
  added_at TEXT DEFAULT (datetime('now','localtime'))
);

CREATE TABLE IF NOT EXISTS analysis_cache (
  code TEXT PRIMARY KEY,
  report TEXT,
  created_at TEXT DEFAULT (datetime('now','localtime'))
);

CREATE TABLE IF NOT EXISTS analysis_history (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  code TEXT NOT NULL,
  stock_name TEXT DEFAULT '',
  focus TEXT DEFAULT '',
  report TEXT,
  sources TEXT DEFAULT '{}',        -- JSON：大模型参考的数据源（RAG新闻/概况/指标/评级/持仓/提示词）
  created_at TEXT DEFAULT (datetime('now','localtime'))
);
CREATE INDEX IF NOT EXISTS idx_history_code ON analysis_history(code);

CREATE TABLE IF NOT EXISTS strategy_backtests (
  strategy TEXT NOT NULL,
  code TEXT NOT NULL,
  stock_name TEXT DEFAULT '',
  metrics TEXT DEFAULT '{}',       -- JSON：收益/回撤/夏普/胜率等
  equity TEXT DEFAULT '[]',       -- JSON：[{date,value,bh}, ...] 净值曲线
  trades TEXT DEFAULT '[]',       -- JSON：交易明细
  run_at TEXT DEFAULT (datetime('now','localtime')),
  PRIMARY KEY (strategy, code)
);

CREATE TABLE IF NOT EXISTS settings (
  key TEXT PRIMARY KEY,
  value TEXT DEFAULT ''
);

CREATE TABLE IF NOT EXISTS notification_log (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  news_id INTEGER DEFAULT 0,
  title TEXT DEFAULT '',
  category TEXT DEFAULT '',
  sentiment REAL DEFAULT 0,
  importance REAL DEFAULT 0,
  related TEXT DEFAULT '',
  status TEXT DEFAULT 'sent',       -- sent / failed
  message TEXT DEFAULT '',
  sent_at TEXT DEFAULT (datetime('now','localtime'))
);

CREATE TABLE IF NOT EXISTS market_index (
  date TEXT PRIMARY KEY,
  sh REAL DEFAULT 0,   -- 上证指数（点）
  sz REAL DEFAULT 0,   -- 深证成指（点）
  cy REAL DEFAULT 0    -- 创业板指（点）
);
"""


def get_conn():
    conn = getattr(_local, "conn", None)
    if conn is None:
        conn = sqlite3.connect(DB_PATH, timeout=30)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA busy_timeout=30000")
        _local.conn = conn
    return conn


@contextmanager
def db():
    conn = get_conn()
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise


def init_db():
    with db() as conn:
        conn.executescript(SCHEMA)


def query(sql, args=()):
    with db() as conn:
        cur = conn.execute(sql, args)
        return [dict(r) for r in cur.fetchall()]


def query_one(sql, args=()):
    rows = query(sql, args)
    return rows[0] if rows else None


def execute(sql, args=()):
    with db() as conn:
        cur = conn.execute(sql, args)
        return cur.lastrowid


def executemany(sql, seq):
    with db() as conn:
        conn.executemany(sql, seq)


def table_count(name):
    return query_one(f'SELECT COUNT(*) AS c FROM "{name}"')["c"]


def wipe_all():
    """清空业务表（保留结构）+ 重置自增序列，用于重灌数据"""
    for t in ("stock_daily", "inst_ratings", "fund_holdings", "news",
              "institutions", "stocks", "watchlist", "analysis_cache", "analysis_history",
              "market_index", "strategy_backtests", "notification_log"):
        with db() as conn:
            conn.execute(f'DELETE FROM "{t}"')
    with db() as conn:
        conn.execute("DELETE FROM sqlite_sequence")   # 重置 AUTOINCREMENT，保证重灌后 ID 从 1 开始
    # settings 保留（用户配置不清空）