82 lines
2.6 KiB
Python
82 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
|
|
from app.services.sql_analyzer import SQLAnalyzer
|
|
|
|
MAX_ROWS = 10000
|
|
|
|
|
|
class QueryService:
|
|
_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} 查询")
|
|
|
|
# AST-based SQL validation (replaces keyword matching)
|
|
is_safe, error_msg = QueryService._analyzer.validate_readonly(sql)
|
|
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,
|
|
)
|
|
|
|
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()
|