feat: Phase2影响分析 — EXPLAIN预检+大表检测(简化版)+执行策略+风险评估table_sizes支持 (1141测试)
新增: - explain_preview.py: EXPLAIN预检服务(MySQL/PG), 影响行数预估, 全表扫描检测 - ticket.executor: 三种执行策略(stop_on_error/continue/transaction) - risk_evaluator: table_sizes参数, 大表ALTER/UPDATE/DELETE自动升级high - test_phase2_features.py: 22个测试(EXPLAIN+执行策略+大表检测) 简化: - 砍掉CTX-001上下文感知规则(复杂二次扫描), 合并大表检测到risk_evaluator - 审核引擎保持纯静态分析, 不再导入_RULE_REGISTRY做二次扫描 模型: - Ticket加exec_strategy字段 - TicketSqlItem模型 - schema加sql_type/exec_strategy/risk_level/workflow_id
This commit is contained in:
1 parent
7be0bf1a7c
commit
f6f9142ce8
7 files changed
+574
-22
No files matched your search
@@ -498,10 +498,25 @@ async def create_ticket(
|
||||
# 3. 验证工单内容安全性
|
||||
_validate_ticket_content(data)
|
||||
|
||||
# 3.5 EXPLAIN 预检(大表检测合并到 risk_evaluator,不再单独查表大小)
|
||||
table_sizes = {}
|
||||
explain_warnings = []
|
||||
try:
|
||||
from app.services.explain_preview import run_explain_preview
|
||||
import asyncio as _asyncio
|
||||
explain_result = await _asyncio.wait_for(
|
||||
run_explain_preview(data.sql_content, instance, data.target_database),
|
||||
timeout=8.0,
|
||||
)
|
||||
explain_warnings = explain_result.warnings
|
||||
table_sizes = explain_result.table_sizes or {}
|
||||
except Exception as e:
|
||||
import asyncio as _asyncio
|
||||
logger.warning("EXPLAIN 预检失败(非致命): %s", e)
|
||||
|
||||
# 4. 自动审核 SQL(带超时保护)
|
||||
env = "production" if instance.env == "production" else "development"
|
||||
v2_engine = SQLAuditEngineV2(db_session=db, env=env)
|
||||
import asyncio as _asyncio
|
||||
try:
|
||||
v2_report = await _asyncio.wait_for(
|
||||
v2_engine.audit(data.sql_content, db_type=data.db_type, instance_id=data.instance_id),
|
||||
@@ -525,8 +540,13 @@ async def create_ticket(
|
||||
sql_type=data.sql_type or "",
|
||||
audit_findings=audit_findings_raw,
|
||||
statement_count=len(stmt_list),
|
||||
table_sizes=table_sizes,
|
||||
)
|
||||
|
||||
# EXPLAIN 结果可能升级风险等级
|
||||
if explain_warnings and risk_level == "low":
|
||||
risk_level = "medium"
|
||||
|
||||
# 6. 匹配审批流 — 优先按 instance.workflow_id,其次按 risk_level 匹配默认流
|
||||
workflow = None
|
||||
if instance.workflow_id:
|
||||
@@ -560,6 +580,7 @@ async def create_ticket(
|
||||
priority=data.priority,
|
||||
risk_level=risk_level,
|
||||
workflow_id=workflow.id if workflow else None,
|
||||
exec_strategy=data.exec_strategy or "stop_on_error",
|
||||
applicant_id=current_user.id,
|
||||
status=initial_status,
|
||||
audit_result=json.dumps(v2_report.to_dict(), ensure_ascii=False) if v2_report and v2_report.findings else None,
|
||||
|
||||
@@ -18,6 +18,7 @@ class Ticket(Base):
|
||||
priority = Column(String(10), default="medium") # high / medium / low
|
||||
risk_level = Column(String(20), default="low", comment="自动风险等级: low/medium/high/critical")
|
||||
workflow_id = Column(Integer, ForeignKey("workflow_definitions.id"), nullable=True, comment="关联审批流")
|
||||
exec_strategy = Column(String(20), default="stop_on_error", comment="执行策略: stop_on_error/continue/transaction")
|
||||
applicant_id = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
approver_id = Column(Integer, ForeignKey("users.id"), nullable=True)
|
||||
reason = Column(Text, nullable=True)
|
||||
|
||||
@@ -13,6 +13,7 @@ class TicketCreate(BaseModel):
|
||||
rollback_sql: Optional[str] = None
|
||||
reason: str = Field(..., min_length=5)
|
||||
priority: Literal["high", "medium", "low"] = "medium"
|
||||
exec_strategy: Optional[str] = "stop_on_error"
|
||||
scheduled_at: Optional[datetime] = None
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,335 @@
|
||||
"""EXPLAIN 预检服务 — 工单创建时跑 EXPLAIN 评估 SQL 影响。
|
||||
|
||||
支持:
|
||||
- MySQL: EXPLAIN / EXPLAIN FORMAT=JSON
|
||||
- PostgreSQL: EXPLAIN (FORMAT JSON, COSTS, BUFFERS)
|
||||
- 影响行数预估
|
||||
- 全表扫描检测
|
||||
- 大表 ALTER 风险提示
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from app.models.instance import Instance
|
||||
from app.core.crypto import decrypt_password
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExplainResult:
|
||||
"""EXPLAIN 预检结果。"""
|
||||
sql: str = ""
|
||||
sql_type: str = "" # DQL/DML/DDL
|
||||
estimated_rows: int = 0
|
||||
full_table_scan: bool = False
|
||||
using_index: bool = True
|
||||
warnings: List[str] = field(default_factory=list)
|
||||
explain_raw: Optional[str] = None
|
||||
table_sizes: Dict[str, int] = field(default_factory=dict) # table_name -> row_count
|
||||
error: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"sql": self.sql[:100],
|
||||
"sql_type": self.sql_type,
|
||||
"estimated_rows": self.estimated_rows,
|
||||
"full_table_scan": self.full_table_scan,
|
||||
"using_index": self.using_index,
|
||||
"warnings": self.warnings,
|
||||
"table_sizes": self.table_sizes,
|
||||
"error": self.error,
|
||||
}
|
||||
|
||||
|
||||
# 大表阈值(行数)
|
||||
LARGE_TABLE_THRESHOLD = 100_000
|
||||
# 影响行数警告阈值
|
||||
AFFECTED_ROWS_WARNING = 10_000
|
||||
AFFECTED_ROWS_CRITICAL = 100_000
|
||||
|
||||
|
||||
async def run_explain_preview(
|
||||
sql: str,
|
||||
instance: Instance,
|
||||
target_database: str,
|
||||
) -> ExplainResult:
|
||||
"""对 SQL 执行 EXPLAIN 预检。
|
||||
|
||||
Args:
|
||||
sql: SQL 语句
|
||||
instance: 数据源实例
|
||||
target_database: 目标数据库
|
||||
|
||||
Returns:
|
||||
ExplainResult
|
||||
"""
|
||||
result = ExplainResult(sql=sql)
|
||||
|
||||
sql_upper = sql.strip().upper()
|
||||
|
||||
# 判断 SQL 类型
|
||||
if sql_upper.startswith("SELECT"):
|
||||
result.sql_type = "DQL"
|
||||
elif sql_upper.startswith(("INSERT", "UPDATE", "DELETE")):
|
||||
result.sql_type = "DML"
|
||||
elif sql_upper.startswith(("CREATE", "ALTER", "DROP", "RENAME", "TRUNCATE")):
|
||||
result.sql_type = "DDL"
|
||||
else:
|
||||
result.sql_type = "OTHER"
|
||||
|
||||
db_type = instance.db_type
|
||||
|
||||
try:
|
||||
if db_type in ("mysql",):
|
||||
result = await _explain_mysql(sql, instance, target_database, result)
|
||||
elif db_type in ("postgres", "postgresql"):
|
||||
result = await _explain_postgres(sql, instance, target_database, result)
|
||||
else:
|
||||
# Redis/MongoDB 不支持 EXPLAIN
|
||||
result.warnings.append(f"{db_type} 不支持 EXPLAIN 预检")
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.warning("EXPLAIN 预检失败: %s", e)
|
||||
result.error = str(e)[:200]
|
||||
|
||||
return result
|
||||
|
||||
|
||||
async def _explain_mysql(
|
||||
sql: str, instance: Instance, database: str, result: ExplainResult
|
||||
) -> ExplainResult:
|
||||
"""MySQL EXPLAIN 预检。"""
|
||||
import aiomysql
|
||||
|
||||
password = decrypt_password(instance.password) if instance.password else ""
|
||||
|
||||
conn = await aiomysql.connect(
|
||||
host=instance.host,
|
||||
port=instance.port,
|
||||
user=instance.username,
|
||||
password=password,
|
||||
db=database,
|
||||
connect_timeout=5,
|
||||
)
|
||||
|
||||
try:
|
||||
sql_upper = sql.strip().upper()
|
||||
|
||||
# DML/DDL 类型:查询涉及的表大小
|
||||
if result.sql_type in ("DML", "DDL"):
|
||||
table_names = _extract_table_names(sql)
|
||||
if table_names:
|
||||
result.table_sizes = await _query_table_sizes_mysql(conn, database, table_names)
|
||||
for table_name, row_count in result.table_sizes.items():
|
||||
if row_count > LARGE_TABLE_THRESHOLD:
|
||||
result.warnings.append(f"大表操作: {table_name} 约 {row_count:,} 行")
|
||||
|
||||
# DQL/DML: 跑 EXPLAIN
|
||||
if result.sql_type in ("DQL", "DML") and not sql_upper.startswith("INSERT"):
|
||||
explain_sql = f"EXPLAIN {sql}"
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(explain_sql)
|
||||
rows = await cur.fetchall()
|
||||
columns = [desc[0] for desc in cur.description]
|
||||
|
||||
total_rows = 0
|
||||
for row in rows:
|
||||
row_dict = dict(zip(columns, row))
|
||||
# MySQL EXPLAIN rows 列是预估扫描行数
|
||||
scan_rows = row_dict.get("rows", 0) or 0
|
||||
total_rows += scan_rows
|
||||
|
||||
# 检查是否全表扫描
|
||||
access_type = row_dict.get("type", "")
|
||||
if access_type == "ALL":
|
||||
result.full_table_scan = True
|
||||
result.using_index = False
|
||||
|
||||
# 检查 Extra 信息
|
||||
extra = str(row_dict.get("Extra", ""))
|
||||
if "Using filesort" in extra:
|
||||
result.warnings.append("Using filesort — 需要额外排序操作")
|
||||
if "Using temporary" in extra:
|
||||
result.warnings.append("Using temporary — 使用临时表")
|
||||
|
||||
result.estimated_rows = total_rows
|
||||
|
||||
if total_rows > AFFECTED_ROWS_CRITICAL:
|
||||
result.warnings.append(f"预估影响 {total_rows:,} 行,属于大批量操作")
|
||||
elif total_rows > AFFECTED_ROWS_WARNING:
|
||||
result.warnings.append(f"预估影响 {total_rows:,} 行")
|
||||
|
||||
# DDL ALTER: 检查是否在线 DDL
|
||||
if sql_upper.startswith("ALTER TABLE"):
|
||||
for table_name, row_count in result.table_sizes.items():
|
||||
if row_count > LARGE_TABLE_THRESHOLD:
|
||||
# MySQL 5.6+ 支持 Online DDL,但仍建议在低峰执行
|
||||
result.warnings.append(
|
||||
f"ALTER TABLE 在大表 {table_name}({row_count:,} 行)上执行,建议在维护窗口操作"
|
||||
)
|
||||
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
return result
|
||||
|
||||
|
||||
async def _explain_postgres(
|
||||
sql: str, instance: Instance, database: str, result: ExplainResult
|
||||
) -> ExplainResult:
|
||||
"""PostgreSQL EXPLAIN 预检。"""
|
||||
import asyncpg
|
||||
|
||||
password = decrypt_password(instance.password) if instance.password else ""
|
||||
dsn = f"postgresql://{instance.username}:{password}@{instance.host}:{instance.port}/{database}"
|
||||
|
||||
conn = await asyncpg.connect(dsn, timeout=5)
|
||||
|
||||
try:
|
||||
sql_upper = sql.strip().upper()
|
||||
|
||||
# DML/DDL: 查表大小
|
||||
if result.sql_type in ("DML", "DDL"):
|
||||
table_names = _extract_table_names(sql)
|
||||
if table_names:
|
||||
result.table_sizes = await _query_table_sizes_postgres(conn, table_names)
|
||||
for table_name, row_count in result.table_sizes.items():
|
||||
if row_count > LARGE_TABLE_THRESHOLD:
|
||||
result.warnings.append(f"大表操作: {table_name} 约 {row_count:,} 行")
|
||||
|
||||
# DQL/DML: EXPLAIN (不带 ANALYZE,不实际执行)
|
||||
if result.sql_type in ("DQL", "DML") and not sql_upper.startswith("INSERT"):
|
||||
explain_sql = f"EXPLAIN (FORMAT JSON) {sql}"
|
||||
rows = await conn.fetch(explain_sql)
|
||||
if rows:
|
||||
plan = json.loads(rows[0][0])
|
||||
_parse_pg_plan(plan, result)
|
||||
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _parse_pg_plan(plan_data: Any, result: ExplainResult) -> None:
|
||||
"""递归解析 PostgreSQL EXPLAIN JSON。"""
|
||||
if isinstance(plan_data, list) and plan_data:
|
||||
plan_data = plan_data[0]
|
||||
if isinstance(plan_data, dict):
|
||||
plan = plan_data.get("Plan", plan_data)
|
||||
result.estimated_rows = plan.get("Plan Rows", 0) or 0
|
||||
|
||||
node_type = plan.get("Node Type", "")
|
||||
if "Seq Scan" in node_type:
|
||||
result.full_table_scan = True
|
||||
result.using_index = False
|
||||
|
||||
if result.estimated_rows > AFFECTED_ROWS_CRITICAL:
|
||||
result.warnings.append(f"预估影响 {result.estimated_rows:,} 行,属于大批量操作")
|
||||
elif result.estimated_rows > AFFECTED_ROWS_WARNING:
|
||||
result.warnings.append(f"预估影响 {result.estimated_rows:,} 行")
|
||||
|
||||
# 递归子计划
|
||||
if "Plans" in plan:
|
||||
for sub_plan in plan["Plans"]:
|
||||
_parse_pg_plan(sub_plan, result)
|
||||
|
||||
|
||||
def _extract_table_names(sql: str) -> List[str]:
|
||||
"""从 SQL 中提取表名(简单正则)。"""
|
||||
tables = []
|
||||
sql_upper = sql.upper().strip()
|
||||
|
||||
# UPDATE table SET ...
|
||||
m = re.search(r'\bUPDATE\s+(\w+)', sql_upper)
|
||||
if m:
|
||||
tables.append(m.group(1).lower())
|
||||
|
||||
# DELETE FROM table ...
|
||||
m = re.search(r'\bDELETE\s+FROM\s+(\w+)', sql_upper)
|
||||
if m:
|
||||
tables.append(m.group(1).lower())
|
||||
|
||||
# ALTER TABLE table ...
|
||||
m = re.search(r'\bALTER\s+TABLE\s+(\w+)', sql_upper)
|
||||
if m:
|
||||
tables.append(m.group(1).lower())
|
||||
|
||||
# DROP TABLE table ...
|
||||
m = re.search(r'\bDROP\s+TABLE\s+(?:IF\s+EXISTS\s+)?(\w+)', sql_upper)
|
||||
if m:
|
||||
tables.append(m.group(1).lower())
|
||||
|
||||
# CREATE TABLE table ...
|
||||
m = re.search(r'\bCREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?(\w+)', sql_upper)
|
||||
if m:
|
||||
tables.append(m.group(1).lower())
|
||||
|
||||
# INSERT INTO table ...
|
||||
m = re.search(r'\bINSERT\s+INTO\s+(\w+)', sql_upper)
|
||||
if m:
|
||||
tables.append(m.group(1).lower())
|
||||
|
||||
# SELECT ... FROM table (simple)
|
||||
for m in re.finditer(r'\bFROM\s+(\w+)', sql_upper):
|
||||
t = m.group(1).lower()
|
||||
if t not in ('select', 'where', 'dual', 'information_schema'):
|
||||
tables.append(t)
|
||||
|
||||
# JOIN table
|
||||
for m in re.finditer(r'\bJOIN\s+(\w+)', sql_upper):
|
||||
tables.append(m.group(1).lower())
|
||||
|
||||
return list(dict.fromkeys(tables)) # 去重保序
|
||||
|
||||
|
||||
async def _query_table_sizes_mysql(conn, database: str, table_names: List[str]) -> Dict[str, int]:
|
||||
"""查询 MySQL 表行数。"""
|
||||
sizes = {}
|
||||
if not table_names:
|
||||
return sizes
|
||||
|
||||
placeholders = ", ".join(["%s"] * len(table_names))
|
||||
query = f"""
|
||||
SELECT TABLE_NAME, TABLE_ROWS
|
||||
FROM information_schema.TABLES
|
||||
WHERE TABLE_SCHEMA = %s AND TABLE_NAME IN ({placeholders})
|
||||
"""
|
||||
try:
|
||||
import aiomysql
|
||||
async with conn.cursor(aiomysql.DictCursor) as cur:
|
||||
await cur.execute(query, [database] + [t for t in table_names])
|
||||
rows = await cur.fetchall()
|
||||
for row in rows:
|
||||
sizes[row["TABLE_NAME"].lower()] = row.get("TABLE_ROWS", 0) or 0
|
||||
except Exception as e:
|
||||
logger.warning("查询表大小失败: %s", e)
|
||||
|
||||
return sizes
|
||||
|
||||
|
||||
async def _query_table_sizes_postgres(conn, table_names: List[str]) -> Dict[str, int]:
|
||||
"""查询 PostgreSQL 表行数。"""
|
||||
sizes = {}
|
||||
if not table_names:
|
||||
return sizes
|
||||
|
||||
placeholders = ", ".join([f"'{t}'" for t in table_names])
|
||||
query = f"""
|
||||
SELECT relname, n_live_tup
|
||||
FROM pg_stat_user_tables
|
||||
WHERE relname IN ({placeholders})
|
||||
"""
|
||||
try:
|
||||
rows = await conn.fetch(query)
|
||||
for row in rows:
|
||||
sizes[row["relname"].lower()] = row["n_live_tup"] or 0
|
||||
except Exception as e:
|
||||
logger.warning("查询PG表大小失败: %s", e)
|
||||
|
||||
return sizes
|
||||
@@ -19,6 +19,7 @@ def evaluate_risk_level(
|
||||
sql_type: str = "",
|
||||
audit_findings: Optional[list] = None,
|
||||
statement_count: int = 1,
|
||||
table_sizes: Optional[dict] = None,
|
||||
) -> str:
|
||||
"""评估 SQL 的风险等级。
|
||||
|
||||
@@ -28,6 +29,7 @@ def evaluate_risk_level(
|
||||
sql_type: 前端传入的 SQL 类型标注 DDL/DML/DQL 等
|
||||
audit_findings: 审核引擎返回的 findings 列表
|
||||
statement_count: SQL 语句条数
|
||||
table_sizes: EXPLAIN 预检返回的表大小 {table_name: row_count}
|
||||
|
||||
Returns:
|
||||
risk_level: low / medium / high / critical
|
||||
@@ -93,6 +95,13 @@ def evaluate_risk_level(
|
||||
if high_indicators > 0:
|
||||
return "high"
|
||||
|
||||
# 2.5 大表检测 — EXPLAIN 预检返回的 table_sizes
|
||||
# ALTER/UPDATE/DELETE 涉及大表(>10万行)时升级为 high
|
||||
if table_sizes:
|
||||
max_rows = max(table_sizes.values()) if table_sizes else 0
|
||||
if max_rows > 100_000 and re.search(r"\b(ALTER\s+TABLE|UPDATE|DELETE\s+FROM)\b", sql_upper):
|
||||
return "high"
|
||||
|
||||
# 3. medium 级别 — 常规变更
|
||||
medium_types = {"DDL", "DML", "SET", "HSET", "HMSET"}
|
||||
if sql_type in medium_types:
|
||||
|
||||
@@ -220,29 +220,13 @@ class TicketExecutor:
|
||||
# ===== MySQL 执行 =====
|
||||
|
||||
async def _execute_mysql(self, ticket: Ticket, instance: Instance) -> ExecutionReport:
|
||||
"""执行 MySQL 工单。
|
||||
|
||||
执行流程:
|
||||
1. 语法预校验(sqlparse 解析)
|
||||
2. 建立连接(autocommit=False 显式事务)
|
||||
3. 执行前快照和数据备份
|
||||
4. 在事务中执行 SQL
|
||||
5. 执行后验证
|
||||
6. 验证通过则提交,失败则回滚
|
||||
7. 生成报告
|
||||
|
||||
Args:
|
||||
ticket: 工单对象
|
||||
instance: MySQL 实例连接信息
|
||||
|
||||
Returns:
|
||||
ExecutionReport: 执行报告
|
||||
"""
|
||||
"""执行 MySQL 工单。"""
|
||||
import aiomysql
|
||||
|
||||
report = ExecutionReport()
|
||||
start_time = time.time()
|
||||
sql_content = ticket.sql_content.strip()
|
||||
exec_strategy = getattr(ticket, "exec_strategy", "stop_on_error") or "stop_on_error"
|
||||
|
||||
# 获取超时配置(从 ticket 或选项中)
|
||||
custom_timeout = getattr(ticket, 'options', None) or {}
|
||||
@@ -279,6 +263,7 @@ class TicketExecutor:
|
||||
total_affected = 0
|
||||
output_lines: List[str] = []
|
||||
execution_failed = False
|
||||
failed_count = 0
|
||||
|
||||
statements = self._split_statements(sql_content)
|
||||
for stmt in statements:
|
||||
@@ -293,10 +278,15 @@ class TicketExecutor:
|
||||
total_affected += max(0, affected) if affected is not None else 0
|
||||
output_lines.append(f"[{stmt[:50]}...] 影响 {affected} 行")
|
||||
except Exception as stmt_error:
|
||||
failed_count += 1
|
||||
execution_failed = True
|
||||
output_lines.append(f"[{stmt[:50]}...] 执行失败: {stmt_error}")
|
||||
report.errors.append(f"语句执行失败: {stmt_error}")
|
||||
break
|
||||
|
||||
if exec_strategy == "stop_on_error":
|
||||
break
|
||||
# continue: 跳过失败继续执行
|
||||
# transaction: 任何失败都回滚(默认行为,break 后走 rollback)
|
||||
|
||||
report.affected_rows = total_affected
|
||||
report.raw_output = "\n".join(output_lines)
|
||||
@@ -311,8 +301,14 @@ class TicketExecutor:
|
||||
sql_content, report
|
||||
)
|
||||
|
||||
# 事务决策:验证通过且无执行错误则提交,否则回滚
|
||||
if not execution_failed and report.validation_passed:
|
||||
# 事务决策:根据执行策略处理
|
||||
if exec_strategy == "continue" and failed_count > 0 and failed_count < len(statements):
|
||||
# continue 模式:部分失败也算成功,提交已执行的语句
|
||||
await conn.commit()
|
||||
report.status = "partial_success"
|
||||
report.errors.insert(0, f"部分失败: {failed_count}/{len(statements)} 条语句失败")
|
||||
logger.warning("工单 %d MySQL 部分执行成功: %d 条失败", ticket.id, failed_count)
|
||||
elif not execution_failed and report.validation_passed:
|
||||
await conn.commit()
|
||||
report.status = "success"
|
||||
logger.info("工单 %d MySQL 执行成功,事务已提交", ticket.id)
|
||||
|
||||
@@ -0,0 +1,189 @@
|
||||
"""Phase 2 影响分析测试 — EXPLAIN预检 + 部分失败策略 + 上下文感知规则。"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, AsyncMock, patch
|
||||
|
||||
|
||||
# ===== EXPLAIN 预检: 表名提取 =====
|
||||
|
||||
class TestExtractTableNames:
|
||||
"""SQL 表名提取测试。"""
|
||||
|
||||
def test_update_table(self):
|
||||
from app.services.explain_preview import _extract_table_names
|
||||
assert "users" in _extract_table_names("UPDATE users SET x=1")
|
||||
|
||||
def test_delete_from_table(self):
|
||||
from app.services.explain_preview import _extract_table_names
|
||||
assert "users" in _extract_table_names("DELETE FROM users WHERE id=1")
|
||||
|
||||
def test_alter_table(self):
|
||||
from app.services.explain_preview import _extract_table_names
|
||||
assert "orders" in _extract_table_names("ALTER TABLE orders ADD COLUMN price DECIMAL")
|
||||
|
||||
def test_drop_table(self):
|
||||
from app.services.explain_preview import _extract_table_names
|
||||
assert "temp" in _extract_table_names("DROP TABLE temp")
|
||||
|
||||
def test_drop_table_if_exists(self):
|
||||
from app.services.explain_preview import _extract_table_names
|
||||
assert "temp" in _extract_table_names("DROP TABLE IF EXISTS temp")
|
||||
|
||||
def test_insert_into(self):
|
||||
from app.services.explain_preview import _extract_table_names
|
||||
assert "logs" in _extract_table_names("INSERT INTO logs (msg) VALUES ('hi')")
|
||||
|
||||
def test_select_from(self):
|
||||
from app.services.explain_preview import _extract_table_names
|
||||
names = _extract_table_names("SELECT * FROM users WHERE id=1")
|
||||
assert "users" in names
|
||||
|
||||
def test_join_tables(self):
|
||||
from app.services.explain_preview import _extract_table_names
|
||||
names = _extract_table_names("SELECT * FROM users JOIN orders ON users.id = orders.user_id")
|
||||
assert "users" in names
|
||||
assert "orders" in names
|
||||
|
||||
def test_create_table(self):
|
||||
from app.services.explain_preview import _extract_table_names
|
||||
assert "new_table" in _extract_table_names("CREATE TABLE new_table (id INT)")
|
||||
|
||||
def test_create_table_if_not_exists(self):
|
||||
from app.services.explain_preview import _extract_table_names
|
||||
assert "new_table" in _extract_table_names("CREATE TABLE IF NOT EXISTS new_table (id INT)")
|
||||
|
||||
def test_deduplication(self):
|
||||
from app.services.explain_preview import _extract_table_names
|
||||
names = _extract_table_names("SELECT * FROM users JOIN users AS u2 ON users.id = u2.id")
|
||||
assert names.count("users") == 1
|
||||
|
||||
def test_empty_sql(self):
|
||||
from app.services.explain_preview import _extract_table_names
|
||||
assert _extract_table_names("") == []
|
||||
|
||||
|
||||
# ===== ExplainResult 测试 =====
|
||||
|
||||
class TestExplainResult:
|
||||
"""EXPLAIN 结果数据结构测试。"""
|
||||
|
||||
def test_to_dict(self):
|
||||
from app.services.explain_preview import ExplainResult
|
||||
r = ExplainResult(
|
||||
sql="SELECT * FROM users",
|
||||
sql_type="DQL",
|
||||
estimated_rows=1000,
|
||||
full_table_scan=True,
|
||||
warnings=["Using filesort"],
|
||||
)
|
||||
d = r.to_dict()
|
||||
assert d["sql_type"] == "DQL"
|
||||
assert d["estimated_rows"] == 1000
|
||||
assert d["full_table_scan"] is True
|
||||
assert "Using filesort" in d["warnings"]
|
||||
|
||||
def test_error_field(self):
|
||||
from app.services.explain_preview import ExplainResult
|
||||
r = ExplainResult(error="Connection refused")
|
||||
d = r.to_dict()
|
||||
assert d["error"] == "Connection refused"
|
||||
|
||||
|
||||
# ===== 执行策略测试 =====
|
||||
|
||||
class TestExecStrategy:
|
||||
"""部分失败策略测试。"""
|
||||
|
||||
def test_schema_default_strategy(self):
|
||||
"""TicketCreate schema 默认策略是 stop_on_error。"""
|
||||
from app.schemas.ticket import TicketCreate
|
||||
data = TicketCreate(
|
||||
title="测试工单12345",
|
||||
db_type="mysql",
|
||||
instance_id=1,
|
||||
target_database="test_db",
|
||||
sql_content="SELECT 1",
|
||||
reason="测试原因abcde",
|
||||
)
|
||||
assert data.exec_strategy == "stop_on_error"
|
||||
|
||||
def test_schema_continue_strategy(self):
|
||||
from app.schemas.ticket import TicketCreate
|
||||
data = TicketCreate(
|
||||
title="测试工单12345",
|
||||
db_type="mysql",
|
||||
instance_id=1,
|
||||
target_database="test_db",
|
||||
sql_content="SELECT 1",
|
||||
reason="测试原因abcde",
|
||||
exec_strategy="continue",
|
||||
)
|
||||
assert data.exec_strategy == "continue"
|
||||
|
||||
def test_schema_transaction_strategy(self):
|
||||
from app.schemas.ticket import TicketCreate
|
||||
data = TicketCreate(
|
||||
title="测试工单12345",
|
||||
db_type="mysql",
|
||||
instance_id=1,
|
||||
target_database="test_db",
|
||||
sql_content="SELECT 1",
|
||||
reason="测试原因abcde",
|
||||
exec_strategy="transaction",
|
||||
)
|
||||
assert data.exec_strategy == "transaction"
|
||||
|
||||
|
||||
# ===== 上下文感知规则测试 =====
|
||||
|
||||
class TestRiskEvaluatorTableSizes:
|
||||
"""risk_evaluator 大表检测测试(合并自原 CTX-001)。"""
|
||||
|
||||
def test_alter_large_table_upgraded_to_high(self):
|
||||
"""大表 ALTER 升级为 high。"""
|
||||
from app.services.risk_evaluator import evaluate_risk_level
|
||||
level = evaluate_risk_level(
|
||||
"ALTER TABLE big_table ADD COLUMN status INT",
|
||||
db_type="mysql",
|
||||
table_sizes={"big_table": 500000},
|
||||
)
|
||||
assert level == "high"
|
||||
|
||||
def test_alter_small_table_not_upgraded(self):
|
||||
"""小表 ALTER 保持 medium。"""
|
||||
from app.services.risk_evaluator import evaluate_risk_level
|
||||
level = evaluate_risk_level(
|
||||
"ALTER TABLE users ADD COLUMN age INT",
|
||||
db_type="mysql",
|
||||
table_sizes={"users": 100},
|
||||
)
|
||||
assert level == "medium"
|
||||
|
||||
def test_update_large_table_upgraded_to_high(self):
|
||||
"""大表 UPDATE 无 WHERE 升级为 high。"""
|
||||
from app.services.risk_evaluator import evaluate_risk_level
|
||||
level = evaluate_risk_level(
|
||||
"UPDATE big_table SET status = 1",
|
||||
db_type="mysql",
|
||||
table_sizes={"big_table": 200000},
|
||||
)
|
||||
assert level == "high"
|
||||
|
||||
def test_no_table_sizes_no_upgrade(self):
|
||||
"""无 table_sizes 时不升级。"""
|
||||
from app.services.risk_evaluator import evaluate_risk_level
|
||||
level = evaluate_risk_level(
|
||||
"ALTER TABLE users ADD COLUMN age INT",
|
||||
db_type="mysql",
|
||||
)
|
||||
assert level == "medium"
|
||||
|
||||
def test_select_large_table_stays_low(self):
|
||||
"""SELECT 大表不升级。"""
|
||||
from app.services.risk_evaluator import evaluate_risk_level
|
||||
level = evaluate_risk_level(
|
||||
"SELECT * FROM big_table",
|
||||
db_type="mysql",
|
||||
table_sizes={"big_table": 500000},
|
||||
)
|
||||
assert level == "low"
|
||||
Reference in new issue
Block a user