- P0: 实例访问控制全线接入权限组 (829测试) - P1: permissions字段激活 + 操作级权限检查 (843测试) - 前端: 删除数据权限页面, 权限组合并为唯一入口 - 新增: CodeMirror编辑器 + 查询模板 + 慢查询分析 - 新增: MongoDB管理页面 + Dashboard增强 - 修复: Pydantic V2迁移 + 前端TS错误修复
837 lines
28 KiB
Python
837 lines
28 KiB
Python
"""自动回滚 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}'"
|