Files
opsxsql/backend/app/services/query_service.py
T
logan 0aa4595c90 fix: 查询模板占位符替换 + _auto_add_limit 跳过非SELECT语句 + 默认用户密码修正
修复:
- getTemplateSQL() 支持 {table}/{db}/{collection} 占位符自动替换
- _auto_add_limit() 仅对 SELECT 语句添加 LIMIT, SHOW/DESC/EXPLAIN 跳过
- seed_data 默认用户密码从 '***' 改为 'admin123'
- DB 中已有用户密码已重置为 admin123

测试: 75 passed (新增2个 SHOW/DESC LIMIT 测试)
2026-05-24 15:32:51 +08:00

1262 lines
42 KiB
Python
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.
"""统一多数据库查询服务 — 支持 MySQL、PostgreSQL、Redis、MongoDB 的只读查询。
核心功能:
1. 统一入口: 根据 db_type 自动路由到对应查询处理器
2. 安全验证:
- SQL 数据库: 使用 SQLAnalyzer 验证只读 + 增强注入防护
- Redis: 使用 READONLY_COMMANDS 白名单 + 子命令验证
- MongoDB: 使用 READONLY_OPERATIONS 白名单 + pipeline 阶段验证
3. 结果脱敏: 集成 DataMasker 对敏感字段脱敏
4. 查询限制: 超时控制、行数限制、内存保护
5. 审计日志: 完整的查询日志记录
使用方式:
result = await QueryService.execute_query(instance, database, query, user_id)
"""
import asyncio
import json
import logging
import re
import time
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Tuple
from app.models.instance import Instance
from app.schemas.query import UnifiedQueryResult
from app.services.db_connector import DBConnector
from app.services.redis_connector import RedisConnector
from app.services.sql_analyzer import SQLAnalyzer
logger = logging.getLogger(__name__)
# ===== 配置常量 =====
# 默认最大返回行数
MAX_ROWS = 10000
# 默认查询超时(秒)
QUERY_TIMEOUT = 30
# 最大查询长度(字符)
MAX_QUERY_LENGTH = 10000
# 最大结果集内存限制(MB)
MAX_RESULT_MEMORY_MB = 50
# ===== Redis 只读命令白名单 =====
REDIS_READONLY_COMMANDS = {
# 字符串操作
"GET", "MGET", "STRLEN", "GETRANGE", "GETBIT", "SUBSTR",
# 哈希操作
"HGET", "HGETALL", "HKEYS", "HVALS", "HLEN", "HEXISTS", "HMGET", "HSTRLEN",
# 列表操作
"LRANGE", "LLEN", "LINDEX", "LPOS",
# 集合操作
"SMEMBERS", "SCARD", "SISMEMBER", "SRANDMEMBER", "SUNION", "SINTER", "SDIFF",
# 有序集合操作
"ZRANGE", "ZREVRANGE", "ZRANGEBYSCORE", "ZREVRANGEBYSCORE",
"ZCARD", "ZSCORE", "ZRANK", "ZREVRANK", "ZCOUNT", "ZLEXCOUNT",
"ZRANGEBYLEX", "ZREVRANGEBYLEX",
# 通用操作
"EXISTS", "TTL", "PTTL", "TYPE", "DBSIZE", "KEYS",
"SCAN", "HSCAN", "SSCAN", "ZSCAN",
"INFO", "CONFIG", "MEMORY", "OBJECT", "SLOWLOG", "CLIENT", "CLUSTER",
# 位图/地理位置/HyperLogLog(只读)
"BITCOUNT", "BITPOS", "GEOPOS", "GEODIST", "GEOHASH", "GEORADIUS_RO",
"GEORADIUSBYMEMBER_RO", "PFCOUNT",
# 其他
"PING", "ECHO", "TIME", "LASTSAVE", "ROLE",
}
# 需要特殊处理的子命令: CONFIG GET / MEMORY USAGE / OBJECT ENCODING 等
REDIS_SUBCOMMAND_ALLOWLIST = {
("CONFIG", "GET"),
("CONFIG", "RESETSTAT"),
("MEMORY", "USAGE"),
("MEMORY", "STATS"),
("MEMORY", "DOCTOR"),
("OBJECT", "ENCODING"),
("OBJECT", "IDLETIME"),
("OBJECT", "REFCOUNT"),
("OBJECT", "FREQ"),
("CLIENT", "LIST"),
("CLIENT", "INFO"),
("CLUSTER", "INFO"),
("CLUSTER", "NODES"),
("CLUSTER", "SLOTS"),
("SLOWLOG", "GET"),
("SLOWLOG", "LEN"),
}
# Redis 危险命令黑名单(即使被误分类为只读也禁止)
REDIS_COMMAND_BLACKLIST = {
"FLUSHDB", "FLUSHALL", "SHUTDOWN", "DEBUG", "SAVE", "BGSAVE",
"BGREWRITEAOF", "SLAVEOF", "REPLICAOF", "SYNC", "PSYNC",
"CONFIG", "ACL", "MODULE", "COMMAND",
}
# ===== MongoDB 只读操作白名单 =====
MONGODB_READONLY_OPS = {
"find", "findOne", "countDocuments", "estimatedDocumentCount",
"distinct", "aggregate",
}
MONGODB_WRITE_OPS = {
"insertOne", "insertMany", "updateOne", "updateMany",
"replaceOne", "deleteOne", "deleteMany", "drop",
"createIndex", "dropIndex", "renameCollection",
"findOneAndUpdate", "findOneAndReplace", "findOneAndDelete",
"bulkWrite",
}
# MongoDB aggregate 危险阶段黑名单
MONGODB_AGGREGATE_BLACKLIST = {
"$out", "$merge", "$currentOp", "$listLocalSessions",
}
# ===== SQL 注入防护模式 =====
# 危险的 SQL 注释模式(用于注释掉合法查询尾部)
SQL_COMMENT_PATTERNS = [
r"--\s*$", # 行尾注释
r"/\*.*\*/", # 块注释
r"#\s*$", # MySQL 风格注释
]
# 危险的 UNION 注入模式
SQL_UNION_INJECTION_PATTERN = re.compile(
r"\bUNION\s+(ALL\s+)?SELECT\b", re.IGNORECASE
)
# 堆叠查询检测(多条语句)
SQL_STACKED_PATTERN = re.compile(r"\b(INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|TRUNCATE|EXEC|EXECUTE)\b", re.IGNORECASE)
# 时间延迟注入检测
SQL_TIMING_ATTACK_PATTERN = re.compile(
r"\b(BENCHMARK|SLEEP|PG_SLEEP|WAITFOR\s+DELAY|WAITFOR\s+TIME)\s*\(",
re.IGNORECASE,
)
class QuerySecurityError(ValueError):
"""查询安全验证失败异常。"""
pass
class QueryTimeoutError(ValueError):
"""查询超时异常。"""
pass
class QueryResultTooLargeError(ValueError):
"""查询结果过大异常。"""
pass
# ===== 查询日志记录器 =====
class QueryLogger:
"""查询操作日志记录器。
记录每次查询的完整上下文,用于安全审计和问题追溯。
"""
@staticmethod
def log_query(
user_id: int,
db_type: str,
database: Optional[str],
query: str,
elapsed_ms: int,
row_count: int,
is_success: bool,
error_msg: Optional[str] = None,
client_ip: Optional[str] = None,
) -> None:
"""记录查询日志。
Args:
user_id: 查询用户ID
db_type: 数据库类型
database: 目标数据库
query: 查询内容(已脱敏)
elapsed_ms: 执行耗时
row_count: 返回行数
is_success: 是否成功
error_msg: 错误信息
client_ip: 客户端IP
"""
log_data = {
"event": "query_executed",
"user_id": user_id,
"db_type": db_type,
"database": database,
"query_preview": query[:500] if query else "",
"query_length": len(query) if query else 0,
"elapsed_ms": elapsed_ms,
"row_count": row_count,
"is_success": is_success,
"client_ip": client_ip,
}
if error_msg:
log_data["error"] = error_msg[:500]
if is_success:
logger.info("Query executed: %s", json.dumps(log_data, ensure_ascii=False, default=str))
else:
logger.warning("Query failed: %s", json.dumps(log_data, ensure_ascii=False, default=str))
@staticmethod
def log_security_event(
event_type: str,
user_id: int,
db_type: str,
query: str,
reason: str,
client_ip: Optional[str] = None,
) -> None:
"""记录安全事件日志。
Args:
event_type: 事件类型(injection_attempt, unauthorized_command 等)
user_id: 用户ID
db_type: 数据库类型
query: 查询内容
reason: 阻止原因
client_ip: 客户端IP
"""
log_data = {
"event": "security_event",
"event_type": event_type,
"user_id": user_id,
"db_type": db_type,
"query_preview": query[:500] if query else "",
"reason": reason,
"client_ip": client_ip,
}
logger.warning("Security event: %s", json.dumps(log_data, ensure_ascii=False, default=str))
# ===== 查询处理器基类 =====
class BaseQueryHandler(ABC):
"""查询处理器基类。
所有数据库类型的查询处理器必须继承此基类,
并实现 execute 和 validate_readonly 方法。
"""
@abstractmethod
async def execute(
self,
instance: Instance,
database: Optional[str],
query: str,
user_id: int,
options: Optional[Dict[str, Any]] = None,
) -> UnifiedQueryResult:
"""执行查询并返回统一格式的结果。
Args:
instance: 数据源实例配置
database: 目标数据库名/Redis DB号/MongoDB collection
query: 查询语句(SQL/Redis命令/MongoDB查询JSON)
user_id: 查询用户ID
options: 额外选项(如 limit, offset 等)
Returns:
UnifiedQueryResult: 统一格式的查询结果
"""
pass
@abstractmethod
def validate_readonly(self, query: str) -> Tuple[bool, str]:
"""验证查询是否为只读。
Args:
query: 查询语句
Returns:
(is_safe, error_message): 是否安全及错误信息
"""
pass
def _validate_query_length(self, query: str) -> None:
"""验证查询长度是否在允许范围内。
Args:
query: 查询语句
Raises:
QuerySecurityError: 查询过长
"""
if not query or not query.strip():
raise QuerySecurityError("查询语句不能为空")
if len(query) > MAX_QUERY_LENGTH:
raise QuerySecurityError(
f"查询语句过长({len(query)} 字符),最大允许 {MAX_QUERY_LENGTH} 字符"
)
def _check_result_size(self, rows: List[Dict[str, Any]]) -> bool:
"""检查结果集大小是否超过内存限制。
Args:
rows: 结果行列表
Returns:
是否超过限制
"""
if not rows:
return False
# 估算内存使用(粗略)
try:
size_bytes = len(json.dumps(rows, default=str).encode("utf-8"))
size_mb = size_bytes / (1024 * 1024)
return size_mb > MAX_RESULT_MEMORY_MB
except (OverflowError, MemoryError):
return True
# ===== SQL 查询处理器(MySQL / PostgreSQL 共用基类) =====
class SQLQueryHandlerBase(BaseQueryHandler):
"""SQL 查询处理器基类(MySQL 和 PostgreSQL 共用逻辑)。"""
_analyzer = SQLAnalyzer()
def validate_readonly(self, query: str) -> Tuple[bool, str]:
"""使用 SQLAnalyzer 验证 SQL 是否为只读。
Args:
query: SQL 查询语句
Returns:
(is_safe, error_message)
"""
is_safe, error_msg = self._analyzer.validate_readonly(query)
return is_safe, error_msg
def _enhanced_security_check(self, query: str, user_id: int) -> None:
"""增强安全验证 — 检测 SQL 注入攻击模式。
检测项:
1. 查询长度限制
2. 堆叠查询(多条语句)
3. UNION 注入
4. 时间延迟攻击
5. 注释攻击
6. 危险的子查询
Args:
query: SQL 查询语句
user_id: 查询用户ID
Raises:
QuerySecurityError: 检测到安全问题
"""
# 1. 长度验证
self._validate_query_length(query)
upper_query = query.upper().strip()
# 2. 检测堆叠查询(尝试执行多条语句)
stripped = query.strip().rstrip(";").strip()
in_string = False
string_char = None
for ch in stripped:
if ch in ("'", '"') and not in_string:
in_string = True
string_char = ch
elif ch == string_char and in_string:
# 检查转义
string_char = None
in_string = False
elif ch == ";" and not in_string:
raise QuerySecurityError(
"检测到多条 SQL 语句(分号分隔),可能存在堆叠查询注入风险"
)
# 3. 检测 UNION 注入(只允许以 SELECT 或 WITH 或 EXPLAIN 开头)
allowed_prefixes = ("SELECT", "WITH", "EXPLAIN", "SHOW", "DESCRIBE", "DESC")
first_word = upper_query.split()[0] if upper_query else ""
if first_word not in allowed_prefixes:
# 进一步检查是否是 EXPLAIN SELECT 或 SHOW ...
if not any(upper_query.startswith(p) for p in allowed_prefixes):
raise QuerySecurityError(
f"查询必须以 SELECT/WITH/EXPLAIN/SHOW 开头,当前以 '{first_word}' 开头"
)
# 4. 检测时间延迟攻击
if SQL_TIMING_ATTACK_PATTERN.search(query):
QueryLogger.log_security_event(
event_type="timing_attack_detected",
user_id=user_id,
db_type="sql",
query=query,
reason="检测到时间延迟函数(SLEEP/BENCHMARK等)",
)
raise QuerySecurityError("查询包含时间延迟函数,已被安全系统拦截")
# 5. 检测 INTO OUTFILE/DUMPFILE(文件写入)
if "INTO OUTFILE" in upper_query or "INTO DUMPFILE" in upper_query:
QueryLogger.log_security_event(
event_type="file_write_attempt",
user_id=user_id,
db_type="sql",
query=query,
reason="检测到 INTO OUTFILE/DUMPFILE 文件写入",
)
raise QuerySecurityError("SELECT INTO OUTFILE/DUMPFILE 被禁止")
# 6. 检测危险的存储过程调用
if re.search(r"\b(Exec|Execute|Sp_|Xp_)", query, re.IGNORECASE):
raise QuerySecurityError("存储过程调用在查询模式中被禁止")
# 不需要添加 LIMIT 的语句前缀
_NO_LIMIT_PREFIXES = ("SHOW", "DESC", "DESCRIBE", "EXPLAIN", "WITH")
def _auto_add_limit(self, sql: str, max_rows: int = MAX_ROWS) -> str:
"""如果 SQL 没有 LIMIT 则自动添加。仅对 SELECT 语句生效。
Args:
sql: 原始 SQL
max_rows: 最大返回行数
Returns:
添加了 LIMIT 的 SQL
"""
upper_sql = sql.strip().upper()
# SHOW / DESC / EXPLAIN 等语句不支持 LIMIT(或语义不同),跳过
if any(upper_sql.startswith(p) for p in self._NO_LIMIT_PREFIXES):
return sql
# 仅对 SELECT 语句添加 LIMIT
if not upper_sql.startswith("SELECT"):
return sql
if not self._has_limit_clause(upper_sql):
sql = sql.rstrip(";") + f" LIMIT {max_rows}"
return sql
@staticmethod
def _has_limit_clause(sql_upper: str) -> bool:
"""安全地检测 SQL 是否已包含 LIMIT 子句。
使用简单的词法分析,避免字符串内容的干扰。
Args:
sql_upper: 大写的 SQL 语句
Returns:
是否包含 LIMIT 子句
"""
# 简单但有效的检测:在字符串外查找 LIMIT 关键字
tokens = re.split(r"('[^']*'|'[^']*)", sql_upper)
# 偶数索引的 token 在字符串外
for i in range(0, len(tokens), 2):
if re.search(r"\bLIMIT\s+\d+", tokens[i]):
return True
return False
class MySQLQueryHandler(SQLQueryHandlerBase):
"""MySQL 查询处理器。
使用 DBConnector._query_mysql 执行查询,
SQLAnalyzer 验证只读性,自动添加 LIMIT,超时控制。
"""
async def execute(
self,
instance: Instance,
database: Optional[str],
query: str,
user_id: int,
options: Optional[Dict[str, Any]] = None,
) -> UnifiedQueryResult:
"""执行 MySQL 只读查询。
Args:
instance: MySQL 实例配置
database: 目标数据库名
query: SQL 查询语句
user_id: 查询用户ID
options: 额外选项
Returns:
UnifiedQueryResult: 统一格式的查询结果
Raises:
QuerySecurityError: 查询不安全
QueryTimeoutError: 查询超时
"""
options = options or {}
custom_timeout = options.get("timeout", QUERY_TIMEOUT)
custom_max_rows = options.get("limit", MAX_ROWS)
custom_max_rows = min(custom_max_rows, MAX_ROWS) # 不允许超过全局上限
# 1. 增强安全验证
self._enhanced_security_check(query, user_id)
# 2. 验证只读
is_safe, error_msg = self.validate_readonly(query)
if not is_safe:
raise QuerySecurityError(error_msg)
# 3. 自动添加 LIMIT
query = self._auto_add_limit(query, custom_max_rows)
# 4. 执行查询
start = time.time()
rows: List[Dict[str, Any]] = []
columns: List[str] = []
try:
rows, columns = await asyncio.wait_for(
DBConnector._query_mysql(instance, database, query),
timeout=custom_timeout,
)
except asyncio.TimeoutError:
QueryLogger.log_query(
user_id=user_id,
db_type="mysql",
database=database,
query=query,
elapsed_ms=int((time.time() - start) * 1000),
row_count=0,
is_success=False,
error_msg=f"查询超时({custom_timeout}秒)",
)
raise QueryTimeoutError(f"查询超时({custom_timeout}秒)")
elapsed = int((time.time() - start) * 1000)
row_count = len(rows) if rows else 0
# 5. 检查结果集大小
truncated = row_count >= custom_max_rows
if self._check_result_size(rows if rows else []):
truncated = True
# 截断结果以保护内存
rows = rows[:custom_max_rows] if rows else []
logger.warning("查询结果超过内存限制 %dMB,已截断", MAX_RESULT_MEMORY_MB)
# 6. 记录查询日志
QueryLogger.log_query(
user_id=user_id,
db_type="mysql",
database=database,
query=query,
elapsed_ms=elapsed,
row_count=row_count,
is_success=True,
)
return UnifiedQueryResult(
columns=columns,
rows=rows if rows else [],
raw_result=None,
total=row_count,
elapsed_ms=elapsed,
truncated=truncated,
db_type="mysql",
query_type="SELECT",
)
class PostgresQueryHandler(SQLQueryHandlerBase):
"""PostgreSQL 查询处理器。
使用 DBConnector._query_postgres 执行查询,
SQLAnalyzer 验证只读性(SQL 标准兼容),自动添加 LIMIT,超时控制。
"""
async def execute(
self,
instance: Instance,
database: Optional[str],
query: str,
user_id: int,
options: Optional[Dict[str, Any]] = None,
) -> UnifiedQueryResult:
"""执行 PostgreSQL 只读查询。
Args:
instance: PostgreSQL 实例配置
database: 目标数据库名
query: SQL 查询语句
user_id: 查询用户ID
options: 额外选项
Returns:
UnifiedQueryResult: 统一格式的查询结果
Raises:
QuerySecurityError: 查询不安全
QueryTimeoutError: 查询超时
"""
options = options or {}
custom_timeout = options.get("timeout", QUERY_TIMEOUT)
custom_max_rows = options.get("limit", MAX_ROWS)
custom_max_rows = min(custom_max_rows, MAX_ROWS)
# 1. 增强安全验证
self._enhanced_security_check(query, user_id)
# 2. 验证只读
is_safe, error_msg = self.validate_readonly(query)
if not is_safe:
raise QuerySecurityError(error_msg)
# 3. 自动添加 LIMIT
query = self._auto_add_limit(query, custom_max_rows)
# 4. 执行查询
start = time.time()
rows: List[Dict[str, Any]] = []
columns: List[str] = []
try:
rows, columns = await asyncio.wait_for(
DBConnector._query_postgres(instance, database, query),
timeout=custom_timeout,
)
except asyncio.TimeoutError:
QueryLogger.log_query(
user_id=user_id,
db_type="postgres",
database=database,
query=query,
elapsed_ms=int((time.time() - start) * 1000),
row_count=0,
is_success=False,
error_msg=f"查询超时({custom_timeout}秒)",
)
raise QueryTimeoutError(f"查询超时({custom_timeout}秒)")
elapsed = int((time.time() - start) * 1000)
row_count = len(rows) if rows else 0
# 5. 检查结果集大小
truncated = row_count >= custom_max_rows
if self._check_result_size(rows if rows else []):
truncated = True
rows = rows[:custom_max_rows] if rows else []
logger.warning("查询结果超过内存限制 %dMB,已截断", MAX_RESULT_MEMORY_MB)
# 6. 记录查询日志
QueryLogger.log_query(
user_id=user_id,
db_type="postgres",
database=database,
query=query,
elapsed_ms=elapsed,
row_count=row_count,
is_success=True,
)
return UnifiedQueryResult(
columns=columns,
rows=rows if rows else [],
raw_result=None,
total=row_count,
elapsed_ms=elapsed,
truncated=truncated,
db_type="postgres",
query_type="SELECT",
)
# ===== Redis 查询处理器 =====
class RedisQueryHandler(BaseQueryHandler):
"""Redis 查询处理器。
使用 READONLY_COMMANDS 白名单 + BLACKLIST 验证命令安全性,
RedisConnector.execute_command 执行命令,
返回格式化为表格/JSON 的结果。
"""
def validate_readonly(self, query: str) -> Tuple[bool, str]:
"""验证 Redis 命令是否在只读白名单中。
增强验证:
1. 命令黑名单检查
2. 子命令白名单检查
3. 参数注入防护(检测管道符、shell注入)
Args:
query: Redis 命令字符串,如 "GET mykey"
Returns:
(is_safe, error_message)
"""
if not query or not query.strip():
return False, "Redis 命令不能为空"
# 长度验证
if len(query) > MAX_QUERY_LENGTH:
return False, f"Redis 命令过长({len(query)} 字符),最大允许 {MAX_QUERY_LENGTH}"
# 检测 shell 注入(管道符、重定向、命令替换)
shell_injection_chars = ["|", ";", "&", "$", "`", "$(", ">", "<"]
for char in shell_injection_chars:
if char in query:
return False, f"Redis 命令包含非法字符 '{char}',可能存在注入风险"
parts = query.strip().split()
if not parts:
return False, "Redis 命令不能为空"
cmd = parts[0].upper()
# 检查命令黑名单(最严格的拦截)
if cmd in REDIS_COMMAND_BLACKLIST:
# 但允许 CONFIG GET 等白名单子命令
if len(parts) >= 2:
sub = parts[1].upper()
subcmd_key = (cmd, sub)
if subcmd_key in REDIS_SUBCOMMAND_ALLOWLIST:
return True, ""
return False, f"Redis 命令 '{cmd}' 已被列入黑名单,禁止执行"
# 检查子命令(如 CONFIG GET, MEMORY USAGE)
if len(parts) >= 2:
sub = parts[1].upper()
subcmd_key = (cmd, sub)
# 如果是已知的多词命令,检查是否在允许列表中
if cmd in ("CONFIG", "MEMORY", "OBJECT", "CLIENT", "CLUSTER", "SLOWLOG"):
if subcmd_key not in REDIS_SUBCOMMAND_ALLOWLIST:
return False, f"Redis 命令 {cmd} {sub} 不在只读白名单中"
return True, ""
if cmd not in REDIS_READONLY_COMMANDS:
return False, f"Redis 命令 '{cmd}' 不在只读白名单中,禁止执行"
return True, ""
async def execute(
self,
instance: Instance,
database: Optional[str],
query: str,
user_id: int,
options: Optional[Dict[str, Any]] = None,
) -> UnifiedQueryResult:
"""执行 Redis 只读命令。
Args:
instance: Redis 实例配置
database: Redis DB 号(字符串,如 "0")
query: Redis 命令字符串
user_id: 查询用户ID
options: 额外选项
Returns:
UnifiedQueryResult: 统一格式的查询结果
Raises:
QuerySecurityError: 命令不安全
QueryTimeoutError: 命令超时
"""
options = options or {}
custom_timeout = options.get("timeout", QUERY_TIMEOUT)
# 1. 验证只读
is_safe, error_msg = self.validate_readonly(query)
if not is_safe:
QueryLogger.log_security_event(
event_type="unauthorized_redis_command",
user_id=user_id,
db_type="redis",
query=query,
reason=error_msg,
)
raise QuerySecurityError(error_msg)
# 2. 解析 DB 号
try:
db = int(database) if database else 0
if db < 0 or db > 15:
raise ValueError("Redis DB 号必须在 0-15 之间")
except (ValueError, TypeError):
raise QuerySecurityError("Redis DB 号必须是 0-15 之间的整数")
# 3. 执行命令
start = time.time()
try:
result_str = await asyncio.wait_for(
RedisConnector.execute_command(instance, db, query),
timeout=custom_timeout,
)
except asyncio.TimeoutError:
QueryLogger.log_query(
user_id=user_id,
db_type="redis",
database=f"DB{db}",
query=query,
elapsed_ms=int((time.time() - start) * 1000),
row_count=0,
is_success=False,
error_msg=f"Redis 命令超时({custom_timeout}秒)",
)
raise QueryTimeoutError(f"Redis 命令超时({custom_timeout}秒)")
except Exception as e:
raise QuerySecurityError(f"Redis 命令执行失败: {str(e)}")
elapsed = int((time.time() - start) * 1000)
# 4. 格式化结果为表格形式
parts = query.strip().split()
cmd = parts[0].upper() if parts else ""
columns, rows = self._format_result(cmd, result_str)
# 5. 记录查询日志
QueryLogger.log_query(
user_id=user_id,
db_type="redis",
database=f"DB{db}",
query=query,
elapsed_ms=elapsed,
row_count=len(rows),
is_success=True,
)
return UnifiedQueryResult(
columns=columns,
rows=rows,
raw_result=result_str,
total=len(rows),
elapsed_ms=elapsed,
truncated=False,
db_type="redis",
query_type=cmd,
)
def _format_result(self, cmd: str, result: str) -> Tuple[List[str], List[Dict[str, Any]]]:
"""将 Redis 原始结果格式化为表格。
Args:
cmd: Redis 命令名
result: 原始结果字符串
Returns:
(columns, rows): 列名列表和行数据列表
"""
if not result or result == "(nil)":
return ["result"], [{"result": "(nil)"}]
lines = result.split("\n")
# 对于返回列表的命令(KEYS, HKEYS, HVALS, SMEMBERS 等)
list_commands = {"KEYS", "HKEYS", "HVALS", "SMEMBERS", "LRANGE", "SINTER", "SUNION", "SDIFF"}
if cmd in list_commands and len(lines) > 1:
return ["value"], [{"value": line} for line in lines if line.strip()]
# 对于 INFO 命令,返回 key-value 表格
if cmd == "INFO":
rows = []
for line in lines:
line = line.strip()
if line and not line.startswith("#") and ":" in line:
k, v = line.split(":", 1)
rows.append({"key": k.strip(), "value": v.strip()})
if rows:
return ["key", "value"], rows
# 对于 CONFIG GET,返回 key-value 表格
if cmd == "CONFIG":
rows = []
for i in range(0, len(lines), 2):
if i + 1 < len(lines):
rows.append({"key": lines[i].strip(), "value": lines[i + 1].strip()})
elif lines[i].strip():
rows.append({"key": lines[i].strip(), "value": ""})
if rows:
return ["key", "value"], rows
# 对于 HGETALL,返回 field-value 表格
if cmd == "HGETALL":
rows = []
for i in range(0, len(lines), 2):
if i + 1 < len(lines):
rows.append({"field": lines[i].strip(), "value": lines[i + 1].strip()})
if rows:
return ["field", "value"], rows
# 默认: 返回单行结果
return ["result"], [{"result": result}]
# ===== MongoDB 查询处理器 =====
class MongoQueryHandler(BaseQueryHandler):
"""MongoDB 查询处理器。
使用 READONLY_OPERATIONS 白名单验证操作安全性,
motor 异步驱动执行查询,
增强的注入防护和结果大小限制。
"""
def validate_readonly(self, query: str) -> Tuple[bool, str]:
"""验证 MongoDB 操作是否在只读白名单中。
增强验证:
1. 严格的 JSON 解析
2. 操作类型白名单
3. aggregate pipeline 阶段检查
4. 禁止 $where 和 JavaScript 执行
Args:
query: MongoDB 查询 JSON 字符串
Returns:
(is_safe, error_message)
"""
if not query or not query.strip():
return False, "MongoDB 查询不能为空"
# 长度验证
if len(query) > MAX_QUERY_LENGTH:
return False, f"MongoDB 查询过长({len(query)} 字符),最大允许 {MAX_QUERY_LENGTH}"
try:
query_obj = json.loads(query)
except json.JSONDecodeError as e:
return False, f"MongoDB 查询 JSON 格式错误: {str(e)}"
if not isinstance(query_obj, dict):
return False, "MongoDB 查询必须是 JSON 对象"
operation = query_obj.get("operation", "")
if not operation:
return False, "MongoDB 查询必须指定 'operation' 字段"
if operation in MONGODB_WRITE_OPS:
return False, f"MongoDB 操作 '{operation}' 是写操作,禁止执行"
if operation not in MONGODB_READONLY_OPS:
return False, f"MongoDB 操作 '{operation}' 不在允许的操作列表中"
# 对 aggregate 进行额外检查: 禁止 $out, $merge 等写阶段
if operation == "aggregate":
pipeline = query_obj.get("pipeline", [])
if not isinstance(pipeline, list):
return False, "aggregate 的 pipeline 必须是数组"
for idx, stage in enumerate(pipeline):
if not isinstance(stage, dict):
return False, f"aggregate pipeline 第 {idx} 阶段必须是对象"
stage_name = list(stage.keys())[0] if stage else ""
if stage_name in MONGODB_AGGREGATE_BLACKLIST:
return False, f"aggregate 管道中的 '{stage_name}' 阶段是写操作,禁止执行"
# 禁止 $where 子句(JavaScript 执行风险)
filter_doc = query_obj.get("filter", {})
if isinstance(filter_doc, dict) and "$where" in filter_doc:
return False, "MongoDB 查询禁止使用 $where 子句(JavaScript 执行风险)"
# 禁止 $function, $accumulator 等可执行 JavaScript 的操作符
js_operators = {"$function", "$accumulator"}
query_str = json.dumps(query_obj)
for op in js_operators:
if op in query_str:
return False, f"MongoDB 查询禁止使用 '{op}' 操作符(JavaScript 执行风险)"
return True, ""
async def execute(
self,
instance: Instance,
database: Optional[str],
query: str,
user_id: int,
options: Optional[Dict[str, Any]] = None,
) -> UnifiedQueryResult:
"""执行 MongoDB 只读查询。
Args:
instance: MongoDB 实例配置
database: MongoDB 数据库名
query: MongoDB 查询 JSON 字符串
user_id: 查询用户ID
options: 额外选项
Returns:
UnifiedQueryResult: 统一格式的查询结果
Raises:
QuerySecurityError: 操作不安全
QueryTimeoutError: 查询超时
"""
options = options or {}
custom_timeout = options.get("timeout", QUERY_TIMEOUT)
custom_max_rows = options.get("limit", MAX_ROWS)
custom_max_rows = min(custom_max_rows, MAX_ROWS)
# 1. 验证只读
is_safe, error_msg = self.validate_readonly(query)
if not is_safe:
QueryLogger.log_security_event(
event_type="unauthorized_mongo_operation",
user_id=user_id,
db_type="mongo",
query=query,
reason=error_msg,
)
raise QuerySecurityError(error_msg)
# 2. 解析查询 JSON
try:
query_obj = json.loads(query)
except json.JSONDecodeError as e:
raise QuerySecurityError(f"MongoDB 查询 JSON 格式错误: {str(e)}")
collection_name = query_obj.get("collection", database or "")
operation = query_obj.get("operation", "")
filter_doc = query_obj.get("filter", {})
projection = query_obj.get("projection", None)
limit_val = min(query_obj.get("limit", custom_max_rows), custom_max_rows)
skip_val = query_obj.get("skip", 0)
sort_doc = query_obj.get("sort", None)
pipeline = query_obj.get("pipeline", [])
if not collection_name:
raise QuerySecurityError("MongoDB 查询必须指定 'collection' 字段或在 database 参数中提供")
# 验证 collection 名称安全(防止 NoSQL 注入)
if not re.match(r"^[\w]+$", collection_name):
raise QuerySecurityError("MongoDB collection 名称只能包含字母数字下划线")
# 3. 导入 MongoConnector(延迟导入避免循环依赖)
from app.services.mongo_connector import MongoConnector
start = time.time()
rows: List[Dict[str, Any]] = []
columns: List[str] = []
total = 0
try:
target_db = database or instance.database or ""
if operation == "find":
docs, total = await asyncio.wait_for(
MongoConnector.execute_find(
instance, target_db,
collection_name, filter_doc, projection,
limit_val, skip_val, sort_doc,
),
timeout=custom_timeout,
)
columns = list(docs[0].keys()) if docs else []
# 将 _id 转换为字符串以便 JSON 序列化
for doc in docs:
if "_id" in doc:
doc["_id"] = str(doc["_id"])
rows = docs
elif operation == "findOne":
doc = await asyncio.wait_for(
MongoConnector.execute_find_one(
instance, target_db,
collection_name, filter_doc, projection,
),
timeout=custom_timeout,
)
if doc:
if "_id" in doc:
doc["_id"] = str(doc["_id"])
columns = list(doc.keys())
rows = [doc]
total = 1
else:
columns, rows, total = [], [], 0
elif operation == "countDocuments":
count = await asyncio.wait_for(
MongoConnector.execute_count(
instance, target_db,
collection_name, filter_doc,
),
timeout=custom_timeout,
)
columns = ["count"]
rows = [{"count": count}]
total = 1
elif operation == "estimatedDocumentCount":
count = await asyncio.wait_for(
MongoConnector.execute_estimated_count(
instance, target_db,
collection_name,
),
timeout=custom_timeout,
)
columns = ["count"]
rows = [{"count": count}]
total = 1
elif operation == "distinct":
field = query_obj.get("field", "")
if not field:
raise QuerySecurityError("distinct 操作必须指定 'field' 字段")
# 验证字段名安全
if not re.match(r"^[\w.]+$", field):
raise QuerySecurityError("字段名只能包含字母数字下划线和点号")
values = await asyncio.wait_for(
MongoConnector.execute_distinct(
instance, target_db,
collection_name, field, filter_doc,
),
timeout=custom_timeout,
)
columns = ["value"]
rows = [{"value": str(v)} for v in values]
total = len(values)
elif operation == "aggregate":
docs = await asyncio.wait_for(
MongoConnector.execute_aggregate(
instance, target_db,
collection_name, pipeline,
),
timeout=custom_timeout,
)
columns = list(docs[0].keys()) if docs else []
for doc in docs:
if "_id" in doc:
doc["_id"] = str(doc["_id"])
rows = docs
total = len(docs)
else:
raise QuerySecurityError(f"不支持的 MongoDB 操作: {operation}")
except asyncio.TimeoutError:
QueryLogger.log_query(
user_id=user_id,
db_type="mongo",
database=database,
query=query,
elapsed_ms=int((time.time() - start) * 1000),
row_count=0,
is_success=False,
error_msg=f"MongoDB 查询超时({custom_timeout}秒)",
)
raise QueryTimeoutError(f"MongoDB 查询超时({custom_timeout}秒)")
except ValueError:
raise
except Exception as e:
raise QuerySecurityError(f"MongoDB 查询执行失败: {str(e)}")
elapsed = int((time.time() - start) * 1000)
# 4. 检查结果集大小
truncated = total >= custom_max_rows
if self._check_result_size(rows):
truncated = True
rows = rows[:custom_max_rows]
logger.warning("MongoDB 查询结果超过内存限制 %dMB,已截断", MAX_RESULT_MEMORY_MB)
# 5. 记录查询日志
QueryLogger.log_query(
user_id=user_id,
db_type="mongo",
database=database,
query=query,
elapsed_ms=elapsed,
row_count=total,
is_success=True,
)
return UnifiedQueryResult(
columns=columns,
rows=rows,
raw_result={"operation": operation, "collection": collection_name},
total=total,
elapsed_ms=elapsed,
truncated=truncated,
db_type="mongo",
query_type=operation,
)
# ===== 统一查询服务入口 =====
class QueryService:
"""统一多数据库查询服务入口。
根据实例的 db_type 自动路由到对应的查询处理器,
支持 MySQL、PostgreSQL、Redis、MongoDB 四种数据库。
Examples:
>>> result = await QueryService.execute_query(instance, database, query, user_id)
>>> print(result.rows)
"""
_handlers: Dict[str, BaseQueryHandler] = {
"mysql": MySQLQueryHandler(),
"postgres": PostgresQueryHandler(),
"postgresql": PostgresQueryHandler(),
"redis": RedisQueryHandler(),
"mongo": MongoQueryHandler(),
}
@staticmethod
async def execute_query(
instance: Instance,
database: Optional[str],
query: str,
user_id: int,
options: Optional[Dict[str, Any]] = None,
) -> UnifiedQueryResult:
"""统一查询入口。
根据 instance.db_type 自动路由到对应查询处理器,
执行安全验证后返回统一格式的结果。
Args:
instance: 数据源实例配置
database: 目标数据库/Redis DB号/MongoDB collection
query: 查询语句(SQL/Redis命令/MongoDB查询JSON)
user_id: 查询用户ID
options: 额外选项(如 limit, offset, timeout 等)
Returns:
UnifiedQueryResult: 统一格式的查询结果
Raises:
QuerySecurityError: 不支持的数据库类型或查询不安全
QueryTimeoutError: 查询超时
"""
db_type = instance.db_type.lower()
handler = QueryService._handlers.get(db_type)
if handler is None:
raise QuerySecurityError(f"暂不支持 {instance.db_type} 查询")
return await handler.execute(instance, database, query, user_id, options)
@staticmethod
async def execute_dml(
instance: Instance, database: str, sql: str
) -> str:
"""执行 DML 语句(INSERT/UPDATE/DELETE)。
目前仅支持 SQL 数据库(MySQL、PostgreSQL)。
Redis 和 MongoDB 的变更操作通过各自连接器处理。
Args:
instance: 数据源实例配置
database: 目标数据库名
sql: DML SQL 语句
Returns:
str: 执行结果描述
Raises:
QuerySecurityError: 不支持的数据库类型或执行失败
"""
return await DBConnector.execute_dml(instance, database, sql)