update by wensicheng on 0629
This commit is contained in:
tao.chen
@@ -1,80 +1,137 @@
|
||||
"""PySpark SQL Guardrails — validator and safe execution wrapper.
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Dependency-free Spark SQL allowlist guard.
|
||||
|
||||
This module is the canonical implementation referenced by
|
||||
.claude/skills/pyspark-sql-guardrails/SKILL.md. Do not redefine the
|
||||
regex inline at call sites; do not write "lighter" versions. Always
|
||||
import this module and route `spark.sql(...)` through `safe_spark_sql`.
|
||||
|
||||
The validator is pure-Python and has no Spark / I/O dependencies. It
|
||||
is safe to import in unit tests, CI, notebooks, and ad-hoc scripts.
|
||||
The guard answers one question only: is this SQL statement type safe enough
|
||||
to show or pass to spark.sql? It does not prove business correctness.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
|
||||
_FORBIDDEN_SQL = re.compile(
|
||||
r"\b(delete|truncate|drop|alter|create|replace|merge|update|msck|refresh|analyze|cache|uncache|vacuum|optimize)\b",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_LEADING_COMMENTS = re.compile(r"\A\s*(?:--[^\n]*\n|/\*.*?\*/\s*)*", re.DOTALL)
|
||||
_STRING_LITERAL = re.compile(r"'(?:[^']|'')*'|\"(?:[^\"]|\"\")*\"")
|
||||
ALLOWED_FIRST = {"select", "with", "insert"}
|
||||
FORBIDDEN = {
|
||||
"delete", "truncate", "drop", "alter", "create", "replace", "merge", "update",
|
||||
"msck", "refresh", "analyze", "cache", "uncache", "vacuum", "optimize",
|
||||
"call", "grant", "revoke",
|
||||
}
|
||||
|
||||
|
||||
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``).
|
||||
@dataclass(frozen=True)
|
||||
class GuardResult:
|
||||
status: str
|
||||
first_keyword: str
|
||||
reason: str = ""
|
||||
|
||||
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 as_dict(self) -> dict[str, str]:
|
||||
return {
|
||||
"status": self.status,
|
||||
"first_keyword": self.first_keyword,
|
||||
"reason": self.reason,
|
||||
}
|
||||
|
||||
|
||||
def _mask(sql: str) -> str:
|
||||
"""Mask strings, comments, and backtick identifiers before token checks."""
|
||||
out: list[str] = []
|
||||
i = 0
|
||||
state = "normal"
|
||||
quote = ""
|
||||
while i < len(sql):
|
||||
ch = sql[i]
|
||||
nxt = sql[i + 1] if i + 1 < len(sql) else ""
|
||||
if state == "normal":
|
||||
if ch in {"'", '"', "`"}:
|
||||
state = "quote"
|
||||
quote = ch
|
||||
out.append(" ")
|
||||
i += 1
|
||||
elif ch == "-" and nxt == "-":
|
||||
state = "line_comment"
|
||||
out.extend(" ")
|
||||
i += 2
|
||||
elif ch == "/" and nxt == "*":
|
||||
state = "block_comment"
|
||||
out.extend(" ")
|
||||
i += 2
|
||||
else:
|
||||
out.append(ch)
|
||||
i += 1
|
||||
elif state == "quote":
|
||||
out.append(" ")
|
||||
if ch == quote:
|
||||
if nxt == quote and quote in {"'", '"'}:
|
||||
out.append(" ")
|
||||
i += 2
|
||||
else:
|
||||
state = "normal"
|
||||
i += 1
|
||||
elif ch == "\\" and nxt:
|
||||
out.append(" ")
|
||||
i += 2
|
||||
else:
|
||||
i += 1
|
||||
elif state == "line_comment":
|
||||
out.append("\n" if ch == "\n" else " ")
|
||||
if ch == "\n":
|
||||
state = "normal"
|
||||
i += 1
|
||||
else:
|
||||
out.append("\n" if ch == "\n" else " ")
|
||||
if ch == "*" and nxt == "/":
|
||||
out.append(" ")
|
||||
state = "normal"
|
||||
i += 2
|
||||
else:
|
||||
i += 1
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def _body(masked: str) -> str:
|
||||
text = masked.strip()
|
||||
return text[:-1].strip() if text.endswith(";") else text
|
||||
|
||||
|
||||
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 ""
|
||||
match = re.search(r"\b([A-Za-z_][\w]*)\b", _body(_mask(sql)))
|
||||
return match.group(1).lower() if match else "unknown"
|
||||
|
||||
|
||||
def assert_select_or_insert(sql: str) -> str:
|
||||
"""Return the SQL body if it is a single SELECT/INSERT, else raise ValueError.
|
||||
def validate_sql(sql: str, *, allow_overwrite: bool = False) -> GuardResult:
|
||||
if not sql or not sql.strip():
|
||||
return GuardResult("FAIL", "unknown", "SQL is empty")
|
||||
|
||||
Hard-allowlist. No I/O, no logging side effects, no Spark dependency.
|
||||
Safe to call in unit tests, CI, notebooks, and ad-hoc scripts.
|
||||
"""
|
||||
text = sql.strip()
|
||||
if not text:
|
||||
raise ValueError("SQL is empty")
|
||||
masked_body = _body(_mask(sql))
|
||||
first = first_keyword(sql)
|
||||
|
||||
body = text[:-1].strip() if text.endswith(";") else text
|
||||
if ";" in body:
|
||||
raise ValueError("Multiple SQL statements are not allowed")
|
||||
if ";" in masked_body:
|
||||
return GuardResult("FAIL", first, "Multiple SQL statements are not allowed")
|
||||
if first not in ALLOWED_FIRST:
|
||||
return GuardResult("FAIL", first, f"Only SELECT, WITH, or INSERT is allowed; got {first!r}")
|
||||
|
||||
normalized = _LEADING_COMMENTS.sub("", body).lstrip()
|
||||
first = normalized.split(None, 1)[0].lower() if normalized else ""
|
||||
lowered = masked_body.lower()
|
||||
found = sorted(word for word in FORBIDDEN if re.search(rf"\b{word}\b", lowered))
|
||||
if found:
|
||||
return GuardResult("FAIL", first, "Forbidden keyword found: " + ", ".join(found))
|
||||
if first == "with" and not re.search(r"\bselect\b", lowered):
|
||||
return GuardResult("FAIL", first, "WITH must contain SELECT")
|
||||
if first == "insert":
|
||||
if not re.search(r"\bselect\b", lowered):
|
||||
return GuardResult("FAIL", first, "INSERT must be based on SELECT")
|
||||
if re.search(r"\binsert\s+overwrite\b", lowered) and not allow_overwrite:
|
||||
return GuardResult("FAIL", first, "INSERT OVERWRITE requires --allow-overwrite")
|
||||
|
||||
if first not in {"select", "with", "insert"}:
|
||||
raise ValueError(f"Only SELECT and INSERT SQL are allowed, got {first!r}")
|
||||
|
||||
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):
|
||||
raise ValueError("WITH statements must be SELECT queries")
|
||||
|
||||
return body
|
||||
return GuardResult("PASS", first, "")
|
||||
|
||||
|
||||
def safe_spark_sql(spark, sql: str):
|
||||
"""Validate then execute. The ONLY sanctioned path to spark.sql()."""
|
||||
return spark.sql(assert_select_or_insert(sql))
|
||||
def assert_allowed_sql(sql: str, *, allow_overwrite: bool = False) -> str:
|
||||
result = validate_sql(sql, allow_overwrite=allow_overwrite)
|
||||
if result.status != "PASS":
|
||||
raise ValueError(result.reason)
|
||||
return sql.strip().rstrip(";")
|
||||
|
||||
|
||||
def safe_spark_sql(spark, sql: str, *, allow_overwrite: bool = False):
|
||||
return spark.sql(assert_allowed_sql(sql, allow_overwrite=allow_overwrite))
|
||||
|
||||
Reference in New Issue
Block a user