logan
4ac2528f8e
后端审核规则修复: - Redis SEC-001: SET/DEL等常规命令不再误报AUTH要求，改为仅对管理级命令WARNING - Redis SEC-003: 降级为INFO，排除已由CMD-001拦截的FLUSHDB/FLUSHALL - Redis key提取: 排除CONFIG/ACL等管理命令，不再把子命令参数误认为key名 - MySQL SEC-006: 新增GRANT ALL/REVOKE ALL全局权限拦截 - MySQL SEC-007: 新增LOAD DATA/LOAD XML拦截 - MySQL SEC-008: 新增SLEEP()/BENCHMARK()时间盲注拦截 - MySQL SEC-009: 新增INTO OUTFILE/DUMPFILE文件写入拦截 前端执行结果结构化: - TicketDetailPage执行结果从原始JSON改为结构化展示 - 解析JSON展示: 状态徽章+影响行数+耗时+验证状态 - Errors/Warnings逐条渲染: 红色错误框+黄色警告框+具体消息 - 删除废弃的executeStats/renderExecuteStats 其他: 审核规则6分组/慢查询移除/查询数据重命名/AuditLogs修复 测试: 1040 passed, TSC零错误
346 lines
14 KiB
Python
346 lines
14 KiB
Python
"""Database connector — unified interface for MySQL/PostgreSQL/SQLite.
|
|
|
|
Provides real database connections for query execution and connection testing.
|
|
Replaces mock aiomysql calls with actual driver connections.
|
|
"""
|
|
import asyncio
|
|
import time
|
|
import asyncio
|
|
import sqlite3
|
|
from typing import Optional, List, Dict, Any, Tuple
|
|
|
|
from app.models.instance import Instance
|
|
from app.services.sql_analyzer import SQLAnalyzer
|
|
|
|
MAX_ROWS = 10000
|
|
QUERY_TIMEOUT = 30
|
|
DML_TIMEOUT = 60
|
|
|
|
|
|
class ConnectionResult:
|
|
"""Result of a connection test."""
|
|
def __init__(self, success: bool, message: str, version: str = None, latency_ms: int = 0):
|
|
self.success = success
|
|
self.message = message
|
|
self.version = version
|
|
self.latency_ms = latency_ms
|
|
|
|
def to_dict(self):
|
|
return {
|
|
"success": self.success,
|
|
"message": self.message,
|
|
"version": self.version,
|
|
"latency_ms": self.latency_ms,
|
|
}
|
|
|
|
|
|
class QueryResult:
|
|
"""Result of a SQL query execution."""
|
|
def __init__(self, columns: List[str], rows: List[Dict[str, Any]],
|
|
total: int, elapsed_ms: int, truncated: bool = False):
|
|
self.columns = columns
|
|
self.rows = rows
|
|
self.total = total
|
|
self.elapsed_ms = elapsed_ms
|
|
self.truncated = truncated
|
|
|
|
|
|
class DBConnector:
|
|
"""Unified database connector supporting MySQL, PostgreSQL, SQLite."""
|
|
|
|
_analyzer = SQLAnalyzer()
|
|
|
|
@staticmethod
|
|
async def test_connection(instance: Instance) -> ConnectionResult:
|
|
"""Test database connectivity and return version info."""
|
|
start = time.time()
|
|
|
|
try:
|
|
if instance.db_type == "mysql":
|
|
return await DBConnector._test_mysql(instance, start)
|
|
elif instance.db_type in ("postgres", "postgresql"):
|
|
return await DBConnector._test_postgres(instance, start)
|
|
elif instance.db_type == "sqlite":
|
|
return await DBConnector._test_sqlite(instance, start)
|
|
elif instance.db_type == "redis":
|
|
return await DBConnector._test_redis(instance, start)
|
|
elif instance.db_type == "mongodb":
|
|
return await DBConnector._test_mongodb(instance, start)
|
|
else:
|
|
return ConnectionResult(False, f"暂不支持 {instance.db_type} 类型数据源")
|
|
except Exception as e:
|
|
elapsed = int((time.time() - start) * 1000)
|
|
return ConnectionResult(False, f"连接失败: {str(e)}", latency_ms=elapsed)
|
|
|
|
|
|
@staticmethod
|
|
async def _test_redis(instance: 'Instance', start: float) -> ConnectionResult:
|
|
"""Test Redis connectivity using synchronous redis client in thread."""
|
|
try:
|
|
import redis as redis_lib
|
|
|
|
def _connect():
|
|
client = redis_lib.Redis(
|
|
host=instance.host,
|
|
port=instance.port,
|
|
password=instance.password if instance.password and instance.password != "default" else None,
|
|
db=0,
|
|
socket_connect_timeout=5,
|
|
socket_timeout=5,
|
|
)
|
|
info = client.info("server")
|
|
version = info.get("redis_version", "unknown")
|
|
client.close()
|
|
return version
|
|
|
|
version = await asyncio.to_thread(_connect)
|
|
elapsed = int((time.time() - start) * 1000)
|
|
return ConnectionResult(True, f"Redis {version} 连接成功", version=version, latency_ms=elapsed)
|
|
except Exception as e:
|
|
elapsed = int((time.time() - start) * 1000)
|
|
return ConnectionResult(False, f"Redis 连接失败: {str(e)}", latency_ms=elapsed)
|
|
|
|
@staticmethod
|
|
async def _test_mongodb(instance: 'Instance', start: float) -> ConnectionResult:
|
|
"""Test MongoDB connectivity using pymongo in thread."""
|
|
try:
|
|
from pymongo import MongoClient
|
|
|
|
def _connect():
|
|
if instance.username and instance.username != "default":
|
|
uri = f"mongodb://{instance.username}:{instance.password}@{instance.host}:{instance.port}/"
|
|
else:
|
|
uri = f"mongodb://{instance.host}:{instance.port}/"
|
|
client = MongoClient(uri, serverSelectionTimeoutMS=5000, connectTimeoutMS=5000)
|
|
info = client.server_info()
|
|
version = info.get("version", "unknown")
|
|
client.close()
|
|
return version
|
|
|
|
version = await asyncio.to_thread(_connect)
|
|
elapsed = int((time.time() - start) * 1000)
|
|
return ConnectionResult(True, f"MongoDB {version} 连接成功", version=version, latency_ms=elapsed)
|
|
except Exception as e:
|
|
elapsed = int((time.time() - start) * 1000)
|
|
return ConnectionResult(False, f"MongoDB 连接失败: {str(e)}", latency_ms=elapsed)
|
|
|
|
@staticmethod
|
|
async def execute_query(
|
|
instance: Instance, database: Optional[str], sql: str, user_id: int
|
|
) -> 'QueryResult':
|
|
"""Execute a read-only SQL query on the target database."""
|
|
# Validate SQL is read-only
|
|
is_safe, error_msg = DBConnector._analyzer.validate_readonly(sql)
|
|
if not is_safe:
|
|
raise ValueError(error_msg)
|
|
|
|
# Auto-add LIMIT
|
|
upper_sql = sql.strip().upper()
|
|
import re
|
|
if not re.search(r"LIMIT\s+\d+", upper_sql, re.IGNORECASE):
|
|
sql = sql.rstrip(";") + f" LIMIT {MAX_ROWS};"
|
|
|
|
start = time.time()
|
|
|
|
if instance.db_type == "mysql":
|
|
rows, columns = await DBConnector._query_mysql(instance, database, sql)
|
|
elif instance.db_type in ("postgres", "postgresql"):
|
|
rows, columns = await DBConnector._query_postgres(instance, database, sql)
|
|
elif instance.db_type == "sqlite":
|
|
rows, columns = await DBConnector._query_sqlite(instance, sql)
|
|
else:
|
|
raise ValueError(f"暂不支持 {instance.db_type} 查询")
|
|
|
|
elapsed = int((time.time() - start) * 1000)
|
|
return QueryResult(
|
|
columns=columns,
|
|
rows=rows,
|
|
total=len(rows),
|
|
elapsed_ms=elapsed,
|
|
truncated=len(rows) >= MAX_ROWS,
|
|
)
|
|
|
|
@staticmethod
|
|
async def execute_dml(instance: Instance, database: str, sql: str) -> str:
|
|
"""Execute a DML statement (INSERT/UPDATE/DELETE)."""
|
|
if instance.db_type == "mysql":
|
|
return await DBConnector._dml_mysql(instance, database, sql)
|
|
elif instance.db_type in ("postgres", "postgresql"):
|
|
return await DBConnector._dml_postgres(instance, database, sql)
|
|
elif instance.db_type == "sqlite":
|
|
return await DBConnector._dml_sqlite(instance, sql)
|
|
else:
|
|
raise ValueError(f"暂不支持 {instance.db_type} 执行")
|
|
|
|
# ===== MySQL =====
|
|
|
|
@staticmethod
|
|
async def _test_mysql(instance: Instance, start: float) -> ConnectionResult:
|
|
import aiomysql
|
|
conn = await asyncio.wait_for(
|
|
aiomysql.connect(
|
|
host=instance.host, port=instance.port,
|
|
user=instance.username, password=instance.password,
|
|
db=instance.database,
|
|
),
|
|
timeout=10,
|
|
)
|
|
try:
|
|
async with conn.cursor() as cur:
|
|
await cur.execute("SELECT VERSION()")
|
|
version = (await cur.fetchone())[0]
|
|
elapsed = int((time.time() - start) * 1000)
|
|
return ConnectionResult(True, "连接成功", version=version, latency_ms=elapsed)
|
|
finally:
|
|
conn.close()
|
|
|
|
@staticmethod
|
|
async def _query_mysql(instance: Instance, database: Optional[str], sql: str) -> Tuple[List[Dict], List[str]]:
|
|
import aiomysql
|
|
conn = await aiomysql.connect(
|
|
host=instance.host, port=instance.port,
|
|
user=instance.username, password=instance.password,
|
|
db=database or instance.database,
|
|
)
|
|
try:
|
|
async with conn.cursor(aiomysql.DictCursor) as cur:
|
|
await asyncio.wait_for(cur.execute(sql), timeout=QUERY_TIMEOUT)
|
|
rows = await cur.fetchall()
|
|
columns = [desc[0] for desc in cur.description] if cur.description else []
|
|
return rows, columns
|
|
except asyncio.TimeoutError:
|
|
raise ValueError("查询超时（30秒）")
|
|
finally:
|
|
conn.close()
|
|
|
|
@staticmethod
|
|
async def _dml_mysql(instance: Instance, database: str, sql: str) -> str:
|
|
import aiomysql
|
|
conn = await aiomysql.connect(
|
|
host=instance.host, port=instance.port,
|
|
user=instance.username, password=instance.password,
|
|
db=database, autocommit=True,
|
|
)
|
|
try:
|
|
async with conn.cursor() as cur:
|
|
await asyncio.wait_for(cur.execute(sql), timeout=DML_TIMEOUT)
|
|
affected = cur.rowcount
|
|
return f"执行成功，影响 {affected} 行"
|
|
except asyncio.TimeoutError:
|
|
raise ValueError("执行超时（60秒）")
|
|
finally:
|
|
conn.close()
|
|
|
|
# ===== PostgreSQL =====
|
|
|
|
@staticmethod
|
|
async def _test_postgres(instance: Instance, start: float) -> ConnectionResult:
|
|
import asyncpg
|
|
dsn = f"postgresql://{instance.username}:{instance.password}@{instance.host}:{instance.port}/{instance.database or 'postgres'}"
|
|
conn = await asyncio.wait_for(asyncpg.connect(dsn), timeout=10)
|
|
try:
|
|
version = await conn.fetchval("SELECT version()")
|
|
elapsed = int((time.time() - start) * 1000)
|
|
# version string is long, take first part
|
|
ver_short = version.split(",")[0] if version else "unknown"
|
|
return ConnectionResult(True, "连接成功", version=ver_short, latency_ms=elapsed)
|
|
finally:
|
|
await conn.close()
|
|
|
|
@staticmethod
|
|
async def _query_postgres(instance: Instance, database: Optional[str], sql: str) -> Tuple[List[Dict], List[str]]:
|
|
import asyncpg
|
|
dsn = f"postgresql://{instance.username}:{instance.password}@{instance.host}:{instance.port}/{database or instance.database or 'postgres'}"
|
|
conn = await asyncpg.connect(dsn)
|
|
try:
|
|
# asyncpg returns Records, convert to dicts
|
|
stmt = await asyncio.wait_for(conn.fetch(sql), timeout=QUERY_TIMEOUT)
|
|
rows = [dict(r) for r in stmt]
|
|
columns = list(stmt[0].keys()) if stmt else []
|
|
return rows, columns
|
|
except asyncio.TimeoutError:
|
|
raise ValueError("查询超时（30秒）")
|
|
finally:
|
|
await conn.close()
|
|
|
|
@staticmethod
|
|
async def _dml_postgres(instance: Instance, database: str, sql: str) -> str:
|
|
import asyncpg
|
|
dsn = f"postgresql://{instance.username}:{instance.password}@{instance.host}:{instance.port}/{database}"
|
|
conn = await asyncpg.connect(dsn)
|
|
try:
|
|
status = await asyncio.wait_for(conn.execute(sql), timeout=DML_TIMEOUT)
|
|
# status like "UPDATE 3"
|
|
parts = status.split()
|
|
affected = int(parts[1]) if len(parts) >= 2 else 0
|
|
return f"执行成功，影响 {affected} 行"
|
|
except asyncio.TimeoutError:
|
|
raise ValueError("执行超时（60秒）")
|
|
finally:
|
|
await conn.close()
|
|
|
|
# ===== SQLite (remote instance) =====
|
|
|
|
@staticmethod
|
|
async def _test_sqlite(instance: Instance, start: float) -> ConnectionResult:
|
|
# SQLite uses host as file path
|
|
db_path = instance.host
|
|
loop = asyncio.get_event_loop()
|
|
try:
|
|
def _connect():
|
|
conn = sqlite3.connect(db_path)
|
|
cur = conn.execute("SELECT sqlite_version()")
|
|
ver = cur.fetchone()[0]
|
|
conn.close()
|
|
return ver
|
|
version = await loop.run_in_executor(None, _connect)
|
|
elapsed = int((time.time() - start) * 1000)
|
|
return ConnectionResult(True, "连接成功", version=f"SQLite {version}", latency_ms=elapsed)
|
|
except Exception as e:
|
|
elapsed = int((time.time() - start) * 1000)
|
|
return ConnectionResult(False, f"连接失败: {e}", latency_ms=elapsed)
|
|
|
|
@staticmethod
|
|
async def _query_sqlite(instance: Instance, sql: str) -> Tuple[List[Dict], List[str]]:
|
|
db_path = instance.host
|
|
loop = asyncio.get_event_loop()
|
|
|
|
def _query():
|
|
conn = sqlite3.connect(db_path)
|
|
conn.row_factory = sqlite3.Row
|
|
cur = conn.execute(sql)
|
|
rows = [dict(r) for r in cur.fetchall()]
|
|
columns = [desc[0] for desc in cur.description] if cur.description else []
|
|
conn.close()
|
|
return rows, columns
|
|
|
|
try:
|
|
return await asyncio.wait_for(
|
|
loop.run_in_executor(None, _query),
|
|
timeout=QUERY_TIMEOUT,
|
|
)
|
|
except asyncio.TimeoutError:
|
|
raise ValueError("查询超时（30秒）")
|
|
|
|
@staticmethod
|
|
async def _dml_sqlite(instance: Instance, sql: str) -> str:
|
|
db_path = instance.host
|
|
loop = asyncio.get_event_loop()
|
|
|
|
def _execute():
|
|
conn = sqlite3.connect(db_path)
|
|
cur = conn.execute(sql)
|
|
affected = cur.rowcount
|
|
conn.commit()
|
|
conn.close()
|
|
return affected
|
|
|
|
try:
|
|
affected = await asyncio.wait_for(
|
|
loop.run_in_executor(None, _execute),
|
|
timeout=DML_TIMEOUT,
|
|
)
|
|
return f"执行成功，影响 {affected} 行"
|
|
except asyncio.TimeoutError:
|
|
raise ValueError("执行超时（60秒）")
|