| |
| import json |
| import math |
| import os |
| import re |
| import sqlite3 |
| from contextlib import contextmanager |
|
|
| |
| 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() |
|
|
