Files
huihuiSquare/backend/app/core/database.py
stefanfeng cd07776914 feat: 多项功能更新
- 日志时间改为北京时间(TZ=Asia/Shanghai)
- 评论达上限后继续执行点赞/收藏/转发
- 用户信息同步改用 PATCH /v2/users/current
- 一键登出全部功能
- 一键登出全部前端按钮
- update.sh 一键更新脚本
2026-03-31 10:29:26 +08:00

73 lines
2.2 KiB
Python
Executable File
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.
"""数据库连接管理"""
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