forked from xiaohei/taiji-AI-PAD
491 lines
16 KiB
Python
491 lines
16 KiB
Python
"""
|
||
数据库配置和连接管理
|
||
"""
|
||
|
||
import asyncio
|
||
import ssl
|
||
from typing import AsyncGenerator
|
||
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
|
||
from sqlalchemy.orm import sessionmaker
|
||
from sqlalchemy import text
|
||
from sqlalchemy.engine import make_url
|
||
import logging
|
||
from urllib.parse import urlparse, parse_qs, urlencode, urlunparse
|
||
|
||
from config import settings
|
||
from models import Base
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def prepare_database_url(url: str) -> str:
|
||
"""
|
||
处理数据库 URL，移除 asyncpg 不支持的参数（如 sslmode）
|
||
"""
|
||
if not url or "asyncpg" not in url:
|
||
return url
|
||
|
||
parsed = urlparse(url)
|
||
query_params = parse_qs(parsed.query)
|
||
|
||
# 移除 sslmode 参数（asyncpg 不支持，需要用 ssl connect_args 替代）
|
||
if "sslmode" in query_params:
|
||
del query_params["sslmode"]
|
||
|
||
# 重新构建查询字符串
|
||
new_query = urlencode(query_params, doseq=True)
|
||
new_parsed = parsed._replace(query=new_query)
|
||
|
||
return urlunparse(new_parsed)
|
||
|
||
|
||
# 处理数据库 URL
|
||
database_url = prepare_database_url(settings.database_url)
|
||
database_url_obj = make_url(database_url)
|
||
|
||
engine_kwargs = {
|
||
"echo": settings.debug,
|
||
"pool_pre_ping": True,
|
||
}
|
||
|
||
if database_url_obj.get_backend_name().startswith("sqlite"):
|
||
engine_kwargs["connect_args"] = {"check_same_thread": False}
|
||
else:
|
||
engine_kwargs.update({
|
||
"pool_size": 5,
|
||
"max_overflow": 10,
|
||
"pool_recycle": 1800,
|
||
"pool_timeout": 30,
|
||
})
|
||
|
||
# Azure Database for PostgreSQL 或任何包含 sslmode 的连接都需要 TLS
|
||
if (database_url_obj.host and database_url_obj.host.endswith("postgres.database.azure.com")) or \
|
||
"sslmode" in settings.database_url:
|
||
# Azure Database for PostgreSQL requires TLS; provide a default SSL context.
|
||
ssl_context = ssl.create_default_context()
|
||
existing_connect_args = engine_kwargs.get("connect_args") or {}
|
||
existing_connect_args["ssl"] = ssl_context
|
||
engine_kwargs["connect_args"] = existing_connect_args
|
||
|
||
# 创建异步数据库引擎
|
||
engine = create_async_engine(database_url, **engine_kwargs)
|
||
|
||
# 判断当前是否使用SQLite后端
|
||
def is_sqlite_backend() -> bool:
|
||
backend = engine.url.get_backend_name()
|
||
return backend.startswith("sqlite")
|
||
|
||
# 创建异步会话工厂
|
||
AsyncSessionLocal = async_sessionmaker(
|
||
engine,
|
||
class_=AsyncSession,
|
||
expire_on_commit=False
|
||
)
|
||
|
||
|
||
async def get_db() -> AsyncGenerator[AsyncSession, None]:
|
||
"""获取数据库会话的依赖注入函数"""
|
||
async with AsyncSessionLocal() as session:
|
||
try:
|
||
yield session
|
||
except Exception as e:
|
||
# HTTPException是正常的业务异常，不应该被当作数据库错误
|
||
from fastapi import HTTPException
|
||
if isinstance(e, HTTPException):
|
||
raise
|
||
logger.error(f"数据库会话错误: {e}")
|
||
await session.rollback()
|
||
raise
|
||
finally:
|
||
await session.close()
|
||
|
||
|
||
async def init_db():
|
||
"""初始化数据库"""
|
||
try:
|
||
# 创建所有表
|
||
async with engine.begin() as conn:
|
||
await conn.run_sync(Base.metadata.create_all)
|
||
|
||
logger.info("数据库表创建成功")
|
||
|
||
# 尝试创建初始数据，失败不阻止启动
|
||
try:
|
||
await create_initial_data()
|
||
except Exception as init_err:
|
||
logger.warning(f"创建初始数据失败（服务将继续启动）: {init_err}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"数据库初始化失败: {e}")
|
||
raise
|
||
|
||
|
||
async def create_initial_data():
|
||
"""创建初始数据"""
|
||
try:
|
||
async with AsyncSessionLocal() as session:
|
||
# 检查是否已有数据
|
||
result = await session.execute(text("SELECT COUNT(*) FROM users"))
|
||
user_count = result.scalar()
|
||
|
||
if user_count == 0:
|
||
# 创建默认管理员用户
|
||
from models import User
|
||
from passlib.context import CryptContext
|
||
|
||
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
|
||
|
||
admin_user = User(
|
||
name="系统管理员", # 必填字段
|
||
username="admin",
|
||
email="admin@taiji-ai.com",
|
||
password_hash=pwd_context.hash("admin123"), # 必填字段
|
||
hashed_password=pwd_context.hash("admin123"),
|
||
full_name="系统管理员",
|
||
role="super_admin", # 设置为超级管理员
|
||
is_active=True,
|
||
is_admin=True,
|
||
status="active",
|
||
)
|
||
|
||
session.add(admin_user)
|
||
await session.commit()
|
||
|
||
logger.info("默认管理员用户创建成功")
|
||
|
||
# 创建示例工具
|
||
await create_sample_tools(session)
|
||
|
||
except Exception as e:
|
||
logger.error(f"创建初始数据失败: {e}")
|
||
raise
|
||
|
||
|
||
async def create_sample_tools(session: AsyncSession):
|
||
"""创建示例工具"""
|
||
try:
|
||
from models import Tool
|
||
|
||
# 检查是否已有工具
|
||
result = await session.execute(text("SELECT COUNT(*) FROM tools"))
|
||
tool_count = result.scalar()
|
||
|
||
if tool_count == 0:
|
||
# 创建示例工具
|
||
sample_tools = [
|
||
{
|
||
"name": "web_search",
|
||
"description": "网络搜索工具",
|
||
"category": "api",
|
||
"schema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"query": {
|
||
"type": "string",
|
||
"description": "搜索查询"
|
||
},
|
||
"limit": {
|
||
"type": "integer",
|
||
"description": "结果数量限制",
|
||
"default": 10
|
||
}
|
||
},
|
||
"required": ["query"]
|
||
},
|
||
"endpoint": "https://api.example.com/search",
|
||
"method": "POST",
|
||
"auth_type": "api_key",
|
||
"rate_limit": 100,
|
||
"cost_per_call": 0.01,
|
||
"is_public": True
|
||
},
|
||
{
|
||
"name": "text_completion",
|
||
"description": "文本补全工具",
|
||
"category": "llm",
|
||
"schema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"prompt": {
|
||
"type": "string",
|
||
"description": "输入提示"
|
||
},
|
||
"max_tokens": {
|
||
"type": "integer",
|
||
"description": "最大token数",
|
||
"default": 150
|
||
},
|
||
"temperature": {
|
||
"type": "number",
|
||
"description": "温度参数",
|
||
"default": 0.7
|
||
}
|
||
},
|
||
"required": ["prompt"]
|
||
},
|
||
"rate_limit": 60,
|
||
"cost_per_call": 0.05,
|
||
"is_public": True
|
||
},
|
||
{
|
||
"name": "weather_api",
|
||
"description": "天气查询API",
|
||
"category": "api",
|
||
"schema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"city": {
|
||
"type": "string",
|
||
"description": "城市名称"
|
||
},
|
||
"unit": {
|
||
"type": "string",
|
||
"enum": ["celsius", "fahrenheit"],
|
||
"default": "celsius"
|
||
}
|
||
},
|
||
"required": ["city"]
|
||
},
|
||
"endpoint": "https://api.openweathermap.org/data/2.5/weather",
|
||
"method": "GET",
|
||
"auth_type": "api_key",
|
||
"rate_limit": 1000,
|
||
"cost_per_call": 0.001,
|
||
"is_public": True
|
||
}
|
||
]
|
||
|
||
for tool_data in sample_tools:
|
||
tool = Tool(**tool_data)
|
||
session.add(tool)
|
||
|
||
await session.commit()
|
||
logger.info(f"创建了 {len(sample_tools)} 个示例工具")
|
||
|
||
except Exception as e:
|
||
logger.error(f"创建示例工具失败: {e}")
|
||
raise
|
||
|
||
|
||
async def check_db_connection():
|
||
"""检查数据库连接"""
|
||
try:
|
||
async with AsyncSessionLocal() as session:
|
||
await session.execute(text("SELECT 1"))
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"数据库连接检查失败: {e}")
|
||
return False
|
||
|
||
|
||
async def get_db_stats():
|
||
"""获取数据库统计信息"""
|
||
try:
|
||
async with AsyncSessionLocal() as session:
|
||
stats = {}
|
||
|
||
# 获取各表的记录数
|
||
tables = ["users", "agents", "tools", "sessions", "executions", "billing"]
|
||
|
||
for table in tables:
|
||
result = await session.execute(text(f"SELECT COUNT(*) FROM {table}"))
|
||
stats[table] = result.scalar()
|
||
|
||
return stats
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取数据库统计失败: {e}")
|
||
return {}
|
||
|
||
|
||
async def cleanup_old_records():
|
||
"""清理旧记录"""
|
||
try:
|
||
async with AsyncSessionLocal() as session:
|
||
# 清理超过30天的执行记录
|
||
if is_sqlite_backend():
|
||
result = await session.execute(text("""
|
||
DELETE FROM executions
|
||
WHERE created_at < datetime('now', '-30 days')
|
||
"""))
|
||
else:
|
||
result = await session.execute(text("""
|
||
DELETE FROM executions
|
||
WHERE created_at < NOW() - INTERVAL '30 days'
|
||
"""))
|
||
|
||
deleted_executions = result.rowcount
|
||
|
||
# 清理超过7天的会话记录
|
||
if is_sqlite_backend():
|
||
result = await session.execute(text("""
|
||
DELETE FROM sessions
|
||
WHERE created_at < datetime('now', '-7 days')
|
||
AND status != 'active'
|
||
"""))
|
||
else:
|
||
result = await session.execute(text("""
|
||
DELETE FROM sessions
|
||
WHERE created_at < NOW() - INTERVAL '7 days'
|
||
AND status != 'active'
|
||
"""))
|
||
|
||
deleted_sessions = result.rowcount
|
||
|
||
await session.commit()
|
||
|
||
logger.info(f"清理完成: 删除了 {deleted_executions} 条执行记录, {deleted_sessions} 条会话记录")
|
||
|
||
return {
|
||
"deleted_executions": deleted_executions,
|
||
"deleted_sessions": deleted_sessions
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"清理旧记录失败: {e}")
|
||
return {}
|
||
|
||
|
||
async def backup_db():
|
||
"""数据库备份"""
|
||
try:
|
||
import subprocess
|
||
from datetime import datetime
|
||
import os
|
||
import shutil
|
||
|
||
# 生成备份文件名
|
||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
backup_dir = "/app/backups"
|
||
os.makedirs(backup_dir, exist_ok=True)
|
||
|
||
if is_sqlite_backend():
|
||
db_path = engine.url.database
|
||
if not db_path:
|
||
raise ValueError("SQLite数据库路径为空，无法备份")
|
||
db_path = os.path.abspath(db_path)
|
||
backup_file = os.path.join(backup_dir, f"mcp_sqlite_backup_{timestamp}.db")
|
||
shutil.copy2(db_path, backup_file)
|
||
logger.info(f"SQLite数据库备份成功: {backup_file}")
|
||
return backup_file
|
||
|
||
backup_file = f"{backup_dir}/taiji_db_backup_{timestamp}.sql"
|
||
# 执行pg_dump命令
|
||
cmd = [
|
||
"pg_dump",
|
||
settings.database_url.replace("postgresql+asyncpg://", "postgresql://"),
|
||
"-f", backup_file
|
||
]
|
||
|
||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||
|
||
if result.returncode == 0:
|
||
logger.info(f"数据库备份成功: {backup_file}")
|
||
return backup_file
|
||
else:
|
||
logger.error(f"数据库备份失败: {result.stderr}")
|
||
return None
|
||
|
||
except Exception as e:
|
||
logger.error(f"数据库备份异常: {e}")
|
||
return None
|
||
|
||
|
||
async def close_db():
|
||
"""关闭数据库连接"""
|
||
try:
|
||
await engine.dispose()
|
||
logger.info("数据库连接已关闭")
|
||
except Exception as e:
|
||
logger.error(f"关闭数据库连接失败: {e}")
|
||
|
||
|
||
# 数据库事件处理
|
||
async def on_startup():
|
||
"""应用启动时的数据库操作"""
|
||
await init_db()
|
||
|
||
|
||
async def on_shutdown():
|
||
"""应用关闭时的数据库操作"""
|
||
await close_db()
|
||
|
||
|
||
# 定期清理任务
|
||
async def periodic_cleanup():
|
||
"""定期清理任务"""
|
||
while True:
|
||
try:
|
||
await asyncio.sleep(3600) # 每小时执行一次
|
||
await cleanup_old_records()
|
||
except Exception as e:
|
||
logger.error(f"定期清理任务异常: {e}")
|
||
|
||
|
||
# 数据库迁移辅助函数
|
||
async def migrate_db():
|
||
"""数据库迁移（简化版本）"""
|
||
try:
|
||
# 这里可以添加数据迁移逻辑
|
||
# 在生产环境中应该使用Alembic进行数据库版本管理
|
||
logger.info("数据库迁移检查完成")
|
||
|
||
except Exception as e:
|
||
logger.error(f"数据库迁移失败: {e}")
|
||
raise
|
||
|
||
|
||
# 性能优化
|
||
async def optimize_db():
|
||
"""数据库性能优化"""
|
||
try:
|
||
async with AsyncSessionLocal() as session:
|
||
# 更新表统计信息
|
||
await session.execute(text("ANALYZE;"))
|
||
|
||
# 重建索引（如果需要）
|
||
# await session.execute(text("REINDEX DATABASE taiji_db;"))
|
||
|
||
await session.commit()
|
||
|
||
logger.info("数据库优化完成")
|
||
|
||
except Exception as e:
|
||
logger.error(f"数据库优化失败: {e}")
|
||
|
||
|
||
# 健康检查
|
||
async def health_check() -> dict:
|
||
"""数据库健康检查"""
|
||
health_info = {
|
||
"database": "unknown",
|
||
"connection_pool": "unknown",
|
||
"stats": {}
|
||
}
|
||
|
||
try:
|
||
# 检查连接
|
||
if await check_db_connection():
|
||
health_info["database"] = "healthy"
|
||
else:
|
||
health_info["database"] = "unhealthy"
|
||
|
||
# 检查连接池状态
|
||
pool = engine.pool
|
||
health_info["connection_pool"] = {
|
||
"size": pool.size(),
|
||
"checked_in": pool.checkedin(),
|
||
"checked_out": pool.checkedout()
|
||
}
|
||
|
||
# 获取统计信息
|
||
health_info["stats"] = await get_db_stats()
|
||
|
||
except Exception as e:
|
||
logger.error(f"数据库健康检查失败: {e}")
|
||
health_info["database"] = "error"
|
||
health_info["error"] = str(e)
|
||
|
||
return health_info
|
||
|