Files
opsxsql/backend/app/services/query_service.py
T
loganandClaude Opus 4.7 439155c92e Initial commit: SQL audit platform v0.1.0
FastAPI backend with auth, user, instance, ticket, query, dashboard APIs.
React + TypeScript frontend with mock data UI (10 pages).
SQLite + async SQLAlchemy, JWT auth, role-based access control.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-22 11:11:57 +08:00

79 lines
2.6 KiB
Python

import asyncio
import re
import time
from typing import Optional
from app.models.instance import Instance
from app.schemas.query import QueryResult
FORBIDDEN_KEYWORDS = ["DROP", "TRUNCATE", "DELETE", "UPDATE", "INSERT", "CREATE", "ALTER", "GRANT", "REVOKE"]
MAX_ROWS = 10000
class QueryService:
@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} 查询")
upper_sql = sql.strip().upper()
for kw in FORBIDDEN_KEYWORDS:
if upper_sql.startswith(kw) or f" {kw} " in upper_sql:
raise ValueError(f"查询模式禁止执行 {kw} 操作")
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,
)
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,
)
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} 执行")
import aiomysql
conn = await aiomysql.connect(
host=instance.host, port=instance.port,
user=instance.username, password=instance.password,
db=database,
autocommit=True,
)
try:
async with conn.cursor() as cur:
await asyncio.wait_for(cur.execute(sql), timeout=60)
affected = cur.rowcount
return f"执行成功,影响 {affected} 行"
except asyncio.TimeoutError:
raise ValueError("执行超时(60秒)")
finally:
conn.close()