1. "first keyword: --"  cosmetic bug
2. re error
This commit is contained in:
tao.chen
2026-06-22 11:31:12 +08:00
parent 6a6df9abf8
commit 87795131a3
3 changed files with 34 additions and 3 deletions
@@ -16,6 +16,34 @@ _FORBIDDEN_SQL = re.compile(
re.IGNORECASE,
)
_LEADING_COMMENTS = re.compile(r"\A\s*(?:--[^\n]*\n|/\*.*?\*/\s*)*", re.DOTALL)
_STRING_LITERAL = re.compile(r"'(?:[^']|'')*'|\"(?:[^\"]|\"\")*\"")
def _strip_string_literals(sql: str) -> str:
"""Replace quoted string literals with empty strings so that the
forbidden-keyword regex does not false-positive on text content
inside quotes (e.g. ``SELECT 'drop' AS note``).
Spark uses single quotes for strings and double quotes for delimited
identifiers; both are replaced because the forbidden-keyword set is
verb-shaped (DROP, DELETE, ...) and never a legitimate identifier.
The doubled-quote escape (``''`` / ``""``) is handled inside the regex.
"""
return _STRING_LITERAL.sub("''", sql)
def first_keyword(sql: str) -> str:
"""Return the lowercased first real SQL keyword after stripping a
leading comment and a trailing semicolon. Pure inspection helper —
does not raise and does not enforce the allowlist.
Intended for CLI reporting so ``first keyword:`` reflects what the
validator actually saw, not the leading comment marker (``--``).
"""
text = sql.strip()
body = text[:-1].strip() if text.endswith(";") else text
normalized = _LEADING_COMMENTS.sub("", body).lstrip()
return normalized.split(None, 1)[0].lower() if normalized else ""
def assert_select_or_insert(sql: str) -> str:
@@ -38,7 +66,7 @@ def assert_select_or_insert(sql: str) -> str:
if first not in {"select", "with", "insert"}:
raise ValueError(f"Only SELECT and INSERT SQL are allowed, got {first!r}")
if _FORBIDDEN_SQL.search(normalized):
if _FORBIDDEN_SQL.search(_strip_string_literals(normalized)):
raise ValueError("Forbidden SQL keyword found")
if first == "with" and not re.search(r"\bselect\b", normalized, re.IGNORECASE):