astock-agent/backend/app/db/database.py
2026-04-24 09:00:46 +08:00

112 lines
6.3 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""SQLAlchemy 异步数据库配置"""
from contextlib import asynccontextmanager
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy import event
from app.config import settings
# SQLite 异步需要 aiosqlite
db_url = settings.database_url.replace("sqlite:///", "sqlite+aiosqlite:///")
engine = create_async_engine(db_url, echo=False)
# 启用 WAL 模式:允许读写并发,写操作不阻塞读操作
@event.listens_for(engine.sync_engine, "connect")
def _set_sqlite_pragma(dbapi_connection, connection_record):
cursor = dbapi_connection.cursor()
cursor.execute("PRAGMA journal_mode=WAL")
cursor.execute("PRAGMA busy_timeout=5000") # 写锁最多等待 5 秒
cursor.close()
async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
@asynccontextmanager
async def get_db():
session = async_session()
try:
yield session
finally:
await session.close()
async def init_db():
"""创建所有表,并补充新增列"""
from app.db.tables import metadata
async with engine.begin() as conn:
await conn.run_sync(metadata.create_all)
# 补充新增列SQLite ALTER TABLE ADD COLUMN已存在会忽略
for col_sql in [
"ALTER TABLE recommendations ADD COLUMN supply_demand_score REAL DEFAULT 0",
"ALTER TABLE recommendations ADD COLUMN price_action_score REAL DEFAULT 0",
"ALTER TABLE recommendations ADD COLUMN position_score REAL",
"ALTER TABLE recommendations ADD COLUMN valuation_score REAL",
"ALTER TABLE recommendations ADD COLUMN risk_note TEXT DEFAULT ''",
"ALTER TABLE recommendations ADD COLUMN action_plan TEXT DEFAULT '观察'",
"ALTER TABLE recommendations ADD COLUMN trigger_condition TEXT DEFAULT ''",
"ALTER TABLE recommendations ADD COLUMN invalidation_condition TEXT DEFAULT ''",
"ALTER TABLE recommendations ADD COLUMN suggested_position_pct REAL DEFAULT 0",
"ALTER TABLE recommendations ADD COLUMN review_after_days INTEGER DEFAULT 3",
"ALTER TABLE recommendations ADD COLUMN lifecycle_status TEXT DEFAULT 'candidate'",
"ALTER TABLE recommendations ADD COLUMN data_freshness TEXT DEFAULT ''",
"ALTER TABLE recommendations ADD COLUMN llm_analysis TEXT DEFAULT ''",
"ALTER TABLE recommendations ADD COLUMN strategy TEXT DEFAULT 'momentum'",
"ALTER TABLE recommendations ADD COLUMN llm_score REAL",
"ALTER TABLE market_temperature ADD COLUMN max_streak INTEGER",
"ALTER TABLE market_temperature ADD COLUMN broken_rate REAL",
"ALTER TABLE recommendations ADD COLUMN entry_signal_type TEXT DEFAULT 'none'",
"ALTER TABLE recommendations ADD COLUMN entry_timing TEXT DEFAULT ''",
"ALTER TABLE recommendations ADD COLUMN recall_tags TEXT DEFAULT '[]'",
"ALTER TABLE recommendations ADD COLUMN prefilter_decision TEXT DEFAULT ''",
"ALTER TABLE recommendations ADD COLUMN prefilter_reason TEXT DEFAULT ''",
"ALTER TABLE recommendations ADD COLUMN focus_points TEXT DEFAULT '[]'",
"ALTER TABLE sector_heat ADD COLUMN stage TEXT",
"ALTER TABLE sector_heat ADD COLUMN board_type TEXT DEFAULT 'theme'",
"ALTER TABLE sector_heat ADD COLUMN theme_id TEXT DEFAULT ''",
"ALTER TABLE sector_heat ADD COLUMN theme_name TEXT DEFAULT ''",
"ALTER TABLE sector_heat ADD COLUMN theme_aliases TEXT DEFAULT '[]'",
"ALTER TABLE sector_heat ADD COLUMN days_continuous INTEGER",
"ALTER TABLE sector_heat ADD COLUMN member_count INTEGER",
"ALTER TABLE sector_heat ADD COLUMN leading_stocks TEXT",
"ALTER TABLE sector_heat ADD COLUMN pct_trend TEXT",
"ALTER TABLE sector_heat ADD COLUMN turnover_avg REAL",
"ALTER TABLE sector_heat ADD COLUMN main_force_ratio REAL",
"ALTER TABLE recommendation_tracking ADD COLUMN max_price REAL",
"ALTER TABLE recommendation_tracking ADD COLUMN min_price REAL",
"ALTER TABLE recommendation_tracking ADD COLUMN max_return_pct REAL",
"ALTER TABLE recommendation_tracking ADD COLUMN max_drawdown_pct REAL",
"ALTER TABLE recommendation_tracking ADD COLUMN days_since_recommendation INTEGER DEFAULT 0",
"ALTER TABLE recommendation_tracking ADD COLUMN close_reason TEXT DEFAULT ''",
"ALTER TABLE recommendation_tracking ADD COLUMN review_note TEXT DEFAULT ''",
"ALTER TABLE users ADD COLUMN email TEXT",
"ALTER TABLE users ADD COLUMN invite_code_used TEXT DEFAULT ''",
"ALTER TABLE stock_diagnoses ADD COLUMN diagnosis_mode TEXT DEFAULT 'entry'",
"ALTER TABLE user_watchlists ADD COLUMN note TEXT DEFAULT ''",
"ALTER TABLE user_watchlists ADD COLUMN watch_group TEXT DEFAULT 'observe'",
"ALTER TABLE user_watchlists ADD COLUMN cost_price REAL",
"ALTER TABLE user_watchlists ADD COLUMN is_active BOOLEAN DEFAULT 1",
"ALTER TABLE user_watchlists ADD COLUMN updated_at DATETIME DEFAULT CURRENT_TIMESTAMP",
"ALTER TABLE watchlist_analyses ADD COLUMN conclusion TEXT DEFAULT '观察'",
"ALTER TABLE watchlist_analyses ADD COLUMN advice TEXT DEFAULT ''",
"ALTER TABLE watchlist_analyses ADD COLUMN trigger_condition TEXT DEFAULT ''",
"ALTER TABLE watchlist_analyses ADD COLUMN risk_note TEXT DEFAULT ''",
"ALTER TABLE watchlist_analyses ADD COLUMN summary TEXT DEFAULT ''",
"ALTER TABLE watchlist_analyses ADD COLUMN full_analysis TEXT DEFAULT ''",
"ALTER TABLE watchlist_analyses ADD COLUMN score_reference REAL DEFAULT 0",
"ALTER TABLE watchlist_analyses ADD COLUMN analysis_mode TEXT DEFAULT 'scheduled'",
]:
try:
await conn.execute(
__import__("sqlalchemy").text(col_sql)
)
except Exception:
pass # 列已存在,忽略
try:
await conn.execute(
__import__("sqlalchemy").text(
"UPDATE users SET email = username WHERE email IS NULL OR email = ''"
)
)
except Exception:
pass