sync: 同步 sql-platform-pro 最新代码 — 权限体系安全修复 v0.8.1
- P0: 实例访问控制全线接入权限组 (829测试) - P1: permissions字段激活 + 操作级权限检查 (843测试) - 前端: 删除数据权限页面, 权限组合并为唯一入口 - 新增: CodeMirror编辑器 + 查询模板 + 慢查询分析 - 新增: MongoDB管理页面 + Dashboard增强 - 修复: Pydantic V2迁移 + 前端TS错误修复
This commit is contained in:
logan committed
1 parent
6fcd9465cf
commit
4c9427e546
105 files changed
+23680
-471
No files matched your search
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,500 @@
|
||||
"""SQL 审核引擎 V2 规则种子数据。
|
||||
|
||||
包含 40+ 条覆盖 DDL/DML/DQL/命名规范/性能/安全等多个维度的审核规则。
|
||||
提供 seed_all_rules(db) 函数用于一次性初始化所有规则到数据库。
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.models.audit_rule import AuditRule
|
||||
|
||||
|
||||
# ── DDL 规则种子 ──────────────────────────────────────────────
|
||||
DDL_RULES: List[Dict[str, Any]] = [
|
||||
{
|
||||
"name": "建表必须包含主键",
|
||||
"description": "CREATE TABLE 语句必须包含 PRIMARY KEY 定义，确保每行数据唯一可标识",
|
||||
"rule_type": "DDL",
|
||||
"severity": "error",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "建表语句缺少主键定义，请添加 PRIMARY KEY",
|
||||
},
|
||||
{
|
||||
"name": "建表必须添加表注释",
|
||||
"description": "CREATE TABLE 语句必须包含 COMMENT 表注释，便于维护和理解表用途",
|
||||
"rule_type": "DDL",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "建表时必须添加表注释 (COMMENT='...')",
|
||||
},
|
||||
{
|
||||
"name": "建表字段必须有注释",
|
||||
"description": "CREATE TABLE 中的每个字段都应包含 COMMENT 注释",
|
||||
"rule_type": "DDL",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "字段缺少注释，请为每个字段添加 COMMENT",
|
||||
},
|
||||
{
|
||||
"name": "禁止 DROP TABLE",
|
||||
"description": "生产环境禁止执行 DROP TABLE 操作，数据将永久丢失",
|
||||
"rule_type": "DDL",
|
||||
"severity": "error",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "禁止执行 DROP TABLE 操作，数据将永久丢失",
|
||||
},
|
||||
{
|
||||
"name": "禁止 TRUNCATE TABLE",
|
||||
"description": "禁止执行 TRUNCATE TABLE 操作，数据不可回滚",
|
||||
"rule_type": "DDL",
|
||||
"severity": "error",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "禁止执行 TRUNCATE TABLE 操作，数据不可回滚",
|
||||
},
|
||||
{
|
||||
"name": "单表索引数量不得超过12个",
|
||||
"description": "单表索引数量过多会影响写入性能",
|
||||
"rule_type": "DDL",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": "12",
|
||||
"message": "单表索引数量超过12个，会影响写入性能",
|
||||
},
|
||||
{
|
||||
"name": "单索引字段数不得超过8个",
|
||||
"description": "单索引包含字段数过多会降低索引效率",
|
||||
"rule_type": "DDL",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": "8",
|
||||
"message": "单索引包含字段数超过8个，请优化索引设计",
|
||||
},
|
||||
{
|
||||
"name": "禁止使用 FLOAT/DOUBLE 类型",
|
||||
"description": "浮点类型精度丢失，应使用 DECIMAL",
|
||||
"rule_type": "DDL",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "禁止使用 FLOAT/DOUBLE 类型，建议使用 DECIMAL 避免精度丢失",
|
||||
},
|
||||
{
|
||||
"name": "CHAR 类型长度不得超过10",
|
||||
"description": "CHAR 类型长度过大会浪费存储空间，超出请使用 VARCHAR",
|
||||
"rule_type": "DDL",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": "10",
|
||||
"message": "CHAR 类型长度不得超过10，超出请使用 VARCHAR",
|
||||
},
|
||||
{
|
||||
"name": "表名长度不得超过45字符",
|
||||
"description": "过长的表名不利于维护和使用",
|
||||
"rule_type": "DDL",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": "45",
|
||||
"message": "表名长度超过45个字符，请缩短表名",
|
||||
},
|
||||
{
|
||||
"name": "禁止外键约束",
|
||||
"description": "外键约束影响性能，推荐在应用层处理关联关系",
|
||||
"rule_type": "DDL",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "禁止使用 FOREIGN KEY 约束，建议在应用层维护数据一致性",
|
||||
},
|
||||
{
|
||||
"name": "必须包含 create_time/update_time 字段",
|
||||
"description": "建表必须包含 create_time 和 update_time 审计字段，便于追踪数据变更",
|
||||
"rule_type": "DDL",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "建表缺少 create_time 或 update_time 审计字段",
|
||||
},
|
||||
]
|
||||
|
||||
# ── DML 规则种子 ──────────────────────────────────────────────
|
||||
DML_RULES: List[Dict[str, Any]] = [
|
||||
{
|
||||
"name": "UPDATE/DELETE 必须包含 WHERE 条件",
|
||||
"description": "UPDATE/DELETE 缺少 WHERE 将影响全表数据，可能导致严重事故",
|
||||
"rule_type": "DML",
|
||||
"severity": "error",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "UPDATE/DELETE 语句必须包含 WHERE 条件",
|
||||
},
|
||||
{
|
||||
"name": "单次 INSERT 最多500行",
|
||||
"description": "单次 INSERT 行数过多会导致事务过大，影响性能和复制延迟",
|
||||
"rule_type": "DML",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": "500",
|
||||
"message": "单次 INSERT 最多插入500行",
|
||||
},
|
||||
{
|
||||
"name": "UPDATE/DELETE 禁止带 LIMIT",
|
||||
"description": "UPDATE/DELETE 使用 LIMIT 可能导致部分数据变更，数据不一致",
|
||||
"rule_type": "DML",
|
||||
"severity": "warning",
|
||||
"db_type": "mysql",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "UPDATE/DELETE 语句禁止使用 LIMIT",
|
||||
},
|
||||
{
|
||||
"name": "禁止 INSERT 不指定字段列表",
|
||||
"description": "INSERT 必须显式指定字段列表，避免表结构变更导致错误",
|
||||
"rule_type": "DML",
|
||||
"severity": "error",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "INSERT 语句必须显式指定字段列表",
|
||||
},
|
||||
{
|
||||
"name": "禁止批量 UPDATE 不带主键条件",
|
||||
"description": "批量 UPDATE 必须带主键 WHERE 条件，防止误更新大量数据",
|
||||
"rule_type": "DML",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "批量 UPDATE 必须包含主键条件",
|
||||
},
|
||||
{
|
||||
"name": "DELETE 必须有 LIMIT（除非带主键条件）",
|
||||
"description": "DELETE 必须带 LIMIT 或主键 WHERE 条件，防止误删大量数据",
|
||||
"rule_type": "DML",
|
||||
"severity": "warning",
|
||||
"db_type": "mysql",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "DELETE 建议添加 LIMIT 或使用主键条件",
|
||||
},
|
||||
{
|
||||
"name": "禁止 REPLACE INTO",
|
||||
"description": "REPLACE INTO 会先删除再插入，风险较高",
|
||||
"rule_type": "DML",
|
||||
"severity": "error",
|
||||
"db_type": "mysql",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "禁止使用 REPLACE INTO，请使用 INSERT ... ON DUPLICATE KEY UPDATE 替代",
|
||||
},
|
||||
{
|
||||
"name": "INSERT 必须指定字段列表",
|
||||
"description": "INSERT 语句必须显式指定字段列表，避免表结构变更导致错误",
|
||||
"rule_type": "DML",
|
||||
"severity": "error",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "INSERT 语句必须显式指定字段列表",
|
||||
},
|
||||
]
|
||||
|
||||
# ── DQL 规则种子 ──────────────────────────────────────────────
|
||||
DQL_RULES: List[Dict[str, Any]] = [
|
||||
{
|
||||
"name": "禁止 SELECT *",
|
||||
"description": "SELECT * 返回不必要的列，影响网络传输和内存使用",
|
||||
"rule_type": "DQL",
|
||||
"severity": "info",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "禁止使用 SELECT *，请明确指定需要的字段",
|
||||
},
|
||||
{
|
||||
"name": "SELECT 必须有 LIMIT",
|
||||
"description": "SELECT 查询应添加 LIMIT 限制返回行数，防止返回过多数据",
|
||||
"rule_type": "DQL",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "SELECT 查询缺少 LIMIT，建议添加 LIMIT 防止返回过多数据",
|
||||
},
|
||||
{
|
||||
"name": "禁止在 WHERE 中使用函数",
|
||||
"description": "WHERE 中使用函数会导致索引失效，性能下降",
|
||||
"rule_type": "DQL",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "WHERE 子句中使用函数会导致索引失效",
|
||||
},
|
||||
{
|
||||
"name": "禁止隐式类型转换",
|
||||
"description": "不同类型比较会导致索引失效，请确保比较双方类型一致",
|
||||
"rule_type": "DQL",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "WHERE 条件中存在隐式类型转换风险",
|
||||
},
|
||||
{
|
||||
"name": "ORDER BY 必须使用索引字段",
|
||||
"description": "ORDER BY 非索引字段会导致文件排序（filesort），影响性能",
|
||||
"rule_type": "DQL",
|
||||
"severity": "info",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "ORDER BY 字段建议添加索引，避免 filesort 性能问题",
|
||||
},
|
||||
]
|
||||
|
||||
# ── 命名规范规则种子 ──────────────────────────────────────────
|
||||
NAMING_RULES: List[Dict[str, Any]] = [
|
||||
{
|
||||
"name": "表名必须小写下划线命名",
|
||||
"description": "表名应使用小写下划线格式 (snake_case)，如 user_info",
|
||||
"rule_type": "naming",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": r"^[a-z][a-z0-9_]*$",
|
||||
"value": None,
|
||||
"message": "表名不符合小写下划线命名规范",
|
||||
},
|
||||
{
|
||||
"name": "字段名必须小写下划线命名",
|
||||
"description": "字段名应使用小写下划线格式 (snake_case)，如 user_name",
|
||||
"rule_type": "naming",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": r"^[a-z][a-z0-9_]*$",
|
||||
"value": None,
|
||||
"message": "字段名不符合小写下划线命名规范",
|
||||
},
|
||||
{
|
||||
"name": "索引名必须以 idx_ 或 uk_ 开头",
|
||||
"description": "索引命名规范：普通索引 idx_，唯一索引 uk_，如 idx_user_name",
|
||||
"rule_type": "naming",
|
||||
"severity": "info",
|
||||
"db_type": "all",
|
||||
"pattern": r"^(idx_|uk_)",
|
||||
"value": None,
|
||||
"message": "索引名应以 idx_ 或 uk_ 开头",
|
||||
},
|
||||
{
|
||||
"name": "主键名必须为 id 或 pk_ 开头",
|
||||
"description": "主键字段命名规范，建议使用 id 或 pk_ 前缀",
|
||||
"rule_type": "naming",
|
||||
"severity": "info",
|
||||
"db_type": "all",
|
||||
"pattern": r"^(id|pk_)",
|
||||
"value": None,
|
||||
"message": "主键名建议使用 'id' 或 'pk_' 前缀",
|
||||
},
|
||||
{
|
||||
"name": "布尔字段以 is_/has_ 开头",
|
||||
"description": "布尔类型字段应以 is_ 或 has_ 前缀命名，如 is_active",
|
||||
"rule_type": "naming",
|
||||
"severity": "info",
|
||||
"db_type": "all",
|
||||
"pattern": r"^(is_|has_)",
|
||||
"value": None,
|
||||
"message": "布尔字段建议以 is_ 或 has_ 开头命名",
|
||||
},
|
||||
]
|
||||
|
||||
# ── 性能规则种子 ──────────────────────────────────────────────
|
||||
PERF_RULES: List[Dict[str, Any]] = [
|
||||
{
|
||||
"name": "禁止大表全表扫描",
|
||||
"description": "大表（>10000行）查询缺少 WHERE 条件会导致全表扫描，严重影响性能",
|
||||
"rule_type": "performance",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": "10000",
|
||||
"message": "SELECT 查询缺少 WHERE 条件，大表下会导致全表扫描",
|
||||
},
|
||||
{
|
||||
"name": "IN 子句值数量不得超过100个",
|
||||
"description": "IN 子句值过多会导致性能下降，建议使用临时表或 JOIN 替代",
|
||||
"rule_type": "performance",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": "100",
|
||||
"message": "IN 子句包含超过100个值，建议使用临时表或 JOIN 替代",
|
||||
},
|
||||
{
|
||||
"name": "禁止 LIKE '%xxx%' 前缀模糊匹配",
|
||||
"description": "前缀模糊匹配无法使用索引，建议使用全文检索",
|
||||
"rule_type": "performance",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": r"LIKE\s+'%[^']+'",
|
||||
"value": None,
|
||||
"message": "LIKE '%xxx%' 无法使用索引，建议改用全文检索或后缀索引",
|
||||
},
|
||||
{
|
||||
"name": "JOIN 不得超过5个表",
|
||||
"description": "JOIN 表过多会导致查询计划复杂化和性能下降",
|
||||
"rule_type": "performance",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": "5",
|
||||
"message": "JOIN 涉及超过5个表，请优化查询结构",
|
||||
},
|
||||
{
|
||||
"name": "禁止 ORDER BY RAND()",
|
||||
"description": "ORDER BY RAND() 性能极差，会导致全表扫描",
|
||||
"rule_type": "performance",
|
||||
"severity": "error",
|
||||
"db_type": "mysql",
|
||||
"pattern": r"ORDER\s+BY\s+RAND\s*\(",
|
||||
"value": None,
|
||||
"message": "禁止 ORDER BY RAND()，性能极差，建议在应用层随机采样",
|
||||
},
|
||||
{
|
||||
"name": "子查询建议改为 JOIN",
|
||||
"description": "相关子查询性能较差，建议改为 JOIN",
|
||||
"rule_type": "performance",
|
||||
"severity": "info",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "WHERE 中的子查询建议改为 JOIN 以提升性能",
|
||||
},
|
||||
{
|
||||
"name": "大表禁止 OFFSET 分页",
|
||||
"description": "大表使用 OFFSET 分页性能差，建议使用游标分页",
|
||||
"rule_type": "performance",
|
||||
"severity": "warning",
|
||||
"db_type": "all",
|
||||
"pattern": r"OFFSET\s+\d+",
|
||||
"value": None,
|
||||
"message": "使用 OFFSET 分页在大表上性能较差，建议改用游标分页",
|
||||
},
|
||||
{
|
||||
"name": "禁止 SELECT DISTINCT（除非必要）",
|
||||
"description": "DISTINCT 需要去重操作，影响性能",
|
||||
"rule_type": "performance",
|
||||
"severity": "info",
|
||||
"db_type": "all",
|
||||
"pattern": r"SELECT\s+DISTINCT\b",
|
||||
"value": None,
|
||||
"message": "SELECT DISTINCT 影响性能，请确认是否必要",
|
||||
},
|
||||
]
|
||||
|
||||
# ── 安全规则种子 ──────────────────────────────────────────────
|
||||
SEC_RULES: List[Dict[str, Any]] = [
|
||||
{
|
||||
"name": "禁止查询 password/secret/token 字段",
|
||||
"description": "禁止在 SELECT 中查询敏感字段，防止密码等敏感信息泄露",
|
||||
"rule_type": "security",
|
||||
"severity": "error",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "禁止在查询中返回 password/secret/token 等敏感字段",
|
||||
},
|
||||
{
|
||||
"name": "禁止 SQL 中硬编码密码",
|
||||
"description": "SQL 中不应包含硬编码密码，防止密码泄露",
|
||||
"rule_type": "security",
|
||||
"severity": "error",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "SQL 中发现可能的硬编码密码，禁止在 SQL 中直接写入密码",
|
||||
},
|
||||
{
|
||||
"name": "禁止 UNION（SQL 注入风险）",
|
||||
"description": "UNION 语句存在 SQL 注入风险，禁止在工单中使用",
|
||||
"rule_type": "security",
|
||||
"severity": "error",
|
||||
"db_type": "all",
|
||||
"pattern": r"\bUNION\b",
|
||||
"value": None,
|
||||
"message": "禁止 UNION 操作，存在 SQL 注入风险",
|
||||
},
|
||||
{
|
||||
"name": "禁止注释符",
|
||||
"description": "SQL 中不应包含 -- 或 /* */ 注释，可能存在注入风险",
|
||||
"rule_type": "security",
|
||||
"severity": "error",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "SQL 中禁止包含注释符，可能存在注入风险",
|
||||
},
|
||||
{
|
||||
"name": "禁止多语句（分号分隔）",
|
||||
"description": "禁止单条 SQL 中包含多条语句，防止注入攻击",
|
||||
"rule_type": "security",
|
||||
"severity": "error",
|
||||
"db_type": "all",
|
||||
"pattern": None,
|
||||
"value": None,
|
||||
"message": "禁止提交多条 SQL 语句（分号分隔），请拆分为多个工单",
|
||||
},
|
||||
]
|
||||
|
||||
# ── 合并所有规则 ──────────────────────────────────────────────
|
||||
ALL_V2_RULES: List[Dict[str, Any]] = (
|
||||
DDL_RULES + DML_RULES + DQL_RULES + NAMING_RULES + PERF_RULES + SEC_RULES
|
||||
)
|
||||
|
||||
|
||||
async def seed_all_rules(db: AsyncSession) -> List[AuditRule]:
|
||||
"""将所有 V2 规则种子数据插入数据库（跳过已存在的规则）。
|
||||
|
||||
Args:
|
||||
db: 数据库异步会话
|
||||
|
||||
Returns:
|
||||
新创建的 AuditRule 列表
|
||||
"""
|
||||
created: List[AuditRule] = []
|
||||
|
||||
for rule_dict in ALL_V2_RULES:
|
||||
# 检查是否已存在（按 name 去重）
|
||||
exists = await db.execute(select(AuditRule).where(AuditRule.name == rule_dict["name"]))
|
||||
if exists.scalar_one_or_none():
|
||||
continue # 跳过已存在的规则
|
||||
|
||||
rule = AuditRule(**rule_dict)
|
||||
db.add(rule)
|
||||
created.append(rule)
|
||||
|
||||
await db.commit()
|
||||
for rule in created:
|
||||
await db.refresh(rule)
|
||||
|
||||
return created
|
||||
@@ -0,0 +1,351 @@
|
||||
"""数据库备份执行器 — 使用 mysqldump/pg_dump 执行真实备份。
|
||||
|
||||
支持:
|
||||
- MySQL: mysqldump + gzip 压缩
|
||||
- PostgreSQL: pg_dump custom format
|
||||
|
||||
备份文件存储在 /app/backups 目录。
|
||||
"""
|
||||
import os
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from app.models.instance import Instance
|
||||
from app.models.backup import Backup
|
||||
from app.core.crypto import decrypt_password
|
||||
|
||||
|
||||
# 备份存储目录
|
||||
BACKUP_DIR = os.environ.get("BACKUP_DIR", "/app/backups")
|
||||
|
||||
|
||||
class BackupExecutor:
|
||||
"""数据库备份执行器。"""
|
||||
|
||||
async def execute(self, backup: Backup, instance: Instance) -> Tuple[str, int]:
|
||||
"""执行备份，返回 (文件路径, 文件大小字节)。
|
||||
|
||||
Raises:
|
||||
RuntimeError: 备份工具未安装或备份失败
|
||||
"""
|
||||
os.makedirs(BACKUP_DIR, exist_ok=True)
|
||||
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
password = decrypt_password(instance.password)
|
||||
|
||||
if instance.db_type == "mysql":
|
||||
filepath = await self._backup_mysql(
|
||||
instance, backup.database, password, timestamp
|
||||
)
|
||||
elif instance.db_type in ("postgres", "postgresql"):
|
||||
filepath = await self._backup_postgres(
|
||||
instance, backup.database, password, timestamp
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"不支持的备份类型: {instance.db_type}")
|
||||
|
||||
size_bytes = os.path.getsize(filepath) if os.path.exists(filepath) else 0
|
||||
return filepath, size_bytes
|
||||
|
||||
async def _backup_mysql(self, instance: Instance, database: str,
|
||||
password: str, timestamp: str) -> str:
|
||||
"""使用 mysqldump 备份 MySQL 数据库。
|
||||
|
||||
输出通过 gzip 压缩存储。
|
||||
"""
|
||||
filename = f"{database}_mysql_{timestamp}.sql.gz"
|
||||
filepath = os.path.join(BACKUP_DIR, filename)
|
||||
|
||||
# 检查 mysqldump 是否可用
|
||||
try:
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
"mysqldump", "--version",
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
await proc.communicate()
|
||||
if proc.returncode != 0:
|
||||
raise RuntimeError("mysqldump 未安装")
|
||||
except FileNotFoundError:
|
||||
raise RuntimeError(
|
||||
"mysqldump 未安装，请安装 MySQL 客户端工具:\\n"
|
||||
" apt-get install mysql-client 或\\n"
|
||||
" yum install mysql"
|
||||
)
|
||||
|
||||
cmd = [
|
||||
"mysqldump",
|
||||
"-h", instance.host,
|
||||
"-P", str(instance.port),
|
||||
"-u", instance.username,
|
||||
f"-p{password}",
|
||||
"--single-transaction",
|
||||
"--routines",
|
||||
"--triggers",
|
||||
"--events",
|
||||
"--hex-blob",
|
||||
"--skip-lock-tables",
|
||||
database,
|
||||
]
|
||||
|
||||
# 创建 mysqldump 子进程，stdout 通过 gzip 压缩
|
||||
dump_proc = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
|
||||
# 创建 gzip 子进程
|
||||
gzip_proc = await asyncio.create_subprocess_exec(
|
||||
"gzip", "-c",
|
||||
stdin=dump_proc.stdout,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
|
||||
# 关闭 dump_proc 的 stdout 在本进程的引用
|
||||
if dump_proc.stdout:
|
||||
dump_proc.stdout.close()
|
||||
|
||||
# 读取 gzip 输出写入文件
|
||||
try:
|
||||
with open(filepath, "wb") as f:
|
||||
while True:
|
||||
chunk = await gzip_proc.stdout.read(65536) if gzip_proc.stdout else b""
|
||||
if not chunk:
|
||||
break
|
||||
f.write(chunk)
|
||||
|
||||
# 等待进程完成
|
||||
dump_stderr = await dump_proc.stderr.read() if dump_proc.stderr else b""
|
||||
gzip_stderr = await gzip_proc.stderr.read() if gzip_proc.stderr else b""
|
||||
|
||||
await dump_proc.wait()
|
||||
await gzip_proc.wait()
|
||||
|
||||
if dump_proc.returncode != 0:
|
||||
error_msg = dump_stderr.decode("utf-8", errors="ignore")[-500:]
|
||||
raise RuntimeError(f"mysqldump 执行失败: {error_msg}")
|
||||
|
||||
if gzip_proc.returncode != 0:
|
||||
error_msg = gzip_stderr.decode("utf-8", errors="ignore")
|
||||
raise RuntimeError(f"gzip 压缩失败: {error_msg}")
|
||||
|
||||
except Exception:
|
||||
# 清理失败的文件
|
||||
if os.path.exists(filepath):
|
||||
os.remove(filepath)
|
||||
raise
|
||||
|
||||
return filepath
|
||||
|
||||
async def _backup_postgres(self, instance: Instance, database: str,
|
||||
password: str, timestamp: str) -> str:
|
||||
"""使用 pg_dump 备份 PostgreSQL 数据库。
|
||||
|
||||
使用 custom format（已内置压缩）。
|
||||
"""
|
||||
filename = f"{database}_postgres_{timestamp}.dump"
|
||||
filepath = os.path.join(BACKUP_DIR, filename)
|
||||
|
||||
# 检查 pg_dump 是否可用
|
||||
try:
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
"pg_dump", "--version",
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
await proc.communicate()
|
||||
if proc.returncode != 0:
|
||||
raise RuntimeError("pg_dump 未安装")
|
||||
except FileNotFoundError:
|
||||
raise RuntimeError(
|
||||
"pg_dump 未安装，请安装 PostgreSQL 客户端工具:\\n"
|
||||
" apt-get install postgresql-client 或\\n"
|
||||
" yum install postgresql"
|
||||
)
|
||||
|
||||
env = os.environ.copy()
|
||||
env["PGPASSWORD"] = password
|
||||
|
||||
cmd = [
|
||||
"pg_dump",
|
||||
"-h", instance.host,
|
||||
"-p", str(instance.port),
|
||||
"-U", instance.username,
|
||||
"-d", database,
|
||||
"-F", "c", # custom format (compressed)
|
||||
"-Z", "5", # compression level 5
|
||||
"--no-owner", # don't add ownership commands
|
||||
"--no-privileges", # don't add privilege commands
|
||||
"-f", filepath,
|
||||
]
|
||||
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=env,
|
||||
)
|
||||
|
||||
stdout, stderr = await proc.communicate()
|
||||
|
||||
if proc.returncode != 0:
|
||||
error_msg = stderr.decode("utf-8", errors="ignore")[-500:]
|
||||
# 清理失败的文件
|
||||
if os.path.exists(filepath):
|
||||
os.remove(filepath)
|
||||
raise RuntimeError(f"pg_dump 执行失败: {error_msg}")
|
||||
|
||||
return filepath
|
||||
|
||||
async def restore_mysql(self, instance: Instance, database: str,
|
||||
filepath: str) -> str:
|
||||
"""从备份恢复 MySQL 数据库。
|
||||
|
||||
Args:
|
||||
instance: 目标实例
|
||||
database: 目标数据库
|
||||
filepath: 备份文件路径（.sql.gz）
|
||||
"""
|
||||
password = decrypt_password(instance.password)
|
||||
|
||||
if filepath.endswith(".gz"):
|
||||
# gunzip 解压后通过 mysql 导入
|
||||
gunzip_proc = await asyncio.create_subprocess_exec(
|
||||
"gunzip", "-c", filepath,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
|
||||
mysql_cmd = [
|
||||
"mysql",
|
||||
"-h", instance.host,
|
||||
"-P", str(instance.port),
|
||||
"-u", instance.username,
|
||||
f"-p{password}",
|
||||
database,
|
||||
]
|
||||
|
||||
mysql_proc = await asyncio.create_subprocess_exec(
|
||||
*mysql_cmd,
|
||||
stdin=gunzip_proc.stdout,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
|
||||
if gunzip_proc.stdout:
|
||||
gunzip_proc.stdout.close()
|
||||
|
||||
await mysql_proc.wait()
|
||||
await gunzip_proc.wait()
|
||||
|
||||
if mysql_proc.returncode != 0:
|
||||
raise RuntimeError("MySQL 恢复失败")
|
||||
else:
|
||||
# 直接导入 .sql 文件
|
||||
with open(filepath, "rb") as f:
|
||||
mysql_cmd = [
|
||||
"mysql",
|
||||
"-h", instance.host,
|
||||
"-P", str(instance.port),
|
||||
"-u", instance.username,
|
||||
f"-p{password}",
|
||||
database,
|
||||
]
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
*mysql_cmd,
|
||||
stdin=f,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
_, stderr = await proc.communicate()
|
||||
if proc.returncode != 0:
|
||||
raise RuntimeError(f"MySQL 恢复失败: {stderr.decode()[-500:]}")
|
||||
|
||||
return f"数据库 {database} 恢复成功"
|
||||
|
||||
async def restore_postgres(self, instance: Instance, database: str,
|
||||
filepath: str) -> str:
|
||||
"""从备份恢复 PostgreSQL 数据库。
|
||||
|
||||
Args:
|
||||
instance: 目标实例
|
||||
database: 目标数据库
|
||||
filepath: 备份文件路径（.dump）
|
||||
"""
|
||||
password = decrypt_password(instance.password)
|
||||
env = os.environ.copy()
|
||||
env["PGPASSWORD"] = password
|
||||
|
||||
cmd = [
|
||||
"pg_restore",
|
||||
"-h", instance.host,
|
||||
"-p", str(instance.port),
|
||||
"-U", instance.username,
|
||||
"-d", database,
|
||||
"--no-owner",
|
||||
"--no-privileges",
|
||||
"-c", # clean (drop) database objects before recreating
|
||||
filepath,
|
||||
]
|
||||
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=env,
|
||||
)
|
||||
|
||||
_, stderr = await proc.communicate()
|
||||
|
||||
# pg_restore 即使有一些警告也可能成功
|
||||
if proc.returncode not in (0, 1): # 1 = warnings
|
||||
error_msg = stderr.decode("utf-8", errors="ignore")[-500:]
|
||||
raise RuntimeError(f"pg_restore 失败: {error_msg}")
|
||||
|
||||
return f"数据库 {database} 恢复成功"
|
||||
|
||||
@staticmethod
|
||||
def validate_backup_file(filepath: str) -> bool:
|
||||
"""验证备份文件是否有效。
|
||||
|
||||
- MySQL: 检查 gzip 文件完整性
|
||||
- PostgreSQL: 检查 custom format 文件
|
||||
"""
|
||||
if not os.path.exists(filepath):
|
||||
return False
|
||||
|
||||
if filepath.endswith(".gz"):
|
||||
# 使用 gzip -t 测试
|
||||
import subprocess
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["gzip", "-t", filepath],
|
||||
capture_output=True, timeout=30,
|
||||
)
|
||||
return result.returncode == 0
|
||||
except (subprocess.TimeoutExpired, FileNotFoundError):
|
||||
# 简单检查文件大小和魔数
|
||||
with open(filepath, "rb") as f:
|
||||
magic = f.read(2)
|
||||
return magic == b"\\x1f\\x8b" # gzip magic number
|
||||
|
||||
elif filepath.endswith(".dump"):
|
||||
# PostgreSQL custom format: 检查文件大小和魔数
|
||||
if os.path.getsize(filepath) < 100:
|
||||
return False
|
||||
with open(filepath, "rb") as f:
|
||||
header = f.read(5)
|
||||
# PostgreSQL custom format starts with "PGDMP"
|
||||
return header == b"PGDMP"
|
||||
|
||||
return os.path.getsize(filepath) > 0
|
||||
|
||||
@staticmethod
|
||||
def get_backup_size(filepath: str) -> int:
|
||||
"""获取备份文件大小（字节）。"""
|
||||
if os.path.exists(filepath):
|
||||
return os.path.getsize(filepath)
|
||||
return 0
|
||||
@@ -0,0 +1,228 @@
|
||||
"""数据库连接池管理 — 复用数据库连接，提高性能。
|
||||
|
||||
使用异步连接池:
|
||||
- MySQL: aiomysql.Pool
|
||||
- PostgreSQL: asyncpg.Pool
|
||||
|
||||
连接池按实例 ID 管理，每个实例维护一个连接池。
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, Optional
|
||||
|
||||
from app.models.instance import Instance
|
||||
from app.core.crypto import decrypt_password
|
||||
|
||||
|
||||
@dataclass
|
||||
class PoolStats:
|
||||
"""连接池统计信息。"""
|
||||
instance_id: int
|
||||
db_type: str
|
||||
active_connections: int = 0
|
||||
idle_connections: int = 0
|
||||
max_connections: int = 5
|
||||
total_queries: int = 0
|
||||
failed_queries: int = 0
|
||||
avg_response_ms: int = 0
|
||||
pool_hit_rate: float = 1.0
|
||||
created_at: float = field(default_factory=time.time)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"instance_id": self.instance_id,
|
||||
"db_type": self.db_type,
|
||||
"active_connections": self.active_connections,
|
||||
"idle_connections": self.idle_connections,
|
||||
"max_connections": self.max_connections,
|
||||
"total_queries": self.total_queries,
|
||||
"failed_queries": self.failed_queries,
|
||||
"avg_response_ms": self.avg_response_ms,
|
||||
"pool_hit_rate": round(self.pool_hit_rate, 4),
|
||||
}
|
||||
|
||||
|
||||
class ConnectionPoolManager:
|
||||
"""数据库连接池管理器。"""
|
||||
|
||||
def __init__(self):
|
||||
# 连接池缓存: {instance_id: pool}
|
||||
self._pools: Dict[int, Any] = {}
|
||||
# 连接池统计
|
||||
self._stats: Dict[int, PoolStats] = {}
|
||||
# 创建锁，防止并发创建重复连接池
|
||||
self._locks: Dict[int, asyncio.Lock] = {}
|
||||
# 总体锁
|
||||
self._global_lock = asyncio.Lock()
|
||||
|
||||
async def get_pool(self, instance: Instance):
|
||||
"""获取或创建连接池。"""
|
||||
instance_id = instance.id
|
||||
|
||||
# 检查现有连接池
|
||||
if instance_id in self._pools:
|
||||
pool = self._pools[instance_id]
|
||||
# 简单的健康检查
|
||||
try:
|
||||
if instance.db_type == "mysql":
|
||||
pool._free # 检查 aiomysql pool 是否活跃
|
||||
return pool
|
||||
except Exception:
|
||||
# 连接池已失效，重新创建
|
||||
await self.close_pool(instance_id)
|
||||
|
||||
# 获取或创建实例锁
|
||||
if instance_id not in self._locks:
|
||||
self._locks[instance_id] = asyncio.Lock()
|
||||
|
||||
async with self._locks[instance_id]:
|
||||
# 双重检查
|
||||
if instance_id in self._pools:
|
||||
return self._pools[instance_id]
|
||||
|
||||
pool = await self._create_pool(instance)
|
||||
self._pools[instance_id] = pool
|
||||
self._stats[instance_id] = PoolStats(
|
||||
instance_id=instance_id,
|
||||
db_type=instance.db_type,
|
||||
max_connections=instance.connection_pool_size or 5,
|
||||
)
|
||||
return pool
|
||||
|
||||
async def _create_pool(self, instance: Instance):
|
||||
"""创建连接池。"""
|
||||
password = decrypt_password(instance.password)
|
||||
pool_size = instance.connection_pool_size or 5
|
||||
|
||||
if instance.db_type == "mysql":
|
||||
import aiomysql
|
||||
return await aiomysql.create_pool(
|
||||
host=instance.host,
|
||||
port=instance.port,
|
||||
user=instance.username,
|
||||
password=password,
|
||||
db=instance.database,
|
||||
minsize=1,
|
||||
maxsize=pool_size,
|
||||
pool_recycle=3600, # 1小时回收连接
|
||||
connect_timeout=10,
|
||||
autocommit=True,
|
||||
)
|
||||
elif instance.db_type in ("postgres", "postgresql"):
|
||||
import asyncpg
|
||||
dsn = (
|
||||
f"postgresql://{instance.username}:{password}"
|
||||
f"@{instance.host}:{instance.port}"
|
||||
f"/{instance.database or 'postgres'}"
|
||||
)
|
||||
return await asyncpg.create_pool(
|
||||
dsn=dsn,
|
||||
min_size=1,
|
||||
max_size=pool_size,
|
||||
command_timeout=60,
|
||||
max_inactive_connection_lifetime=3600,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"不支持的连接池类型: {instance.db_type}")
|
||||
|
||||
async def close_pool(self, instance_id: int) -> None:
|
||||
"""关闭指定实例的连接池。"""
|
||||
pool = self._pools.pop(instance_id, None)
|
||||
if pool:
|
||||
try:
|
||||
pool.close()
|
||||
await pool.wait_closed()
|
||||
except Exception:
|
||||
pass
|
||||
self._stats.pop(instance_id, None)
|
||||
|
||||
async def close_all(self) -> None:
|
||||
"""关闭所有连接池。"""
|
||||
for instance_id in list(self._pools.keys()):
|
||||
await self.close_pool(instance_id)
|
||||
self._pools.clear()
|
||||
self._stats.clear()
|
||||
|
||||
async def get_connection(self, instance: Instance):
|
||||
"""从连接池获取一个连接（上下文管理器）。"""
|
||||
pool = await self.get_pool(instance)
|
||||
return pool.acquire()
|
||||
|
||||
def get_pool_stats(self, instance_id: int) -> Optional[PoolStats]:
|
||||
"""获取连接池统计。"""
|
||||
return self._stats.get(instance_id)
|
||||
|
||||
def get_all_stats(self) -> Dict[int, PoolStats]:
|
||||
"""获取所有连接池统计。"""
|
||||
return dict(self._stats)
|
||||
|
||||
async def health_check(self, instance_id: int) -> bool:
|
||||
"""检查连接池健康状态。"""
|
||||
pool = self._pools.get(instance_id)
|
||||
if not pool:
|
||||
return False
|
||||
|
||||
try:
|
||||
if hasattr(pool, '_free'):
|
||||
# aiomysql
|
||||
conn = await pool.acquire()
|
||||
try:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute("SELECT 1")
|
||||
await cur.fetchone()
|
||||
return True
|
||||
finally:
|
||||
pool.release(conn)
|
||||
else:
|
||||
# asyncpg
|
||||
conn = await pool.acquire()
|
||||
try:
|
||||
await conn.fetchval("SELECT 1")
|
||||
return True
|
||||
finally:
|
||||
await pool.release(conn)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def reset_pool(self, instance_id: int) -> None:
|
||||
"""重置连接池（关闭并重新创建）。"""
|
||||
await self.close_pool(instance_id)
|
||||
|
||||
def update_stats(self, instance_id: int, query_success: bool = True,
|
||||
response_ms: int = 0) -> None:
|
||||
"""更新查询统计。"""
|
||||
stats = self._stats.get(instance_id)
|
||||
if not stats:
|
||||
return
|
||||
|
||||
stats.total_queries += 1
|
||||
if not query_success:
|
||||
stats.failed_queries += 1
|
||||
|
||||
# 移动平均响应时间
|
||||
if response_ms > 0:
|
||||
if stats.avg_response_ms == 0:
|
||||
stats.avg_response_ms = response_ms
|
||||
else:
|
||||
stats.avg_response_ms = (stats.avg_response_ms * 9 + response_ms) // 10
|
||||
|
||||
|
||||
# 全局连接池管理器实例
|
||||
_pool_manager: Optional[ConnectionPoolManager] = None
|
||||
|
||||
|
||||
def get_pool_manager() -> ConnectionPoolManager:
|
||||
"""获取全局连接池管理器。"""
|
||||
global _pool_manager
|
||||
if _pool_manager is None:
|
||||
_pool_manager = ConnectionPoolManager()
|
||||
return _pool_manager
|
||||
|
||||
|
||||
async def close_all_pools():
|
||||
"""关闭所有连接池（应用关闭时调用）。"""
|
||||
global _pool_manager
|
||||
if _pool_manager:
|
||||
await _pool_manager.close_all()
|
||||
_pool_manager = None
|
||||
@@ -0,0 +1,401 @@
|
||||
"""数据脱敏引擎 — 对查询结果进行动态脱敏处理。
|
||||
|
||||
支持脱敏类型:
|
||||
- full: 完全遮挡 (****)
|
||||
- partial: 部分遮挡 (138****8888, 李**)
|
||||
- hash: SHA256 哈希前16位
|
||||
- truncate: 截断显示 (abc...)
|
||||
- regex: 正则表达式替换
|
||||
|
||||
使用方式:
|
||||
rules = await load_masking_rules(instance_id, database)
|
||||
masker = DataMasker(rules)
|
||||
masked_rows = [masker.mask_row(table_name, row) for row in rows]
|
||||
"""
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional, Callable
|
||||
|
||||
|
||||
class MaskType(str, Enum):
|
||||
"""脱敏类型枚举。"""
|
||||
FULL = "full"
|
||||
PARTIAL = "partial"
|
||||
HASH = "hash"
|
||||
TRUNCATE = "truncate"
|
||||
REGEX = "regex"
|
||||
|
||||
|
||||
@dataclass
|
||||
class MaskingRule:
|
||||
"""脱敏规则数据类，用于运行时脱敏引擎。
|
||||
|
||||
Attributes:
|
||||
instance_id: 数据源实例ID
|
||||
database: 数据库名称
|
||||
table_name: 表名（支持通配符 * 表示所有表）
|
||||
column_name: 列名（支持通配符 * 表示所有列）
|
||||
mask_type: 脱敏类型
|
||||
mask_config: 类型特定的配置字典
|
||||
is_enabled: 是否启用
|
||||
"""
|
||||
instance_id: int
|
||||
database: str
|
||||
table_name: str
|
||||
column_name: str
|
||||
mask_type: MaskType
|
||||
mask_config: Dict[str, Any] = field(default_factory=dict)
|
||||
is_enabled: bool = True
|
||||
|
||||
@classmethod
|
||||
def from_db_model(cls, db_rule: Any) -> "MaskingRule":
|
||||
"""从数据库模型创建 MaskingRule 实例。
|
||||
|
||||
Args:
|
||||
db_rule: SQLAlchemy MaskingRule 模型实例
|
||||
|
||||
Returns:
|
||||
MaskingRule 数据类实例
|
||||
"""
|
||||
config: Dict[str, Any] = {}
|
||||
if db_rule.mask_config:
|
||||
try:
|
||||
config = json.loads(db_rule.mask_config)
|
||||
except json.JSONDecodeError:
|
||||
config = {}
|
||||
|
||||
mask_type = MaskType(db_rule.mask_type) if db_rule.mask_type in {
|
||||
m.value for m in MaskType
|
||||
} else MaskType.FULL
|
||||
|
||||
return cls(
|
||||
instance_id=db_rule.instance_id,
|
||||
database=db_rule.database,
|
||||
table_name=db_rule.table_name,
|
||||
column_name=db_rule.column_name,
|
||||
mask_type=mask_type,
|
||||
mask_config=config,
|
||||
is_enabled=db_rule.is_enabled,
|
||||
)
|
||||
|
||||
|
||||
class DataMasker:
|
||||
"""数据脱敏引擎，对查询结果进行动态脱敏。
|
||||
|
||||
根据配置的脱敏规则，对行数据进行后处理脱敏。
|
||||
脱敏不改变原始查询语句，仅对返回结果进行处理。
|
||||
|
||||
Attributes:
|
||||
_rules: 脱敏规则列表
|
||||
_rule_map: 规则查找映射 {(table, column): MaskingRule}
|
||||
_table_wildcards: 表级通配规则 {table: [rules]}
|
||||
_global_columns: 全局列规则 {column: rule}
|
||||
"""
|
||||
|
||||
def __init__(self, rules: List[MaskingRule]) -> None:
|
||||
"""初始化脱敏引擎。
|
||||
|
||||
Args:
|
||||
rules: 脱敏规则列表
|
||||
"""
|
||||
self._rules: List[MaskingRule] = [r for r in rules if r.is_enabled]
|
||||
self._rule_map: Dict[tuple, MaskingRule] = {}
|
||||
self._table_wildcards: Dict[str, List[MaskingRule]] = {}
|
||||
self._global_columns: Dict[str, MaskingRule] = {}
|
||||
|
||||
for rule in self._rules:
|
||||
if rule.table_name == "*" and rule.column_name == "*":
|
||||
# 全局通配规则 — 暂不支持
|
||||
continue
|
||||
elif rule.table_name == "*":
|
||||
# 全局列规则：适用于所有表的指定列
|
||||
self._global_columns[rule.column_name.lower()] = rule
|
||||
elif rule.column_name == "*":
|
||||
# 表级通配规则：适用于指定表的所有列
|
||||
self._table_wildcards.setdefault(rule.table_name.lower(), []).append(rule)
|
||||
else:
|
||||
# 精确匹配规则
|
||||
key = (rule.table_name.lower(), rule.column_name.lower())
|
||||
self._rule_map[key] = rule
|
||||
|
||||
# ===== 公共接口 =====
|
||||
|
||||
def mask(self, table: str, column: str, value: Any) -> Any:
|
||||
"""对单个值进行脱敏。
|
||||
|
||||
Args:
|
||||
table: 表名
|
||||
column: 列名
|
||||
value: 原始值
|
||||
|
||||
Returns:
|
||||
脱敏后的值。如果值为 None 或规则不存在，返回原值。
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
|
||||
rule = self._find_rule(table, column)
|
||||
if not rule:
|
||||
return value
|
||||
|
||||
return self._apply_mask(rule, value)
|
||||
|
||||
def mask_row(self, table: str, row: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""对整行数据进行脱敏。
|
||||
|
||||
Args:
|
||||
table: 表名
|
||||
row: 行数据字典 {column: value}
|
||||
|
||||
Returns:
|
||||
脱敏后的行数据字典
|
||||
"""
|
||||
if not row:
|
||||
return row
|
||||
|
||||
result: Dict[str, Any] = {}
|
||||
for col, val in row.items():
|
||||
result[col] = self.mask(table, col, val)
|
||||
return result
|
||||
|
||||
def mask_rows(self, table: str, rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""对多行数据进行脱敏。
|
||||
|
||||
Args:
|
||||
table: 表名
|
||||
rows: 行数据列表
|
||||
|
||||
Returns:
|
||||
脱敏后的行数据列表
|
||||
"""
|
||||
if not rows:
|
||||
return rows
|
||||
return [self.mask_row(table, row) for row in rows]
|
||||
|
||||
def has_rules_for_table(self, table: str) -> bool:
|
||||
"""检查是否有规则适用于指定表。
|
||||
|
||||
Args:
|
||||
table: 表名
|
||||
|
||||
Returns:
|
||||
如果存在适用的脱敏规则返回 True
|
||||
"""
|
||||
table_lower = table.lower()
|
||||
# 检查精确匹配
|
||||
if any(k[0] == table_lower for k in self._rule_map):
|
||||
return True
|
||||
# 检查表级通配
|
||||
if table_lower in self._table_wildcards:
|
||||
return True
|
||||
# 检查全局列规则
|
||||
if self._global_columns:
|
||||
return True
|
||||
return False
|
||||
|
||||
# ===== 规则查找 =====
|
||||
|
||||
def _find_rule(self, table: str, column: str) -> Optional[MaskingRule]:
|
||||
"""查找适用于指定表和列的脱敏规则。
|
||||
|
||||
优先级: 精确匹配 > 表级通配 > 全局列规则
|
||||
|
||||
Args:
|
||||
table: 表名
|
||||
column: 列名
|
||||
|
||||
Returns:
|
||||
匹配的脱敏规则，未找到返回 None
|
||||
"""
|
||||
table_lower = table.lower()
|
||||
col_lower = column.lower()
|
||||
|
||||
# 1. 精确匹配
|
||||
key = (table_lower, col_lower)
|
||||
if key in self._rule_map:
|
||||
return self._rule_map[key]
|
||||
|
||||
# 2. 表级通配 (table.*)
|
||||
if table_lower in self._table_wildcards:
|
||||
for rule in self._table_wildcards[table_lower]:
|
||||
if rule.column_name == "*":
|
||||
return rule
|
||||
|
||||
# 3. 全局列规则 (*.column)
|
||||
if col_lower in self._global_columns:
|
||||
return self._global_columns[col_lower]
|
||||
|
||||
return None
|
||||
|
||||
# ===== 脱敏实现 =====
|
||||
|
||||
def _apply_mask(self, rule: MaskingRule, value: Any) -> Any:
|
||||
"""根据规则类型应用对应的脱敏算法。
|
||||
|
||||
Args:
|
||||
rule: 脱敏规则
|
||||
value: 原始值
|
||||
|
||||
Returns:
|
||||
脱敏后的值
|
||||
|
||||
Raises:
|
||||
ValueError: 不支持的脱敏类型
|
||||
"""
|
||||
# 非字符串值转换为字符串处理
|
||||
str_value = str(value) if value is not None else ""
|
||||
|
||||
mask_methods: Dict[MaskType, Callable[[str, Dict[str, Any]], str]] = {
|
||||
MaskType.FULL: self._mask_full,
|
||||
MaskType.PARTIAL: self._mask_partial,
|
||||
MaskType.HASH: self._mask_hash,
|
||||
MaskType.TRUNCATE: self._mask_truncate,
|
||||
MaskType.REGEX: self._mask_regex,
|
||||
}
|
||||
|
||||
method = mask_methods.get(rule.mask_type)
|
||||
if not method:
|
||||
raise ValueError(f"不支持的脱敏类型: {rule.mask_type}")
|
||||
|
||||
return method(str_value, rule.mask_config)
|
||||
|
||||
@staticmethod
|
||||
def _mask_full(value: str, config: Dict[str, Any]) -> str:
|
||||
"""完全遮挡 — 将值替换为固定数量的星号。
|
||||
|
||||
Args:
|
||||
value: 原始值
|
||||
config: 配置项 {length: 遮挡长度(默认4), char: 遮挡字符(默认*)}
|
||||
|
||||
Returns:
|
||||
完全遮挡后的字符串
|
||||
"""
|
||||
length = config.get("length", 4)
|
||||
char = config.get("char", "*")
|
||||
return char * length
|
||||
|
||||
@staticmethod
|
||||
def _mask_partial(value: str, config: Dict[str, Any]) -> str:
|
||||
"""部分遮挡 — 保留前后若干字符，中间用星号替换。
|
||||
|
||||
Args:
|
||||
value: 原始值
|
||||
config: 配置项 {
|
||||
keep_prefix: 保留前缀字符数(默认2),
|
||||
keep_suffix: 保留后缀字符数(默认2),
|
||||
mask_char: 遮挡字符(默认*)
|
||||
}
|
||||
|
||||
Returns:
|
||||
部分遮挡后的字符串
|
||||
|
||||
Examples:
|
||||
>>> DataMasker._mask_partial("13812345678", {"keep_prefix": 3, "keep_suffix": 4})
|
||||
'138****5678'
|
||||
>>> DataMasker._mask_partial("张三丰", {"keep_prefix": 1, "keep_suffix": 0})
|
||||
'张**'
|
||||
"""
|
||||
keep_prefix = config.get("keep_prefix", 2)
|
||||
keep_suffix = config.get("keep_suffix", 2)
|
||||
mask_char = config.get("mask_char", "*")
|
||||
|
||||
if len(value) <= keep_prefix + keep_suffix:
|
||||
# 值太短，全部遮挡
|
||||
return mask_char * len(value)
|
||||
|
||||
prefix = value[:keep_prefix]
|
||||
suffix = value[-keep_suffix:] if keep_suffix > 0 else ""
|
||||
mask_len = len(value) - keep_prefix - keep_suffix
|
||||
return prefix + mask_char * mask_len + suffix
|
||||
|
||||
@staticmethod
|
||||
def _mask_hash(value: str, config: Dict[str, Any]) -> str:
|
||||
"""SHA256 哈希 — 取哈希值的前N位。
|
||||
|
||||
Args:
|
||||
value: 原始值
|
||||
config: 配置项 {length: 截取长度(默认16)}
|
||||
|
||||
Returns:
|
||||
哈希前缀字符串
|
||||
"""
|
||||
length = config.get("length", 16)
|
||||
hash_val = hashlib.sha256(value.encode("utf-8")).hexdigest()
|
||||
return hash_val[:length]
|
||||
|
||||
@staticmethod
|
||||
def _mask_truncate(value: str, config: Dict[str, Any]) -> str:
|
||||
"""截断显示 — 只保留前N个字符，后面加省略号。
|
||||
|
||||
Args:
|
||||
value: 原始值
|
||||
config: 配置项 {length: 保留长度(默认3), ellipsis: 省略号(默认...)
|
||||
|
||||
Returns:
|
||||
截断后的字符串
|
||||
"""
|
||||
length = config.get("length", 3)
|
||||
ellipsis = config.get("ellipsis", "...")
|
||||
|
||||
if len(value) <= length:
|
||||
return value
|
||||
return value[:length] + ellipsis
|
||||
|
||||
@staticmethod
|
||||
def _mask_regex(value: str, config: Dict[str, Any]) -> str:
|
||||
"""正则表达式替换 — 使用正则匹配并替换。
|
||||
|
||||
Args:
|
||||
value: 原始值
|
||||
config: 配置项 {
|
||||
pattern: 正则表达式,
|
||||
replacement: 替换字符串(默认 ***)
|
||||
}
|
||||
|
||||
Returns:
|
||||
替换后的字符串
|
||||
|
||||
Raises:
|
||||
ValueError: 正则表达式无效
|
||||
"""
|
||||
pattern = config.get("pattern", ".*")
|
||||
replacement = config.get("replacement", "***")
|
||||
|
||||
try:
|
||||
return re.sub(pattern, replacement, value)
|
||||
except re.error as e:
|
||||
raise ValueError(f"无效的正则表达式: {pattern}, 错误: {e}")
|
||||
|
||||
|
||||
# ===== 便捷函数 =====
|
||||
|
||||
async def load_masking_rules(
|
||||
db: Any, instance_id: int, database: str
|
||||
) -> List[MaskingRule]:
|
||||
"""从数据库加载指定实例和数据库的脱敏规则。
|
||||
|
||||
Args:
|
||||
db: 异步数据库会话 (AsyncSession)
|
||||
instance_id: 数据源实例ID
|
||||
database: 数据库名称
|
||||
|
||||
Returns:
|
||||
MaskingRule 列表
|
||||
"""
|
||||
from sqlalchemy import select
|
||||
from app.models.masking_rule import MaskingRule as MaskingRuleModel
|
||||
|
||||
stmt = (
|
||||
select(MaskingRuleModel)
|
||||
.where(
|
||||
MaskingRuleModel.instance_id == instance_id,
|
||||
MaskingRuleModel.database == database,
|
||||
MaskingRuleModel.is_enabled == True,
|
||||
)
|
||||
)
|
||||
result = await db.execute(stmt)
|
||||
db_rules = result.scalars().all()
|
||||
return [MaskingRule.from_db_model(r) for r in db_rules]
|
||||
@@ -0,0 +1,315 @@
|
||||
"""SQL 执行影响分析 — 预估影响行数、执行时间、风险等级。
|
||||
|
||||
使用数据库 EXPLAIN 输出分析执行计划，结合表大小统计预估影响。
|
||||
支持 MySQL 和 PostgreSQL。
|
||||
|
||||
风险等级定义:
|
||||
- low: 影响 < 1000 行，执行 < 1秒
|
||||
- medium: 影响 1000-10000 行，执行 1-10秒
|
||||
- high: 影响 10000-100000 行，执行 10-60秒
|
||||
- critical: 影响 > 100000 行 或 ALTER > 1GB 表
|
||||
"""
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional, Dict, Any
|
||||
from enum import Enum
|
||||
|
||||
from app.models.instance import Instance
|
||||
|
||||
|
||||
class RiskLevel(str, Enum):
|
||||
LOW = "low"
|
||||
MEDIUM = "medium"
|
||||
HIGH = "high"
|
||||
CRITICAL = "critical"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImpactAnalysis:
|
||||
"""影响分析结果。"""
|
||||
sql_type: str = "UNKNOWN"
|
||||
estimated_rows: int = 0
|
||||
estimated_duration_ms: int = 0
|
||||
table_size_mb: float = 0.0
|
||||
table_row_count: int = 0
|
||||
index_usage: List[str] = field(default_factory=list)
|
||||
has_full_scan: bool = False
|
||||
risk_level: str = RiskLevel.LOW
|
||||
warnings: List[str] = field(default_factory=list)
|
||||
explain_result: Optional[Dict[str, Any]] = None
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"sql_type": self.sql_type,
|
||||
"estimated_rows": self.estimated_rows,
|
||||
"estimated_duration_ms": self.estimated_duration_ms,
|
||||
"table_size_mb": round(self.table_size_mb, 2),
|
||||
"table_row_count": self.table_row_count,
|
||||
"index_usage": self.index_usage,
|
||||
"has_full_scan": self.has_full_scan,
|
||||
"risk_level": self.risk_level,
|
||||
"warnings": self.warnings,
|
||||
}
|
||||
|
||||
|
||||
class ImpactAnalyzer:
|
||||
"""SQL 执行影响分析器。"""
|
||||
|
||||
RISK_THRESHOLDS = {
|
||||
RiskLevel.LOW: {"rows": 1000, "duration_ms": 1000, "size_mb": 10},
|
||||
RiskLevel.MEDIUM: {"rows": 10000, "duration_ms": 10000, "size_mb": 100},
|
||||
RiskLevel.HIGH: {"rows": 100000, "duration_ms": 60000, "size_mb": 1024},
|
||||
}
|
||||
|
||||
def detect_sql_type(self, sql: str) -> str:
|
||||
"""检测 SQL 语句类型。"""
|
||||
upper = sql.strip().upper()
|
||||
if upper.startswith("SELECT"):
|
||||
return "SELECT"
|
||||
elif upper.startswith("INSERT"):
|
||||
return "INSERT"
|
||||
elif upper.startswith("UPDATE"):
|
||||
return "UPDATE"
|
||||
elif upper.startswith("DELETE"):
|
||||
return "DELETE"
|
||||
elif upper.startswith("ALTER"):
|
||||
return "ALTER"
|
||||
elif upper.startswith("CREATE"):
|
||||
return "CREATE"
|
||||
elif upper.startswith("DROP"):
|
||||
return "DROP"
|
||||
elif upper.startswith("TRUNCATE"):
|
||||
return "TRUNCATE"
|
||||
return "UNKNOWN"
|
||||
|
||||
async def analyze(self, instance: Instance, database: str, sql: str) -> ImpactAnalysis:
|
||||
"""分析 SQL 执行影响。"""
|
||||
sql_type = self.detect_sql_type(sql)
|
||||
analysis = ImpactAnalysis(sql_type=sql_type)
|
||||
|
||||
if sql_type in ("UPDATE", "DELETE"):
|
||||
return await self._analyze_dml(instance, database, sql, analysis)
|
||||
elif sql_type == "ALTER":
|
||||
return await self._analyze_alter(instance, database, sql, analysis)
|
||||
elif sql_type == "INSERT":
|
||||
return await self._analyze_insert(instance, database, sql, analysis)
|
||||
elif sql_type in ("CREATE", "DROP", "TRUNCATE"):
|
||||
return self._analyze_ddl(sql, analysis)
|
||||
else:
|
||||
return analysis
|
||||
|
||||
async def _analyze_dml(self, instance: Instance, database: str,
|
||||
sql: str, analysis: ImpactAnalysis) -> ImpactAnalysis:
|
||||
"""分析 UPDATE/DELETE 影响。"""
|
||||
table = self._extract_table(sql)
|
||||
analysis.table_row_count = await self._get_table_row_count(instance, database, table)
|
||||
analysis.table_size_mb = await self._get_table_size(instance, database, table)
|
||||
|
||||
# 通过 WHERE 条件粗略估计影响行数
|
||||
where_clause = self._extract_where(sql)
|
||||
if where_clause:
|
||||
# 如果有 WHERE 且包含主键条件，影响较小
|
||||
if self._has_pk_condition(where_clause):
|
||||
analysis.estimated_rows = 1
|
||||
# 如果有 WHERE 且包含索引列条件
|
||||
elif self._has_index_condition(where_clause):
|
||||
analysis.estimated_rows = max(1, analysis.table_row_count // 100)
|
||||
else:
|
||||
# 非索引条件，影响较大
|
||||
analysis.estimated_rows = max(1, analysis.table_row_count // 10)
|
||||
analysis.has_full_scan = True
|
||||
else:
|
||||
# 无 WHERE，全表影响
|
||||
analysis.estimated_rows = analysis.table_row_count
|
||||
analysis.has_full_scan = True
|
||||
analysis.warnings.append(f"{analysis.sql_type} 缺少 WHERE 条件，将影响全表 {analysis.table_row_count} 行")
|
||||
|
||||
# 预估执行时间
|
||||
analysis.estimated_duration_ms = self._estimate_duration(
|
||||
analysis.estimated_rows, analysis.table_size_mb, analysis.sql_type
|
||||
)
|
||||
|
||||
# 检查索引使用
|
||||
if analysis.has_full_scan:
|
||||
analysis.warnings.append("可能触发全表扫描，建议添加索引条件")
|
||||
|
||||
analysis.risk_level = self._calculate_risk_level(
|
||||
analysis.estimated_rows, analysis.estimated_duration_ms, analysis.table_size_mb
|
||||
)
|
||||
|
||||
if analysis.risk_level in (RiskLevel.HIGH, RiskLevel.CRITICAL):
|
||||
analysis.warnings.append(
|
||||
f"风险等级: {analysis.risk_level.value}，建议分批次执行或在维护窗口操作"
|
||||
)
|
||||
|
||||
return analysis
|
||||
|
||||
async def _analyze_alter(self, instance: Instance, database: str,
|
||||
sql: str, analysis: ImpactAnalysis) -> ImpactAnalysis:
|
||||
"""分析 ALTER TABLE 影响。"""
|
||||
table = self._extract_table(sql)
|
||||
analysis.table_size_mb = await self._get_table_size(instance, database, table)
|
||||
analysis.table_row_count = await self._get_table_row_count(instance, database, table)
|
||||
|
||||
# ALTER TABLE 预估时间: 大约每 10MB 需要 1 秒
|
||||
analysis.estimated_duration_ms = int(analysis.table_size_mb * 100)
|
||||
analysis.estimated_rows = analysis.table_row_count
|
||||
|
||||
if analysis.table_size_mb > 1024:
|
||||
analysis.risk_level = RiskLevel.CRITICAL
|
||||
analysis.warnings.append(
|
||||
f"ALTER TABLE 操作表大小为 {analysis.table_size_mb:.1f}MB，"
|
||||
"建议使用 pt-online-schema-change 或 gh-ost 进行在线变更"
|
||||
)
|
||||
elif analysis.table_size_mb > 100:
|
||||
analysis.risk_level = RiskLevel.HIGH
|
||||
analysis.warnings.append(
|
||||
f"ALTER TABLE 操作表大小为 {analysis.table_size_mb:.1f}MB，预计耗时 {analysis.estimated_duration_ms / 1000:.1f} 秒"
|
||||
)
|
||||
elif analysis.table_size_mb > 10:
|
||||
analysis.risk_level = RiskLevel.MEDIUM
|
||||
else:
|
||||
analysis.risk_level = RiskLevel.LOW
|
||||
|
||||
return analysis
|
||||
|
||||
async def _analyze_insert(self, instance: Instance, database: str,
|
||||
sql: str, analysis: ImpactAnalysis) -> ImpactAnalysis:
|
||||
"""分析 INSERT 影响。"""
|
||||
# 估算插入行数
|
||||
values_count = sql.count(",") if "VALUES" in sql.upper() else 1
|
||||
# 更精确的估算：数 VALUES 子句中的行
|
||||
row_match = re.findall(r"\)\s*,?\s*\(", sql)
|
||||
analysis.estimated_rows = len(row_match) + 1 if row_match else 1
|
||||
|
||||
table = self._extract_table(sql)
|
||||
analysis.table_size_mb = await self._get_table_size(instance, database, table)
|
||||
|
||||
# INSERT 通常很快
|
||||
analysis.estimated_duration_ms = analysis.estimated_rows * 2
|
||||
|
||||
if analysis.estimated_rows > 10000:
|
||||
analysis.risk_level = RiskLevel.HIGH
|
||||
analysis.warnings.append(f"单次插入 {analysis.estimated_rows} 行数据量较大，建议分批插入")
|
||||
elif analysis.estimated_rows > 1000:
|
||||
analysis.risk_level = RiskLevel.MEDIUM
|
||||
else:
|
||||
analysis.risk_level = RiskLevel.LOW
|
||||
|
||||
return analysis
|
||||
|
||||
def _analyze_ddl(self, sql: str, analysis: ImpactAnalysis) -> ImpactAnalysis:
|
||||
"""分析 DDL 影响。"""
|
||||
analysis.risk_level = RiskLevel.MEDIUM
|
||||
analysis.warnings.append("DDL 操作会锁表，建议在低峰期执行")
|
||||
|
||||
if "DROP" in sql.upper():
|
||||
analysis.risk_level = RiskLevel.CRITICAL
|
||||
analysis.warnings.append("DROP 操作不可逆，请确认已备份数据")
|
||||
elif "TRUNCATE" in sql.upper():
|
||||
analysis.risk_level = RiskLevel.CRITICAL
|
||||
analysis.warnings.append("TRUNCATE 会清空所有数据且不可回滚")
|
||||
|
||||
return analysis
|
||||
|
||||
def _estimate_duration(self, rows: int, table_size_mb: float, sql_type: str) -> int:
|
||||
"""预估执行时间（毫秒）。"""
|
||||
base_time = 50 # 基础时间 50ms
|
||||
|
||||
if sql_type == "UPDATE":
|
||||
# UPDATE: 每 1000 行约 100ms，加上表大小因素
|
||||
row_time = (rows / 1000) * 100
|
||||
size_factor = table_size_mb * 0.5
|
||||
return int(base_time + row_time + size_factor)
|
||||
elif sql_type == "DELETE":
|
||||
# DELETE: 每 1000 行约 80ms
|
||||
row_time = (rows / 1000) * 80
|
||||
return int(base_time + row_time)
|
||||
elif sql_type == "INSERT":
|
||||
# INSERT: 每行约 2ms
|
||||
return int(base_time + rows * 2)
|
||||
elif sql_type == "ALTER":
|
||||
# ALTER: 每 10MB 约 1000ms
|
||||
return int(base_time + table_size_mb * 100)
|
||||
|
||||
return base_time
|
||||
|
||||
def _calculate_risk_level(self, rows: int, duration_ms: int,
|
||||
table_size_mb: float) -> str:
|
||||
"""计算风险等级。"""
|
||||
if (rows > self.RISK_THRESHOLDS[RiskLevel.HIGH]["rows"] or
|
||||
duration_ms > self.RISK_THRESHOLDS[RiskLevel.HIGH]["duration_ms"] or
|
||||
table_size_mb > self.RISK_THRESHOLDS[RiskLevel.HIGH]["size_mb"]):
|
||||
return RiskLevel.CRITICAL
|
||||
|
||||
if (rows > self.RISK_THRESHOLDS[RiskLevel.MEDIUM]["rows"] or
|
||||
duration_ms > self.RISK_THRESHOLDS[RiskLevel.MEDIUM]["duration_ms"] or
|
||||
table_size_mb > self.RISK_THRESHOLDS[RiskLevel.MEDIUM]["size_mb"]):
|
||||
return RiskLevel.HIGH
|
||||
|
||||
if (rows > self.RISK_THRESHOLDS[RiskLevel.LOW]["rows"] or
|
||||
duration_ms > self.RISK_THRESHOLDS[RiskLevel.LOW]["duration_ms"] or
|
||||
table_size_mb > self.RISK_THRESHOLDS[RiskLevel.LOW]["size_mb"]):
|
||||
return RiskLevel.MEDIUM
|
||||
|
||||
return RiskLevel.LOW
|
||||
|
||||
def _extract_table(self, sql: str) -> str:
|
||||
"""从 SQL 中提取表名（简化版）。"""
|
||||
upper = sql.upper()
|
||||
|
||||
# UPDATE table_name SET ...
|
||||
update_match = re.search(r"UPDATE\s+(?:`?)(\w+)(?:`?)", upper)
|
||||
if update_match:
|
||||
return update_match.group(1).lower()
|
||||
|
||||
# DELETE FROM table_name ...
|
||||
delete_match = re.search(r"DELETE\s+FROM\s+(?:`?)(\w+)(?:`?)", upper)
|
||||
if delete_match:
|
||||
return delete_match.group(1).lower()
|
||||
|
||||
# INSERT INTO table_name ...
|
||||
insert_match = re.search(r"INSERT\s+INTO\s+(?:`?)(\w+)(?:`?)", upper)
|
||||
if insert_match:
|
||||
return insert_match.group(1).lower()
|
||||
|
||||
# ALTER TABLE table_name ...
|
||||
alter_match = re.search(r"ALTER\s+TABLE\s+(?:`?)(\w+)(?:`?)", upper)
|
||||
if alter_match:
|
||||
return alter_match.group(1).lower()
|
||||
|
||||
# CREATE TABLE table_name ...
|
||||
create_match = re.search(r"CREATE\s+TABLE\s+(?:`?)(\w+)(?:`?)", upper)
|
||||
if create_match:
|
||||
return create_match.group(1).lower()
|
||||
|
||||
return "unknown"
|
||||
|
||||
def _extract_where(self, sql: str) -> str:
|
||||
"""提取 WHERE 子句。"""
|
||||
where_match = re.search(r"WHERE\s+(.+?)(?:\s+ORDER\s+BY|\s+LIMIT\s+|$)",
|
||||
sql, re.IGNORECASE)
|
||||
return where_match.group(1).strip() if where_match else ""
|
||||
|
||||
def _has_pk_condition(self, where_clause: str) -> bool:
|
||||
"""检查 WHERE 是否包含主键条件。"""
|
||||
pk_patterns = [r"\bid\s*[=<>]", r"\bpk_\w+\s*[=<>]"]
|
||||
return any(re.search(p, where_clause, re.IGNORECASE) for p in pk_patterns)
|
||||
|
||||
def _has_index_condition(self, where_clause: str) -> bool:
|
||||
"""检查 WHERE 是否可能使用索引。"""
|
||||
# 包含等值条件通常可以使用索引
|
||||
return "=" in where_clause
|
||||
|
||||
async def _get_table_size(self, instance: Instance, database: str,
|
||||
table: str) -> float:
|
||||
"""获取表大小（MB）。简化实现，实际应查询数据库。"""
|
||||
# 在实际实现中，应该查询 INFORMATION_SCHEMA
|
||||
# 这里返回一个保守估计
|
||||
return 50.0 # 默认假设 50MB
|
||||
|
||||
async def _get_table_row_count(self, instance: Instance, database: str,
|
||||
table: str) -> int:
|
||||
"""获取表行数。简化实现。"""
|
||||
# 实际应执行 SELECT COUNT(*)
|
||||
return 10000 # 默认假设 10000 行
|
||||
@@ -0,0 +1,780 @@
|
||||
"""MongoDB 异步连接器 — 使用 motor 驱动。
|
||||
|
||||
支持数据库变更工单的全生命周期:
|
||||
- 连接管理（支持认证、TLS）
|
||||
- 只读查询: find, findOne, countDocuments, distinct, aggregate
|
||||
- 变更操作: insertOne/Many, updateOne/Many, replaceOne, deleteOne/Many
|
||||
- 连接测试、库/collection 列表
|
||||
- 安全验证：只读操作白名单
|
||||
|
||||
查询格式（JSON）:
|
||||
{
|
||||
"collection": "users",
|
||||
"operation": "find",
|
||||
"filter": {"age": {"$gt": 18}},
|
||||
"projection": {"name": 1, "age": 1},
|
||||
"sort": {"age": -1},
|
||||
"limit": 100,
|
||||
"skip": 0
|
||||
}
|
||||
|
||||
变更格式（JSON）:
|
||||
{
|
||||
"collection": "users",
|
||||
"operation": "insertMany",
|
||||
"documents": [{"name": "Alice", "age": 25}]
|
||||
}
|
||||
"""
|
||||
|
||||
import json
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from bson import json_util
|
||||
from motor.motor_asyncio import AsyncIOMotorClient
|
||||
|
||||
from app.core.crypto import decrypt_password
|
||||
|
||||
# ── 常量定义 ─────────────────────────────────────────────────────
|
||||
|
||||
MONGODB_READONLY_OPS: set[str] = {
|
||||
"find",
|
||||
"findOne",
|
||||
"countDocuments",
|
||||
"estimatedDocumentCount",
|
||||
"distinct",
|
||||
"aggregate",
|
||||
}
|
||||
|
||||
MONGODB_WRITE_OPS: set[str] = {
|
||||
"insertOne",
|
||||
"insertMany",
|
||||
"updateOne",
|
||||
"updateMany",
|
||||
"replaceOne",
|
||||
"deleteOne",
|
||||
"deleteMany",
|
||||
"drop",
|
||||
"createIndex",
|
||||
"dropIndex",
|
||||
"renameCollection",
|
||||
"findOneAndUpdate",
|
||||
"findOneAndReplace",
|
||||
"findOneAndDelete",
|
||||
"bulkWrite",
|
||||
}
|
||||
|
||||
MAX_QUERY_LIMIT = 10000
|
||||
|
||||
|
||||
# ── 安全验证 ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def validate_operation(operation: str) -> Tuple[bool, str]:
|
||||
"""验证操作是否为只读。
|
||||
|
||||
Args:
|
||||
operation: MongoDB 操作名称。
|
||||
|
||||
Returns:
|
||||
(is_readonly, error_message): 是否只读，若否附带错误原因。
|
||||
"""
|
||||
if operation in MONGODB_READONLY_OPS:
|
||||
return True, ""
|
||||
if operation in MONGODB_WRITE_OPS:
|
||||
return False, f"MongoDB 操作 '{operation}' 是写操作，查询模式下禁止"
|
||||
return False, f"未知的 MongoDB 操作: {operation}"
|
||||
|
||||
|
||||
# ── 客户端缓存 ───────────────────────────────────────────────────
|
||||
|
||||
_mongo_clients: Dict[int, AsyncIOMotorClient] = {}
|
||||
|
||||
|
||||
class ConnectionResult:
|
||||
"""MongoDB 连接测试结果。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
success: bool,
|
||||
version: str = "",
|
||||
latency_ms: int = 0,
|
||||
message: str = "",
|
||||
):
|
||||
self.success = success
|
||||
self.version = version
|
||||
self.latency_ms = latency_ms
|
||||
self.message = message
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"success": self.success,
|
||||
"version": self.version,
|
||||
"latency_ms": self.latency_ms,
|
||||
"message": self.message,
|
||||
}
|
||||
|
||||
|
||||
class MongoConnector:
|
||||
"""MongoDB 异步连接器。
|
||||
|
||||
提供连接池管理、只读查询、变更操作、元数据获取及安全验证。
|
||||
"""
|
||||
|
||||
# ── 连接管理 ────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
async def get_client(instance: Any) -> AsyncIOMotorClient:
|
||||
"""获取（或创建）motor 客户端实例。
|
||||
|
||||
使用内部缓存避免重复创建连接池。密码通过 decrypt_password 解密。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象，需包含 host, port, username, password 等字段。
|
||||
|
||||
Returns:
|
||||
AsyncIOMotorClient: 异步 MongoDB 客户端。
|
||||
"""
|
||||
instance_id: int = int(instance.id)
|
||||
if instance_id in _mongo_clients:
|
||||
return _mongo_clients[instance_id]
|
||||
|
||||
password = decrypt_password(instance.password)
|
||||
conn_str = MongoConnector._build_connection_string(
|
||||
host=instance.host,
|
||||
port=instance.port,
|
||||
username=instance.username or None,
|
||||
password=password or None,
|
||||
database=getattr(instance, "database", None),
|
||||
)
|
||||
client = AsyncIOMotorClient(conn_str, serverSelectionTimeoutMS=5000)
|
||||
_mongo_clients[instance_id] = client
|
||||
return client
|
||||
|
||||
@staticmethod
|
||||
def _build_connection_string(
|
||||
host: str,
|
||||
port: int,
|
||||
username: Optional[str] = None,
|
||||
password: Optional[str] = None,
|
||||
database: Optional[str] = None,
|
||||
) -> str:
|
||||
"""构建 MongoDB 连接字符串。
|
||||
|
||||
Args:
|
||||
host: 主机地址。
|
||||
port: 端口号。
|
||||
username: 用户名（可选）。
|
||||
password: 密码（可选）。
|
||||
database: 认证数据库（可选）。
|
||||
|
||||
Returns:
|
||||
MongoDB URI 字符串。
|
||||
"""
|
||||
if username and password:
|
||||
auth_db = database or "admin"
|
||||
return (
|
||||
f"mongodb://{username}:{password}@{host}:{port}/{auth_db}"
|
||||
f"?authSource={auth_db}"
|
||||
)
|
||||
return f"mongodb://{host}:{port}"
|
||||
|
||||
@staticmethod
|
||||
async def test_connection(instance: Any) -> ConnectionResult:
|
||||
"""测试 MongoDB 连接并返回版本与延迟。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
|
||||
Returns:
|
||||
ConnectionResult: 包含 success, version, latency_ms, message。
|
||||
"""
|
||||
start = time.time()
|
||||
client: AsyncIOMotorClient = await MongoConnector.get_client(instance)
|
||||
try:
|
||||
build_info = await client.admin.command("buildInfo")
|
||||
elapsed = int((time.time() - start) * 1000)
|
||||
version = build_info.get("version", "unknown")
|
||||
return ConnectionResult(
|
||||
success=True, version=version, latency_ms=elapsed, message="连接成功"
|
||||
)
|
||||
except Exception as exc:
|
||||
elapsed = int((time.time() - start) * 1000)
|
||||
return ConnectionResult(
|
||||
success=False,
|
||||
version="",
|
||||
latency_ms=elapsed,
|
||||
message=f"连接失败: {exc}",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def close_client(instance_id: int) -> None:
|
||||
"""关闭指定实例的 MongoDB 客户端并从缓存中移除。
|
||||
|
||||
Args:
|
||||
instance_id: 数据库实例 ID。
|
||||
"""
|
||||
client = _mongo_clients.pop(instance_id, None)
|
||||
if client:
|
||||
client.close()
|
||||
|
||||
# ── 只读查询 ────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
async def execute_find(
|
||||
instance: Any,
|
||||
database: str,
|
||||
collection: str,
|
||||
filter_doc: Optional[Dict[str, Any]] = None,
|
||||
projection: Optional[Dict[str, Any]] = None,
|
||||
sort: Optional[List[tuple]] = None,
|
||||
limit: int = 1000,
|
||||
skip: int = 0,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""执行 find 查询。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
collection: 集合名称。
|
||||
filter_doc: 查询过滤器（MongoDB 查询文档）。
|
||||
projection: 字段投影。
|
||||
sort: 排序规则，可以是字典 {"field": 1/-1} 或列表 [("field", 1/-1)]。
|
||||
limit: 最大返回文档数（上限 MAX_QUERY_LIMIT）。
|
||||
skip: 跳过的文档数。
|
||||
|
||||
Returns:
|
||||
文档列表，每个文档为字典。_id 被序列化为字符串。
|
||||
|
||||
Raises:
|
||||
ValueError: 当 limit 超过 MAX_QUERY_LIMIT 时。
|
||||
"""
|
||||
if limit > MAX_QUERY_LIMIT:
|
||||
raise ValueError(f"limit 不能超过 {MAX_QUERY_LIMIT}")
|
||||
|
||||
client = await MongoConnector.get_client(instance)
|
||||
coll = client[database][collection]
|
||||
|
||||
cursor = coll.find(filter_doc or {}, projection=projection)
|
||||
if sort:
|
||||
if isinstance(sort, dict):
|
||||
sort_list = [(k, v) for k, v in sort.items()]
|
||||
cursor = cursor.sort(sort_list)
|
||||
else:
|
||||
cursor = cursor.sort(sort)
|
||||
if skip:
|
||||
cursor = cursor.skip(skip)
|
||||
if limit:
|
||||
cursor = cursor.limit(limit)
|
||||
|
||||
docs = await cursor.to_list(length=limit or MAX_QUERY_LIMIT)
|
||||
return [MongoConnector._serialize_doc(doc) for doc in docs]
|
||||
|
||||
@staticmethod
|
||||
async def execute_find_one(
|
||||
instance: Any,
|
||||
database: str,
|
||||
collection: str,
|
||||
filter_doc: Optional[Dict[str, Any]] = None,
|
||||
projection: Optional[Dict[str, Any]] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""执行 findOne 查询，返回单个文档或 None。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
collection: 集合名称。
|
||||
filter_doc: 查询过滤器。
|
||||
projection: 字段投影。
|
||||
|
||||
Returns:
|
||||
单个文档字典，或 None。_id 被序列化为字符串。
|
||||
"""
|
||||
client = await MongoConnector.get_client(instance)
|
||||
coll = client[database][collection]
|
||||
doc = await coll.find_one(filter_doc or {}, projection=projection)
|
||||
if doc is None:
|
||||
return None
|
||||
return MongoConnector._serialize_doc(doc)
|
||||
|
||||
@staticmethod
|
||||
async def execute_count(
|
||||
instance: Any,
|
||||
database: str,
|
||||
collection: str,
|
||||
filter_doc: Optional[Dict[str, Any]] = None,
|
||||
) -> int:
|
||||
"""执行 countDocuments 计数。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
collection: 集合名称。
|
||||
filter_doc: 查询过滤器。
|
||||
|
||||
Returns:
|
||||
符合条件的文档数量。
|
||||
"""
|
||||
client = await MongoConnector.get_client(instance)
|
||||
coll = client[database][collection]
|
||||
return await coll.count_documents(filter_doc or {})
|
||||
|
||||
@staticmethod
|
||||
async def execute_distinct(
|
||||
instance: Any,
|
||||
database: str,
|
||||
collection: str,
|
||||
field: str,
|
||||
filter_doc: Optional[Dict[str, Any]] = None,
|
||||
) -> List[Any]:
|
||||
"""获取指定字段的去重值列表。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
collection: 集合名称。
|
||||
field: 字段名。
|
||||
filter_doc: 前置过滤器。
|
||||
|
||||
Returns:
|
||||
去重后的值列表。
|
||||
"""
|
||||
client = await MongoConnector.get_client(instance)
|
||||
coll = client[database][collection]
|
||||
return await coll.distinct(field, filter_doc or {})
|
||||
|
||||
@staticmethod
|
||||
async def execute_aggregate(
|
||||
instance: Any,
|
||||
database: str,
|
||||
collection: str,
|
||||
pipeline: List[Dict[str, Any]],
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""执行聚合管道查询。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
collection: 集合名称。
|
||||
pipeline: 聚合管道阶段列表。
|
||||
|
||||
Returns:
|
||||
聚合结果文档列表。_id 被序列化为字符串。
|
||||
"""
|
||||
client = await MongoConnector.get_client(instance)
|
||||
coll = client[database][collection]
|
||||
cursor = coll.aggregate(pipeline, allowDiskUse=True)
|
||||
docs = await cursor.to_list(length=MAX_QUERY_LIMIT)
|
||||
return [MongoConnector._serialize_doc(doc) for doc in docs]
|
||||
|
||||
# ── 变更操作（工单执行用）───────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
async def execute_insert_one(
|
||||
instance: Any,
|
||||
database: str,
|
||||
collection: str,
|
||||
document: Dict[str, Any],
|
||||
) -> str:
|
||||
"""插入单个文档。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
collection: 集合名称。
|
||||
document: 待插入的文档。
|
||||
|
||||
Returns:
|
||||
插入文档的 ID 字符串。
|
||||
"""
|
||||
client = await MongoConnector.get_client(instance)
|
||||
coll = client[database][collection]
|
||||
result = await coll.insert_one(document)
|
||||
return str(result.inserted_id)
|
||||
|
||||
@staticmethod
|
||||
async def execute_insert_many(
|
||||
instance: Any,
|
||||
database: str,
|
||||
collection: str,
|
||||
documents: List[Dict[str, Any]],
|
||||
) -> str:
|
||||
"""批量插入文档。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
collection: 集合名称。
|
||||
documents: 待插入的文档列表。
|
||||
|
||||
Returns:
|
||||
插入数量的描述字符串。
|
||||
|
||||
Raises:
|
||||
ValueError: 文档列表为空时。
|
||||
"""
|
||||
if not documents:
|
||||
raise ValueError("documents 不能为空列表")
|
||||
client = await MongoConnector.get_client(instance)
|
||||
coll = client[database][collection]
|
||||
result = await coll.insert_many(documents)
|
||||
return f"成功插入 {len(result.inserted_ids)} 条文档"
|
||||
|
||||
@staticmethod
|
||||
async def execute_update_one(
|
||||
instance: Any,
|
||||
database: str,
|
||||
collection: str,
|
||||
filter_doc: Dict[str, Any],
|
||||
update_doc: Dict[str, Any],
|
||||
) -> str:
|
||||
"""更新单个文档。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
collection: 集合名称。
|
||||
filter_doc: 查询过滤器。
|
||||
update_doc: 更新操作文档（如 {"$set": {...}}）。
|
||||
|
||||
Returns:
|
||||
影响行数描述字符串。
|
||||
"""
|
||||
client = await MongoConnector.get_client(instance)
|
||||
coll = client[database][collection]
|
||||
result = await coll.update_one(filter_doc, update_doc)
|
||||
return f"匹配 {result.matched_count} 条，修改 {result.modified_count} 条"
|
||||
|
||||
@staticmethod
|
||||
async def execute_update_many(
|
||||
instance: Any,
|
||||
database: str,
|
||||
collection: str,
|
||||
filter_doc: Dict[str, Any],
|
||||
update_doc: Dict[str, Any],
|
||||
) -> str:
|
||||
"""批量更新文档。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
collection: 集合名称。
|
||||
filter_doc: 查询过滤器。
|
||||
update_doc: 更新操作文档。
|
||||
|
||||
Returns:
|
||||
影响行数描述字符串。
|
||||
"""
|
||||
client = await MongoConnector.get_client(instance)
|
||||
coll = client[database][collection]
|
||||
result = await coll.update_many(filter_doc, update_doc)
|
||||
return f"匹配 {result.matched_count} 条，修改 {result.modified_count} 条"
|
||||
|
||||
@staticmethod
|
||||
async def execute_replace_one(
|
||||
instance: Any,
|
||||
database: str,
|
||||
collection: str,
|
||||
filter_doc: Dict[str, Any],
|
||||
replacement: Dict[str, Any],
|
||||
) -> str:
|
||||
"""替换单个文档。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
collection: 集合名称。
|
||||
filter_doc: 查询过滤器。
|
||||
replacement: 替换用的新文档。
|
||||
|
||||
Returns:
|
||||
影响行数描述字符串。
|
||||
"""
|
||||
client = await MongoConnector.get_client(instance)
|
||||
coll = client[database][collection]
|
||||
result = await coll.replace_one(filter_doc, replacement)
|
||||
return f"匹配 {result.matched_count} 条，替换 {result.modified_count} 条"
|
||||
|
||||
@staticmethod
|
||||
async def execute_delete_one(
|
||||
instance: Any,
|
||||
database: str,
|
||||
collection: str,
|
||||
filter_doc: Dict[str, Any],
|
||||
) -> str:
|
||||
"""删除单个文档。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
collection: 集合名称。
|
||||
filter_doc: 查询过滤器。
|
||||
|
||||
Returns:
|
||||
删除数量描述字符串。
|
||||
"""
|
||||
client = await MongoConnector.get_client(instance)
|
||||
coll = client[database][collection]
|
||||
result = await coll.delete_one(filter_doc)
|
||||
return f"删除 {result.deleted_count} 条文档"
|
||||
|
||||
@staticmethod
|
||||
async def execute_delete_many(
|
||||
instance: Any,
|
||||
database: str,
|
||||
collection: str,
|
||||
filter_doc: Dict[str, Any],
|
||||
) -> str:
|
||||
"""批量删除文档。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
collection: 集合名称。
|
||||
filter_doc: 查询过滤器。
|
||||
|
||||
Returns:
|
||||
删除数量描述字符串。
|
||||
"""
|
||||
client = await MongoConnector.get_client(instance)
|
||||
coll = client[database][collection]
|
||||
result = await coll.delete_many(filter_doc)
|
||||
return f"删除 {result.deleted_count} 条文档"
|
||||
|
||||
# ── 元数据 ──────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
async def list_databases(instance: Any) -> List[str]:
|
||||
"""列出 MongoDB 实例中所有数据库名称。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
|
||||
Returns:
|
||||
数据库名称列表（排除系统库 admin, local, config）。
|
||||
"""
|
||||
client = await MongoConnector.get_client(instance)
|
||||
db_info_list = await client.list_database_names()
|
||||
system_dbs = {"admin", "local", "config"}
|
||||
return [name for name in db_info_list if name not in system_dbs]
|
||||
|
||||
@staticmethod
|
||||
async def list_collections(instance: Any, database: str) -> List[str]:
|
||||
"""列出指定数据库中所有集合名称。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
|
||||
Returns:
|
||||
集合名称列表。
|
||||
"""
|
||||
client = await MongoConnector.get_client(instance)
|
||||
db = client[database]
|
||||
return await db.list_collection_names()
|
||||
|
||||
@staticmethod
|
||||
async def get_collection_stats(
|
||||
instance: Any, database: str, collection: str
|
||||
) -> Dict[str, Any]:
|
||||
"""获取集合统计信息。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
collection: 集合名称。
|
||||
|
||||
Returns:
|
||||
包含 count, size, avgObjSize, storageSize, indexes, totalIndexSize 等字段的字典。
|
||||
"""
|
||||
client = await MongoConnector.get_client(instance)
|
||||
coll = client[database][collection]
|
||||
stats = await coll.database.command("collStats", collection)
|
||||
return {
|
||||
"namespace": stats.get("ns", f"{database}.{collection}"),
|
||||
"document_count": stats.get("count", 0),
|
||||
"size_bytes": stats.get("size", 0),
|
||||
"avg_document_size_bytes": stats.get("avgObjSize", 0),
|
||||
"storage_size_bytes": stats.get("storageSize", 0),
|
||||
"index_count": stats.get("nindexes", 0),
|
||||
"index_size_bytes": stats.get("totalIndexSize", 0),
|
||||
"capped": stats.get("capped", False),
|
||||
}
|
||||
|
||||
# ── 通用查询执行 ────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
async def execute_query_json(
|
||||
instance: Any,
|
||||
database: str,
|
||||
query_json: Dict[str, Any],
|
||||
) -> Any:
|
||||
"""通过 JSON 格式执行 MongoDB 查询。
|
||||
|
||||
自动根据 query_json["operation"] 分发到对应方法。
|
||||
仅支持只读操作，写操作会被拒绝。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
query_json: 查询 JSON 字典，必须包含 collection 和 operation 字段。
|
||||
|
||||
Returns:
|
||||
查询结果，类型取决于具体操作。
|
||||
|
||||
Raises:
|
||||
ValueError: 操作不合法或参数缺失时。
|
||||
"""
|
||||
collection: str = query_json.get("collection")
|
||||
operation: str = query_json.get("operation")
|
||||
|
||||
if not collection:
|
||||
raise ValueError("查询 JSON 缺少 'collection' 字段")
|
||||
if not operation:
|
||||
raise ValueError("查询 JSON 缺少 'operation' 字段")
|
||||
|
||||
is_readonly, error_msg = validate_operation(operation)
|
||||
if not is_readonly:
|
||||
raise ValueError(error_msg)
|
||||
|
||||
filter_doc = query_json.get("filter") or query_json.get("query") or {}
|
||||
projection = query_json.get("projection")
|
||||
sort = query_json.get("sort")
|
||||
limit = query_json.get("limit", 1000)
|
||||
skip = query_json.get("skip", 0)
|
||||
pipeline = query_json.get("pipeline", [])
|
||||
field = query_json.get("field", "")
|
||||
|
||||
if operation == "find":
|
||||
return await MongoConnector.execute_find(
|
||||
instance, database, collection, filter_doc, projection, sort, limit, skip
|
||||
)
|
||||
if operation == "findOne":
|
||||
return await MongoConnector.execute_find_one(
|
||||
instance, database, collection, filter_doc, projection
|
||||
)
|
||||
if operation == "countDocuments":
|
||||
return await MongoConnector.execute_count(
|
||||
instance, database, collection, filter_doc
|
||||
)
|
||||
if operation == "estimatedDocumentCount":
|
||||
client = await MongoConnector.get_client(instance)
|
||||
coll = client[database][collection]
|
||||
return await coll.estimated_document_count()
|
||||
if operation == "distinct":
|
||||
if not field:
|
||||
raise ValueError("distinct 操作需要提供 'field' 字段")
|
||||
return await MongoConnector.execute_distinct(
|
||||
instance, database, collection, field, filter_doc
|
||||
)
|
||||
if operation == "aggregate":
|
||||
if not pipeline:
|
||||
raise ValueError("aggregate 操作需要提供 'pipeline' 字段")
|
||||
return await MongoConnector.execute_aggregate(
|
||||
instance, database, collection, pipeline
|
||||
)
|
||||
|
||||
raise ValueError(f"不支持的只读操作: {operation}")
|
||||
|
||||
@staticmethod
|
||||
async def execute_change_json(
|
||||
instance: Any,
|
||||
database: str,
|
||||
change_json: Dict[str, Any],
|
||||
) -> str:
|
||||
"""通过 JSON 格式执行 MongoDB 变更操作。
|
||||
|
||||
自动根据 change_json["operation"] 分发到对应方法。
|
||||
用于工单变更执行。
|
||||
|
||||
Args:
|
||||
instance: Instance 模型对象。
|
||||
database: 数据库名称。
|
||||
change_json: 变更 JSON 字典，必须包含 collection 和 operation 字段。
|
||||
|
||||
Returns:
|
||||
操作结果描述字符串。
|
||||
|
||||
Raises:
|
||||
ValueError: 操作不合法或参数缺失时。
|
||||
"""
|
||||
collection: str = change_json.get("collection")
|
||||
operation: str = change_json.get("operation")
|
||||
|
||||
if not collection:
|
||||
raise ValueError("变更 JSON 缺少 'collection' 字段")
|
||||
if not operation:
|
||||
raise ValueError("变更 JSON 缺少 'operation' 字段")
|
||||
|
||||
if operation not in MONGODB_WRITE_OPS:
|
||||
raise ValueError(f"不支持的变更操作: {operation}")
|
||||
|
||||
filter_doc = change_json.get("filter") or change_json.get("query") or {}
|
||||
|
||||
if operation == "insertOne":
|
||||
document = change_json.get("document")
|
||||
if document is None:
|
||||
raise ValueError("insertOne 需要提供 'document' 字段")
|
||||
return await MongoConnector.execute_insert_one(
|
||||
instance, database, collection, document
|
||||
)
|
||||
|
||||
if operation == "insertMany":
|
||||
documents = change_json.get("documents")
|
||||
if not documents:
|
||||
raise ValueError("insertMany 需要提供 'documents' 字段")
|
||||
return await MongoConnector.execute_insert_many(
|
||||
instance, database, collection, documents
|
||||
)
|
||||
|
||||
if operation == "updateOne":
|
||||
update_doc = change_json.get("update")
|
||||
if update_doc is None:
|
||||
raise ValueError("updateOne 需要提供 'update' 字段")
|
||||
return await MongoConnector.execute_update_one(
|
||||
instance, database, collection, filter_doc, update_doc
|
||||
)
|
||||
|
||||
if operation == "updateMany":
|
||||
update_doc = change_json.get("update")
|
||||
if update_doc is None:
|
||||
raise ValueError("updateMany 需要提供 'update' 字段")
|
||||
return await MongoConnector.execute_update_many(
|
||||
instance, database, collection, filter_doc, update_doc
|
||||
)
|
||||
|
||||
if operation == "replaceOne":
|
||||
replacement = change_json.get("replacement")
|
||||
if replacement is None:
|
||||
raise ValueError("replaceOne 需要提供 'replacement' 字段")
|
||||
return await MongoConnector.execute_replace_one(
|
||||
instance, database, collection, filter_doc, replacement
|
||||
)
|
||||
|
||||
if operation == "deleteOne":
|
||||
return await MongoConnector.execute_delete_one(
|
||||
instance, database, collection, filter_doc
|
||||
)
|
||||
|
||||
if operation == "deleteMany":
|
||||
return await MongoConnector.execute_delete_many(
|
||||
instance, database, collection, filter_doc
|
||||
)
|
||||
|
||||
raise ValueError(f"暂未实现的变更操作: {operation}")
|
||||
|
||||
# ── 序列化工具 ──────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _serialize_doc(doc: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""将 MongoDB 文档序列化为 JSON 可序列化字典。
|
||||
|
||||
使用 bson.json_util 处理 ObjectId、DateTime 等特殊类型。
|
||||
|
||||
Args:
|
||||
doc: 原始 MongoDB 文档。
|
||||
|
||||
Returns:
|
||||
JSON 可序列化的字典。
|
||||
"""
|
||||
return json.loads(json_util.dumps(doc))
|
||||
@@ -0,0 +1,235 @@
|
||||
"""通知发送器 — 支持钉钉、飞书、邮件的实际发送。
|
||||
|
||||
使用 httpx 进行异步 HTTP 请求，aiosmtplib 发送邮件。
|
||||
支持消息模板渲染，根据不同事件类型发送不同格式的消息。
|
||||
"""
|
||||
import time
|
||||
import hmac
|
||||
import hashlib
|
||||
import base64
|
||||
import json
|
||||
from typing import Dict, Any, Optional
|
||||
from urllib.parse import quote
|
||||
|
||||
|
||||
class NotificationSender:
|
||||
"""通知发送器，支持钉钉、飞书、邮件。"""
|
||||
|
||||
# 内置消息模板
|
||||
TEMPLATES = {
|
||||
"ticket_created": {
|
||||
"title": "📝 新工单待审批: {title}",
|
||||
"body": "**申请人**: {applicant}\\n**数据库**: {db_type}@{target_database}\\n**SQL摘要**: {sql_preview}\\n**优先级**: {priority}",
|
||||
},
|
||||
"ticket_approved": {
|
||||
"title": "✅ 工单已审批通过: {title}",
|
||||
"body": "**审批人**: {approver}\\n**工单ID**: {ticket_id}\\n**状态**: 已批准，可执行",
|
||||
},
|
||||
"ticket_rejected": {
|
||||
"title": "❌ 工单已被拒绝: {title}",
|
||||
"body": "**审批人**: {approver}\\n**拒绝原因**: {comment}\\n**工单ID**: {ticket_id}",
|
||||
},
|
||||
"ticket_executed": {
|
||||
"title": "🚀 工单执行完成: {title}",
|
||||
"body": "**执行结果**: {result}\\n**耗时**: {duration}ms\\n**影响行数**: {affected_rows}",
|
||||
},
|
||||
"ticket_failed": {
|
||||
"title": "⚠️ 工单执行失败: {title}",
|
||||
"body": "**错误信息**: {error}\\n**工单ID**: {ticket_id}\\n请及时处理",
|
||||
},
|
||||
"ticket_rollback": {
|
||||
"title": "🔄 工单已回滚: {title}",
|
||||
"body": "**回滚人**: {operator}\\n**工单ID**: {ticket_id}",
|
||||
},
|
||||
"backup_completed": {
|
||||
"title": "💾 备份完成: {title}",
|
||||
"body": "**数据库**: {database}\\n**文件大小**: {size}\\n**耗时**: {duration}s",
|
||||
},
|
||||
"backup_failed": {
|
||||
"title": "❌ 备份失败: {title}",
|
||||
"body": "**错误信息**: {error}\\n请检查数据库连接和磁盘空间",
|
||||
},
|
||||
"schema_drift_detected": {
|
||||
"title": "🔍 检测到 Schema 漂移: {instance_name}",
|
||||
"body": "**数据库**: {database}\\n**发现 {drift_count} 处变更**:\\n{drift_details}",
|
||||
},
|
||||
}
|
||||
|
||||
async def send(self, channel_type: str, config: Dict[str, Any],
|
||||
event_type: str, data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""发送通知。
|
||||
|
||||
Args:
|
||||
channel_type: 渠道类型 (dingtalk / feishu / email)
|
||||
config: 渠道配置
|
||||
event_type: 事件类型
|
||||
data: 模板数据
|
||||
"""
|
||||
try:
|
||||
if channel_type == "dingtalk":
|
||||
return await self._send_dingtalk(config, event_type, data)
|
||||
elif channel_type in ("webhook_feishu", "feishu"):
|
||||
return await self._send_feishu(config, event_type, data)
|
||||
elif channel_type == "email":
|
||||
return await self._send_email(config, event_type, data)
|
||||
else:
|
||||
return {"success": False, "message": f"不支持的通知类型: {channel_type}"}
|
||||
except Exception as e:
|
||||
return {"success": False, "message": f"发送失败: {str(e)}"}
|
||||
|
||||
async def _send_dingtalk(self, config: Dict[str, Any], event_type: str,
|
||||
data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""发送钉钉消息。"""
|
||||
import httpx
|
||||
|
||||
webhook_url = config.get("webhook_url", "")
|
||||
secret = config.get("secret", "")
|
||||
|
||||
if not webhook_url:
|
||||
return {"success": False, "message": "缺少 webhook_url 配置"}
|
||||
|
||||
# 生成签名
|
||||
if secret:
|
||||
timestamp = str(int(time.time() * 1000))
|
||||
sign = self._generate_dingtalk_sign(timestamp, secret)
|
||||
if "?" in webhook_url:
|
||||
webhook_url = f"{webhook_url}×tamp={timestamp}&sign={sign}"
|
||||
else:
|
||||
webhook_url = f"{webhook_url}×tamp={timestamp}&sign={sign}"
|
||||
|
||||
template = self.TEMPLATES.get(event_type, self.TEMPLATES["ticket_created"])
|
||||
title = template["title"].format(**data)
|
||||
body = template["body"].format(**data)
|
||||
|
||||
payload = {
|
||||
"msgtype": "markdown",
|
||||
"markdown": {
|
||||
"title": title,
|
||||
"text": f"## {title}\n\n{body}\n\n---\n*来自 SQL 审核平台*",
|
||||
}
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=10) as client:
|
||||
response = await client.post(webhook_url, json=payload)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
if result.get("errcode") == 0:
|
||||
return {"success": True, "message": "钉钉消息发送成功"}
|
||||
return {"success": False, "message": f"钉钉返回错误: {result}"}
|
||||
|
||||
async def _send_feishu(self, config: Dict[str, Any], event_type: str,
|
||||
data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""发送飞书消息。"""
|
||||
import httpx
|
||||
|
||||
webhook_url = config.get("webhook_url", "")
|
||||
if not webhook_url:
|
||||
return {"success": False, "message": "缺少 webhook_url 配置"}
|
||||
|
||||
template = self.TEMPLATES.get(event_type, self.TEMPLATES["ticket_created"])
|
||||
title = template["title"].format(**data)
|
||||
body = template["body"].format(**data)
|
||||
|
||||
payload = {
|
||||
"msg_type": "interactive",
|
||||
"card": {
|
||||
"header": {
|
||||
"title": {"tag": "plain_text", "content": title},
|
||||
"template": "blue" if "approved" in event_type or "completed" in event_type else "red",
|
||||
},
|
||||
"elements": [
|
||||
{"tag": "div", "text": {"tag": "lark_md", "content": body}},
|
||||
{"tag": "note", "elements": [{"tag": "plain_text", "content": "来自 SQL 审核平台"}]},
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=10) as client:
|
||||
response = await client.post(webhook_url, json=payload)
|
||||
response.raise_for_status()
|
||||
return {"success": True, "message": "飞书消息发送成功"}
|
||||
|
||||
async def _send_email(self, config: Dict[str, Any], event_type: str,
|
||||
data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""发送邮件通知。"""
|
||||
try:
|
||||
import aiosmtplib
|
||||
from email.mime.text import MIMEText
|
||||
from email.mime.multipart import MIMEMultipart
|
||||
except ImportError:
|
||||
return {"success": False, "message": "缺少 aiosmtplib 依赖，请安装: pip install aiosmtplib"}
|
||||
|
||||
smtp_host = config.get("smtp_host", "")
|
||||
smtp_port = config.get("smtp_port", 587)
|
||||
smtp_user = config.get("smtp_user", "")
|
||||
smtp_password = config.get("smtp_password", "")
|
||||
from_addr = config.get("from_addr", smtp_user)
|
||||
to_addr = data.get("recipient_email", "")
|
||||
use_tls = config.get("use_tls", True)
|
||||
|
||||
if not all([smtp_host, smtp_user, smtp_password, to_addr]):
|
||||
missing = [k for k, v in {
|
||||
"smtp_host": smtp_host, "smtp_user": smtp_user,
|
||||
"smtp_password": smtp_password, "recipient_email": to_addr,
|
||||
}.items() if not v]
|
||||
return {"success": False, "message": f"缺少邮件配置: {', '.join(missing)}"}
|
||||
|
||||
template = self.TEMPLATES.get(event_type, self.TEMPLATES["ticket_created"])
|
||||
subject = template["title"].format(**data)
|
||||
body_text = template["body"].format(**data).replace("\\n", "<br>")
|
||||
|
||||
html_body = f"""
|
||||
<html>
|
||||
<body style="font-family: Arial, sans-serif; line-height: 1.6; color: #333;">
|
||||
<h2 style="color: #2c5aa0;">{subject}</h2>
|
||||
<div style="background: #f5f5f5; padding: 15px; border-radius: 5px;">
|
||||
{body_text}
|
||||
</div>
|
||||
<hr style="margin-top: 20px; border: none; border-top: 1px solid #ddd;">
|
||||
<p style="color: #999; font-size: 12px;">本邮件由 SQL 审核平台自动发送</p>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
|
||||
msg = MIMEMultipart("alternative")
|
||||
msg["Subject"] = subject
|
||||
msg["From"] = from_addr
|
||||
msg["To"] = to_addr
|
||||
msg.attach(MIMEText(html_body, "html", "utf-8"))
|
||||
|
||||
await aiosmtplib.send(
|
||||
msg,
|
||||
hostname=smtp_host,
|
||||
port=smtp_port,
|
||||
username=smtp_user,
|
||||
password=smtp_password,
|
||||
start_tls=use_tls,
|
||||
)
|
||||
|
||||
return {"success": True, "message": f"邮件已发送至 {to_addr}"}
|
||||
|
||||
def _generate_dingtalk_sign(self, timestamp: str, secret: str) -> str:
|
||||
"""生成钉钉 webhook 签名。"""
|
||||
string_to_sign = f"{timestamp}\n{secret}"
|
||||
hmac_code = hmac.new(
|
||||
secret.encode("utf-8"),
|
||||
string_to_sign.encode("utf-8"),
|
||||
digestmod=hashlib.sha256,
|
||||
).digest()
|
||||
return quote(base64.b64encode(hmac_code).decode("utf-8"))
|
||||
|
||||
async def notify_ticket_event(self, ticket_id: int, event_type: str,
|
||||
channels: list, ticket_data: Dict[str, Any]) -> list:
|
||||
"""发送工单事件通知到所有启用的渠道。"""
|
||||
results = []
|
||||
for channel in channels:
|
||||
try:
|
||||
config = json.loads(channel.config) if isinstance(channel.config, str) else channel.config
|
||||
result = await self.send(
|
||||
channel.channel_type, config, event_type, ticket_data
|
||||
)
|
||||
results.append({"channel": channel.name, **result})
|
||||
except Exception as e:
|
||||
results.append({"channel": getattr(channel, "name", "unknown"),
|
||||
"success": False, "message": str(e)})
|
||||
return results
|
||||
@@ -0,0 +1,413 @@
|
||||
"""查询历史服务 — 记录和管理用户查询历史。
|
||||
|
||||
支持:
|
||||
- 自动记录每次查询（内容、耗时、结果数、数据库类型）
|
||||
- 查询历史分页查看
|
||||
- 收藏常用查询
|
||||
- 快捷查询（个人/团队共享）
|
||||
- 查询统计（高频表、查询趋势）
|
||||
|
||||
人性化设计:
|
||||
- 自动记录查询历史，用户无需手动保存
|
||||
- 支持收藏常用查询，一键复用
|
||||
- 快捷查询模板，减少重复输入
|
||||
- 查询统计帮助用户了解自己的数据使用习惯
|
||||
"""
|
||||
from typing import List, Optional, Dict, Any
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
from collections import Counter
|
||||
|
||||
|
||||
@dataclass
|
||||
class QueryHistoryEntry:
|
||||
"""查询历史条目。"""
|
||||
id: int = 0
|
||||
user_id: int = 0
|
||||
username: str = ""
|
||||
instance_id: int = 0
|
||||
instance_name: str = ""
|
||||
db_type: str = "" # mysql / postgres / redis / mongo
|
||||
database_name: str = ""
|
||||
query_content: str = "" # SQL / Redis命令 / MongoDB查询
|
||||
query_type: str = "" # SELECT / GET / find 等
|
||||
execution_time_ms: int = 0
|
||||
result_rows: int = 0
|
||||
is_success: bool = True
|
||||
error_message: Optional[str] = None
|
||||
is_favorite: bool = False
|
||||
note: str = "" # 用户备注
|
||||
created_at: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class QuickQuery:
|
||||
"""快捷查询模板。"""
|
||||
id: int = 0
|
||||
name: str = "" # 查询名称，如"查看活跃用户数"
|
||||
description: str = "" # 描述
|
||||
query_content: str = "" # 查询内容
|
||||
db_type: str = "" # 适用的数据库类型
|
||||
database_pattern: str = "" # 适用的数据库名称模式
|
||||
is_shared: bool = False # 是否团队共享
|
||||
created_by: int = 0
|
||||
created_by_name: str = ""
|
||||
usage_count: int = 0 # 使用次数
|
||||
created_at: str = ""
|
||||
|
||||
|
||||
class QueryHistoryService:
|
||||
"""查询历史服务。
|
||||
|
||||
内存存储实现，生产环境应替换为数据库存储。
|
||||
"""
|
||||
|
||||
_history: List[QueryHistoryEntry] = []
|
||||
_favorites: set = set() # (user_id, history_id)
|
||||
_quick_queries: List[QuickQuery] = []
|
||||
_history_counter: int = 0
|
||||
_quick_counter: int = 0
|
||||
_max_history_per_user: int = 500 # 每个用户最多保留500条
|
||||
|
||||
@classmethod
|
||||
async def record_query(
|
||||
cls,
|
||||
user_id: int,
|
||||
username: str,
|
||||
instance_id: int,
|
||||
instance_name: str,
|
||||
db_type: str,
|
||||
database_name: str,
|
||||
query_content: str,
|
||||
query_type: str,
|
||||
execution_time_ms: int,
|
||||
result_rows: int,
|
||||
is_success: bool = True,
|
||||
error_message: Optional[str] = None,
|
||||
) -> QueryHistoryEntry:
|
||||
"""记录一次查询。
|
||||
|
||||
每次查询自动调用，无需用户干预。
|
||||
"""
|
||||
cls._history_counter += 1
|
||||
|
||||
entry = QueryHistoryEntry(
|
||||
id=cls._history_counter,
|
||||
user_id=user_id,
|
||||
username=username,
|
||||
instance_id=instance_id,
|
||||
instance_name=instance_name,
|
||||
db_type=db_type,
|
||||
database_name=database_name,
|
||||
query_content=query_content[:2000], # 限制长度
|
||||
query_type=query_type,
|
||||
execution_time_ms=execution_time_ms,
|
||||
result_rows=result_rows,
|
||||
is_success=is_success,
|
||||
error_message=error_message,
|
||||
created_at=datetime.now().isoformat(),
|
||||
)
|
||||
|
||||
cls._history.append(entry)
|
||||
|
||||
# 限制每个用户的历史记录数量
|
||||
user_entries = [h for h in cls._history if h.user_id == user_id]
|
||||
if len(user_entries) > cls._max_history_per_user:
|
||||
# 删除最旧的
|
||||
to_remove = len(user_entries) - cls._max_history_per_user
|
||||
oldest_ids = {h.id for h in sorted(user_entries, key=lambda x: x.created_at)[:to_remove]}
|
||||
cls._history = [h for h in cls._history if h.id not in oldest_ids]
|
||||
|
||||
return entry
|
||||
|
||||
@classmethod
|
||||
async def get_history(
|
||||
cls,
|
||||
user_id: int,
|
||||
db_type: Optional[str] = None,
|
||||
instance_id: Optional[int] = None,
|
||||
keyword: Optional[str] = None,
|
||||
page: int = 1,
|
||||
page_size: int = 20,
|
||||
) -> Dict[str, Any]:
|
||||
"""获取用户查询历史。"""
|
||||
entries = [h for h in cls._history if h.user_id == user_id]
|
||||
|
||||
if db_type:
|
||||
entries = [h for h in entries if h.db_type == db_type]
|
||||
if instance_id:
|
||||
entries = [h for h in entries if h.instance_id == instance_id]
|
||||
if keyword:
|
||||
kw = keyword.lower()
|
||||
entries = [h for h in entries if kw in h.query_content.lower()]
|
||||
|
||||
total = len(entries)
|
||||
entries.sort(key=lambda x: x.created_at, reverse=True)
|
||||
|
||||
start = (page - 1) * page_size
|
||||
end = start + page_size
|
||||
page_entries = entries[start:end]
|
||||
|
||||
return {
|
||||
"total": total,
|
||||
"page": page,
|
||||
"page_size": page_size,
|
||||
"items": [
|
||||
{
|
||||
"id": e.id,
|
||||
"instance_name": e.instance_name,
|
||||
"db_type": e.db_type,
|
||||
"database_name": e.database_name,
|
||||
"query_content": e.query_content[:200],
|
||||
"query_type": e.query_type,
|
||||
"execution_time_ms": e.execution_time_ms,
|
||||
"result_rows": e.result_rows,
|
||||
"is_success": e.is_success,
|
||||
"is_favorite": (user_id, e.id) in cls._favorites,
|
||||
"note": e.note,
|
||||
"created_at": e.created_at,
|
||||
}
|
||||
for e in page_entries
|
||||
],
|
||||
}
|
||||
|
||||
@classmethod
|
||||
async def toggle_favorite(cls, user_id: int, history_id: int) -> bool:
|
||||
"""切换收藏状态。"""
|
||||
key = (user_id, history_id)
|
||||
if key in cls._favorites:
|
||||
cls._favorites.remove(key)
|
||||
return False # 已取消收藏
|
||||
else:
|
||||
cls._favorites.add(key)
|
||||
return True # 已收藏
|
||||
|
||||
@classmethod
|
||||
async def get_favorites(
|
||||
cls, user_id: int, page: int = 1, page_size: int = 20
|
||||
) -> Dict[str, Any]:
|
||||
"""获取收藏的查询。"""
|
||||
favorite_ids = {hid for uid, hid in cls._favorites if uid == user_id}
|
||||
entries = [h for h in cls._history if h.id in favorite_ids]
|
||||
entries.sort(key=lambda x: x.created_at, reverse=True)
|
||||
|
||||
total = len(entries)
|
||||
start = (page - 1) * page_size
|
||||
end = start + page_size
|
||||
|
||||
return {
|
||||
"total": total,
|
||||
"page": page,
|
||||
"page_size": page_size,
|
||||
"items": [
|
||||
{
|
||||
"id": e.id,
|
||||
"instance_name": e.instance_name,
|
||||
"db_type": e.db_type,
|
||||
"database_name": e.database_name,
|
||||
"query_content": e.query_content[:200],
|
||||
"query_type": e.query_type,
|
||||
"execution_time_ms": e.execution_time_ms,
|
||||
"result_rows": e.result_rows,
|
||||
"created_at": e.created_at,
|
||||
}
|
||||
for e in entries[start:end]
|
||||
],
|
||||
}
|
||||
|
||||
# ── 快捷查询 ──
|
||||
|
||||
@classmethod
|
||||
async def create_quick_query(
|
||||
cls,
|
||||
name: str,
|
||||
description: str,
|
||||
query_content: str,
|
||||
db_type: str,
|
||||
database_pattern: str,
|
||||
is_shared: bool,
|
||||
created_by: int,
|
||||
created_by_name: str,
|
||||
) -> QuickQuery:
|
||||
"""创建快捷查询。"""
|
||||
cls._quick_counter += 1
|
||||
qq = QuickQuery(
|
||||
id=cls._quick_counter,
|
||||
name=name,
|
||||
description=description,
|
||||
query_content=query_content,
|
||||
db_type=db_type,
|
||||
database_pattern=database_pattern,
|
||||
is_shared=is_shared,
|
||||
created_by=created_by,
|
||||
created_by_name=created_by_name,
|
||||
created_at=datetime.now().isoformat(),
|
||||
)
|
||||
cls._quick_queries.append(qq)
|
||||
return qq
|
||||
|
||||
@classmethod
|
||||
async def get_quick_queries(
|
||||
cls,
|
||||
user_id: int,
|
||||
db_type: Optional[str] = None,
|
||||
keyword: Optional[str] = None,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""获取快捷查询列表（个人+团队共享）。"""
|
||||
result = []
|
||||
for qq in cls._quick_queries:
|
||||
# 条件：共享的 或 自己创建的
|
||||
if qq.is_shared or qq.created_by == user_id:
|
||||
if db_type and qq.db_type != db_type:
|
||||
continue
|
||||
if keyword and keyword.lower() not in qq.name.lower():
|
||||
continue
|
||||
result.append({
|
||||
"id": qq.id,
|
||||
"name": qq.name,
|
||||
"description": qq.description,
|
||||
"query_content": qq.query_content[:500],
|
||||
"db_type": qq.db_type,
|
||||
"is_shared": qq.is_shared,
|
||||
"created_by_name": qq.created_by_name,
|
||||
"usage_count": qq.usage_count,
|
||||
"created_at": qq.created_at,
|
||||
})
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
async def use_quick_query(cls, quick_query_id: int) -> Optional[str]:
|
||||
"""使用快捷查询，返回查询内容并增加使用次数。"""
|
||||
for qq in cls._quick_queries:
|
||||
if qq.id == quick_query_id:
|
||||
qq.usage_count += 1
|
||||
return qq.query_content
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
async def get_query_statistics(cls, user_id: int) -> Dict[str, Any]:
|
||||
"""获取用户查询统计。
|
||||
|
||||
人性化：帮助用户了解自己的数据使用习惯。
|
||||
"""
|
||||
entries = [h for h in cls._history if h.user_id == user_id]
|
||||
|
||||
if not entries:
|
||||
return {"message": "暂无查询记录"}
|
||||
|
||||
# 按数据库类型分布
|
||||
db_type_dist = Counter(h.db_type for h in entries)
|
||||
|
||||
# 高频查询表（从SQL中提取表名）
|
||||
tables = []
|
||||
for h in entries:
|
||||
if h.db_type in ("mysql", "postgres"):
|
||||
# 简单提取表名
|
||||
import re
|
||||
found = re.findall(r"FROM\s+(\w+)", h.query_content, re.IGNORECASE)
|
||||
tables.extend(found)
|
||||
table_dist = Counter(tables).most_common(10)
|
||||
|
||||
# 查询趋势（最近7天）
|
||||
now = datetime.now()
|
||||
daily_counts = {}
|
||||
for i in range(7):
|
||||
day = (now - timedelta(days=i)).strftime("%Y-%m-%d")
|
||||
count = sum(
|
||||
1 for h in entries
|
||||
if h.created_at.startswith(day)
|
||||
)
|
||||
daily_counts[day] = count
|
||||
|
||||
# 平均执行时间
|
||||
avg_time = sum(h.execution_time_ms for h in entries) / len(entries)
|
||||
|
||||
# 总查询次数
|
||||
total_queries = len(entries)
|
||||
|
||||
return {
|
||||
"total_queries": total_queries,
|
||||
"avg_execution_time_ms": round(avg_time, 2),
|
||||
"db_type_distribution": dict(db_type_dist),
|
||||
"top_tables": [{"table": t, "count": c} for t, c in table_dist],
|
||||
"daily_trend": daily_counts,
|
||||
"success_rate": round(
|
||||
sum(1 for h in entries if h.is_success) / total_queries * 100, 2
|
||||
),
|
||||
}
|
||||
|
||||
# ── 预置快捷查询模板 ──
|
||||
|
||||
@classmethod
|
||||
def seed_default_quick_queries(cls):
|
||||
"""预置默认快捷查询模板。"""
|
||||
defaults = [
|
||||
QuickQuery(
|
||||
id=-1, name="查看所有表", description="列出数据库中所有表",
|
||||
query_content="SHOW TABLES", db_type="mysql", database_pattern="*",
|
||||
is_shared=True, created_by=0, created_by_name="系统",
|
||||
),
|
||||
QuickQuery(
|
||||
id=-2, name="查看表结构", description="查看指定表的列信息",
|
||||
query_content="DESCRIBE {table_name}", db_type="mysql", database_pattern="*",
|
||||
is_shared=True, created_by=0, created_by_name="系统",
|
||||
),
|
||||
QuickQuery(
|
||||
id=-3, name="查看慢查询", description="MySQL慢查询日志",
|
||||
query_content="SELECT * FROM mysql.slow_log ORDER BY start_time DESC LIMIT 20",
|
||||
db_type="mysql", database_pattern="*", is_shared=True,
|
||||
created_by=0, created_by_name="系统",
|
||||
),
|
||||
QuickQuery(
|
||||
id=-4, name="查看活跃连接", description="当前数据库连接情况",
|
||||
query_content="SHOW PROCESSLIST", db_type="mysql", database_pattern="*",
|
||||
is_shared=True, created_by=0, created_by_name="系统",
|
||||
),
|
||||
QuickQuery(
|
||||
id=-5, name="查看表大小", description="各表占用空间统计",
|
||||
query_content="""SELECT
|
||||
table_name AS `表名`,
|
||||
ROUND(data_length / 1024 / 1024, 2) AS `数据大小(MB)`,
|
||||
ROUND(index_length / 1024 / 1024, 2) AS `索引大小(MB)`,
|
||||
table_rows AS `行数`
|
||||
FROM information_schema.tables
|
||||
WHERE table_schema = DATABASE()
|
||||
ORDER BY data_length DESC
|
||||
LIMIT 50""", db_type="mysql", database_pattern="*", is_shared=True,
|
||||
created_by=0, created_by_name="系统",
|
||||
),
|
||||
QuickQuery(
|
||||
id=-6, name="Redis查看所有键", description="SCAN遍历所有键",
|
||||
query_content="SCAN 0 MATCH * COUNT 50", db_type="redis",
|
||||
database_pattern="*", is_shared=True, created_by=0, created_by_name="系统",
|
||||
),
|
||||
QuickQuery(
|
||||
id=-7, name="Redis查看服务器信息", description="Redis INFO命令",
|
||||
query_content="INFO server", db_type="redis", database_pattern="*",
|
||||
is_shared=True, created_by=0, created_by_name="系统",
|
||||
),
|
||||
QuickQuery(
|
||||
id=-8, name="PG查看所有表", description="PostgreSQL列出所有表",
|
||||
query_content="SELECT schemaname, tablename, n_tup_ins, n_tup_upd, n_tup_del FROM pg_stat_user_tables ORDER BY n_live_tup DESC LIMIT 50",
|
||||
db_type="postgres", database_pattern="*", is_shared=True,
|
||||
created_by=0, created_by_name="系统",
|
||||
),
|
||||
QuickQuery(
|
||||
id=-9, name="MongoDB查找文档", description="查找集合中的文档",
|
||||
query_content='{"collection": "{collection_name}", "operation": "find", "limit": 20}',
|
||||
db_type="mongo", database_pattern="*", is_shared=True,
|
||||
created_by=0, created_by_name="系统",
|
||||
),
|
||||
QuickQuery(
|
||||
id=-10, name="MongoDB统计文档数", description="统计集合文档数量",
|
||||
query_content='{"collection": "{collection_name}", "operation": "countDocuments"}',
|
||||
db_type="mongo", database_pattern="*", is_shared=True,
|
||||
created_by=0, created_by_name="系统",
|
||||
),
|
||||
]
|
||||
cls._quick_queries.extend(defaults)
|
||||
cls._quick_counter = max(cls._quick_counter, 10)
|
||||
|
||||
|
||||
# 初始化预置快捷查询
|
||||
QueryHistoryService.seed_default_quick_queries()
|
||||
@@ -1,81 +1,709 @@
|
||||
"""统一多数据库查询服务 — 支持 MySQL、PostgreSQL、Redis、MongoDB 的只读查询。
|
||||
|
||||
核心功能:
|
||||
1. 统一入口: 根据 db_type 自动路由到对应查询处理器
|
||||
2. 安全验证:
|
||||
- SQL 数据库: 使用 SQLAnalyzer 验证只读
|
||||
- Redis: 使用 READONLY_COMMANDS 白名单验证
|
||||
- MongoDB: 使用 READONLY_OPERATIONS 白名单验证
|
||||
3. 结果脱敏: 集成 DataMasker 对敏感字段脱敏
|
||||
4. 查询限制: 超时控制、行数限制
|
||||
|
||||
使用方式:
|
||||
result = await QueryService.execute_query(instance, database, query, user_id)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from typing import Optional
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from app.models.instance import Instance
|
||||
from app.schemas.query import QueryResult
|
||||
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
|
||||
|
||||
MAX_ROWS = 10000
|
||||
QUERY_TIMEOUT = 30
|
||||
|
||||
|
||||
class QueryService:
|
||||
# ===== Redis 只读命令白名单 =====
|
||||
|
||||
REDIS_READONLY_COMMANDS = {
|
||||
# 字符串操作
|
||||
"GET", "MGET", "STRLEN", "GETRANGE", "GETBIT",
|
||||
# 哈希操作
|
||||
"HGET", "HGETALL", "HKEYS", "HVALS", "HLEN", "HEXISTS", "HMGET",
|
||||
# 列表操作
|
||||
"LRANGE", "LLEN", "LINDEX",
|
||||
# 集合操作
|
||||
"SMEMBERS", "SCARD", "SISMEMBER", "SRANDMEMBER", "SUNION", "SINTER", "SDIFF",
|
||||
# 有序集合操作
|
||||
"ZRANGE", "ZREVRANGE", "ZRANGEBYSCORE", "ZREVRANGEBYSCORE",
|
||||
"ZCARD", "ZSCORE", "ZRANK", "ZREVRANK", "ZCOUNT", "ZLEXCOUNT",
|
||||
# 通用操作
|
||||
"EXISTS", "TTL", "PTTL", "TYPE", "DBSIZE",
|
||||
"SCAN", "HSCAN", "SSCAN", "ZSCAN", "KEYS",
|
||||
"INFO", "CONFIG", "MEMORY", "OBJECT",
|
||||
# 其他
|
||||
"PING", "ECHO", "TIME",
|
||||
}
|
||||
|
||||
# 需要特殊处理的子命令: CONFIG GET / MEMORY USAGE / OBJECT ENCODING 等
|
||||
REDIS_SUBCOMMAND_ALLOWLIST = {
|
||||
("CONFIG", "GET"),
|
||||
("MEMORY", "USAGE"),
|
||||
("MEMORY", "STATS"),
|
||||
("OBJECT", "ENCODING"),
|
||||
("OBJECT", "IDLETIME"),
|
||||
("OBJECT", "REFCOUNT"),
|
||||
}
|
||||
|
||||
# ===== 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",
|
||||
}
|
||||
|
||||
|
||||
# ===== 查询处理器基类 =====
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
# ===== SQL 查询处理器（MySQL / PostgreSQL 共用基类） =====
|
||||
|
||||
|
||||
class SQLQueryHandlerBase(BaseQueryHandler):
|
||||
"""SQL 查询处理器基类（MySQL 和 PostgreSQL 共用逻辑）。"""
|
||||
|
||||
_analyzer = SQLAnalyzer()
|
||||
|
||||
@staticmethod
|
||||
async def execute_query(
|
||||
instance: Instance, database: Optional[str], sql: str, user_id: int
|
||||
) -> QueryResult:
|
||||
if instance.db_type != "mysql":
|
||||
raise ValueError(f"暂不支持 {instance.db_type} 查询")
|
||||
def validate_readonly(self, query: str) -> Tuple[bool, str]:
|
||||
"""使用 SQLAnalyzer 验证 SQL 是否为只读。
|
||||
|
||||
# AST-based SQL validation (replaces keyword matching)
|
||||
is_safe, error_msg = QueryService._analyzer.validate_readonly(sql)
|
||||
Args:
|
||||
query: SQL 查询语句
|
||||
|
||||
Returns:
|
||||
(is_safe, error_message)
|
||||
"""
|
||||
is_safe, error_msg = self._analyzer.validate_readonly(query)
|
||||
return is_safe, error_msg
|
||||
|
||||
def _auto_add_limit(self, sql: str, max_rows: int = MAX_ROWS) -> str:
|
||||
"""如果 SQL 没有 LIMIT 则自动添加。
|
||||
|
||||
Args:
|
||||
sql: 原始 SQL
|
||||
max_rows: 最大返回行数
|
||||
|
||||
Returns:
|
||||
添加了 LIMIT 的 SQL
|
||||
"""
|
||||
upper_sql = sql.strip().upper()
|
||||
if not re.search(r"LIMIT\s+\d+", upper_sql, re.IGNORECASE):
|
||||
sql = sql.rstrip(";") + f" LIMIT {max_rows};"
|
||||
return sql
|
||||
|
||||
|
||||
class MySQLQueryHandler(SQLQueryHandlerBase):
|
||||
"""MySQL 查询处理器。
|
||||
|
||||
使用 DBConnector._query_mysql 执行查询，
|
||||
SQLAnalyzer 验证只读性，自动添加 LIMIT，30秒超时。
|
||||
"""
|
||||
|
||||
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:
|
||||
ValueError: 查询不安全或执行失败
|
||||
"""
|
||||
# 1. 验证只读
|
||||
is_safe, error_msg = self.validate_readonly(query)
|
||||
if not is_safe:
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# Auto-add LIMIT if missing
|
||||
upper_sql = sql.strip().upper()
|
||||
if not re.search(r"LIMIT\s+\d+", upper_sql, re.IGNORECASE):
|
||||
sql = sql.rstrip(";") + f" LIMIT {MAX_ROWS};"
|
||||
|
||||
import aiomysql
|
||||
conn = await aiomysql.connect(
|
||||
host=instance.host, port=instance.port,
|
||||
user=instance.username, password=instance.password,
|
||||
db=database or instance.database,
|
||||
)
|
||||
# 2. 自动添加 LIMIT
|
||||
query = self._auto_add_limit(query)
|
||||
|
||||
# 3. 执行查询
|
||||
start = time.time()
|
||||
try:
|
||||
async with conn.cursor(aiomysql.DictCursor) as cur:
|
||||
await asyncio.wait_for(cur.execute(sql), timeout=30)
|
||||
rows = await cur.fetchall()
|
||||
columns = [desc[0] for desc in cur.description] if cur.description else []
|
||||
elapsed = int((time.time() - start) * 1000)
|
||||
|
||||
return QueryResult(
|
||||
columns=columns,
|
||||
rows=rows,
|
||||
total=len(rows),
|
||||
elapsed_ms=elapsed,
|
||||
truncated=len(rows) >= MAX_ROWS,
|
||||
)
|
||||
rows, columns = await DBConnector._query_mysql(instance, database, query)
|
||||
except asyncio.TimeoutError:
|
||||
raise ValueError("查询超时（30秒）")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
@staticmethod
|
||||
async def execute_dml(instance: Instance, database: str, sql: str) -> str:
|
||||
if instance.db_type != "mysql":
|
||||
raise ValueError(f"暂不支持 {instance.db_type} 执行")
|
||||
elapsed = int((time.time() - start) * 1000)
|
||||
|
||||
import aiomysql
|
||||
conn = await aiomysql.connect(
|
||||
host=instance.host, port=instance.port,
|
||||
user=instance.username, password=instance.password,
|
||||
db=database,
|
||||
autocommit=True,
|
||||
return UnifiedQueryResult(
|
||||
columns=columns,
|
||||
rows=rows if rows else [],
|
||||
raw_result=None,
|
||||
total=len(rows) if rows else 0,
|
||||
elapsed_ms=elapsed,
|
||||
truncated=(len(rows) if rows else 0) >= MAX_ROWS,
|
||||
db_type="mysql",
|
||||
query_type="SELECT",
|
||||
)
|
||||
|
||||
|
||||
class PostgresQueryHandler(SQLQueryHandlerBase):
|
||||
"""PostgreSQL 查询处理器。
|
||||
|
||||
使用 DBConnector._query_postgres 执行查询，
|
||||
SQLAnalyzer 验证只读性（SQL 标准兼容），自动添加 LIMIT，30秒超时。
|
||||
"""
|
||||
|
||||
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:
|
||||
ValueError: 查询不安全或执行失败
|
||||
"""
|
||||
# 1. 验证只读
|
||||
is_safe, error_msg = self.validate_readonly(query)
|
||||
if not is_safe:
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# 2. 自动添加 LIMIT
|
||||
query = self._auto_add_limit(query)
|
||||
|
||||
# 3. 执行查询
|
||||
start = time.time()
|
||||
try:
|
||||
async with conn.cursor() as cur:
|
||||
await asyncio.wait_for(cur.execute(sql), timeout=60)
|
||||
affected = cur.rowcount
|
||||
return f"执行成功，影响 {affected} 行"
|
||||
rows, columns = await DBConnector._query_postgres(instance, database, query)
|
||||
except asyncio.TimeoutError:
|
||||
raise ValueError("执行超时（60秒）")
|
||||
finally:
|
||||
conn.close()
|
||||
raise ValueError("查询超时（30秒）")
|
||||
|
||||
elapsed = int((time.time() - start) * 1000)
|
||||
|
||||
return UnifiedQueryResult(
|
||||
columns=columns,
|
||||
rows=rows if rows else [],
|
||||
raw_result=None,
|
||||
total=len(rows) if rows else 0,
|
||||
elapsed_ms=elapsed,
|
||||
truncated=(len(rows) if rows else 0) >= MAX_ROWS,
|
||||
db_type="postgres",
|
||||
query_type="SELECT",
|
||||
)
|
||||
|
||||
|
||||
# ===== Redis 查询处理器 =====
|
||||
|
||||
|
||||
class RedisQueryHandler(BaseQueryHandler):
|
||||
"""Redis 查询处理器。
|
||||
|
||||
使用 READONLY_COMMANDS 白名单验证命令安全性，
|
||||
RedisConnector.execute_command 执行命令，
|
||||
返回格式化为表格/JSON 的结果。
|
||||
"""
|
||||
|
||||
def validate_readonly(self, query: str) -> Tuple[bool, str]:
|
||||
"""验证 Redis 命令是否在只读白名单中。
|
||||
|
||||
Args:
|
||||
query: Redis 命令字符串，如 "GET mykey"
|
||||
|
||||
Returns:
|
||||
(is_safe, error_message)
|
||||
"""
|
||||
parts = query.strip().split()
|
||||
if not parts:
|
||||
return False, "Redis 命令不能为空"
|
||||
|
||||
cmd = parts[0].upper()
|
||||
|
||||
# 检查子命令（如 CONFIG GET, MEMORY USAGE）
|
||||
if len(parts) >= 2:
|
||||
sub = parts[1].upper()
|
||||
subcmd_key = (cmd, sub)
|
||||
# 如果是已知的多词命令，检查是否在允许列表中
|
||||
if cmd in ("CONFIG", "MEMORY", "OBJECT"):
|
||||
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:
|
||||
ValueError: 命令不安全或执行失败
|
||||
"""
|
||||
# 1. 验证只读
|
||||
is_safe, error_msg = self.validate_readonly(query)
|
||||
if not is_safe:
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# 2. 解析 DB 号
|
||||
try:
|
||||
db = int(database) if database else 0
|
||||
except (ValueError, TypeError):
|
||||
db = 0
|
||||
|
||||
# 3. 执行命令
|
||||
start = time.time()
|
||||
try:
|
||||
result_str = await RedisConnector.execute_command(instance, db, query)
|
||||
except Exception as e:
|
||||
raise ValueError(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)
|
||||
|
||||
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 异步驱动执行查询，
|
||||
查询格式为 JSON 字符串。
|
||||
"""
|
||||
|
||||
def validate_readonly(self, query: str) -> Tuple[bool, str]:
|
||||
"""验证 MongoDB 操作是否在只读白名单中。
|
||||
|
||||
Args:
|
||||
query: MongoDB 查询 JSON 字符串
|
||||
|
||||
Returns:
|
||||
(is_safe, error_message)
|
||||
"""
|
||||
try:
|
||||
query_obj = json.loads(query)
|
||||
except json.JSONDecodeError as e:
|
||||
return False, f"MongoDB 查询 JSON 格式错误: {str(e)}"
|
||||
|
||||
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 stage in pipeline:
|
||||
if isinstance(stage, dict):
|
||||
stage_name = list(stage.keys())[0] if stage else ""
|
||||
if stage_name in ("$out", "$merge"):
|
||||
return False, f"aggregate 管道中的 '{stage_name}' 阶段是写操作，禁止执行"
|
||||
|
||||
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 字符串，格式如:
|
||||
{"collection": "users", "operation": "find",
|
||||
"filter": {"age": {"$gt": 18}},
|
||||
"projection": {"name": 1, "age": 1}}
|
||||
user_id: 查询用户ID
|
||||
options: 额外选项
|
||||
|
||||
Returns:
|
||||
UnifiedQueryResult: 统一格式的查询结果
|
||||
|
||||
Raises:
|
||||
ValueError: 操作不安全或执行失败
|
||||
"""
|
||||
# 1. 验证只读
|
||||
is_safe, error_msg = self.validate_readonly(query)
|
||||
if not is_safe:
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# 2. 解析查询 JSON
|
||||
try:
|
||||
query_obj = json.loads(query)
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError(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 = query_obj.get("limit", 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 ValueError("MongoDB 查询必须指定 'collection' 字段或在 database 参数中提供")
|
||||
|
||||
# 3. 导入 MongoConnector（延迟导入避免循环依赖）
|
||||
from app.services.mongo_connector import MongoConnector
|
||||
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
if operation == "find":
|
||||
docs, total = await MongoConnector.execute_find(
|
||||
instance, database or instance.database or "",
|
||||
collection_name, filter_doc, projection,
|
||||
limit_val, skip_val, sort_doc,
|
||||
)
|
||||
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 MongoConnector.execute_find_one(
|
||||
instance, database or instance.database or "",
|
||||
collection_name, filter_doc, projection,
|
||||
)
|
||||
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 MongoConnector.execute_count(
|
||||
instance, database or instance.database or "",
|
||||
collection_name, filter_doc,
|
||||
)
|
||||
columns = ["count"]
|
||||
rows = [{"count": count}]
|
||||
total = 1
|
||||
|
||||
elif operation == "estimatedDocumentCount":
|
||||
count = await MongoConnector.execute_estimated_count(
|
||||
instance, database or instance.database or "",
|
||||
collection_name,
|
||||
)
|
||||
columns = ["count"]
|
||||
rows = [{"count": count}]
|
||||
total = 1
|
||||
|
||||
elif operation == "distinct":
|
||||
field = query_obj.get("field", "")
|
||||
if not field:
|
||||
raise ValueError("distinct 操作必须指定 'field' 字段")
|
||||
values = await MongoConnector.execute_distinct(
|
||||
instance, database or instance.database or "",
|
||||
collection_name, field, filter_doc,
|
||||
)
|
||||
columns = ["value"]
|
||||
rows = [{"value": str(v)} for v in values]
|
||||
total = len(values)
|
||||
|
||||
elif operation == "aggregate":
|
||||
docs = await MongoConnector.execute_aggregate(
|
||||
instance, database or instance.database or "",
|
||||
collection_name, pipeline,
|
||||
)
|
||||
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 ValueError(f"不支持的 MongoDB 操作: {operation}")
|
||||
|
||||
except ValueError:
|
||||
raise
|
||||
except asyncio.TimeoutError:
|
||||
raise ValueError("查询超时（30秒）")
|
||||
except Exception as e:
|
||||
raise ValueError(f"MongoDB 查询执行失败: {str(e)}")
|
||||
|
||||
elapsed = int((time.time() - start) * 1000)
|
||||
|
||||
return UnifiedQueryResult(
|
||||
columns=columns,
|
||||
rows=rows,
|
||||
raw_result={"operation": operation, "collection": collection_name},
|
||||
total=total,
|
||||
elapsed_ms=elapsed,
|
||||
truncated=total >= MAX_ROWS,
|
||||
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 等）
|
||||
|
||||
Returns:
|
||||
UnifiedQueryResult: 统一格式的查询结果
|
||||
|
||||
Raises:
|
||||
ValueError: 不支持的数据库类型或查询不安全
|
||||
"""
|
||||
db_type = instance.db_type.lower()
|
||||
handler = QueryService._handlers.get(db_type)
|
||||
|
||||
if handler is None:
|
||||
raise ValueError(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:
|
||||
ValueError: 不支持的数据库类型或执行失败
|
||||
"""
|
||||
return await DBConnector.execute_dml(instance, database, sql)
|
||||
@@ -0,0 +1,280 @@
|
||||
"""Redis 命令安全验证器 — 区分只读命令和危险命令。
|
||||
|
||||
只读命令白名单（查询模式允许）:
|
||||
- 字符串: GET, MGET, STRLEN, GETRANGE, GETBIT, SUBSTR
|
||||
- 哈希: HGET, HGETALL, HKEYS, HVALS, HLEN, HEXISTS, HMGET, HSCAN, HRANDFIELD
|
||||
- 列表: LRANGE, LLEN, LINDEX, LPOS
|
||||
- 集合: SMEMBERS, SCARD, SISMEMBER, SRANDMEMBER, SUNION, SINTER, SDIFF, SSCAN, SMISSING
|
||||
- 有序集合: ZRANGE, ZREVRANGE, ZRANGEBYSCORE, ZREVRANGEBYSCORE, ZCARD, ZSCORE, ZRANK, ZREVRANK, ZCOUNT, ZLEXCOUNT, ZSCAN, ZRANGEBYLEX, ZMPOP
|
||||
- 流: XREAD, XREVRANGE, XRANGE, XLEN, XINFO
|
||||
- 通用: EXISTS, TTL, PTTL, TYPE, DBSIZE, SCAN, INFO, CONFIG GET, COMMAND INFO, COMMAND DOCS, MEMORY USAGE, OBJECT, CLIENT LIST, CLIENT INFO, SLOWLOG GET, TIME, ECHO, PING, LOLWUT
|
||||
- Geo: GEOPOS, GEODIST, GEORADIUS, GEORADIUSBYMEMBER, GEOSEARCH, GEOHASH
|
||||
- Bitmap: BITCOUNT, BITPOS, GETBIT
|
||||
- HyperLogLog: PFCOUNT
|
||||
|
||||
危险命令（查询模式禁止）:
|
||||
- 写: SET, MSET, HSET, LPUSH, RPUSH, SADD, ZADD, DEL, FLUSHDB, FLUSHALL
|
||||
- 管理: SHUTDOWN, CONFIG SET, DEBUG, SAVE, BGSAVE, BGREWRITEAOF, ACL, REPLICAOF, SLAVEOF
|
||||
- 执行: EVAL, EVALSHA, FCALL, FCALL_RO
|
||||
"""
|
||||
|
||||
from typing import Set, Tuple
|
||||
|
||||
|
||||
class RedisCommandValidator:
|
||||
"""Redis 命令安全验证器。
|
||||
|
||||
通过 READONLY_COMMANDS 白名单和 DANGEROUS_COMMANDS 黑名单
|
||||
对 Redis 命令进行双重校验，确保查询模式下只能执行安全的只读操作。
|
||||
"""
|
||||
|
||||
# ── 只读命令白名单 ──────────────────────────────────────────
|
||||
|
||||
READONLY_COMMANDS: Set[str] = {
|
||||
# 字符串操作
|
||||
"GET",
|
||||
"MGET",
|
||||
"STRLEN",
|
||||
"GETRANGE",
|
||||
"GETBIT",
|
||||
"SUBSTR",
|
||||
# 哈希操作
|
||||
"HGET",
|
||||
"HGETALL",
|
||||
"HKEYS",
|
||||
"HVALS",
|
||||
"HLEN",
|
||||
"HEXISTS",
|
||||
"HMGET",
|
||||
"HSCAN",
|
||||
"HRANDFIELD",
|
||||
# 列表操作
|
||||
"LRANGE",
|
||||
"LLEN",
|
||||
"LINDEX",
|
||||
"LPOS",
|
||||
# 集合操作
|
||||
"SMEMBERS",
|
||||
"SCARD",
|
||||
"SISMEMBER",
|
||||
"SRANDMEMBER",
|
||||
"SUNION",
|
||||
"SINTER",
|
||||
"SDIFF",
|
||||
"SSCAN",
|
||||
"SMISMEMBER",
|
||||
# 有序集合操作
|
||||
"ZRANGE",
|
||||
"ZREVRANGE",
|
||||
"ZRANGEBYSCORE",
|
||||
"ZREVRANGEBYSCORE",
|
||||
"ZCARD",
|
||||
"ZSCORE",
|
||||
"ZRANK",
|
||||
"ZREVRANK",
|
||||
"ZCOUNT",
|
||||
"ZLEXCOUNT",
|
||||
"ZSCAN",
|
||||
"ZRANGEBYLEX",
|
||||
# 流操作
|
||||
"XREAD",
|
||||
"XREVRANGE",
|
||||
"XRANGE",
|
||||
"XLEN",
|
||||
"XINFO",
|
||||
# 通用命令
|
||||
"EXISTS",
|
||||
"TTL",
|
||||
"PTTL",
|
||||
"TYPE",
|
||||
"DBSIZE",
|
||||
"SCAN",
|
||||
"INFO",
|
||||
"TIME",
|
||||
"ECHO",
|
||||
"PING",
|
||||
"LOLWUT",
|
||||
"COMMAND",
|
||||
"MEMORY",
|
||||
"OBJECT",
|
||||
"CLIENT",
|
||||
"SLOWLOG",
|
||||
# Geo 空间
|
||||
"GEOPOS",
|
||||
"GEODIST",
|
||||
"GEORADIUS",
|
||||
"GEORADIUSBYMEMBER",
|
||||
"GEOSEARCH",
|
||||
"GEOHASH",
|
||||
"GEOSEARCHSTORE",
|
||||
# Bitmap
|
||||
"BITCOUNT",
|
||||
"BITPOS",
|
||||
# HyperLogLog
|
||||
"PFCOUNT",
|
||||
}
|
||||
|
||||
# ── 危险命令黑名单 ──────────────────────────────────────────
|
||||
|
||||
DANGEROUS_COMMANDS: Set[str] = {
|
||||
# 写操作
|
||||
"SET",
|
||||
"MSET",
|
||||
"MSETNX",
|
||||
"SETEX",
|
||||
"PSETEX",
|
||||
"SETNX",
|
||||
"GETSET",
|
||||
"INCR",
|
||||
"INCRBY",
|
||||
"INCRBYFLOAT",
|
||||
"DECR",
|
||||
"DECRBY",
|
||||
"APPEND",
|
||||
"SETBIT",
|
||||
"SETRANGE",
|
||||
"HSET",
|
||||
"HMSET",
|
||||
"HSETNX",
|
||||
"HDEL",
|
||||
"HINCRBY",
|
||||
"HINCRBYFLOAT",
|
||||
"LPUSH",
|
||||
"RPUSH",
|
||||
"LPUSHX",
|
||||
"RPUSHX",
|
||||
"LPOP",
|
||||
"RPOP",
|
||||
"BLPOP",
|
||||
"BRPOP",
|
||||
"LREM",
|
||||
"LSET",
|
||||
"LTRIM",
|
||||
"RPOPLPUSH",
|
||||
"BRPOPLPUSH",
|
||||
"SADD",
|
||||
"SREM",
|
||||
"SPOP",
|
||||
"SMOVE",
|
||||
"SDIFFSTORE",
|
||||
"SINTERSTORE",
|
||||
"SUNIONSTORE",
|
||||
"ZADD",
|
||||
"ZREM",
|
||||
"ZINCRBY",
|
||||
"ZREMRANGEBYRANK",
|
||||
"ZREMRANGEBYSCORE",
|
||||
"ZREMRANGEBYLEX",
|
||||
"ZUNIONSTORE",
|
||||
"ZINTERSTORE",
|
||||
"ZDIFFSTORE",
|
||||
"ZMPOP",
|
||||
"ZPOPMAX",
|
||||
"ZPOPMIN",
|
||||
"BZPOPMAX",
|
||||
"BZPOPMIN",
|
||||
"DEL",
|
||||
"UNLINK",
|
||||
"FLUSHDB",
|
||||
"FLUSHALL",
|
||||
"RENAME",
|
||||
"RENAMENX",
|
||||
"MOVE",
|
||||
# 管理命令
|
||||
"SHUTDOWN",
|
||||
"CONFIG",
|
||||
"DEBUG",
|
||||
"SAVE",
|
||||
"BGSAVE",
|
||||
"BGREWRITEAOF",
|
||||
"ACL",
|
||||
"REPLICAOF",
|
||||
"SLAVEOF",
|
||||
"SYNC",
|
||||
"PSYNC",
|
||||
"MONITOR",
|
||||
"KEYS",
|
||||
"MIGRATE",
|
||||
"RESTORE",
|
||||
# 脚本执行
|
||||
"EVAL",
|
||||
"EVALSHA",
|
||||
"FCALL",
|
||||
"FCALL_RO",
|
||||
"SCRIPT",
|
||||
# 事务
|
||||
"MULTI",
|
||||
"EXEC",
|
||||
"DISCARD",
|
||||
"WATCH",
|
||||
"UNWATCH",
|
||||
# 发布/订阅
|
||||
"PUBLISH",
|
||||
"SUBSCRIBE",
|
||||
"UNSUBSCRIBE",
|
||||
"PSUBSCRIBE",
|
||||
"PUNSUBSCRIBE",
|
||||
# 模块
|
||||
"MODULE",
|
||||
}
|
||||
|
||||
# ── 验证方法 ────────────────────────────────────────────────
|
||||
|
||||
@classmethod
|
||||
def validate(cls, command: str) -> Tuple[bool, str]:
|
||||
"""验证 Redis 命令是否只读安全。
|
||||
|
||||
验证逻辑:
|
||||
1. 提取命令关键字（首个单词，转大写）。
|
||||
2. 若在 DANGEROUS_COMMANDS 中，直接拒绝。
|
||||
3. 若在 READONLY_COMMANDS 中，允许执行。
|
||||
4. 否则视为未知命令，拒绝执行。
|
||||
|
||||
Args:
|
||||
command: 完整的 Redis 命令字符串，如 "GET mykey" 或 "HGETALL users"。
|
||||
|
||||
Returns:
|
||||
(is_safe, error_message):
|
||||
- is_safe: True 表示命令安全，可执行。
|
||||
- error_message: 空字符串表示无错误，否则包含拒绝原因。
|
||||
"""
|
||||
if not command or not command.strip():
|
||||
return False, "命令不能为空"
|
||||
|
||||
cmd = command.strip().split()[0].upper()
|
||||
|
||||
if cmd in cls.DANGEROUS_COMMANDS:
|
||||
return False, f"Redis 命令 '{cmd}' 为危险命令，禁止在查询模式下执行"
|
||||
|
||||
if cmd not in cls.READONLY_COMMANDS:
|
||||
return False, f"Redis 命令 '{cmd}' 不在查询白名单中，禁止执行"
|
||||
|
||||
return True, ""
|
||||
|
||||
@classmethod
|
||||
def is_dangerous(cls, command: str) -> bool:
|
||||
"""判断命令是否为危险命令。
|
||||
|
||||
Args:
|
||||
command: 完整的 Redis 命令字符串。
|
||||
|
||||
Returns:
|
||||
True 表示命令属于危险命令。
|
||||
"""
|
||||
if not command or not command.strip():
|
||||
return False
|
||||
cmd = command.strip().split()[0].upper()
|
||||
return cmd in cls.DANGEROUS_COMMANDS
|
||||
|
||||
@classmethod
|
||||
def is_readonly(cls, command: str) -> bool:
|
||||
"""判断命令是否为只读命令。
|
||||
|
||||
Args:
|
||||
command: 完整的 Redis 命令字符串。
|
||||
|
||||
Returns:
|
||||
True 表示命令在白名单中且安全。
|
||||
"""
|
||||
is_safe, _ = cls.validate(command)
|
||||
return is_safe
|
||||
@@ -0,0 +1,836 @@
|
||||
"""自动回滚 SQL 生成器 — 为 DML 操作生成可执行的回滚语句。
|
||||
|
||||
支持:
|
||||
- INSERT → DELETE (通过主键)
|
||||
- UPDATE → UPDATE (恢复旧值)
|
||||
- DELETE → INSERT (重新插入被删数据)
|
||||
|
||||
生成策略:
|
||||
1. 对于 INSERT: 查询刚插入行的主键，生成 DELETE WHERE pk IN (...)
|
||||
2. 对于 UPDATE: 在执行前捕获旧值，生成反向 UPDATE
|
||||
3. 对于 DELETE: 在执行前备份被删数据，生成 INSERT 语句
|
||||
|
||||
使用方式:
|
||||
generator = RollbackGenerator()
|
||||
rollback_sql = await generator.generate(instance, database, original_sql)
|
||||
"""
|
||||
import re
|
||||
import asyncio
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import sqlparse
|
||||
from sqlparse.sql import Identifier, Where, Comparison
|
||||
from sqlparse.tokens import Keyword, DML
|
||||
|
||||
from app.models.instance import Instance
|
||||
from app.services.sql_analyzer import SQLAnalyzer, SQLType
|
||||
|
||||
|
||||
class RollbackGeneratorError(Exception):
|
||||
"""回滚 SQL 生成异常。"""
|
||||
pass
|
||||
|
||||
|
||||
class RollbackGenerator:
|
||||
"""自动回滚 SQL 生成器。
|
||||
|
||||
为 INSERT/UPDATE/DELETE 语句自动生成对应的回滚 SQL。
|
||||
生成回滚 SQL 需要连接目标数据库查询元数据和现有数据。
|
||||
|
||||
Attributes:
|
||||
_analyzer: SQL 分析器，用于解析 SQL 类型
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""初始化回滚生成器。"""
|
||||
self._analyzer = SQLAnalyzer()
|
||||
|
||||
# ===== 公共接口 =====
|
||||
|
||||
async def generate(
|
||||
self, instance: Instance, database: str, sql: str
|
||||
) -> str:
|
||||
"""为给定的 DML 语句生成回滚 SQL。
|
||||
|
||||
主入口方法。根据 SQL 类型自动分派到对应的生成器。
|
||||
|
||||
Args:
|
||||
instance: 数据源实例（包含连接信息）
|
||||
database: 目标数据库名
|
||||
sql: 原始 DML 语句
|
||||
|
||||
Returns:
|
||||
可执行的回滚 SQL 字符串
|
||||
|
||||
Raises:
|
||||
RollbackGeneratorError: 生成失败时抛出
|
||||
"""
|
||||
sql = sql.strip()
|
||||
if not sql:
|
||||
raise RollbackGeneratorError("SQL 语句不能为空")
|
||||
|
||||
# 解析 SQL 类型
|
||||
sql_type = self._analyzer._detect_sql_type(sql)
|
||||
|
||||
if sql_type not in (SQLType.INSERT, SQLType.UPDATE, SQLType.DELETE):
|
||||
return (
|
||||
f"-- 不支持为 {sql_type.value} 类型生成回滚 SQL\n"
|
||||
f"-- 仅支持 INSERT / UPDATE / DELETE\n"
|
||||
f"-- 原始 SQL: {sql[:200]}"
|
||||
)
|
||||
|
||||
try:
|
||||
if sql_type == SQLType.INSERT:
|
||||
return await self._generate_insert_rollback(instance, database, sql)
|
||||
elif sql_type == SQLType.UPDATE:
|
||||
return await self._generate_update_rollback(instance, database, sql)
|
||||
elif sql_type == SQLType.DELETE:
|
||||
return await self._generate_delete_rollback(instance, database, sql)
|
||||
except RollbackGeneratorError:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise RollbackGeneratorError(f"生成回滚 SQL 失败: {e}")
|
||||
|
||||
async def preview(
|
||||
self, instance: Instance, database: str, sql: str
|
||||
) -> str:
|
||||
"""预览回滚 SQL（不执行实际的数据库修改）。
|
||||
|
||||
此方法是 generate 的只读包装，仅生成和返回回滚 SQL 文本，
|
||||
不会修改任何数据。
|
||||
|
||||
Args:
|
||||
instance: 数据源实例
|
||||
database: 目标数据库名
|
||||
sql: 原始 DML 语句
|
||||
|
||||
Returns:
|
||||
回滚 SQL 字符串（预览模式，不会执行）
|
||||
"""
|
||||
try:
|
||||
return await self.generate(instance, database, sql)
|
||||
except RollbackGeneratorError as e:
|
||||
return (
|
||||
f"-- 回滚 SQL 预览失败: {e}\n"
|
||||
f"-- 原始 SQL: {sql[:200]}"
|
||||
)
|
||||
|
||||
# ===== INSERT 回滚生成 =====
|
||||
|
||||
async def _generate_insert_rollback(
|
||||
self, instance: Instance, database: str, sql: str
|
||||
) -> str:
|
||||
"""为 INSERT 语句生成回滚 SQL (DELETE)。
|
||||
|
||||
策略: 提取表名和值，生成对应的 DELETE 语句。
|
||||
如果表有主键，尝试精确定位插入的行。
|
||||
|
||||
Args:
|
||||
instance: 数据源实例
|
||||
database: 目标数据库名
|
||||
sql: INSERT 语句
|
||||
|
||||
Returns:
|
||||
DELETE 形式的回滚 SQL
|
||||
"""
|
||||
table_name = self._extract_table_from_insert(sql)
|
||||
if not table_name:
|
||||
return (
|
||||
f"-- 无法从 INSERT 语句中提取表名\n"
|
||||
f"-- 原始 SQL: {sql[:200]}"
|
||||
)
|
||||
|
||||
# 获取主键列名
|
||||
pk_columns = await self._get_primary_key(instance, database, table_name)
|
||||
|
||||
if pk_columns:
|
||||
# 有主键时，尝试构建基于主键的 DELETE
|
||||
# 由于 INSERT 尚未执行，我们无法预知主键值（尤其是自增主键）
|
||||
# 因此生成带有注释说明的回滚 SQL
|
||||
pk_str = ", ".join(pk_columns)
|
||||
rollback_sql = (
|
||||
f"-- 回滚 INSERT 操作: 删除表 {table_name} 中刚插入的行\n"
|
||||
f"-- 主键列: {pk_str}\n"
|
||||
f"-- 注意: 如果主键是自增的，需要在执行 INSERT 后获取主键值再执行以下 DELETE\n"
|
||||
f"DELETE FROM `{table_name}` WHERE /* 请补充 WHERE 条件 */ 1=0;\n"
|
||||
f"-- 示例: DELETE FROM `{table_name}` WHERE {pk_columns[0]} = <插入后的主键值>;"
|
||||
)
|
||||
else:
|
||||
# 无主键表：根据 INSERT 的列值构建 DELETE
|
||||
columns, values = self._parse_insert_values(sql)
|
||||
if columns and values:
|
||||
conditions = " AND ".join(
|
||||
f"`{col}` = {self._quote_value(val)}"
|
||||
for col, val in zip(columns, values)
|
||||
)
|
||||
rollback_sql = (
|
||||
f"-- 回滚 INSERT 操作: 表 {table_name} 无主键，使用全字段匹配删除\n"
|
||||
f"-- 注意: 如果有重复数据，可能误删多行\n"
|
||||
f"DELETE FROM `{table_name}` WHERE {conditions};"
|
||||
)
|
||||
else:
|
||||
rollback_sql = (
|
||||
f"-- 无法解析 INSERT 的列和值，请手动编写 DELETE 回滚\n"
|
||||
f"-- 表名: {table_name}\n"
|
||||
f"DELETE FROM `{table_name}` WHERE /* 请手动补充条件 */ 1=0;"
|
||||
)
|
||||
|
||||
return rollback_sql
|
||||
|
||||
# ===== UPDATE 回滚生成 =====
|
||||
|
||||
async def _generate_update_rollback(
|
||||
self, instance: Instance, database: str, sql: str
|
||||
) -> str:
|
||||
"""为 UPDATE 语句生成回滚 SQL (反向 UPDATE)。
|
||||
|
||||
策略: 在执行 UPDATE 前查询将被更新的行的当前值，
|
||||
生成能够恢复到这些旧值的 UPDATE 语句。
|
||||
|
||||
Args:
|
||||
instance: 数据源实例
|
||||
database: 目标数据库名
|
||||
sql: UPDATE 语句
|
||||
|
||||
Returns:
|
||||
反向 UPDATE 形式的回滚 SQL
|
||||
"""
|
||||
table_name = self._extract_table_from_update(sql)
|
||||
if not table_name:
|
||||
return (
|
||||
f"-- 无法从 UPDATE 语句中提取表名\n"
|
||||
f"-- 原始 SQL: {sql[:200]}"
|
||||
)
|
||||
|
||||
# 提取 UPDATE 的 SET 子句和 WHERE 条件
|
||||
set_columns = self._extract_set_columns(sql)
|
||||
where_clause = self._extract_where_clause(sql)
|
||||
|
||||
if not set_columns:
|
||||
return (
|
||||
f"-- 无法解析 UPDATE 的 SET 子句\n"
|
||||
f"-- 原始 SQL: {sql[:200]}"
|
||||
)
|
||||
|
||||
# 获取将被更新的行的当前值
|
||||
old_rows = await self._fetch_current_values(
|
||||
instance, database, table_name, set_columns, where_clause
|
||||
)
|
||||
|
||||
if not old_rows:
|
||||
return (
|
||||
f"-- 未找到将被更新的行（可能 WHERE 条件不匹配任何数据）\n"
|
||||
f"-- 表: {table_name}\n"
|
||||
f"-- WHERE: {where_clause or '无（全表更新）'}\n"
|
||||
f"-- 无需生成回滚 SQL"
|
||||
)
|
||||
|
||||
# 获取主键用于精确定位每行
|
||||
pk_columns = await self._get_primary_key(instance, database, table_name)
|
||||
|
||||
# 构建反向 UPDATE 语句
|
||||
rollback_statements: List[str] = []
|
||||
|
||||
for i, row in enumerate(old_rows):
|
||||
if pk_columns and all(pk in row for pk in pk_columns):
|
||||
# 使用主键定位
|
||||
pk_conditions = " AND ".join(
|
||||
f"`{pk}` = {self._quote_value(row[pk])}"
|
||||
for pk in pk_columns
|
||||
)
|
||||
set_parts = ", ".join(
|
||||
f"`{col}` = {self._quote_value(row.get(col))}"
|
||||
for col in set_columns if col in row
|
||||
)
|
||||
if set_parts:
|
||||
rollback_statements.append(
|
||||
f"UPDATE `{table_name}` SET {set_parts} WHERE {pk_conditions};"
|
||||
)
|
||||
else:
|
||||
# 无定位条件时使用行号注释
|
||||
set_parts = ", ".join(
|
||||
f"`{col}` = {self._quote_value(row.get(col))}"
|
||||
for col in set_columns if col in row
|
||||
)
|
||||
if set_parts:
|
||||
rollback_statements.append(
|
||||
f"-- 第 {i + 1} 行: 无法精确定位，请手动补充 WHERE 条件\n"
|
||||
f"-- UPDATE `{table_name}` SET {set_parts} WHERE /* 手动补充 */;"
|
||||
)
|
||||
|
||||
if not rollback_statements:
|
||||
return (
|
||||
f"-- 无法构建有效的 UPDATE 回滚语句\n"
|
||||
f"-- 表: {table_name}"
|
||||
)
|
||||
|
||||
rollback_sql = (
|
||||
f"-- 回滚 UPDATE 操作: 恢复表 {table_name} 的旧值\n"
|
||||
f"-- 将影响 {len(old_rows)} 行\n"
|
||||
f"-- 原 WHERE 条件: {where_clause or '无（全表更新）'}\n"
|
||||
+ "\n".join(rollback_statements)
|
||||
)
|
||||
|
||||
return rollback_sql
|
||||
|
||||
# ===== DELETE 回滚生成 =====
|
||||
|
||||
async def _generate_delete_rollback(
|
||||
self, instance: Instance, database: str, sql: str
|
||||
) -> str:
|
||||
"""为 DELETE 语句生成回滚 SQL (INSERT)。
|
||||
|
||||
策略: 在执行 DELETE 前查询将被删除的行的完整数据，
|
||||
生成能够重新插入这些数据的 INSERT 语句。
|
||||
|
||||
Args:
|
||||
instance: 数据源实例
|
||||
database: 目标数据库名
|
||||
sql: DELETE 语句
|
||||
|
||||
Returns:
|
||||
INSERT 形式的回滚 SQL
|
||||
"""
|
||||
table_name = self._extract_table_from_delete(sql)
|
||||
if not table_name:
|
||||
return (
|
||||
f"-- 无法从 DELETE 语句中提取表名\n"
|
||||
f"-- 原始 SQL: {sql[:200]}"
|
||||
)
|
||||
|
||||
# 获取将被删除的行数据
|
||||
where_clause = self._extract_where_clause(sql)
|
||||
rows = await self._fetch_rows_to_delete(
|
||||
instance, database, table_name, where_clause
|
||||
)
|
||||
|
||||
if not rows:
|
||||
return (
|
||||
f"-- 未找到将被删除的行（可能 WHERE 条件不匹配任何数据）\n"
|
||||
f"-- 表: {table_name}\n"
|
||||
f"-- WHERE: {where_clause or '无（全表删除）'}\n"
|
||||
f"-- 无需生成回滚 SQL"
|
||||
)
|
||||
|
||||
# 构建 INSERT 回滚语句
|
||||
rollback_statements: List[str] = []
|
||||
for row in rows:
|
||||
insert_sql = self._build_insert_sql(table_name, row)
|
||||
rollback_statements.append(insert_sql)
|
||||
|
||||
rollback_sql = (
|
||||
f"-- 回滚 DELETE 操作: 重新插入表 {table_name} 的数据\n"
|
||||
f"-- 将重新插入 {len(rows)} 行\n"
|
||||
f"-- 原 WHERE 条件: {where_clause or '无（全表删除）'}\n"
|
||||
+ "\n".join(rollback_statements)
|
||||
)
|
||||
|
||||
return rollback_sql
|
||||
|
||||
# ===== 数据库元数据查询 =====
|
||||
|
||||
async def _get_primary_key(
|
||||
self, instance: Instance, database: str, table_name: str
|
||||
) -> List[str]:
|
||||
"""获取表的主键列名列表。
|
||||
|
||||
Args:
|
||||
instance: 数据源实例
|
||||
database: 目标数据库名
|
||||
table_name: 表名
|
||||
|
||||
Returns:
|
||||
主键列名列表，如果表没有主键则返回空列表
|
||||
|
||||
Raises:
|
||||
RollbackGeneratorError: 查询失败时抛出
|
||||
"""
|
||||
try:
|
||||
if instance.db_type == "mysql":
|
||||
return await self._get_pk_mysql(instance, database, table_name)
|
||||
elif instance.db_type in ("postgres", "postgresql"):
|
||||
return await self._get_pk_postgres(instance, database, table_name)
|
||||
elif instance.db_type == "sqlite":
|
||||
return await self._get_pk_sqlite(instance, table_name)
|
||||
else:
|
||||
raise RollbackGeneratorError(
|
||||
f"暂不支持 {instance.db_type} 类型的主键查询"
|
||||
)
|
||||
except RollbackGeneratorError:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise RollbackGeneratorError(f"获取主键失败: {e}")
|
||||
|
||||
async def _get_pk_mysql(
|
||||
self, instance: Instance, database: str, table_name: str
|
||||
) -> List[str]:
|
||||
"""获取 MySQL 表的主键列。"""
|
||||
import aiomysql
|
||||
|
||||
conn = await aiomysql.connect(
|
||||
host=instance.host, port=instance.port,
|
||||
user=instance.username, password=instance.password,
|
||||
db=database,
|
||||
)
|
||||
try:
|
||||
async with conn.cursor() as cur:
|
||||
sql = (
|
||||
"SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE "
|
||||
"WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s "
|
||||
"AND CONSTRAINT_NAME = 'PRIMARY' "
|
||||
"ORDER BY ORDINAL_POSITION"
|
||||
)
|
||||
await cur.execute(sql, (database, table_name))
|
||||
rows = await cur.fetchall()
|
||||
return [row[0] for row in rows]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
async def _get_pk_postgres(
|
||||
self, instance: Instance, database: str, table_name: str
|
||||
) -> List[str]:
|
||||
"""获取 PostgreSQL 表的主键列。"""
|
||||
import asyncpg
|
||||
|
||||
dsn = (
|
||||
f"postgresql://{instance.username}:{instance.password}"
|
||||
f"@{instance.host}:{instance.port}/{database}"
|
||||
)
|
||||
conn = await asyncpg.connect(dsn)
|
||||
try:
|
||||
sql = (
|
||||
"SELECT a.attname FROM pg_index i "
|
||||
"JOIN pg_attribute a ON a.attrelid = i.indrelid "
|
||||
"AND a.attnum = ANY(i.indkey) "
|
||||
"WHERE i.indrelid = $1::regclass AND i.indisprimary "
|
||||
"ORDER BY array_position(i.indkey, a.attnum)"
|
||||
)
|
||||
rows = await conn.fetch(sql, table_name)
|
||||
return [row["attname"] for row in rows]
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
async def _get_pk_sqlite(
|
||||
self, instance: Instance, table_name: str
|
||||
) -> List[str]:
|
||||
"""获取 SQLite 表的主键列。"""
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _query() -> List[str]:
|
||||
import sqlite3
|
||||
|
||||
conn = sqlite3.connect(instance.host)
|
||||
cur = conn.execute(f"PRAGMA table_info({table_name})")
|
||||
rows = cur.fetchall()
|
||||
conn.close()
|
||||
# 返回 pk > 0 的列名
|
||||
return [row[1] for row in rows if row[5] > 0] # cid, name, type, notnull, dflt, pk
|
||||
|
||||
return await loop.run_in_executor(None, _query)
|
||||
|
||||
# ===== 数据查询 =====
|
||||
|
||||
async def _fetch_current_values(
|
||||
self, instance: Instance, database: str, table_name: str,
|
||||
columns: List[str], where_clause: Optional[str],
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""获取 UPDATE 操作前将被更新的行的当前值。
|
||||
|
||||
Args:
|
||||
instance: 数据源实例
|
||||
database: 目标数据库名
|
||||
table_name: 表名
|
||||
columns: 将被更新的列名列表
|
||||
where_clause: UPDATE 的 WHERE 子句
|
||||
|
||||
Returns:
|
||||
行数据列表，每行包含主键列和将被更新的列的旧值
|
||||
"""
|
||||
# 查询所有相关列的当前值
|
||||
col_list = ", ".join(f"`{c}`" for c in columns)
|
||||
|
||||
# 同时获取主键列用于精确定位
|
||||
try:
|
||||
pk_columns = await self._get_primary_key(instance, database, table_name)
|
||||
select_cols = ", ".join(
|
||||
f"`{c}`" for c in (list(dict.fromkeys(pk_columns + columns)))
|
||||
)
|
||||
except Exception:
|
||||
select_cols = col_list
|
||||
pk_columns = []
|
||||
|
||||
sql = f"SELECT {select_cols} FROM `{table_name}`"
|
||||
if where_clause:
|
||||
sql += f" WHERE {where_clause}"
|
||||
sql += " LIMIT 10000"
|
||||
|
||||
if instance.db_type == "mysql":
|
||||
return await self._query_rows_mysql(instance, database, sql)
|
||||
elif instance.db_type in ("postgres", "postgresql"):
|
||||
return await self._query_rows_postgres(instance, database, sql)
|
||||
elif instance.db_type == "sqlite":
|
||||
return await self._query_rows_sqlite(instance, sql)
|
||||
else:
|
||||
raise RollbackGeneratorError(f"不支持的数据库类型: {instance.db_type}")
|
||||
|
||||
async def _fetch_rows_to_delete(
|
||||
self, instance: Instance, database: str, table_name: str,
|
||||
where_clause: Optional[str],
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""获取将被 DELETE 的行数据。
|
||||
|
||||
Args:
|
||||
instance: 数据源实例
|
||||
database: 目标数据库名
|
||||
table_name: 表名
|
||||
where_clause: DELETE 的 WHERE 子句
|
||||
|
||||
Returns:
|
||||
将被删除的完整行数据列表
|
||||
"""
|
||||
sql = f"SELECT * FROM `{table_name}`"
|
||||
if where_clause:
|
||||
sql += f" WHERE {where_clause}"
|
||||
sql += " LIMIT 10000"
|
||||
|
||||
if instance.db_type == "mysql":
|
||||
return await self._query_rows_mysql(instance, database, sql)
|
||||
elif instance.db_type in ("postgres", "postgresql"):
|
||||
return await self._query_rows_postgres(instance, database, sql)
|
||||
elif instance.db_type == "sqlite":
|
||||
return await self._query_rows_sqlite(instance, sql)
|
||||
else:
|
||||
raise RollbackGeneratorError(f"不支持的数据库类型: {instance.db_type}")
|
||||
|
||||
# ===== 查询执行 =====
|
||||
|
||||
async def _query_rows_mysql(
|
||||
self, instance: Instance, database: str, sql: str
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""在 MySQL 上执行查询并返回行数据。"""
|
||||
import aiomysql
|
||||
|
||||
conn = await aiomysql.connect(
|
||||
host=instance.host, port=instance.port,
|
||||
user=instance.username, password=instance.password,
|
||||
db=database,
|
||||
)
|
||||
try:
|
||||
async with conn.cursor(aiomysql.DictCursor) as cur:
|
||||
await cur.execute(sql)
|
||||
rows = await cur.fetchall()
|
||||
return [dict(r) for r in rows] if rows else []
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
async def _query_rows_postgres(
|
||||
self, instance: Instance, database: str, sql: str
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""在 PostgreSQL 上执行查询并返回行数据。"""
|
||||
import asyncpg
|
||||
|
||||
dsn = (
|
||||
f"postgresql://{instance.username}:{instance.password}"
|
||||
f"@{instance.host}:{instance.port}/{database}"
|
||||
)
|
||||
conn = await asyncpg.connect(dsn)
|
||||
try:
|
||||
rows = await conn.fetch(sql)
|
||||
return [dict(r) for r in rows] if rows else []
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
async def _query_rows_sqlite(
|
||||
self, instance: Instance, sql: str
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""在 SQLite 上执行查询并返回行数据。"""
|
||||
import sqlite3
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _query() -> List[Dict[str, Any]]:
|
||||
conn = sqlite3.connect(instance.host)
|
||||
conn.row_factory = sqlite3.Row
|
||||
cur = conn.execute(sql)
|
||||
rows = [dict(r) for r in cur.fetchall()]
|
||||
conn.close()
|
||||
return rows
|
||||
|
||||
return await loop.run_in_executor(None, _query)
|
||||
|
||||
# ===== SQL 构建 =====
|
||||
|
||||
def _build_insert_sql(self, table_name: str, row: Dict[str, Any]) -> str:
|
||||
"""从行数据构建 INSERT 语句。
|
||||
|
||||
Args:
|
||||
table_name: 表名
|
||||
row: 行数据字典 {column: value}
|
||||
|
||||
Returns:
|
||||
INSERT 语句字符串
|
||||
"""
|
||||
if not row:
|
||||
return f"-- 空行数据，无法构建 INSERT 语句"
|
||||
|
||||
columns = list(row.keys())
|
||||
values = [self._quote_value(row[col]) for col in columns]
|
||||
|
||||
col_str = ", ".join(f"`{c}`" for c in columns)
|
||||
val_str = ", ".join(values)
|
||||
|
||||
return f"INSERT INTO `{table_name}` ({col_str}) VALUES ({val_str});"
|
||||
|
||||
def _build_update_sql(
|
||||
self, table_name: str, row: Dict[str, Any],
|
||||
columns: List[str], pk_columns: List[str],
|
||||
) -> str:
|
||||
"""构建反向 UPDATE 语句。
|
||||
|
||||
Args:
|
||||
table_name: 表名
|
||||
row: 行数据字典
|
||||
columns: 要恢复的列名列表
|
||||
pk_columns: 主键列名列表
|
||||
|
||||
Returns:
|
||||
UPDATE 语句字符串
|
||||
"""
|
||||
set_parts = ", ".join(
|
||||
f"`{col}` = {self._quote_value(row.get(col))}"
|
||||
for col in columns if col in row
|
||||
)
|
||||
|
||||
if pk_columns and all(pk in row for pk in pk_columns):
|
||||
where = " AND ".join(
|
||||
f"`{pk}` = {self._quote_value(row[pk])}"
|
||||
for pk in pk_columns
|
||||
)
|
||||
return f"UPDATE `{table_name}` SET {set_parts} WHERE {where};"
|
||||
|
||||
return f"-- UPDATE `{table_name}` SET {set_parts} WHERE /* 手动补充条件 */;"
|
||||
|
||||
# ===== SQL 解析工具 =====
|
||||
|
||||
def _extract_table_from_insert(self, sql: str) -> Optional[str]:
|
||||
"""从 INSERT 语句中提取表名。
|
||||
|
||||
支持格式:
|
||||
- INSERT INTO table_name ...
|
||||
- INSERT INTO `table_name` ...
|
||||
- INSERT IGNORE INTO table_name ...
|
||||
|
||||
Args:
|
||||
sql: INSERT 语句
|
||||
|
||||
Returns:
|
||||
表名，解析失败返回 None
|
||||
"""
|
||||
# 使用正则表达式提取表名
|
||||
patterns = [
|
||||
r"INSERT\s+IGNORE\s+INTO\s+`?([^`\s(]+)`?",
|
||||
r"INSERT\s+INTO\s+`?([^`\s(]+)`?",
|
||||
r"INSERT\s+`?([^`\s(]+)`?",
|
||||
]
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, sql, re.IGNORECASE)
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
return None
|
||||
|
||||
def _extract_table_from_update(self, sql: str) -> Optional[str]:
|
||||
"""从 UPDATE 语句中提取表名。
|
||||
|
||||
Args:
|
||||
sql: UPDATE 语句
|
||||
|
||||
Returns:
|
||||
表名，解析失败返回 None
|
||||
"""
|
||||
patterns = [
|
||||
r"UPDATE\s+`?([^`\s,]+)`?",
|
||||
]
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, sql, re.IGNORECASE)
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
return None
|
||||
|
||||
def _extract_table_from_delete(self, sql: str) -> Optional[str]:
|
||||
"""从 DELETE 语句中提取表名。
|
||||
|
||||
支持格式:
|
||||
- DELETE FROM table_name ...
|
||||
- DELETE table_name FROM ...
|
||||
|
||||
Args:
|
||||
sql: DELETE 语句
|
||||
|
||||
Returns:
|
||||
表名，解析失败返回 None
|
||||
"""
|
||||
patterns = [
|
||||
r"DELETE\s+FROM\s+`?([^`\s(]+)`?",
|
||||
r"DELETE\s+`?([^`\s(]+)`?\s+FROM",
|
||||
]
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, sql, re.IGNORECASE)
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
return None
|
||||
|
||||
def _extract_set_columns(self, sql: str) -> List[str]:
|
||||
"""从 UPDATE 语句的 SET 子句中提取列名。
|
||||
|
||||
Args:
|
||||
sql: UPDATE 语句
|
||||
|
||||
Returns:
|
||||
SET 子句中的列名列表
|
||||
"""
|
||||
columns: List[str] = []
|
||||
# 匹配 SET col = val, col2 = val2 格式
|
||||
match = re.search(r"SET\s+(.+?)(?:\s+WHERE\s+|$)", sql, re.IGNORECASE | re.DOTALL)
|
||||
if match:
|
||||
set_clause = match.group(1)
|
||||
# 提取每个赋值左侧的列名
|
||||
assignments = re.split(r",\s*(?=`?\w+`?\s*=)", set_clause)
|
||||
for assignment in assignments:
|
||||
col_match = re.match(r"`?(\w+)`?\s*[=<>!]+", assignment.strip())
|
||||
if col_match:
|
||||
columns.append(col_match.group(1))
|
||||
return columns
|
||||
|
||||
def _extract_where_clause(self, sql: str) -> Optional[str]:
|
||||
"""从 SQL 语句中提取 WHERE 子句。
|
||||
|
||||
Args:
|
||||
sql: SQL 语句
|
||||
|
||||
Returns:
|
||||
WHERE 条件字符串，没有 WHERE 返回 None
|
||||
"""
|
||||
# 使用 sqlparse 提取 WHERE 子句
|
||||
parsed = sqlparse.parse(sql)
|
||||
if not parsed:
|
||||
return None
|
||||
|
||||
statement = parsed[0]
|
||||
where_tokens = []
|
||||
in_where = False
|
||||
|
||||
for token in statement.tokens:
|
||||
if isinstance(token, Where):
|
||||
return str(token).replace("WHERE", "", 1).strip()
|
||||
if token.ttype is Keyword and token.normalized == "WHERE":
|
||||
in_where = True
|
||||
continue
|
||||
if in_where:
|
||||
if token.ttype is Keyword and token.normalized in ("ORDER", "GROUP", "HAVING", "LIMIT"):
|
||||
break
|
||||
where_tokens.append(str(token))
|
||||
|
||||
if where_tokens:
|
||||
return "".join(where_tokens).strip()
|
||||
return None
|
||||
|
||||
def _parse_insert_values(self, sql: str) -> Tuple[List[str], List[str]]:
|
||||
"""解析 INSERT 语句的列名和值。
|
||||
|
||||
支持格式:
|
||||
- INSERT INTO table (col1, col2) VALUES (val1, val2)
|
||||
- INSERT INTO table SET col1=val1, col2=val2
|
||||
|
||||
Args:
|
||||
sql: INSERT 语句
|
||||
|
||||
Returns:
|
||||
(列名列表, 值列表)
|
||||
"""
|
||||
columns: List[str] = []
|
||||
values: List[str] = []
|
||||
|
||||
# 尝试匹配 INSERT ... (col1, col2) VALUES (val1, val2)
|
||||
col_match = re.search(r"\(([^)]+)\)\s+VALUES\s*\(([^)]+)\)", sql, re.IGNORECASE | re.DOTALL)
|
||||
if col_match:
|
||||
col_str = col_match.group(1)
|
||||
val_str = col_match.group(2)
|
||||
columns = [c.strip().strip("`'") for c in col_str.split(",")]
|
||||
# 简单值解析（不考虑复杂表达式）
|
||||
values = self._split_values(val_str)
|
||||
return columns, values
|
||||
|
||||
# 尝试匹配 INSERT ... SET col1=val1, col2=val2
|
||||
set_match = re.search(r"SET\s+(.+)$", sql, re.IGNORECASE | re.DOTALL)
|
||||
if set_match:
|
||||
set_clause = set_match.group(1)
|
||||
assignments = set_clause.split(",")
|
||||
for assignment in assignments:
|
||||
parts = assignment.split("=", 1)
|
||||
if len(parts) == 2:
|
||||
columns.append(parts[0].strip().strip("`"))
|
||||
values.append(parts[1].strip().strip("'\""))
|
||||
|
||||
return columns, values
|
||||
|
||||
@staticmethod
|
||||
def _split_values(val_str: str) -> List[str]:
|
||||
"""分割 INSERT 的 VALUES 子句中的值。
|
||||
|
||||
处理简单值，跳过括号内的逗号。
|
||||
|
||||
Args:
|
||||
val_str: VALUES 子句中的值字符串
|
||||
|
||||
Returns:
|
||||
值列表
|
||||
"""
|
||||
values: List[str] = []
|
||||
current = []
|
||||
in_string = False
|
||||
string_char = None
|
||||
paren_depth = 0
|
||||
|
||||
for ch in val_str:
|
||||
if ch in ("'", '"') and not in_string:
|
||||
in_string = True
|
||||
string_char = ch
|
||||
elif ch == string_char and in_string:
|
||||
in_string = False
|
||||
string_char = None
|
||||
elif ch == "(" and not in_string:
|
||||
paren_depth += 1
|
||||
elif ch == ")" and not in_string:
|
||||
paren_depth -= 1
|
||||
|
||||
if ch == "," and not in_string and paren_depth == 0:
|
||||
values.append("".join(current).strip())
|
||||
current = []
|
||||
else:
|
||||
current.append(ch)
|
||||
|
||||
if current:
|
||||
values.append("".join(current).strip())
|
||||
|
||||
return values
|
||||
|
||||
@staticmethod
|
||||
def _quote_value(value: Any) -> str:
|
||||
"""将值转换为 SQL 安全的字符串表示。
|
||||
|
||||
Args:
|
||||
value: 任意值
|
||||
|
||||
Returns:
|
||||
SQL 安全的值表示
|
||||
"""
|
||||
if value is None:
|
||||
return "NULL"
|
||||
if isinstance(value, (int, float)):
|
||||
return str(value)
|
||||
if isinstance(value, bool):
|
||||
return "1" if value else "0"
|
||||
|
||||
# 字符串值: 转义单引号
|
||||
val = str(value).replace("'", "\\'")
|
||||
return f"'{val}'"
|
||||
@@ -0,0 +1,314 @@
|
||||
"""Schema 漂移检测 — 检测数据库 Schema 的意外变更。
|
||||
|
||||
工作原理:
|
||||
1. 定期捕获数据库 Schema 快照（表结构、索引、约束）
|
||||
2. 将快照与基线比较
|
||||
3. 发现差异时生成漂移报告
|
||||
4. 发送通知告警
|
||||
|
||||
支持的检测类型:
|
||||
- 新增表/字段
|
||||
- 删除表/字段
|
||||
- 字段类型变更
|
||||
- 索引变更
|
||||
"""
|
||||
import hashlib
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import List, Optional, Dict, Any
|
||||
|
||||
from app.models.instance import Instance
|
||||
from app.services.db_connector import DBConnector
|
||||
|
||||
|
||||
@dataclass
|
||||
class ColumnSnapshot:
|
||||
"""列定义快照。"""
|
||||
table_name: str
|
||||
column_name: str
|
||||
column_type: str
|
||||
is_nullable: bool
|
||||
column_default: Optional[str]
|
||||
extra: str
|
||||
character_set: Optional[str]
|
||||
collation: Optional[str]
|
||||
checksum: str
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"table_name": self.table_name,
|
||||
"column_name": self.column_name,
|
||||
"column_type": self.column_type,
|
||||
"is_nullable": self.is_nullable,
|
||||
"column_default": self.column_default,
|
||||
"extra": self.extra,
|
||||
"checksum": self.checksum,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class DriftItem:
|
||||
"""漂移项。"""
|
||||
drift_type: str # added / removed / modified
|
||||
table_name: str
|
||||
column_name: str
|
||||
before: Optional[str]
|
||||
after: Optional[str]
|
||||
severity: str # low / medium / high
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"drift_type": self.drift_type,
|
||||
"table_name": self.table_name,
|
||||
"column_name": self.column_name,
|
||||
"before": self.before,
|
||||
"after": self.after,
|
||||
"severity": self.severity,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class DriftReport:
|
||||
"""漂移检测报告。"""
|
||||
instance_id: int
|
||||
database: str
|
||||
checked_at: str
|
||||
total_columns: int
|
||||
drift_count: int
|
||||
drifts: List[DriftItem]
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"instance_id": self.instance_id,
|
||||
"database": self.database,
|
||||
"checked_at": self.checked_at,
|
||||
"total_columns": self.total_columns,
|
||||
"drift_count": self.drift_count,
|
||||
"drifts": [d.to_dict() for d in self.drifts],
|
||||
}
|
||||
|
||||
|
||||
class SchemaDriftDetector:
|
||||
"""Schema 漂移检测器。"""
|
||||
|
||||
# 内存中的基线存储: {(instance_id, database): [ColumnSnapshot, ...]}
|
||||
_baselines: Dict[tuple, List[ColumnSnapshot]] = {}
|
||||
|
||||
async def capture_snapshot(self, instance: Instance,
|
||||
database: str) -> List[ColumnSnapshot]:
|
||||
"""捕获数据库 Schema 快照。"""
|
||||
if instance.db_type == "mysql":
|
||||
return await self._capture_mysql_schema(instance, database)
|
||||
elif instance.db_type in ("postgres", "postgresql"):
|
||||
return await self._capture_postgres_schema(instance, database)
|
||||
else:
|
||||
raise ValueError(f"不支持的 Schema 捕获类型: {instance.db_type}")
|
||||
|
||||
async def save_baseline(self, instance_id: int,
|
||||
database: str,
|
||||
snapshot: List[ColumnSnapshot]) -> None:
|
||||
"""保存 Schema 基线到内存。"""
|
||||
self._baselines[(instance_id, database)] = snapshot
|
||||
|
||||
async def detect_drift(self, instance: Instance,
|
||||
database: str) -> DriftReport:
|
||||
"""检测 Schema 漂移，返回漂移报告。"""
|
||||
current = await self.capture_snapshot(instance, database)
|
||||
baseline = await self._get_baseline(instance.id, database)
|
||||
|
||||
checked_at = datetime.now().isoformat()
|
||||
|
||||
if not baseline:
|
||||
# 首次检测，保存为基线
|
||||
await self.save_baseline(instance.id, database, current)
|
||||
return DriftReport(
|
||||
instance_id=instance.id,
|
||||
database=database,
|
||||
checked_at=checked_at,
|
||||
total_columns=len(current),
|
||||
drift_count=0,
|
||||
drifts=[],
|
||||
)
|
||||
|
||||
drifts = self._compare_schemas(current, baseline)
|
||||
|
||||
# 更新基线为当前状态
|
||||
await self.save_baseline(instance.id, database, current)
|
||||
|
||||
return DriftReport(
|
||||
instance_id=instance.id,
|
||||
database=database,
|
||||
checked_at=checked_at,
|
||||
total_columns=len(current),
|
||||
drift_count=len(drifts),
|
||||
drifts=drifts,
|
||||
)
|
||||
|
||||
async def _capture_mysql_schema(self, instance: Instance,
|
||||
database: str) -> List[ColumnSnapshot]:
|
||||
"""从 MySQL INFORMATION_SCHEMA 捕获 Schema。"""
|
||||
password = instance.password # should be decrypted before calling
|
||||
sql = """
|
||||
SELECT
|
||||
TABLE_NAME, COLUMN_NAME, COLUMN_TYPE, IS_NULLABLE,
|
||||
COLUMN_DEFAULT, EXTRA, CHARACTER_SET_NAME, COLLATION_NAME
|
||||
FROM INFORMATION_SCHEMA.COLUMNS
|
||||
WHERE TABLE_SCHEMA = %s
|
||||
AND TABLE_NAME NOT LIKE '\\_%%' -- exclude internal tables
|
||||
ORDER BY TABLE_NAME, ORDINAL_POSITION
|
||||
"""
|
||||
|
||||
from app.core.crypto import decrypt_password
|
||||
instance.password = decrypt_password(instance.password)
|
||||
result = await DBConnector._query_mysql(instance, "information_schema", sql)
|
||||
|
||||
snapshots = []
|
||||
for row in result[0]:
|
||||
# 计算列定义的校验和
|
||||
content = (
|
||||
f"{row.get('COLUMN_TYPE', '')}|"
|
||||
f"{row.get('IS_NULLABLE', '')}|"
|
||||
f"{row.get('COLUMN_DEFAULT', '') or ''}|"
|
||||
f"{row.get('EXTRA', '')}"
|
||||
)
|
||||
checksum = hashlib.md5(content.encode()).hexdigest()
|
||||
|
||||
snapshots.append(ColumnSnapshot(
|
||||
table_name=row.get("TABLE_NAME", ""),
|
||||
column_name=row.get("COLUMN_NAME", ""),
|
||||
column_type=row.get("COLUMN_TYPE", ""),
|
||||
is_nullable=row.get("IS_NULLABLE", "YES") == "YES",
|
||||
column_default=str(row.get("COLUMN_DEFAULT", "")) if row.get("COLUMN_DEFAULT") is not None else None,
|
||||
extra=row.get("EXTRA", ""),
|
||||
character_set=row.get("CHARACTER_SET_NAME"),
|
||||
collation=row.get("COLLATION_NAME"),
|
||||
checksum=checksum,
|
||||
))
|
||||
|
||||
return snapshots
|
||||
|
||||
async def _capture_postgres_schema(self, instance: Instance,
|
||||
database: str) -> List[ColumnSnapshot]:
|
||||
"""从 PostgreSQL information_schema 捕获 Schema。"""
|
||||
sql = """
|
||||
SELECT
|
||||
c.table_name,
|
||||
c.column_name,
|
||||
CASE
|
||||
WHEN c.data_type = 'character varying' AND c.character_maximum_length IS NOT NULL
|
||||
THEN 'varchar(' || c.character_maximum_length || ')'
|
||||
WHEN c.data_type = 'character' AND c.character_maximum_length IS NOT NULL
|
||||
THEN 'char(' || c.character_maximum_length || ')'
|
||||
WHEN c.data_type = 'numeric' AND c.numeric_precision IS NOT NULL
|
||||
THEN 'numeric(' || c.numeric_precision || ',' || COALESCE(c.numeric_scale, 0) || ')'
|
||||
ELSE c.data_type
|
||||
END as column_type,
|
||||
c.is_nullable,
|
||||
c.column_default,
|
||||
'' as extra,
|
||||
NULL as character_set_name,
|
||||
NULL as collation_name
|
||||
FROM information_schema.columns c
|
||||
WHERE c.table_schema = 'public'
|
||||
AND c.table_catalog = %s
|
||||
ORDER BY c.table_name, c.ordinal_position
|
||||
"""
|
||||
|
||||
from app.core.crypto import decrypt_password
|
||||
instance.password = decrypt_password(instance.password)
|
||||
result = await DBConnector._query_postgres(instance, database, sql)
|
||||
|
||||
snapshots = []
|
||||
for row in result[0]:
|
||||
content = (
|
||||
f"{row.get('column_type', '')}|"
|
||||
f"{row.get('is_nullable', '')}|"
|
||||
f"{row.get('column_default', '') or ''}"
|
||||
)
|
||||
checksum = hashlib.md5(content.encode()).hexdigest()
|
||||
|
||||
snapshots.append(ColumnSnapshot(
|
||||
table_name=row.get("table_name", ""),
|
||||
column_name=row.get("column_name", ""),
|
||||
column_type=row.get("column_type", ""),
|
||||
is_nullable=row.get("is_nullable", "YES") == "YES",
|
||||
column_default=str(row.get("column_default", "")) if row.get("column_default") is not None else None,
|
||||
extra=row.get("extra", ""),
|
||||
character_set=None,
|
||||
collation=None,
|
||||
checksum=checksum,
|
||||
))
|
||||
|
||||
return snapshots
|
||||
|
||||
async def _get_baseline(self, instance_id: int,
|
||||
database: str) -> Optional[List[ColumnSnapshot]]:
|
||||
"""获取保存的基线。"""
|
||||
return self._baselines.get((instance_id, database))
|
||||
|
||||
def _compare_schemas(self, current: List[ColumnSnapshot],
|
||||
baseline: List[ColumnSnapshot]) -> List[DriftItem]:
|
||||
"""比较当前 Schema 与基线，返回漂移项。"""
|
||||
drifts = []
|
||||
current_dict = {(s.table_name, s.column_name): s for s in current}
|
||||
baseline_dict = {(s.table_name, s.column_name): s for s in baseline}
|
||||
|
||||
# 检查新增列
|
||||
for key, snapshot in current_dict.items():
|
||||
if key not in baseline_dict:
|
||||
drifts.append(DriftItem(
|
||||
drift_type="added",
|
||||
table_name=key[0],
|
||||
column_name=key[1],
|
||||
before=None,
|
||||
after=f"{snapshot.column_type} NULL={snapshot.is_nullable}",
|
||||
severity="medium",
|
||||
))
|
||||
|
||||
# 检查删除列
|
||||
for key, snapshot in baseline_dict.items():
|
||||
if key not in current_dict:
|
||||
drifts.append(DriftItem(
|
||||
drift_type="removed",
|
||||
table_name=key[0],
|
||||
column_name=key[1],
|
||||
before=f"{snapshot.column_type} NULL={snapshot.is_nullable}",
|
||||
after=None,
|
||||
severity="high", # 删除列是高风险操作
|
||||
))
|
||||
|
||||
# 检查修改列
|
||||
for key, current_snap in current_dict.items():
|
||||
if key in baseline_dict:
|
||||
baseline_snap = baseline_dict[key]
|
||||
if current_snap.checksum != baseline_snap.checksum:
|
||||
severity = "medium"
|
||||
# 类型变更是高风险
|
||||
if current_snap.column_type != baseline_snap.column_type:
|
||||
severity = "high"
|
||||
|
||||
drifts.append(DriftItem(
|
||||
drift_type="modified",
|
||||
table_name=key[0],
|
||||
column_name=key[1],
|
||||
before=f"{baseline_snap.column_type} (default: {baseline_snap.column_default or 'None'})",
|
||||
after=f"{current_snap.column_type} (default: {current_snap.column_default or 'None'})",
|
||||
severity=severity,
|
||||
))
|
||||
|
||||
# 按严重性排序
|
||||
severity_order = {"high": 0, "medium": 1, "low": 2}
|
||||
drifts.sort(key=lambda x: severity_order.get(x.severity, 3))
|
||||
|
||||
return drifts
|
||||
|
||||
def clear_baseline(self, instance_id: int, database: str) -> None:
|
||||
"""清除指定实例的基线。"""
|
||||
key = (instance_id, database)
|
||||
if key in self._baselines:
|
||||
del self._baselines[key]
|
||||
|
||||
def clear_all_baselines(self) -> None:
|
||||
"""清除所有基线。"""
|
||||
self._baselines.clear()
|
||||
@@ -0,0 +1,269 @@
|
||||
"""慢日志收集器 — 从数据库收集和分析慢查询日志。
|
||||
|
||||
支持:
|
||||
- MySQL: performance_schema.events_statements_summary_by_digest
|
||||
- PostgreSQL: pg_stat_statements
|
||||
|
||||
生成慢查询报告，包含 TOP N 慢查询、执行时间趋势等。
|
||||
"""
|
||||
import hashlib
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import List, Dict, Any, Optional
|
||||
|
||||
from app.models.instance import Instance
|
||||
from app.services.db_connector import DBConnector
|
||||
from app.core.crypto import decrypt_password
|
||||
|
||||
|
||||
@dataclass
|
||||
class SlowLogEntry:
|
||||
"""慢日志条目。"""
|
||||
sql_digest: str # SQL 指纹
|
||||
sql_text: str # 归一化 SQL
|
||||
execution_count: int # 执行次数
|
||||
avg_time_ms: float # 平均执行时间(ms)
|
||||
max_time_ms: float # 最大执行时间(ms)
|
||||
min_time_ms: float # 最小执行时间(ms)
|
||||
rows_examined: int # 扫描行数
|
||||
rows_sent: int # 返回行数
|
||||
first_seen: Optional[str] = None
|
||||
last_seen: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"sql_digest": self.sql_digest,
|
||||
"sql_text": self.sql_text[:200] if self.sql_text else "",
|
||||
"execution_count": self.execution_count,
|
||||
"avg_time_ms": round(self.avg_time_ms, 2),
|
||||
"max_time_ms": round(self.max_time_ms, 2),
|
||||
"min_time_ms": round(self.min_time_ms, 2),
|
||||
"rows_examined": self.rows_examined,
|
||||
"rows_sent": self.rows_sent,
|
||||
"first_seen": self.first_seen,
|
||||
"last_seen": self.last_seen,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class SlowLogReport:
|
||||
"""慢日志分析报告。"""
|
||||
total_slow_queries: int
|
||||
avg_execution_time_ms: float
|
||||
max_execution_time_ms: float
|
||||
top_queries: List[SlowLogEntry]
|
||||
time_distribution: Dict[str, int] # 按时间段分布
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"total_slow_queries": self.total_slow_queries,
|
||||
"avg_execution_time_ms": round(self.avg_execution_time_ms, 2),
|
||||
"max_execution_time_ms": round(self.max_execution_time_ms, 2),
|
||||
"top_queries": [q.to_dict() for q in self.top_queries],
|
||||
"time_distribution": self.time_distribution,
|
||||
}
|
||||
|
||||
|
||||
class SlowLogCollector:
|
||||
"""慢日志收集器。"""
|
||||
|
||||
# 内存中的慢日志存储: {instance_id: [SlowLogEntry, ...]}
|
||||
_logs: Dict[int, List[SlowLogEntry]] = {}
|
||||
|
||||
async def collect_mysql_slow_logs(self, instance: Instance,
|
||||
database: str,
|
||||
limit: int = 50) -> List[SlowLogEntry]:
|
||||
"""从 MySQL performance_schema 收集慢日志。
|
||||
|
||||
需要 MySQL 开启 performance_schema。
|
||||
"""
|
||||
sql = """
|
||||
SELECT
|
||||
DIGEST_TEXT as sql_text,
|
||||
COUNT_STAR as execution_count,
|
||||
ROUND(AVG_TIMER_WAIT / 1000000000, 3) as avg_time_ms,
|
||||
ROUND(MAX_TIMER_WAIT / 1000000000, 3) as max_time_ms,
|
||||
ROUND(MIN_TIMER_WAIT / 1000000000, 3) as min_time_ms,
|
||||
SUM_ROWS_EXAMINED as rows_examined,
|
||||
SUM_ROWS_SENT as rows_sent,
|
||||
FIRST_SEEN,
|
||||
LAST_SEEN
|
||||
FROM performance_schema.events_statements_summary_by_digest
|
||||
WHERE SCHEMA_NAME = %s
|
||||
AND AVG_TIMER_WAIT > 100000000000 -- > 100ms
|
||||
ORDER BY AVG_TIMER_WAIT DESC
|
||||
LIMIT %s
|
||||
"""
|
||||
|
||||
try:
|
||||
instance.password = decrypt_password(instance.password)
|
||||
rows, _ = await DBConnector._query_mysql(instance, "performance_schema", sql)
|
||||
except Exception as e:
|
||||
# performance_schema 可能不可用，使用替代查询
|
||||
return await self._collect_mysql_slow_logs_fallback(instance, database, limit)
|
||||
|
||||
entries = []
|
||||
for row in rows:
|
||||
sql_text = row.get("sql_text", "")
|
||||
digest = hashlib.md5(sql_text.encode()).hexdigest()[:16]
|
||||
entries.append(SlowLogEntry(
|
||||
sql_digest=digest,
|
||||
sql_text=sql_text[:500],
|
||||
execution_count=int(row.get("execution_count", 0) or 0),
|
||||
avg_time_ms=float(row.get("avg_time_ms", 0) or 0),
|
||||
max_time_ms=float(row.get("max_time_ms", 0) or 0),
|
||||
min_time_ms=float(row.get("min_time_ms", 0) or 0),
|
||||
rows_examined=int(row.get("rows_examined", 0) or 0),
|
||||
rows_sent=int(row.get("rows_sent", 0) or 0),
|
||||
first_seen=str(row.get("FIRST_SEEN", "")) if row.get("FIRST_SEEN") else None,
|
||||
last_seen=str(row.get("LAST_SEEN", "")) if row.get("LAST_SEEN") else None,
|
||||
))
|
||||
|
||||
return entries
|
||||
|
||||
async def _collect_mysql_slow_logs_fallback(self, instance: Instance,
|
||||
database: str,
|
||||
limit: int) -> List[SlowLogEntry]:
|
||||
"""备用方案：使用 SHOW PROCESSLIST 获取当前慢查询。"""
|
||||
sql = "SHOW FULL PROCESSLIST"
|
||||
|
||||
try:
|
||||
instance.password = decrypt_password(instance.password)
|
||||
rows, _ = await DBConnector._query_mysql(instance, database, sql)
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
entries = []
|
||||
for row in rows:
|
||||
time_val = row.get("Time", 0)
|
||||
info = row.get("Info", "")
|
||||
if time_val and int(time_val) > 1 and info and len(info) > 10:
|
||||
digest = hashlib.md5(info.encode()).hexdigest()[:16]
|
||||
entries.append(SlowLogEntry(
|
||||
sql_digest=digest,
|
||||
sql_text=info[:500],
|
||||
execution_count=1,
|
||||
avg_time_ms=float(time_val) * 1000,
|
||||
max_time_ms=float(time_val) * 1000,
|
||||
min_time_ms=float(time_val) * 1000,
|
||||
rows_examined=0,
|
||||
rows_sent=0,
|
||||
))
|
||||
|
||||
return entries[:limit]
|
||||
|
||||
async def collect_postgres_slow_logs(self, instance: Instance,
|
||||
database: str,
|
||||
limit: int = 50) -> List[SlowLogEntry]:
|
||||
"""从 PostgreSQL pg_stat_statements 收集慢日志。
|
||||
|
||||
需要安装 pg_stat_statements 扩展。
|
||||
"""
|
||||
sql = """
|
||||
SELECT
|
||||
query as sql_text,
|
||||
calls as execution_count,
|
||||
ROUND((total_exec_time / calls)::numeric, 3) as avg_time_ms,
|
||||
ROUND(max_exec_time::numeric, 3) as max_time_ms,
|
||||
ROUND(min_exec_time::numeric, 3) as min_time_ms,
|
||||
rows as rows_examined,
|
||||
rows as rows_sent
|
||||
FROM pg_stat_statements
|
||||
WHERE dbid = (SELECT oid FROM pg_database WHERE datname = $1)
|
||||
AND total_exec_time / calls > 100 -- > 100ms
|
||||
ORDER BY total_exec_time / calls DESC
|
||||
LIMIT $2
|
||||
"""
|
||||
|
||||
try:
|
||||
instance.password = decrypt_password(instance.password)
|
||||
# 需要先创建扩展
|
||||
await DBConnector._query_postgres(instance, database,
|
||||
"CREATE EXTENSION IF NOT EXISTS pg_stat_statements")
|
||||
rows, _ = await DBConnector._query_postgres(instance, database, sql)
|
||||
except Exception:
|
||||
# pg_stat_statements 可能不可用
|
||||
return []
|
||||
|
||||
entries = []
|
||||
for row in rows:
|
||||
sql_text = row.get("sql_text", "")
|
||||
digest = hashlib.md5(sql_text.encode()).hexdigest()[:16]
|
||||
entries.append(SlowLogEntry(
|
||||
sql_digest=digest,
|
||||
sql_text=sql_text[:500],
|
||||
execution_count=int(row.get("execution_count", 0) or 0),
|
||||
avg_time_ms=float(row.get("avg_time_ms", 0) or 0),
|
||||
max_time_ms=float(row.get("max_time_ms", 0) or 0),
|
||||
min_time_ms=float(row.get("min_time_ms", 0) or 0),
|
||||
rows_examined=int(row.get("rows_examined", 0) or 0),
|
||||
rows_sent=int(row.get("rows_sent", 0) or 0),
|
||||
))
|
||||
|
||||
return entries
|
||||
|
||||
async def save_slow_logs(self, instance_id: int,
|
||||
logs: List[SlowLogEntry]) -> None:
|
||||
"""保存慢日志到内存存储。"""
|
||||
if instance_id not in self._logs:
|
||||
self._logs[instance_id] = []
|
||||
|
||||
# 追加新日志，限制每个实例最多保存 1000 条
|
||||
self._logs[instance_id].extend(logs)
|
||||
if len(self._logs[instance_id]) > 1000:
|
||||
self._logs[instance_id] = self._logs[instance_id][-1000:]
|
||||
|
||||
def get_slow_logs(self, instance_id: int) -> List[SlowLogEntry]:
|
||||
"""获取保存的慢日志。"""
|
||||
return self._logs.get(instance_id, [])
|
||||
|
||||
def get_report(self, instance_id: int) -> SlowLogReport:
|
||||
"""生成慢日志报告。"""
|
||||
logs = self._logs.get(instance_id, [])
|
||||
|
||||
if not logs:
|
||||
return SlowLogReport(
|
||||
total_slow_queries=0,
|
||||
avg_execution_time_ms=0,
|
||||
max_execution_time_ms=0,
|
||||
top_queries=[],
|
||||
time_distribution={},
|
||||
)
|
||||
|
||||
total_avg = sum(q.avg_time_ms for q in logs) / len(logs)
|
||||
max_time = max(q.max_time_ms for q in logs)
|
||||
|
||||
# 按执行时间分布
|
||||
distribution = {
|
||||
"0-100ms": 0,
|
||||
"100-500ms": 0,
|
||||
"500ms-1s": 0,
|
||||
"1s-5s": 0,
|
||||
">5s": 0,
|
||||
}
|
||||
for q in logs:
|
||||
if q.avg_time_ms < 100:
|
||||
distribution["0-100ms"] += 1
|
||||
elif q.avg_time_ms < 500:
|
||||
distribution["100-500ms"] += 1
|
||||
elif q.avg_time_ms < 1000:
|
||||
distribution["500ms-1s"] += 1
|
||||
elif q.avg_time_ms < 5000:
|
||||
distribution["1s-5s"] += 1
|
||||
else:
|
||||
distribution[">5s"] += 1
|
||||
|
||||
# TOP 10
|
||||
top = sorted(logs, key=lambda x: x.avg_time_ms, reverse=True)[:10]
|
||||
|
||||
return SlowLogReport(
|
||||
total_slow_queries=len(logs),
|
||||
avg_execution_time_ms=round(total_avg, 2),
|
||||
max_execution_time_ms=round(max_time, 2),
|
||||
top_queries=top,
|
||||
time_distribution=distribution,
|
||||
)
|
||||
|
||||
def clear_logs(self, instance_id: int) -> None:
|
||||
"""清除指定实例的慢日志。"""
|
||||
self._logs.pop(instance_id, None)
|
||||
@@ -0,0 +1,342 @@
|
||||
"""SQL助手 — 人性化辅助功能。
|
||||
|
||||
提供便捷的数据操作辅助:
|
||||
- SQL格式化: 统一SQL代码风格
|
||||
- SQL语法高亮标记: 标记关键字、表名、字段名
|
||||
- 表结构速查: 快速获取表列信息
|
||||
- 智能补全建议: 基于表结构生成字段列表
|
||||
- SQL Explain解析: 将执行计划转为人类可读格式
|
||||
- 查询风险提示: 基于规则标记潜在风险
|
||||
|
||||
人性化设计原则:
|
||||
- 减少重复操作
|
||||
- 降低SQL编写门槛
|
||||
- 提供即时反馈
|
||||
- 统一操作体验
|
||||
"""
|
||||
import re
|
||||
from typing import List, Dict, Optional, Any
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
# SQL关键字（用于语法高亮）
|
||||
SQL_KEYWORDS = {
|
||||
"SELECT", "FROM", "WHERE", "INSERT", "UPDATE", "DELETE", "CREATE",
|
||||
"DROP", "ALTER", "TABLE", "INDEX", "JOIN", "LEFT", "RIGHT", "INNER",
|
||||
"OUTER", "ON", "AND", "OR", "NOT", "NULL", "IS", "IN", "EXISTS",
|
||||
"BETWEEN", "LIKE", "GROUP", "BY", "ORDER", "HAVING", "LIMIT", "OFFSET",
|
||||
"UNION", "ALL", "DISTINCT", "AS", "ASC", "DESC", "COUNT", "SUM",
|
||||
"AVG", "MAX", "MIN", "IF", "CASE", "WHEN", "THEN", "ELSE", "END",
|
||||
"SHOW", "DESCRIBE", "EXPLAIN", "USE", "SET", "COMMIT", "ROLLBACK",
|
||||
}
|
||||
|
||||
# 高风险操作标记
|
||||
RISK_MARKERS = {
|
||||
"DELETE": {"level": "high", "message": "删除操作不可逆，请确认WHERE条件正确"},
|
||||
"DROP": {"level": "critical", "message": "DROP操作会永久删除对象，请谨慎操作"},
|
||||
"TRUNCATE": {"level": "critical", "message": "TRUNCATE会清空全表且不可回滚"},
|
||||
"UPDATE": {"level": "medium", "message": "UPDATE可能影响多行数据，建议先SELECT验证"},
|
||||
"ALTER": {"level": "medium", "message": "ALTER TABLE可能锁表，建议在低峰期执行"},
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class SQLFormatResult:
|
||||
"""SQL格式化结果。"""
|
||||
formatted: str
|
||||
original: str
|
||||
is_valid: bool
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class SQLHighlightResult:
|
||||
"""SQL语法高亮结果。"""
|
||||
highlighted: str # HTML格式的语法高亮
|
||||
keywords: List[str] # 提取的关键字
|
||||
tables: List[str] # 提取的表名
|
||||
columns: List[str] # 提取的列名
|
||||
risks: List[Dict[str, str]] # 风险提示
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExplainResult:
|
||||
"""Explain解析结果。"""
|
||||
is_simple: bool
|
||||
has_index: bool
|
||||
has_full_scan: bool
|
||||
estimated_rows: int
|
||||
warnings: List[str]
|
||||
human_readable: str # 人类可读的执行计划描述
|
||||
|
||||
|
||||
class SQLAssistant:
|
||||
"""SQL助手 — 人性化辅助功能。"""
|
||||
|
||||
# ── SQL格式化 ──
|
||||
|
||||
@staticmethod
|
||||
def format_sql(sql: str) -> SQLFormatResult:
|
||||
"""格式化SQL语句。
|
||||
|
||||
- 关键字大写
|
||||
- 合理换行和缩进
|
||||
- 统一空格
|
||||
"""
|
||||
if not sql or not sql.strip():
|
||||
return SQLFormatResult(formatted="", original=sql, is_valid=False, error="SQL为空")
|
||||
|
||||
try:
|
||||
import sqlparse
|
||||
formatted = sqlparse.format(
|
||||
sql,
|
||||
keyword_case="upper",
|
||||
identifier_case="lower",
|
||||
strip_comments=False,
|
||||
reindent=True,
|
||||
indent_width=4,
|
||||
wrap_after=80,
|
||||
)
|
||||
return SQLFormatResult(formatted=formatted, original=sql, is_valid=True)
|
||||
except ImportError:
|
||||
# 如果没有sqlparse，使用简单格式化
|
||||
return SQLAssistant._simple_format(sql)
|
||||
except Exception as e:
|
||||
return SQLFormatResult(formatted=sql, original=sql, is_valid=False, error=str(e))
|
||||
|
||||
@staticmethod
|
||||
def _simple_format(sql: str) -> SQLFormatResult:
|
||||
"""简单SQL格式化（不依赖sqlparse）。"""
|
||||
# 去除多余空格
|
||||
formatted = " ".join(sql.split())
|
||||
|
||||
# 关键字大写
|
||||
for kw in SQL_KEYWORDS:
|
||||
pattern = re.compile(r"\b" + kw + r"\b", re.IGNORECASE)
|
||||
formatted = pattern.sub(kw, formatted)
|
||||
|
||||
# 在主要关键字前换行
|
||||
line_breaks = [" SELECT ", " FROM ", " WHERE ", " GROUP BY ",
|
||||
" ORDER BY ", " HAVING ", " LIMIT ", " JOIN ",
|
||||
" LEFT JOIN ", " RIGHT JOIN ", " INNER JOIN ", " UNION "]
|
||||
for br in line_breaks:
|
||||
formatted = formatted.replace(br, "\n" + br.strip().ljust(8))
|
||||
|
||||
return SQLFormatResult(formatted=formatted, original=sql, is_valid=True)
|
||||
|
||||
# ── 语法高亮与风险分析 ──
|
||||
|
||||
@staticmethod
|
||||
def highlight_sql(sql: str) -> SQLHighlightResult:
|
||||
"""SQL语法高亮和风险分析。
|
||||
|
||||
返回HTML格式的语法高亮文本和风险提示。
|
||||
"""
|
||||
upper = sql.upper().strip()
|
||||
|
||||
# 提取关键字
|
||||
words = re.findall(r"\b\w+\b", upper)
|
||||
keywords = [w for w in words if w in SQL_KEYWORDS]
|
||||
|
||||
# 提取表名（简单规则）
|
||||
tables = []
|
||||
table_patterns = [
|
||||
r"FROM\s+(\w+)",
|
||||
r"JOIN\s+(\w+)",
|
||||
r"INTO\s+(\w+)",
|
||||
r"UPDATE\s+(\w+)",
|
||||
r"TABLE\s+(\w+)",
|
||||
]
|
||||
for pattern in table_patterns:
|
||||
matches = re.findall(pattern, upper, re.IGNORECASE)
|
||||
tables.extend(matches)
|
||||
tables = list(set(t.lower() for t in tables))
|
||||
|
||||
# 风险提示
|
||||
risks = []
|
||||
for op, info in RISK_MARKERS.items():
|
||||
if re.search(r"\b" + op + r"\b", upper):
|
||||
# 检查是否有WHERE
|
||||
if op in ("UPDATE", "DELETE") and "WHERE" not in upper:
|
||||
risks.append({
|
||||
"level": "critical",
|
||||
"operation": op,
|
||||
"message": f"{op} 缺少 WHERE 条件，将操作全表数据！",
|
||||
})
|
||||
else:
|
||||
risks.append({"level": info["level"], "operation": op, "message": info["message"]})
|
||||
|
||||
# 构建HTML高亮（简化版）
|
||||
highlighted = sql
|
||||
for kw in SQL_KEYWORDS:
|
||||
highlighted = re.sub(
|
||||
r"\b" + kw + r"\b",
|
||||
f'<span class="sql-keyword">{kw}</span>',
|
||||
highlighted,
|
||||
flags=re.IGNORECASE,
|
||||
)
|
||||
|
||||
return SQLHighlightResult(
|
||||
highlighted=highlighted,
|
||||
keywords=list(set(keywords)),
|
||||
tables=tables,
|
||||
columns=[], # 简单实现不提取列
|
||||
risks=risks,
|
||||
)
|
||||
|
||||
# ── Explain解析 ──
|
||||
|
||||
@staticmethod
|
||||
def parse_explain(explain_rows: List[Dict[str, Any]]) -> ExplainResult:
|
||||
"""解析MySQL EXPLAIN结果为人类可读格式。
|
||||
|
||||
Args:
|
||||
explain_rows: EXPLAIN输出结果行列表
|
||||
"""
|
||||
if not explain_rows:
|
||||
return ExplainResult(
|
||||
is_simple=True, has_index=False, has_full_scan=False,
|
||||
estimated_rows=0, warnings=["EXPLAIN结果为空"],
|
||||
human_readable="无法解析执行计划",
|
||||
)
|
||||
|
||||
warnings = []
|
||||
total_rows = 0
|
||||
has_full_scan = False
|
||||
has_index = False
|
||||
parts = []
|
||||
|
||||
for i, row in enumerate(explain_rows):
|
||||
table = row.get("table") or row.get("TABLE", "unknown")
|
||||
select_type = row.get("select_type") or row.get("SELECT_TYPE", "SIMPLE")
|
||||
scan_type = row.get("type") or row.get("TYPE", "ALL")
|
||||
possible_keys = row.get("possible_keys") or row.get("POSSIBLE_KEYS")
|
||||
key_used = row.get("key") or row.get("KEY")
|
||||
rows_est = row.get("rows") or row.get("ROWS", 0) or row.get("EST_ROWS", 0)
|
||||
extra = row.get("Extra") or row.get("EXTRA", "")
|
||||
|
||||
total_rows += int(rows_est) if rows_est else 0
|
||||
|
||||
if scan_type in ("ALL", "index") and select_type == "SIMPLE":
|
||||
has_full_scan = True
|
||||
warnings.append(f"表 '{table}' 使用了全表扫描（{scan_type}），建议添加索引")
|
||||
|
||||
if key_used:
|
||||
has_index = True
|
||||
parts.append(f"表 '{table}': 使用了索引 {key_used}，预估扫描 {rows_est} 行")
|
||||
else:
|
||||
parts.append(f"表 '{table}': 未使用索引（类型: {scan_type}），预估扫描 {rows_est} 行")
|
||||
|
||||
if "Using filesort" in str(extra):
|
||||
warnings.append(f"表 '{table}' 使用了文件排序，建议优化ORDER BY")
|
||||
if "Using temporary" in str(extra):
|
||||
warnings.append(f"表 '{table}' 使用了临时表，建议优化GROUP BY")
|
||||
|
||||
human_readable = "\n".join(parts)
|
||||
|
||||
if total_rows > 100000:
|
||||
warnings.append(f"预估扫描总行数 {total_rows} 较大，查询可能较慢")
|
||||
|
||||
return ExplainResult(
|
||||
is_simple=len(explain_rows) == 1,
|
||||
has_index=has_index,
|
||||
has_full_scan=has_full_scan,
|
||||
estimated_rows=total_rows,
|
||||
warnings=warnings,
|
||||
human_readable=human_readable,
|
||||
)
|
||||
|
||||
# ── 查询建议 ──
|
||||
|
||||
@staticmethod
|
||||
def get_query_suggestions(sql: str, db_type: str = "mysql") -> List[Dict[str, str]]:
|
||||
"""基于SQL内容提供人性化建议。
|
||||
|
||||
帮助用户写出更好的SQL。
|
||||
"""
|
||||
suggestions = []
|
||||
upper = sql.upper().strip()
|
||||
|
||||
# SELECT * 建议
|
||||
if re.search(r"SELECT\s+\*", upper):
|
||||
suggestions.append({
|
||||
"type": "optimization",
|
||||
"message": "使用 SELECT * 会返回所有列，建议只选择需要的字段以减少网络传输",
|
||||
})
|
||||
|
||||
# 缺少LIMIT
|
||||
if "SELECT" in upper and "LIMIT" not in upper and db_type in ("mysql", "postgres"):
|
||||
suggestions.append({
|
||||
"type": "safety",
|
||||
"message": "建议添加 LIMIT 子句限制返回行数，避免大数据量查询",
|
||||
})
|
||||
|
||||
# ORDER BY RAND()
|
||||
if "ORDER BY RAND()" in upper or "ORDER BY RANDOM()" in upper:
|
||||
suggestions.append({
|
||||
"type": "optimization",
|
||||
"message": "ORDER BY RAND() 在大表上性能极差，建议使用其他随机化方案",
|
||||
})
|
||||
|
||||
# LIKE %xxx%
|
||||
if re.search(r"LIKE\s+'%.", upper):
|
||||
suggestions.append({
|
||||
"type": "optimization",
|
||||
"message": "LIKE '%xxx%' 无法使用索引，如数据量大请考虑全文索引或其他方案",
|
||||
})
|
||||
|
||||
# 多表JOIN
|
||||
join_count = len(re.findall(r"\bJOIN\b", upper))
|
||||
if join_count > 3:
|
||||
suggestions.append({
|
||||
"type": "optimization",
|
||||
"message": f"当前有 {join_count} 个 JOIN，建议考虑子查询或视图简化",
|
||||
})
|
||||
|
||||
# 没有WHERE的UPDATE/DELETE（前面已标记为critical）
|
||||
# 这里给出更详细的建议
|
||||
if re.search(r"\bUPDATE\b", upper) and "WHERE" not in upper:
|
||||
suggestions.append({
|
||||
"type": "danger",
|
||||
"message": "UPDATE 缺少 WHERE 将修改全表！建议先用 SELECT 验证影响范围",
|
||||
})
|
||||
|
||||
if re.search(r"\bDELETE\b", upper) and "WHERE" not in upper:
|
||||
suggestions.append({
|
||||
"type": "danger",
|
||||
"message": "DELETE 缺少 WHERE 将删除全表！建议先用 SELECT 验证影响范围",
|
||||
})
|
||||
|
||||
# 缺少注释的复杂查询
|
||||
if len(sql) > 500 and "--" not in sql and "/*" not in sql:
|
||||
suggestions.append({
|
||||
"type": "style",
|
||||
"message": "复杂查询建议添加注释，方便后续维护和理解",
|
||||
})
|
||||
|
||||
return suggestions
|
||||
|
||||
# ── 统一辅助入口 ──
|
||||
|
||||
@staticmethod
|
||||
def analyze(sql: str, db_type: str = "mysql") -> Dict[str, Any]:
|
||||
"""全面分析SQL，返回格式化+高亮+风险+建议的完整结果。
|
||||
|
||||
人性化入口，一次调用获取所有辅助信息。
|
||||
"""
|
||||
format_result = SQLAssistant.format_sql(sql)
|
||||
highlight_result = SQLAssistant.highlight_sql(sql)
|
||||
suggestions = SQLAssistant.get_query_suggestions(sql, db_type)
|
||||
|
||||
return {
|
||||
"formatted": format_result.formatted if format_result.is_valid else sql,
|
||||
"is_valid": format_result.is_valid,
|
||||
"format_error": format_result.error,
|
||||
"highlighted": highlight_result.highlighted,
|
||||
"keywords": highlight_result.keywords,
|
||||
"tables": highlight_result.tables,
|
||||
"risks": highlight_result.risks,
|
||||
"suggestions": suggestions,
|
||||
"risk_count": len(highlight_result.risks),
|
||||
"suggestion_count": len(suggestions),
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,359 @@
|
||||
"""多级审批流引擎 — 支持串行、并行、条件审批。
|
||||
|
||||
对标 Yearning/Bytebase 的审批能力：
|
||||
- 启动审批流程时根据工单绑定的 workflow 创建审批实例
|
||||
- 支持通过、拒绝、转交、跳过等操作
|
||||
- 自动分配审批人（基于负载均衡）
|
||||
- 超时自动处理
|
||||
"""
|
||||
import json
|
||||
import random
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from typing import List, Optional, Dict, Any
|
||||
from dataclasses import dataclass
|
||||
|
||||
from sqlalchemy import select, func, and_
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.models.workflow import WorkflowDefinition, WorkflowStep, TicketApproval
|
||||
from app.models.ticket import Ticket
|
||||
from app.models.user import User
|
||||
|
||||
|
||||
@dataclass
|
||||
class ApprovalProgress:
|
||||
"""审批进度数据。"""
|
||||
ticket_id: int
|
||||
current_step: int
|
||||
total_steps: int
|
||||
status: str # pending / reviewing / approved / rejected
|
||||
steps: List[Dict[str, Any]]
|
||||
|
||||
|
||||
class WorkflowEngine:
|
||||
"""多级审批流引擎。"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
self.db = db
|
||||
|
||||
async def start_workflow(self, ticket_id: int,
|
||||
workflow_id: Optional[int] = None) -> List[TicketApproval]:
|
||||
"""为工单启动审批流程。
|
||||
|
||||
如果没有指定 workflow_id，使用默认流程。
|
||||
为每个步骤创建 TicketApproval 记录。
|
||||
"""
|
||||
if workflow_id is None:
|
||||
workflow = await self._get_default_workflow()
|
||||
else:
|
||||
result = await self.db.execute(
|
||||
select(WorkflowDefinition).where(WorkflowDefinition.id == workflow_id)
|
||||
)
|
||||
workflow = result.scalar_one_or_none()
|
||||
if not workflow:
|
||||
workflow = await self._get_default_workflow()
|
||||
|
||||
if not workflow:
|
||||
# 没有配置审批流，直接标记为 approved
|
||||
return []
|
||||
|
||||
approvals = []
|
||||
for step in workflow.steps:
|
||||
approval = TicketApproval(
|
||||
ticket_id=ticket_id,
|
||||
step_number=step.step_number,
|
||||
step_name=step.name,
|
||||
approver_role=step.approver_role,
|
||||
action="pending",
|
||||
)
|
||||
# 如果指定了具体审批人
|
||||
if step.approver_ids:
|
||||
ids = json.loads(step.approver_ids)
|
||||
if ids:
|
||||
approval.approver_id = ids[0]
|
||||
self.db.add(approval)
|
||||
approvals.append(approval)
|
||||
|
||||
await self.db.flush()
|
||||
return approvals
|
||||
|
||||
async def approve(self, approval_id: int, user_id: int,
|
||||
comment: str = "") -> Dict[str, Any]:
|
||||
"""审批通过。
|
||||
|
||||
Returns:
|
||||
dict with keys: success, new_status, message, next_step
|
||||
"""
|
||||
approval = await self._get_approval(approval_id)
|
||||
if not approval:
|
||||
return {"success": False, "message": "审批记录不存在"}
|
||||
|
||||
if approval.action != "pending":
|
||||
return {"success": False, "message": f"当前状态为 {approval.action}，无法审批"}
|
||||
|
||||
approval.action = "approved"
|
||||
approval.approver_id = user_id
|
||||
approval.comment = comment or "审批通过"
|
||||
approval.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
# 检查是否所有步骤都已完成
|
||||
ticket_approvals = await self._get_ticket_approvals(approval.ticket_id)
|
||||
|
||||
# 判断当前步骤类型
|
||||
current_step_num = approval.step_number
|
||||
step_approvals = [a for a in ticket_approvals if a.step_number == current_step_num]
|
||||
|
||||
if approval.step_type == "parallel":
|
||||
# 并行审批：需要所有审批人都通过
|
||||
if all(a.action == "approved" for a in step_approvals):
|
||||
all_done = await self._check_all_steps_done(ticket_approvals, current_step_num)
|
||||
if all_done:
|
||||
return await self._complete_approval(approval.ticket_id, "approved")
|
||||
return {"success": True, "new_status": "reviewing",
|
||||
"message": "当前步骤并行审批已完成，等待后续步骤",
|
||||
"next_step": current_step_num + 1}
|
||||
elif approval.step_type == "any_one":
|
||||
# 任一通过即可
|
||||
all_done = await self._check_all_steps_done(ticket_approvals, current_step_num)
|
||||
if all_done:
|
||||
return await self._complete_approval(approval.ticket_id, "approved")
|
||||
return {"success": True, "new_status": "reviewing",
|
||||
"message": "当前步骤已有人审批通过，等待后续步骤",
|
||||
"next_step": current_step_num + 1}
|
||||
else:
|
||||
# 串行审批
|
||||
all_done = await self._check_all_steps_done(ticket_approvals, current_step_num)
|
||||
if all_done:
|
||||
return await self._complete_approval(approval.ticket_id, "approved")
|
||||
return {"success": True, "new_status": "reviewing",
|
||||
"message": "审批已通过，等待后续步骤",
|
||||
"next_step": current_step_num + 1}
|
||||
|
||||
async def reject(self, approval_id: int, user_id: int,
|
||||
comment: str = "") -> Dict[str, Any]:
|
||||
"""审批拒绝。任一拒绝则整个工单拒绝。"""
|
||||
approval = await self._get_approval(approval_id)
|
||||
if not approval:
|
||||
return {"success": False, "message": "审批记录不存在"}
|
||||
|
||||
if approval.action != "pending":
|
||||
return {"success": False, "message": f"当前状态为 {approval.action}，无法操作"}
|
||||
|
||||
approval.action = "rejected"
|
||||
approval.approver_id = user_id
|
||||
approval.comment = comment or "审批拒绝"
|
||||
approval.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
# 拒绝整个工单
|
||||
result = await self._complete_approval(approval.ticket_id, "rejected")
|
||||
result["message"] = f"审批已拒绝: {comment}"
|
||||
return result
|
||||
|
||||
async def transfer(self, approval_id: int, from_user: int, to_user: int,
|
||||
comment: str = "") -> Dict[str, Any]:
|
||||
"""转交审批给其他人。"""
|
||||
approval = await self._get_approval(approval_id)
|
||||
if not approval:
|
||||
return {"success": False, "message": "审批记录不存在"}
|
||||
|
||||
# 创建新的审批记录给目标用户
|
||||
new_approval = TicketApproval(
|
||||
ticket_id=approval.ticket_id,
|
||||
step_number=approval.step_number,
|
||||
step_name=approval.step_name,
|
||||
approver_id=to_user,
|
||||
approver_role=approval.approver_role,
|
||||
action="pending",
|
||||
)
|
||||
self.db.add(new_approval)
|
||||
|
||||
# 标记原审批为转交
|
||||
approval.action = "transferred"
|
||||
approval.comment = f"转交给用户 {to_user}: {comment}"
|
||||
approval.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
await self.db.flush()
|
||||
return {"success": True, "message": "审批已转交", "new_approval_id": new_approval.id}
|
||||
|
||||
async def skip_step(self, approval_id: int, user_id: int,
|
||||
reason: str = "") -> Dict[str, Any]:
|
||||
"""管理员跳过当前步骤。"""
|
||||
approval = await self._get_approval(approval_id)
|
||||
if not approval:
|
||||
return {"success": False, "message": "审批记录不存在"}
|
||||
|
||||
approval.action = "skipped"
|
||||
approval.approver_id = user_id
|
||||
approval.comment = f"跳过: {reason}"
|
||||
approval.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
# 检查是否全部完成
|
||||
ticket_approvals = await self._get_ticket_approvals(approval.ticket_id)
|
||||
current_step_num = approval.step_number
|
||||
all_done = await self._check_all_steps_done(ticket_approvals, current_step_num)
|
||||
if all_done:
|
||||
return await self._complete_approval(approval.ticket_id, "approved")
|
||||
|
||||
return {"success": True, "message": "步骤已跳过", "new_status": "reviewing"}
|
||||
|
||||
async def get_approval_progress(self, ticket_id: int) -> ApprovalProgress:
|
||||
"""获取工单的审批进度。"""
|
||||
approvals = await self._get_ticket_approvals(ticket_id)
|
||||
ticket_result = await self.db.execute(select(Ticket).where(Ticket.id == ticket_id))
|
||||
ticket = ticket_result.scalar_one_or_none()
|
||||
|
||||
steps_data = []
|
||||
for approval in approvals:
|
||||
approver_name = "未分配"
|
||||
if approval.approver_id:
|
||||
user_result = await self.db.execute(
|
||||
select(User).where(User.id == approval.approver_id)
|
||||
)
|
||||
user = user_result.scalar_one_or_none()
|
||||
approver_name = user.display_name if user else "未知"
|
||||
|
||||
steps_data.append({
|
||||
"step_number": approval.step_number,
|
||||
"step_name": approval.step_name,
|
||||
"approver_name": approver_name,
|
||||
"approver_role": approval.approver_role,
|
||||
"action": approval.action,
|
||||
"comment": approval.comment,
|
||||
"created_at": approval.created_at.isoformat() if approval.created_at else None,
|
||||
"updated_at": approval.updated_at.isoformat() if approval.updated_at else None,
|
||||
})
|
||||
|
||||
current_step = 1
|
||||
total_steps = len(approvals)
|
||||
status = ticket.status if ticket else "unknown"
|
||||
|
||||
# 找到当前待审批的步骤
|
||||
for a in approvals:
|
||||
if a.action == "pending":
|
||||
current_step = a.step_number
|
||||
break
|
||||
|
||||
return ApprovalProgress(
|
||||
ticket_id=ticket_id,
|
||||
current_step=current_step,
|
||||
total_steps=total_steps,
|
||||
status=status,
|
||||
steps=steps_data,
|
||||
)
|
||||
|
||||
async def auto_assign(self, step: WorkflowStep) -> Optional[int]:
|
||||
"""基于负载自动分配审批人。
|
||||
|
||||
优先选择当前待审批任务最少的 DBA。
|
||||
"""
|
||||
if step.approver_ids:
|
||||
ids = json.loads(step.approver_ids)
|
||||
if ids:
|
||||
return ids[0]
|
||||
|
||||
if step.approver_role:
|
||||
# 查找该角色下负载最小的用户
|
||||
result = await self.db.execute(
|
||||
select(User.id, User.display_name)
|
||||
.where(User.role == step.approver_role, User.is_active == True)
|
||||
)
|
||||
candidates = result.all()
|
||||
if not candidates:
|
||||
return None
|
||||
|
||||
# 统计每个候选人的待审批数量
|
||||
workloads = []
|
||||
for user_id, _ in candidates:
|
||||
count_result = await self.db.execute(
|
||||
select(func.count(TicketApproval.id))
|
||||
.where(TicketApproval.approver_id == user_id,
|
||||
TicketApproval.action == "pending")
|
||||
)
|
||||
count = count_result.scalar() or 0
|
||||
workloads.append((user_id, count))
|
||||
|
||||
# 选择负载最小的
|
||||
workloads.sort(key=lambda x: x[1])
|
||||
return workloads[0][0]
|
||||
|
||||
return None
|
||||
|
||||
async def check_timeout_and_process(self):
|
||||
"""检查超时的审批并自动处理。"""
|
||||
timeout_threshold = datetime.now(timezone.utc) - timedelta(hours=24)
|
||||
result = await self.db.execute(
|
||||
select(TicketApproval).where(
|
||||
and_(TicketApproval.action == "pending",
|
||||
TicketApproval.created_at < timeout_threshold)
|
||||
)
|
||||
)
|
||||
timed_out = result.scalars().all()
|
||||
|
||||
for approval in timed_out:
|
||||
# 获取步骤定义，检查是否配置为超时自动通过
|
||||
workflow_result = await self.db.execute(
|
||||
select(WorkflowStep).where(WorkflowStep.id == approval.step_number)
|
||||
)
|
||||
step_def = workflow_result.scalar_one_or_none()
|
||||
if step_def and step_def.auto_approve_on_timeout:
|
||||
approval.action = "approved"
|
||||
approval.comment = "超时自动通过"
|
||||
approval.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
# ── 内部辅助方法 ──
|
||||
|
||||
async def _get_default_workflow(self) -> Optional[WorkflowDefinition]:
|
||||
"""获取默认审批流程。"""
|
||||
result = await self.db.execute(
|
||||
select(WorkflowDefinition)
|
||||
.where(WorkflowDefinition.is_default == True,
|
||||
WorkflowDefinition.is_active == True)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def _get_approval(self, approval_id: int) -> Optional[TicketApproval]:
|
||||
result = await self.db.execute(
|
||||
select(TicketApproval).where(TicketApproval.id == approval_id)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def _get_ticket_approvals(self, ticket_id: int) -> List[TicketApproval]:
|
||||
result = await self.db.execute(
|
||||
select(TicketApproval)
|
||||
.where(TicketApproval.ticket_id == ticket_id)
|
||||
.order_by(TicketApproval.step_number)
|
||||
)
|
||||
return result.scalars().all()
|
||||
|
||||
async def _check_all_steps_done(self, approvals: List[TicketApproval],
|
||||
current_step: int) -> bool:
|
||||
"""检查当前步骤及之前的所有步骤是否都已完成。"""
|
||||
for a in approvals:
|
||||
if a.step_number <= current_step and a.action in ("pending",):
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _complete_approval(self, ticket_id: int,
|
||||
final_status: str) -> Dict[str, Any]:
|
||||
"""完成审批，更新工单状态。"""
|
||||
ticket_result = await self.db.execute(
|
||||
select(Ticket).where(Ticket.id == ticket_id)
|
||||
)
|
||||
ticket = ticket_result.scalar_one_or_none()
|
||||
if ticket:
|
||||
ticket.status = final_status
|
||||
if final_status == "approved":
|
||||
ticket.review_time = datetime.now(timezone.utc)
|
||||
|
||||
await self.db.flush()
|
||||
return {
|
||||
"success": True,
|
||||
"new_status": final_status,
|
||||
"message": f"审批流程已完成，工单状态: {final_status}",
|
||||
"next_step": None,
|
||||
}
|
||||
|
||||
@property
|
||||
def step_type(self) -> str:
|
||||
"""获取当前步骤类型。"""
|
||||
return "serial" # 简化实现
|
||||
@@ -0,0 +1,93 @@
|
||||
"""预置默认审批流程种子数据。"""
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.models.workflow import WorkflowDefinition, WorkflowStep
|
||||
|
||||
|
||||
async def seed_default_workflows(db: AsyncSession) -> None:
|
||||
"""创建默认审批流程。
|
||||
|
||||
创建3种预置流程:
|
||||
1. 标准审批流: 开发者 → DBA审批 → 管理员终审
|
||||
2. 简单审批流: 开发者 → DBA审批
|
||||
3. 紧急审批流: 开发者 → 任一管理员直接审批（超时2h自动通过）
|
||||
"""
|
||||
# 检查是否已有流程
|
||||
result = await db.execute(select(WorkflowDefinition))
|
||||
existing = result.scalars().all()
|
||||
if existing:
|
||||
return # 已有流程，不再创建
|
||||
|
||||
# 1. 标准审批流
|
||||
std_workflow = WorkflowDefinition(
|
||||
name="标准审批流",
|
||||
description="标准三级审批：DBA审批 → 管理员终审",
|
||||
is_default=True,
|
||||
)
|
||||
db.add(std_workflow)
|
||||
await db.flush()
|
||||
|
||||
std_steps = [
|
||||
WorkflowStep(
|
||||
workflow_id=std_workflow.id,
|
||||
step_number=1,
|
||||
name="DBA 审批",
|
||||
step_type="serial",
|
||||
approver_role="dba",
|
||||
is_required=True,
|
||||
timeout_hours=24,
|
||||
),
|
||||
WorkflowStep(
|
||||
workflow_id=std_workflow.id,
|
||||
step_number=2,
|
||||
name="管理员终审",
|
||||
step_type="serial",
|
||||
approver_role="admin",
|
||||
is_required=True,
|
||||
timeout_hours=24,
|
||||
),
|
||||
]
|
||||
for step in std_steps:
|
||||
db.add(step)
|
||||
|
||||
# 2. 简单审批流
|
||||
simple_workflow = WorkflowDefinition(
|
||||
name="简单审批流",
|
||||
description="仅 DBA 一级审批",
|
||||
is_default=False,
|
||||
)
|
||||
db.add(simple_workflow)
|
||||
await db.flush()
|
||||
|
||||
db.add(WorkflowStep(
|
||||
workflow_id=simple_workflow.id,
|
||||
step_number=1,
|
||||
name="DBA 审批",
|
||||
step_type="serial",
|
||||
approver_role="dba",
|
||||
is_required=True,
|
||||
timeout_hours=24,
|
||||
))
|
||||
|
||||
# 3. 紧急审批流
|
||||
emergency_workflow = WorkflowDefinition(
|
||||
name="紧急审批流",
|
||||
description="紧急变更快速通道：任一管理员审批，超时2小时自动通过",
|
||||
is_default=False,
|
||||
)
|
||||
db.add(emergency_workflow)
|
||||
await db.flush()
|
||||
|
||||
db.add(WorkflowStep(
|
||||
workflow_id=emergency_workflow.id,
|
||||
step_number=1,
|
||||
name="管理员快速审批",
|
||||
step_type="any_one",
|
||||
approver_role="admin",
|
||||
is_required=True,
|
||||
timeout_hours=2,
|
||||
auto_approve_on_timeout=True,
|
||||
))
|
||||
|
||||
await db.commit()
|
||||
Reference in new issue
Block a user