51 lines
2.2 KiB
Python
51 lines
2.2 KiB
Python
"""SQLAlchemy 异步数据库配置"""
|
||
|
||
from contextlib import asynccontextmanager
|
||
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
||
from app.config import settings
|
||
|
||
# SQLite 异步需要 aiosqlite
|
||
db_url = settings.database_url.replace("sqlite:///", "sqlite+aiosqlite:///")
|
||
engine = create_async_engine(db_url, echo=False)
|
||
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 position_score REAL",
|
||
"ALTER TABLE recommendations ADD COLUMN valuation_score REAL",
|
||
"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 sector_heat ADD COLUMN stage TEXT",
|
||
"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",
|
||
]:
|
||
try:
|
||
await conn.execute(
|
||
__import__("sqlalchemy").text(col_sql)
|
||
)
|
||
except Exception:
|
||
pass # 列已存在，忽略
|