logan
a9e214ef1a
- AuditRule模型: 新增code/config_json字段，支持规则编码和阈值参数持久化 - audit_engine_v2: _load_rules从DB读config_json覆盖默认config - audit_rule_sync: 启动时自动从_RULE_REGISTRY同步规则到DB(239条) - settings API: update后刷新query_service配置缓存 - query_service: MAX_ROWS/QUERY_TIMEOUT/MAX_QUERY_LENGTH/MAX_RESULT_MEMORY改为动态读取system_settings - ticket_executor: SQL_TIMEOUT/REDIS_TIMEOUT/MONGO_TIMEOUT改为动态获取 - settings_seeds: 新增executor分组(sql/redis/mongo超时) + query_result_max_memory_mb - audit_rule endpoint: 移除无效的AuditEngine缓存清理（每次audit是新实例） - 测试: 1057 passed（修复所有新引入失败，17个历史遗留）
1353 lines
46 KiB
Python
1353 lines
46 KiB
Python
"""统一多数据库查询服务 — 支持 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__)
|
||
|
||
# ===== 配置常量（默认值，优先从 system_settings 读取） =====
|
||
|
||
# 默认最大返回行数
|
||
MAX_ROWS = 10000
|
||
# 默认查询超时（秒）
|
||
QUERY_TIMEOUT = 30
|
||
# 最大查询长度（字符）
|
||
MAX_QUERY_LENGTH = 10000
|
||
# 最大结果集内存限制（MB）
|
||
MAX_RESULT_MEMORY_MB = 50
|
||
|
||
|
||
async def _get_setting(key: str, default: int) -> int:
|
||
"""从 system_settings 读取配置，失败则返回默认值。"""
|
||
try:
|
||
from app.core.database import async_session
|
||
from app.models.system_setting import SystemSetting
|
||
from sqlalchemy import select
|
||
async with async_session() as db:
|
||
result = await db.execute(
|
||
select(SystemSetting.value).where(SystemSetting.key == key)
|
||
)
|
||
val = result.scalar_one_or_none()
|
||
return int(val) if val is not None else default
|
||
except Exception:
|
||
return default
|
||
|
||
|
||
async def get_query_config() -> Dict[str, int]:
|
||
"""批量获取查询相关配置，减少 DB 查询次数。"""
|
||
try:
|
||
from app.core.database import async_session
|
||
from app.models.system_setting import SystemSetting
|
||
from sqlalchemy import select
|
||
async with async_session() as db:
|
||
keys = [
|
||
"query_max_limit", "query_timeout", "query_max_length",
|
||
"query_result_max_memory_mb",
|
||
"executor_sql_timeout", "executor_redis_timeout", "executor_mongo_timeout",
|
||
]
|
||
stmt = select(SystemSetting.key, SystemSetting.value).where(
|
||
SystemSetting.key.in_(keys)
|
||
)
|
||
result = await db.execute(stmt)
|
||
rows = {r[0]: r[1] for r in result.all()}
|
||
return {
|
||
"max_rows": int(rows.get("query_max_limit", MAX_ROWS)),
|
||
"query_timeout": int(rows.get("query_timeout", QUERY_TIMEOUT)),
|
||
"max_query_length": int(rows.get("query_max_length", MAX_QUERY_LENGTH)),
|
||
"max_memory_mb": int(rows.get("query_result_max_memory_mb", MAX_RESULT_MEMORY_MB)),
|
||
"sql_timeout": int(rows.get("executor_sql_timeout", 60)),
|
||
"redis_timeout": int(rows.get("executor_redis_timeout", 30)),
|
||
"mongo_timeout": int(rows.get("executor_mongo_timeout", 60)),
|
||
}
|
||
except Exception:
|
||
return {
|
||
"max_rows": MAX_ROWS,
|
||
"query_timeout": QUERY_TIMEOUT,
|
||
"max_query_length": MAX_QUERY_LENGTH,
|
||
"max_memory_mb": MAX_RESULT_MEMORY_MB,
|
||
"sql_timeout": 60,
|
||
"redis_timeout": 30,
|
||
"mongo_timeout": 60,
|
||
}
|
||
|
||
|
||
# ===== 动态配置缓存 =====
|
||
|
||
_active_config: Optional[Dict[str, int]] = None
|
||
|
||
|
||
async def refresh_config() -> Dict[str, int]:
|
||
"""刷新配置缓存并返回。启动时和 settings 更新后调用。"""
|
||
global _active_config
|
||
_active_config = await get_query_config()
|
||
return _active_config
|
||
|
||
|
||
def current_config() -> Dict[str, int]:
|
||
"""获取当前配置（同步），未初始化则返回默认值。"""
|
||
if _active_config is not None:
|
||
return _active_config
|
||
return {
|
||
"max_rows": MAX_ROWS,
|
||
"query_timeout": QUERY_TIMEOUT,
|
||
"max_query_length": MAX_QUERY_LENGTH,
|
||
"max_memory_mb": MAX_RESULT_MEMORY_MB,
|
||
"sql_timeout": 60,
|
||
"redis_timeout": 30,
|
||
"mongo_timeout": 60,
|
||
}
|
||
|
||
|
||
# ===== 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("查询语句不能为空")
|
||
cfg = current_config()
|
||
if len(query) > cfg["max_query_length"]:
|
||
max_len = cfg["max_query_length"]
|
||
raise QuerySecurityError(
|
||
f"查询语句过长（{len(query)} 字符），最大允许 {max_len} 字符"
|
||
)
|
||
|
||
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 > current_config()["max_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 = -1) -> str:
|
||
"""如果 SQL 没有 LIMIT 则自动添加。仅对 SELECT 语句生效。
|
||
|
||
Args:
|
||
sql: 原始 SQL
|
||
max_rows: 最大返回行数（-1 表示从配置读取）
|
||
|
||
Returns:
|
||
添加了 LIMIT 的 SQL
|
||
"""
|
||
if max_rows < 0:
|
||
max_rows = current_config()["max_rows"]
|
||
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", current_config()["query_timeout"])
|
||
custom_max_rows = options.get("limit", current_config()["max_rows"])
|
||
custom_max_rows = min(custom_max_rows, current_config()["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，已截断", current_config()["max_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", current_config()["query_timeout"])
|
||
custom_max_rows = options.get("limit", current_config()["max_rows"])
|
||
custom_max_rows = min(custom_max_rows, current_config()["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，已截断", current_config()["max_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) > current_config()["max_query_length"]:
|
||
max_len = current_config()["max_query_length"]
|
||
return False, f"Redis 命令过长（{len(query)} 字符），最大允许 {max_len}"
|
||
|
||
# 检测 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", current_config()["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) > current_config()["max_query_length"]:
|
||
max_len = current_config()["max_query_length"]
|
||
return False, f"MongoDB 查询过长（{len(query)} 字符），最大允许 {max_len}"
|
||
|
||
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", current_config()["query_timeout"])
|
||
custom_max_rows = options.get("limit", current_config()["max_rows"])
|
||
custom_max_rows = min(custom_max_rows, current_config()["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，已截断", current_config()["max_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} 查询")
|
||
|
||
# 每次查询刷新配置缓存（轻量：仅查一次 DB）
|
||
await refresh_config()
|
||
|
||
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)
|