375 lines
14 KiB
Python
375 lines
14 KiB
Python
|
"""数据库配置服务"""
|
|||
|
|
from loguru import logger
|
|||
|
|
from typing import List, Dict, Any, Optional
|
|||
|
|
from sqlalchemy.orm import Session
|
|||
|
|
from sqlalchemy import select
|
|||
|
|
from cryptography.fernet import Fernet
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
from ..models.database_config import DatabaseConfig
|
|||
|
|
from utils.util_exceptions import ValidationError, NotFoundError
|
|||
|
|
from .postgresql_tool_manager import get_postgresql_tool
|
|||
|
|
from .mysql_tool_manager import get_mysql_tool
|
|||
|
|
|
|||
|
|
class DatabaseConfigService:
|
|||
|
|
"""数据库配置管理服务"""
|
|||
|
|
|
|||
|
|
def __init__(self, db_session: Session):
|
|||
|
|
self.session = db_session
|
|||
|
|
self.postgresql_tool = get_postgresql_tool()
|
|||
|
|
self.mysql_tool = get_mysql_tool()
|
|||
|
|
# 初始化加密密钥
|
|||
|
|
self.encryption_key = self._get_or_create_encryption_key()
|
|||
|
|
self.cipher = Fernet(self.encryption_key)
|
|||
|
|
def _get_or_create_encryption_key(self) -> bytes:
|
|||
|
|
"""获取或创建加密密钥"""
|
|||
|
|
key_file = "db/db_config_key.key"
|
|||
|
|
if os.path.exists(key_file):
|
|||
|
|
print('find db_config_key')
|
|||
|
|
with open(key_file, 'rb') as f:
|
|||
|
|
return f.read()
|
|||
|
|
|
|||
|
|
else:
|
|||
|
|
print('not find db_config_key')
|
|||
|
|
key = Fernet.generate_key()
|
|||
|
|
with open(key_file, 'wb') as f:
|
|||
|
|
f.write(key)
|
|||
|
|
return key
|
|||
|
|
|
|||
|
|
def _encrypt_password(self, password: str) -> str:
|
|||
|
|
"""加密密码"""
|
|||
|
|
return self.cipher.encrypt(password.encode()).decode()
|
|||
|
|
|
|||
|
|
def _decrypt_password(self, encrypted_password: str) -> str:
|
|||
|
|
"""解密密码"""
|
|||
|
|
return self.cipher.decrypt(encrypted_password.encode()).decode()
|
|||
|
|
|
|||
|
|
async def create_config(self, user_id: int, config_data: Dict[str, Any]) -> DatabaseConfig:
|
|||
|
|
"""创建数据库配置"""
|
|||
|
|
try:
|
|||
|
|
# 验证配置
|
|||
|
|
required_fields = ['name', 'db_type', 'host', 'port', 'database', 'username', 'password']
|
|||
|
|
for field in required_fields:
|
|||
|
|
if field not in config_data:
|
|||
|
|
raise ValidationError(f"缺少必需字段: {field}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 测试连接
|
|||
|
|
test_config = {
|
|||
|
|
'host': config_data['host'],
|
|||
|
|
'port': config_data['port'],
|
|||
|
|
'database': config_data['database'],
|
|||
|
|
'username': config_data['username'],
|
|||
|
|
'password': config_data['password']
|
|||
|
|
}
|
|||
|
|
if 'postgresql' == config_data['db_type']:
|
|||
|
|
test_result = await self.postgresql_tool.execute(
|
|||
|
|
operation="test_connection",
|
|||
|
|
connection_config=test_config
|
|||
|
|
)
|
|||
|
|
if not test_result.success:
|
|||
|
|
raise ValidationError(f"数据库连接测试失败: {test_result.error}")
|
|||
|
|
elif 'mysql' == config_data['db_type']:
|
|||
|
|
test_result = await self.mysql_tool.execute(
|
|||
|
|
operation="test_connection",
|
|||
|
|
connection_config=test_config
|
|||
|
|
)
|
|||
|
|
if not test_result.success:
|
|||
|
|
raise ValidationError(f"数据库连接测试失败: {test_result.error}")
|
|||
|
|
# 如果设置为默认配置，先取消其他默认配置
|
|||
|
|
if config_data.get('is_default', False):
|
|||
|
|
stmt = select(DatabaseConfig).where(
|
|||
|
|
DatabaseConfig.created_by == user_id,
|
|||
|
|
DatabaseConfig.is_default == True
|
|||
|
|
)
|
|||
|
|
result = self.session.execute(stmt)
|
|||
|
|
for config in result.scalars():
|
|||
|
|
config.is_default = False
|
|||
|
|
|
|||
|
|
# 创建配置
|
|||
|
|
db_config = DatabaseConfig(
|
|||
|
|
created_by=user_id,
|
|||
|
|
name=config_data['name'],
|
|||
|
|
db_type=config_data['db_type'],
|
|||
|
|
host=config_data['host'],
|
|||
|
|
port=config_data['port'],
|
|||
|
|
database=config_data['database'],
|
|||
|
|
username=config_data['username'],
|
|||
|
|
password=self._encrypt_password(config_data['password']),
|
|||
|
|
is_active=config_data.get('is_active', True),
|
|||
|
|
is_default=config_data.get('is_default', False),
|
|||
|
|
connection_params=config_data.get('connection_params')
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
self.session.add(db_config)
|
|||
|
|
await self.session.commit()
|
|||
|
|
await self.session.refresh(db_config)
|
|||
|
|
|
|||
|
|
logger.info(f"创建数据库配置成功: {db_config.name} (ID: {db_config.id})")
|
|||
|
|
return db_config
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
await self.session.rollback()
|
|||
|
|
logger.error(f"创建数据库配置失败: {str(e)}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
async def get_user_configs(self, user_id: int, active_only: bool = True) -> List[DatabaseConfig]:
|
|||
|
|
"""获取用户的数据库配置列表"""
|
|||
|
|
stmt = select(DatabaseConfig).where(DatabaseConfig.created_by == user_id)
|
|||
|
|
if active_only:
|
|||
|
|
stmt = stmt.where(DatabaseConfig.is_active == True)
|
|||
|
|
stmt = stmt.order_by(DatabaseConfig.created_at.desc())
|
|||
|
|
return (await self.session.execute(stmt)).scalars().all()
|
|||
|
|
|
|||
|
|
async def get_config_by_id(self, config_id: int, user_id: int) -> Optional[DatabaseConfig]:
|
|||
|
|
"""根据ID获取配置"""
|
|||
|
|
stmt = select(DatabaseConfig).where(
|
|||
|
|
DatabaseConfig.id == config_id,
|
|||
|
|
DatabaseConfig.created_by == user_id
|
|||
|
|
)
|
|||
|
|
return (await self.session.execute(stmt)).scalar_one_or_none()
|
|||
|
|
|
|||
|
|
async def get_default_config(self, user_id: int) -> Optional[DatabaseConfig]:
|
|||
|
|
"""获取用户的默认配置"""
|
|||
|
|
stmt = select(DatabaseConfig).where(
|
|||
|
|
DatabaseConfig.created_by == user_id,
|
|||
|
|
# DatabaseConfig.is_default == True,
|
|||
|
|
DatabaseConfig.is_active == True
|
|||
|
|
)
|
|||
|
|
return (await self.session.execute(stmt)).scalar_one_or_none()
|
|||
|
|
|
|||
|
|
async def test_connection(self, config_id: int, user_id: int) -> Dict[str, Any]:
|
|||
|
|
"""测试数据库连接"""
|
|||
|
|
config = self.get_config_by_id(config_id, user_id)
|
|||
|
|
if not config:
|
|||
|
|
raise NotFoundError("数据库配置不存在")
|
|||
|
|
|
|||
|
|
test_config = {
|
|||
|
|
'host': config.host,
|
|||
|
|
'port': config.port,
|
|||
|
|
'database': config.database,
|
|||
|
|
'username': config.username,
|
|||
|
|
'password': self._decrypt_password(config.password)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
result = await self.postgresql_tool.execute(
|
|||
|
|
operation="test_connection",
|
|||
|
|
connection_config=test_config
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
'success': result.success,
|
|||
|
|
'message': result.result.get('message') if result.success else result.error,
|
|||
|
|
'details': result.result if result.success else None
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
async def connect_and_get_tables(self, config_id: int, user_id: int) -> Dict[str, Any]:
|
|||
|
|
"""连接数据库并获取表列表"""
|
|||
|
|
config = self.get_config_by_id(config_id, user_id)
|
|||
|
|
if not config:
|
|||
|
|
raise NotFoundError("数据库配置不存在")
|
|||
|
|
|
|||
|
|
connection_config = {
|
|||
|
|
'host': config.host,
|
|||
|
|
'port': config.port,
|
|||
|
|
'database': config.database,
|
|||
|
|
'username': config.username,
|
|||
|
|
'password': self._decrypt_password(config.password)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if 'postgresql' == config.db_type:
|
|||
|
|
# 连接数据库
|
|||
|
|
connect_result = await self.postgresql_tool.execute(
|
|||
|
|
operation="connect",
|
|||
|
|
connection_config=connection_config,
|
|||
|
|
user_id=str(user_id)
|
|||
|
|
)
|
|||
|
|
elif 'mysql' == config.db_type:
|
|||
|
|
# 连接数据库
|
|||
|
|
connect_result = await self.mysql_tool.execute(
|
|||
|
|
operation="connect",
|
|||
|
|
connection_config=connection_config,
|
|||
|
|
user_id=str(user_id)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if not connect_result.success:
|
|||
|
|
return {
|
|||
|
|
'success': False,
|
|||
|
|
'message': connect_result.error
|
|||
|
|
}
|
|||
|
|
# 连接信息已保存到PostgreSQLMCPTool的connections中
|
|||
|
|
return {
|
|||
|
|
'success': True,
|
|||
|
|
'data': connect_result.result,
|
|||
|
|
'config_name': config.name
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
async def get_table_data(self, table_name: str, user_id: int, db_type: str, limit: int = 100) -> Dict[str, Any]:
|
|||
|
|
"""获取表数据预览（复用已建立的连接）"""
|
|||
|
|
try:
|
|||
|
|
user_id_str = str(user_id)
|
|||
|
|
|
|||
|
|
# 根据db_type选择相应的数据库工具
|
|||
|
|
if db_type.lower() == 'postgresql':
|
|||
|
|
db_tool = self.postgresql_tool
|
|||
|
|
elif db_type.lower() == 'mysql':
|
|||
|
|
db_tool = self.mysql_tool
|
|||
|
|
else:
|
|||
|
|
return {
|
|||
|
|
'success': False,
|
|||
|
|
'message': f'不支持的数据库类型: {db_type}'
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 检查是否已有连接
|
|||
|
|
if user_id_str not in db_tool.connections:
|
|||
|
|
return {
|
|||
|
|
'success': False,
|
|||
|
|
'message': '数据库连接已断开，请重新连接数据库'
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 直接使用已建立的连接执行查询
|
|||
|
|
sql_query = f"SELECT * FROM {table_name}"
|
|||
|
|
result = await db_tool.execute(
|
|||
|
|
operation="execute_query",
|
|||
|
|
user_id=user_id_str,
|
|||
|
|
sql_query=sql_query,
|
|||
|
|
limit=limit
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if not result.success:
|
|||
|
|
return {
|
|||
|
|
'success': False,
|
|||
|
|
'message': result.error
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
'success': True,
|
|||
|
|
'data': result.result,
|
|||
|
|
'db_type': db_type
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"获取表数据失败: {str(e)}", exc_info=True)
|
|||
|
|
return {
|
|||
|
|
'success': False,
|
|||
|
|
'message': f'获取表数据失败: {str(e)}'
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def disconnect_database(self, user_id: int) -> Dict[str, Any]:
|
|||
|
|
"""断开数据库连接"""
|
|||
|
|
try:
|
|||
|
|
# 从PostgreSQLMCPTool断开连接
|
|||
|
|
self.postgresql_tool.execute(
|
|||
|
|
operation="disconnect",
|
|||
|
|
user_id=str(user_id)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 从本地连接管理中移除
|
|||
|
|
if user_id in self.user_connections:
|
|||
|
|
del self.user_connections[user_id]
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
'success': True,
|
|||
|
|
'message': '数据库连接已断开'
|
|||
|
|
}
|
|||
|
|
except Exception as e:
|
|||
|
|
return {
|
|||
|
|
'success': False,
|
|||
|
|
'message': f'断开连接失败: {str(e)}'
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
async def get_config_by_type(self, user_id: int, db_type: str) -> Optional[DatabaseConfig]:
|
|||
|
|
"""根据数据库类型获取用户配置"""
|
|||
|
|
stmt = select(DatabaseConfig).where(
|
|||
|
|
DatabaseConfig.created_by == user_id,
|
|||
|
|
DatabaseConfig.db_type == db_type,
|
|||
|
|
DatabaseConfig.is_active == True
|
|||
|
|
)
|
|||
|
|
return await self.session.scalar(stmt)
|
|||
|
|
|
|||
|
|
async def create_or_update_config(self, user_id: int, config_data: Dict[str, Any]) -> DatabaseConfig:
|
|||
|
|
"""创建或更新数据库配置（保证db_type唯一性）"""
|
|||
|
|
try:
|
|||
|
|
# 检查是否已存在该类型的配置
|
|||
|
|
existing_config = self.get_config_by_type(user_id, config_data['db_type'])
|
|||
|
|
|
|||
|
|
if existing_config:
|
|||
|
|
# 更新现有配置
|
|||
|
|
for key, value in config_data.items():
|
|||
|
|
if key == 'password':
|
|||
|
|
setattr(existing_config, key, self._encrypt_password(value))
|
|||
|
|
elif hasattr(existing_config, key):
|
|||
|
|
setattr(existing_config, key, value)
|
|||
|
|
|
|||
|
|
await self.session.commit()
|
|||
|
|
await self.session.refresh(existing_config)
|
|||
|
|
logger.info(f"更新数据库配置成功: {existing_config.name} (ID: {existing_config.id})")
|
|||
|
|
return existing_config
|
|||
|
|
else:
|
|||
|
|
# 创建新配置
|
|||
|
|
return await self.create_config(user_id, config_data)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
await self.session.rollback()
|
|||
|
|
logger.error(f"创建或更新数据库配置失败: {str(e)}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def describe_table(self, table_name: str, user_id: int) -> Dict[str, Any]:
|
|||
|
|
"""获取表结构信息（复用已建立的连接）"""
|
|||
|
|
try:
|
|||
|
|
logger.error(f"未实现的逻辑，暂自编 - describe_table: {table_name}")
|
|||
|
|
user_id_str = str(user_id)
|
|||
|
|
|
|||
|
|
# 获取用户默认数据库配置
|
|||
|
|
default_config = self.get_default_config(user_id)
|
|||
|
|
if not default_config:
|
|||
|
|
return {
|
|||
|
|
'success': False,
|
|||
|
|
'message': '未找到默认数据库配置'
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 根据db_type选择相应的数据库工具
|
|||
|
|
if default_config.db_type.lower() == 'postgresql':
|
|||
|
|
db_tool = self.postgresql_tool
|
|||
|
|
elif default_config.db_type.lower() == 'mysql':
|
|||
|
|
db_tool = self.mysql_tool
|
|||
|
|
else:
|
|||
|
|
return {
|
|||
|
|
'success': False,
|
|||
|
|
'message': f'不支持的数据库类型: {default_config.db_type}'
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 检查是否已有连接
|
|||
|
|
if user_id_str not in db_tool.connections:
|
|||
|
|
return {
|
|||
|
|
'success': False,
|
|||
|
|
'message': '数据库连接已断开，请重新连接数据库'
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 使用已建立的连接执行describe_table操作
|
|||
|
|
result = await db_tool.execute(
|
|||
|
|
operation="describe_table",
|
|||
|
|
user_id=user_id_str,
|
|||
|
|
table_name=table_name
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if not result.success:
|
|||
|
|
return {
|
|||
|
|
'success': False,
|
|||
|
|
'message': result.error
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
'success': True,
|
|||
|
|
'data': result.result,
|
|||
|
|
'db_type': default_config.db_type
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"获取表结构失败: {str(e)}", exc_info=True)
|
|||
|
|
return {
|
|||
|
|
'success': False,
|
|||
|
|
'message': f'获取表结构失败: {str(e)}'
|
|||
|
|
}
|
|||
|
|
|