88 lines
3.0 KiB
Python
88 lines
3.0 KiB
Python
|
"""
|
||
|
|
从 DBHub allowed-keywords.ts 等价移植：只读 SQL 判定。
|
||
|
|
参见 dbhub/src/utils/allowed-keywords.ts
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import re
|
||
|
|
from typing import Literal
|
||
|
|
|
||
|
|
from db.dbhub_sql_parser import ConnectorType, strip_comments_and_strings
|
||
|
|
|
||
|
|
ALLOWED_KEYWORDS: dict[ConnectorType, list[str]] = {
|
||
|
|
"postgres": ["select", "with", "explain", "show"],
|
||
|
|
"mysql": ["select", "with", "explain", "show", "describe", "desc"],
|
||
|
|
"mariadb": ["select", "with", "explain", "show", "describe", "desc"],
|
||
|
|
"sqlite": ["select", "with", "explain", "pragma"],
|
||
|
|
"sqlserver": ["select", "with", "explain", "showplan"],
|
||
|
|
}
|
||
|
|
|
||
|
|
_MUTATING = [
|
||
|
|
"insert",
|
||
|
|
"update",
|
||
|
|
"delete",
|
||
|
|
"drop",
|
||
|
|
"alter",
|
||
|
|
"create",
|
||
|
|
"truncate",
|
||
|
|
"merge",
|
||
|
|
"grant",
|
||
|
|
"revoke",
|
||
|
|
"rename",
|
||
|
|
]
|
||
|
|
_mutating_pattern = re.compile(rf"\b(?:{'|'.join(_MUTATING)})\b", re.IGNORECASE)
|
||
|
|
_mutating_pattern_with_replace = re.compile(
|
||
|
|
rf"\b(?:{'|'.join(_MUTATING)}|replace\s+(?:(?:low_priority|delayed)\s+)?into)\b",
|
||
|
|
re.IGNORECASE,
|
||
|
|
)
|
||
|
|
|
||
|
|
_MUTATING_PATTERNS: dict[ConnectorType, re.Pattern[str]] = {
|
||
|
|
"postgres": _mutating_pattern,
|
||
|
|
"mysql": _mutating_pattern_with_replace,
|
||
|
|
"mariadb": _mutating_pattern_with_replace,
|
||
|
|
"sqlite": _mutating_pattern_with_replace,
|
||
|
|
"sqlserver": _mutating_pattern,
|
||
|
|
}
|
||
|
|
|
||
|
|
_SELECT_INTO_PATTERN = re.compile(r"\bselect\b[\s\S]+\binto\b", re.IGNORECASE)
|
||
|
|
|
||
|
|
_EXPLAIN_ANALYZE_PATTERN = re.compile(
|
||
|
|
r"^explain\s+(?:\([^)]*\banalyze\b(?!\s*(?:=\s*)?(?:false|off|0)\b)[^)]*\)|\banalyze\b(?!\s*(?:=\s*)?(?:false|off|0)\b)(?:\s+verbose\b)?)",
|
||
|
|
re.IGNORECASE,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _check_read_only(cleaned_sql: str, connector_type: ConnectorType | str) -> bool:
|
||
|
|
if not cleaned_sql:
|
||
|
|
return False
|
||
|
|
m = re.search(r"\S+", cleaned_sql)
|
||
|
|
first_word = m.group(0) if m else ""
|
||
|
|
keyword_list = ALLOWED_KEYWORDS.get(connector_type, []) # type: ignore[arg-type]
|
||
|
|
if first_word not in keyword_list:
|
||
|
|
return False
|
||
|
|
if first_word == "with":
|
||
|
|
pat = _MUTATING_PATTERNS.get(connector_type, _mutating_pattern) # type: ignore[arg-type]
|
||
|
|
if pat.search(cleaned_sql):
|
||
|
|
return False
|
||
|
|
if first_word in ("select", "with") and _SELECT_INTO_PATTERN.search(cleaned_sql):
|
||
|
|
return False
|
||
|
|
if first_word == "explain":
|
||
|
|
em = _EXPLAIN_ANALYZE_PATTERN.match(cleaned_sql)
|
||
|
|
if em:
|
||
|
|
after_explain = cleaned_sql[em.end() :].strip()
|
||
|
|
if after_explain and not _check_read_only(after_explain, connector_type):
|
||
|
|
return False
|
||
|
|
return True
|
||
|
|
|
||
|
|
|
||
|
|
def is_read_only_sql(sql: str, connector_type: ConnectorType | str) -> bool:
|
||
|
|
"""Check if a SQL query is read-only (DBHub-compatible)."""
|
||
|
|
cleaned = strip_comments_and_strings(sql, connector_type if connector_type in ALLOWED_KEYWORDS else None)
|
||
|
|
cleaned = cleaned.strip().lower()
|
||
|
|
return _check_read_only(cleaned, connector_type)
|
||
|
|
|
||
|
|
|
||
|
|
def allowed_keywords_list(connector_type: ConnectorType | str) -> list[str]:
|
||
|
|
return ALLOWED_KEYWORDS.get(connector_type, []) # type: ignore[arg-type]
|