81 lines
3.1 KiB
Python
81 lines
3.1 KiB
Python
"""PySpark SQL Guardrails — validator and safe execution wrapper.
|
|
|
|
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.
|
|
"""
|
|
|
|
import re
|
|
|
|
_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"'(?:[^']|'')*'|\"(?:[^\"]|\"\")*\"")
|
|
|
|
|
|
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:
|
|
"""Return the SQL body if it is a single SELECT/INSERT, else raise ValueError.
|
|
|
|
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")
|
|
|
|
body = text[:-1].strip() if text.endswith(";") else text
|
|
if ";" in body:
|
|
raise ValueError("Multiple SQL statements are not allowed")
|
|
|
|
normalized = _LEADING_COMMENTS.sub("", body).lstrip()
|
|
first = normalized.split(None, 1)[0].lower() if normalized else ""
|
|
|
|
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
|
|
|
|
|
|
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))
|