0cfc9bf9c8
- 虚拟用户管理（昵称/头像/性别/简介/邮箱同步到目标平台） - AI互动调度（点赞/收藏/评论/转发） - 日志时间改为北京时间 - 评论达上限后继续执行点赞收藏转发 - 一键登出全部功能 - 浅色主题UI
73 lines
2.2 KiB
Python
73 lines
2.2 KiB
Python
"""数据库连接管理"""
|
||
import asyncio
|
||
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
||
from sqlalchemy.orm import DeclarativeBase
|
||
from app.core.config import settings
|
||
from app.core.logger import logger
|
||
|
||
|
||
class Base(DeclarativeBase):
|
||
pass
|
||
|
||
|
||
engine = create_async_engine(
|
||
settings.DATABASE_URL,
|
||
echo=False,
|
||
pool_pre_ping=True,
|
||
pool_recycle=3600,
|
||
pool_size=10,
|
||
max_overflow=20,
|
||
)
|
||
|
||
AsyncSessionLocal = async_sessionmaker(
|
||
engine, class_=AsyncSession, expire_on_commit=False
|
||
)
|
||
|
||
|
||
async def get_db():
|
||
"""获取数据库会话"""
|
||
async with AsyncSessionLocal() as session:
|
||
try:
|
||
yield session
|
||
await session.commit()
|
||
except Exception:
|
||
await session.rollback()
|
||
raise
|
||
finally:
|
||
await session.close()
|
||
|
||
|
||
async def wait_for_db(max_retries: int = 30, interval: int = 2):
|
||
"""等待 MySQL 就绪，最多重试 max_retries 次"""
|
||
for attempt in range(1, max_retries + 1):
|
||
try:
|
||
async with engine.begin() as conn:
|
||
await conn.execute(__import__("sqlalchemy").text("SELECT 1"))
|
||
logger.info(f"✅ 数据库连接成功（第 {attempt} 次尝试）")
|
||
return
|
||
except Exception as e:
|
||
if attempt == max_retries:
|
||
logger.error(f"数据库连接失败，已重试 {max_retries} 次: {e}")
|
||
raise
|
||
logger.warning(f"数据库未就绪，{interval}s 后重试（{attempt}/{max_retries}）: {e}")
|
||
await asyncio.sleep(interval)
|
||
|
||
|
||
async def init_db():
|
||
"""初始化数据库 - 等待 MySQL 就绪并注册所有模型"""
|
||
try:
|
||
# 等待 MySQL 容器真正就绪
|
||
await wait_for_db(max_retries=30, interval=2)
|
||
|
||
# 导入所有模型类，确保 SQLAlchemy ORM 元数据注册
|
||
from app.models import (
|
||
VirtualUser, UserPersonality, InteractionRecord,
|
||
TokenStat, AIModelConfig, SystemConfig, LoginLog
|
||
)
|
||
logger.info("✅ 数据库模型注册成功")
|
||
logger.info("✅ 数据库初始化完成")
|
||
|
||
except Exception as e:
|
||
logger.error(f"数据库初始化失败: {e}")
|
||
raise
|