509 lines
22 KiB
Python
509 lines
22 KiB
Python
from sqlalchemy import create_engine, text, MetaData, Table, Column, inspect
|
||
from sqlalchemy.orm import sessionmaker
|
||
from sqlalchemy.exc import SQLAlchemyError
|
||
from typing import Dict, List, Any, Optional
|
||
import logging
|
||
import oracledb
|
||
from urllib.parse import quote_plus
|
||
|
||
# 配置日志
|
||
logging.basicConfig(level=logging.INFO)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
class DatabaseManager:
|
||
"""数据库管理类，支持多种数据库类型的连接和操作"""
|
||
|
||
def __init__(self):
|
||
self.engines = {}
|
||
self.sessions = {}
|
||
|
||
def create_connection(self, db_type: str, host: str, port: int,
|
||
username: str, password: str, database: str = None, **kwargs) -> str:
|
||
"""创建数据库连接
|
||
|
||
Args:
|
||
db_type: 数据库类型 (mysql, oracle, sqlserver, postgresql)
|
||
host: 数据库主机地址
|
||
port: 数据库端口
|
||
username: 用户名
|
||
password: 密码
|
||
database: 数据库名称
|
||
|
||
Returns:
|
||
connection_id: 连接ID
|
||
"""
|
||
try:
|
||
connection_url = self._build_connection_url(db_type, host, port, username, password, database, **kwargs)
|
||
connection_id = f"{db_type}_{host}_{port}_{database or 'default'}"
|
||
|
||
# 配置引擎参数
|
||
engine_kwargs = {"echo": False, "pool_pre_ping": True}
|
||
if db_type.lower() == 'oracle':
|
||
# Oracle连接池配置
|
||
engine_kwargs.update({
|
||
"pool_size": 5,
|
||
"max_overflow": 10,
|
||
"pool_timeout": 30,
|
||
"pool_recycle": 3600,
|
||
"pool_reset_on_return": "commit"
|
||
})
|
||
|
||
engine = create_engine(connection_url, **engine_kwargs)
|
||
|
||
# 测试连接
|
||
with engine.connect() as conn:
|
||
if db_type.lower() == 'oracle':
|
||
conn.execute(text("SELECT 1 FROM DUAL"))
|
||
else:
|
||
conn.execute(text("SELECT 1"))
|
||
|
||
self.engines[connection_id] = engine
|
||
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||
self.sessions[connection_id] = SessionLocal
|
||
|
||
logger.info(f"成功创建数据库连接: {connection_id}")
|
||
return connection_id
|
||
|
||
except Exception as e:
|
||
logger.error(f"创建数据库连接失败: {str(e)}")
|
||
raise Exception(f"数据库连接失败: {str(e)}")
|
||
|
||
def test_connection(self, db_type: str, host: str, port: int,
|
||
username: str, password: str, database: str = None, **kwargs) -> Dict[str, Any]:
|
||
"""测试数据库是否可连通
|
||
|
||
Args:
|
||
db_type: 数据库类型 (mysql, oracle, sqlserver, postgresql)
|
||
host: 数据库主机地址
|
||
port: 数据库端口
|
||
username: 用户名
|
||
password: 密码
|
||
database: 数据库名称
|
||
|
||
Returns:
|
||
Dict[str, Any]: 测试结果信息，包含连接是否成功与版本信息
|
||
"""
|
||
try:
|
||
# 构建连接URL（内部会对Oracle进行预测试）
|
||
connection_url = self._build_connection_url(db_type, host, port, username, password, database, **kwargs)
|
||
engine = create_engine(connection_url, echo=False, pool_pre_ping=True)
|
||
|
||
server_version = None
|
||
with engine.connect() as conn:
|
||
# 基本连通性测试
|
||
if db_type.lower() == 'oracle':
|
||
conn.execute(text("SELECT 1 FROM DUAL"))
|
||
version_sql = "SELECT BANNER FROM V$VERSION WHERE ROWNUM = 1"
|
||
server_version = conn.execute(text(version_sql)).scalar()
|
||
elif db_type.lower() == 'mysql':
|
||
conn.execute(text("SELECT 1"))
|
||
server_version = conn.execute(text("SELECT VERSION()")).scalar()
|
||
elif db_type.lower() == 'postgresql':
|
||
conn.execute(text("SELECT 1"))
|
||
server_version = conn.execute(text("SELECT version()")).scalar()
|
||
elif db_type.lower() == 'sqlserver':
|
||
conn.execute(text("SELECT 1"))
|
||
server_version = conn.execute(text("SELECT @@VERSION")).scalar()
|
||
else:
|
||
raise ValueError(f"不支持的数据库类型: {db_type}")
|
||
|
||
# 释放临时引擎
|
||
engine.dispose()
|
||
|
||
return {
|
||
"ok": True,
|
||
"db_type": db_type,
|
||
"connection_url": str(connection_url),
|
||
"server_version": server_version
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"连接测试失败: {str(e)}")
|
||
return {
|
||
"ok": False,
|
||
"db_type": db_type,
|
||
"error": str(e)
|
||
}
|
||
|
||
def _build_connection_url(self, db_type: str, host: str, port: int,
|
||
username: str, password: str, database: str = None, **kwargs) -> str:
|
||
"""构建数据库连接URL"""
|
||
# 对用户名和密码进行URL编码，防止特殊字符导致解析错误
|
||
encoded_username = quote_plus(username)
|
||
encoded_password = quote_plus(password)
|
||
|
||
if db_type.lower() == 'mysql':
|
||
db_part = f"/{database}" if database else ""
|
||
return f"mysql+pymysql://{encoded_username}:{encoded_password}@{host}:{port}{db_part}?charset=utf8mb4"
|
||
elif db_type.lower() == 'oracle':
|
||
# Oracle连接格式 - 根据oracledb文档优化
|
||
service_name = database or 'XE'
|
||
|
||
# 先测试Oracle连接以确保参数正确
|
||
try:
|
||
self._test_oracle_connection(
|
||
host=host,
|
||
port=port,
|
||
username=username,
|
||
password=password,
|
||
service_name=service_name,
|
||
**kwargs
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Oracle连接预测试失败: {str(e)}")
|
||
raise e
|
||
|
||
# 使用简化的SQLAlchemy URL格式
|
||
# 根据oracledb文档，SQLAlchemy会自动处理连接参数
|
||
base_url = f"oracle+oracledb://{encoded_username}:{encoded_password}@{host}:{port}/?service_name={service_name}"
|
||
|
||
return base_url
|
||
elif db_type.lower() == 'sqlserver':
|
||
db_part = f"/{database}" if database else ""
|
||
return f"mssql+pymssql://{encoded_username}:{encoded_password}@{host}:{port}{db_part}"
|
||
elif db_type.lower() == 'postgresql':
|
||
db_part = f"/{database}" if database else "/postgres"
|
||
return f"postgresql+psycopg2://{encoded_username}:{encoded_password}@{host}:{port}{db_part}"
|
||
else:
|
||
raise ValueError(f"不支持的数据库类型: {db_type}")
|
||
|
||
def _test_oracle_connection(self, host: str, port: int, username: str, password: str,
|
||
service_name: str = None, **kwargs):
|
||
"""测试Oracle直接连接"""
|
||
service_name = service_name or 'XE'
|
||
|
||
# 尝试多种连接方式
|
||
connection_methods = [
|
||
# 方式1: Easy Connect字符串
|
||
{
|
||
'name': 'Easy Connect',
|
||
'params': {
|
||
'user': username,
|
||
'password': password,
|
||
'dsn': f"{host}:{port}/{service_name}"
|
||
}
|
||
},
|
||
# 方式2: 分离参数
|
||
{
|
||
'name': 'Separate Parameters',
|
||
'params': {
|
||
'user': username,
|
||
'password': password,
|
||
'host': host,
|
||
'port': port,
|
||
'service_name': service_name
|
||
}
|
||
},
|
||
# 方式3: 使用SID
|
||
{
|
||
'name': 'SID Connection',
|
||
'params': {
|
||
'user': username,
|
||
'password': password,
|
||
'dsn': oracledb.makedsn(host, port, sid=service_name)
|
||
}
|
||
}
|
||
]
|
||
|
||
last_error = None
|
||
|
||
for method in connection_methods:
|
||
try:
|
||
logger.info(f"尝试Oracle连接方式: {method['name']}")
|
||
|
||
# 尝试连接
|
||
connection = oracledb.connect(**method['params'])
|
||
|
||
# 测试查询
|
||
cursor = connection.cursor()
|
||
cursor.execute("SELECT 1 FROM DUAL")
|
||
result = cursor.fetchone()
|
||
|
||
# 获取数据库版本信息
|
||
cursor.execute("SELECT BANNER FROM V$VERSION WHERE ROWNUM = 1")
|
||
version = cursor.fetchone()
|
||
|
||
cursor.close()
|
||
connection.close()
|
||
|
||
logger.info(f"Oracle连接成功 ({method['name']}): 查询结果={result}, 版本={version[0] if version else 'Unknown'}")
|
||
return # 连接成功，退出函数
|
||
|
||
except Exception as e:
|
||
last_error = e
|
||
logger.warning(f"Oracle连接方式 '{method['name']}' 失败: {str(e)}")
|
||
continue
|
||
|
||
# 所有连接方式都失败
|
||
error_msg = f"所有Oracle连接方式都失败。最后一个错误: {str(last_error)}"
|
||
logger.error(error_msg)
|
||
raise Exception(error_msg)
|
||
|
||
def get_engine(self, connection_id: str):
|
||
"""获取数据库引擎"""
|
||
if connection_id not in self.engines:
|
||
raise ValueError(f"连接ID不存在: {connection_id}")
|
||
return self.engines[connection_id]
|
||
|
||
def get_session(self, connection_id: str):
|
||
"""获取数据库会话"""
|
||
if connection_id not in self.sessions:
|
||
raise ValueError(f"连接ID不存在: {connection_id}")
|
||
return self.sessions[connection_id]()
|
||
|
||
def execute_query(self, connection_id: str, sql: str, params: Dict = None) -> List[Dict]:
|
||
"""执行查询SQL"""
|
||
try:
|
||
engine = self.get_engine(connection_id)
|
||
with engine.connect() as conn:
|
||
result = conn.execute(text(sql), params or {})
|
||
columns = result.keys()
|
||
rows = result.fetchall()
|
||
return [dict(zip(columns, row)) for row in rows]
|
||
except Exception as e:
|
||
logger.error(f"执行查询失败: {str(e)}")
|
||
raise Exception(f"查询执行失败: {str(e)}")
|
||
|
||
def execute_non_query(self, connection_id: str, sql: str, params: Dict = None) -> int:
|
||
"""执行非查询SQL（INSERT, UPDATE, DELETE）"""
|
||
try:
|
||
engine = self.get_engine(connection_id)
|
||
with engine.connect() as conn:
|
||
with conn.begin():
|
||
result = conn.execute(text(sql), params or {})
|
||
return result.rowcount
|
||
except Exception as e:
|
||
logger.error(f"执行非查询失败: {str(e)}")
|
||
raise Exception(f"非查询执行失败: {str(e)}")
|
||
|
||
def get_database_info(self, connection_id: str) -> Dict:
|
||
"""获取数据库信息"""
|
||
try:
|
||
engine = self.get_engine(connection_id)
|
||
inspector = inspect(engine)
|
||
|
||
# 获取数据库名称
|
||
with engine.connect() as conn:
|
||
db_name_result = conn.execute(text("SELECT DATABASE()" if 'mysql' in str(engine.url)
|
||
else "SELECT CURRENT_DATABASE()" if 'postgresql' in str(engine.url)
|
||
else "SELECT DB_NAME()" if 'mssql' in str(engine.url)
|
||
else "SELECT SYS_CONTEXT('USERENV', 'DB_NAME') FROM DUAL"))
|
||
db_name = db_name_result.scalar()
|
||
|
||
# 获取表列表
|
||
tables = inspector.get_table_names()
|
||
|
||
return {
|
||
"database_name": db_name,
|
||
"tables": tables,
|
||
"table_count": len(tables)
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"获取数据库信息失败: {str(e)}")
|
||
raise Exception(f"获取数据库信息失败: {str(e)}")
|
||
|
||
def get_tables_with_comments(self, connection_id: str) -> List[Dict[str, Any]]:
|
||
"""获取数据库中所有表及其备注信息"""
|
||
try:
|
||
engine = self.get_engine(connection_id)
|
||
url_str = str(engine.url)
|
||
with engine.connect() as conn:
|
||
if 'mysql' in url_str:
|
||
sql = """
|
||
SELECT
|
||
TABLE_NAME AS table_name,
|
||
TABLE_COMMENT AS table_comment
|
||
FROM INFORMATION_SCHEMA.TABLES
|
||
WHERE TABLE_SCHEMA = DATABASE()
|
||
ORDER BY TABLE_NAME
|
||
"""
|
||
return [dict(row) for row in conn.execute(text(sql)).mappings().all()]
|
||
elif 'postgresql' in url_str:
|
||
# 获取当前schema
|
||
schema = conn.execute(text("SELECT current_schema()")).scalar() or 'public'
|
||
sql = """
|
||
SELECT
|
||
c.relname AS table_name,
|
||
obj_description(c.oid) AS table_comment
|
||
FROM pg_class c
|
||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||
WHERE n.nspname = :schema
|
||
AND c.relkind IN ('r','p')
|
||
ORDER BY c.relname
|
||
"""
|
||
return [dict(row) for row in conn.execute(text(sql), {"schema": schema}).mappings().all()]
|
||
elif 'mssql' in url_str:
|
||
sql = """
|
||
SELECT
|
||
t.name AS table_name,
|
||
CAST(ep.value AS NVARCHAR(MAX)) AS table_comment
|
||
FROM sys.tables t
|
||
LEFT JOIN sys.extended_properties ep
|
||
ON ep.major_id = t.object_id
|
||
AND ep.minor_id = 0
|
||
AND ep.name = 'MS_Description'
|
||
ORDER BY t.name
|
||
"""
|
||
return [dict(row) for row in conn.execute(text(sql)).mappings().all()]
|
||
elif 'oracle' in url_str:
|
||
sql = """
|
||
SELECT
|
||
table_name AS table_name,
|
||
comments AS table_comment
|
||
FROM user_tab_comments
|
||
ORDER BY table_name
|
||
"""
|
||
return [dict(row) for row in conn.execute(text(sql)).mappings().all()]
|
||
else:
|
||
raise ValueError("不支持的数据库类型")
|
||
except Exception as e:
|
||
logger.error(f"获取表备注信息失败: {str(e)}")
|
||
raise Exception(f"获取表备注信息失败: {str(e)}")
|
||
|
||
def get_table_info(self, connection_id: str, table_name: str) -> Dict:
|
||
"""获取表信息
|
||
|
||
序列化说明:
|
||
- SQLAlchemy 的列类型对象（如 `VARCHAR`、`NUMBER` 等）不可直接JSON序列化。
|
||
- 本方法将列的 `type` 字段统一转换为字符串表示（例如 `VARCHAR(255)`）。
|
||
- 其他由 Inspector 返回的结构（主键、外键、索引）保持原始可序列化格式。
|
||
"""
|
||
try:
|
||
engine = self.get_engine(connection_id)
|
||
inspector = inspect(engine)
|
||
|
||
# 获取列信息
|
||
raw_columns = inspector.get_columns(table_name)
|
||
columns = []
|
||
for col in raw_columns:
|
||
# 创建可序列化的列字典
|
||
serializable = dict(col)
|
||
col_type = serializable.get('type')
|
||
# 将SQLAlchemy类型对象转换为字符串
|
||
if col_type is not None:
|
||
try:
|
||
serializable['type'] = str(col_type)
|
||
except Exception:
|
||
# 兜底：使用类型名
|
||
serializable['type'] = type(col_type).__name__
|
||
columns.append(serializable)
|
||
|
||
# 获取主键
|
||
primary_keys = inspector.get_pk_constraint(table_name)
|
||
|
||
# 获取外键
|
||
foreign_keys = inspector.get_foreign_keys(table_name)
|
||
|
||
# 获取索引
|
||
indexes = inspector.get_indexes(table_name)
|
||
|
||
return {
|
||
"table_name": table_name,
|
||
"columns": columns,
|
||
"primary_keys": primary_keys,
|
||
"foreign_keys": foreign_keys,
|
||
"indexes": indexes
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"获取表信息失败: {str(e)}")
|
||
raise Exception(f"获取表信息失败: {str(e)}")
|
||
|
||
def get_table_columns(self, connection_id: str, table_name: str) -> List[Dict[str, Any]]:
|
||
"""获取指定表的字段名、类型、备注等信息（兼容多数据库）"""
|
||
try:
|
||
engine = self.get_engine(connection_id)
|
||
url_str = str(engine.url)
|
||
with engine.connect() as conn:
|
||
if 'mysql' in url_str:
|
||
sql = """
|
||
SELECT
|
||
COLUMN_NAME AS column_name,
|
||
DATA_TYPE AS data_type,
|
||
IS_NULLABLE AS is_nullable,
|
||
COLUMN_DEFAULT AS column_default,
|
||
COLUMN_COMMENT AS column_comment,
|
||
CHARACTER_MAXIMUM_LENGTH AS max_length,
|
||
NUMERIC_PRECISION AS numeric_precision,
|
||
NUMERIC_SCALE AS numeric_scale
|
||
FROM INFORMATION_SCHEMA.COLUMNS
|
||
WHERE TABLE_SCHEMA = DATABASE()
|
||
AND TABLE_NAME = :table_name
|
||
ORDER BY ORDINAL_POSITION
|
||
"""
|
||
return [dict(row) for row in conn.execute(text(sql), {"table_name": table_name}).mappings().all()]
|
||
elif 'postgresql' in url_str:
|
||
sql = """
|
||
SELECT
|
||
a.attname AS column_name,
|
||
pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
|
||
NOT a.attnotnull AS is_nullable,
|
||
pg_get_expr(d.adbin, d.adrelid) AS column_default,
|
||
col_description(a.attrelid, a.attnum) AS column_comment
|
||
FROM pg_attribute a
|
||
JOIN pg_class c ON a.attrelid = c.oid
|
||
LEFT JOIN pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum
|
||
WHERE c.relname = :table_name
|
||
AND a.attnum > 0
|
||
AND NOT a.attisdropped
|
||
ORDER BY a.attnum
|
||
"""
|
||
return [dict(row) for row in conn.execute(text(sql), {"table_name": table_name}).mappings().all()]
|
||
elif 'mssql' in url_str:
|
||
sql = """
|
||
SELECT
|
||
c.name AS column_name,
|
||
t.name AS data_type,
|
||
CASE WHEN c.is_nullable = 1 THEN 'YES' ELSE 'NO' END AS is_nullable,
|
||
OBJECT_DEFINITION(c.default_object_id) AS column_default,
|
||
CAST(ep.value AS NVARCHAR(MAX)) AS column_comment,
|
||
COLUMNPROPERTY(c.object_id, c.name, 'charmaxlen') AS max_length
|
||
FROM sys.columns c
|
||
LEFT JOIN sys.types t ON c.user_type_id = t.user_type_id
|
||
LEFT JOIN sys.extended_properties ep
|
||
ON ep.major_id = c.object_id
|
||
AND ep.minor_id = c.column_id
|
||
AND ep.name = 'MS_Description'
|
||
WHERE c.object_id = OBJECT_ID(:table_name)
|
||
ORDER BY c.column_id
|
||
"""
|
||
return [dict(row) for row in conn.execute(text(sql), {"table_name": table_name}).mappings().all()]
|
||
elif 'oracle' in url_str:
|
||
sql = """
|
||
SELECT
|
||
cols.column_name AS column_name,
|
||
cols.data_type AS data_type,
|
||
cols.nullable AS is_nullable,
|
||
cols.data_default AS column_default,
|
||
comm.comments AS column_comment,
|
||
cols.data_length AS max_length
|
||
FROM user_tab_columns cols
|
||
LEFT JOIN user_col_comments comm
|
||
ON comm.table_name = cols.table_name
|
||
AND comm.column_name = cols.column_name
|
||
WHERE cols.table_name = UPPER(:table_name)
|
||
ORDER BY cols.column_id
|
||
"""
|
||
return [dict(row) for row in conn.execute(text(sql), {"table_name": table_name}).mappings().all()]
|
||
else:
|
||
raise ValueError("不支持的数据库类型")
|
||
except Exception as e:
|
||
logger.error(f"获取字段信息失败: {str(e)}")
|
||
raise Exception(f"获取字段信息失败: {str(e)}")
|
||
|
||
def close_connection(self, connection_id: str):
|
||
"""关闭数据库连接"""
|
||
try:
|
||
if connection_id in self.engines:
|
||
self.engines[connection_id].dispose()
|
||
del self.engines[connection_id]
|
||
del self.sessions[connection_id]
|
||
logger.info(f"已关闭数据库连接: {connection_id}")
|
||
except Exception as e:
|
||
logger.error(f"关闭连接失败: {str(e)}")
|
||
|
||
def list_connections(self) -> List[str]:
|
||
"""列出所有活动连接"""
|
||
return list(self.engines.keys())
|
||
|
||
# 全局数据库管理器实例
|
||
db_manager = DatabaseManager()
|