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))
|
||||
|
||||
@@ -1,67 +1,77 @@
|
||||
"""Standard test set for sql_guard.assert_select_or_insert.
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
Run from the scripts/ subdirectory of the skill:
|
||||
cd .claude/skills/pyspark-sql-guardrails/scripts
|
||||
python3 test_sql_guard.py
|
||||
from sql_guard import assert_allowed_sql, validate_sql
|
||||
|
||||
Or from the project root:
|
||||
python3 .claude/skills/pyspark-sql-guardrails/scripts/test_sql_guard.py
|
||||
PASS_CASES = [
|
||||
("plain SELECT", "SELECT 1"),
|
||||
("WITH SELECT", "WITH x AS (SELECT 1 AS a) SELECT a FROM x"),
|
||||
("INSERT INTO SELECT", "INSERT INTO target_table SELECT * FROM source_table"),
|
||||
("trailing semicolon", "SELECT 1;"),
|
||||
("leading line comment", "-- harmless\nSELECT 1"),
|
||||
("leading block comment", "/* harmless */ SELECT 1"),
|
||||
("danger word in string", "SELECT 'drop table is text' AS note"),
|
||||
("danger word in double quote", 'SELECT "delete" AS note'),
|
||||
("danger word in backtick identifier", "SELECT `drop` FROM t"),
|
||||
("semicolon in string", "SELECT ';' AS semi"),
|
||||
("semicolon in comment", "-- ;\nSELECT 1"),
|
||||
]
|
||||
|
||||
Any change to assert_select_or_insert must be followed by running this
|
||||
set: all EXPECT_RAISE cases must raise, all EXPECT_OK cases must return
|
||||
cleanly. Failures are loud (non-zero exit).
|
||||
"""
|
||||
PASS_WITH_OVERWRITE = [
|
||||
("INSERT OVERWRITE when explicitly allowed", "INSERT OVERWRITE target SELECT * FROM source"),
|
||||
]
|
||||
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Allow running this file directly from the scripts/ subdirectory.
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
from sql_guard import assert_select_or_insert
|
||||
FAIL_CASES = [
|
||||
("DROP", "DROP TABLE x"),
|
||||
("UPDATE", "UPDATE t SET a = 1"),
|
||||
("DELETE", "DELETE FROM x"),
|
||||
("TRUNCATE", "TRUNCATE TABLE x"),
|
||||
("ALTER", "ALTER TABLE x ADD COLUMN a int"),
|
||||
("CREATE", "CREATE TABLE x(id int)"),
|
||||
("REPLACE", "CREATE OR REPLACE VIEW v AS SELECT 1"),
|
||||
("MERGE", "MERGE INTO t USING s ON t.id=s.id WHEN MATCHED THEN UPDATE SET a=1"),
|
||||
("MSCK", "MSCK REPAIR TABLE x"),
|
||||
("REFRESH", "REFRESH TABLE x"),
|
||||
("ANALYZE", "ANALYZE TABLE x COMPUTE STATISTICS"),
|
||||
("CACHE", "CACHE TABLE x"),
|
||||
("UNCACHE", "UNCACHE TABLE x"),
|
||||
("VACUUM", "VACUUM x"),
|
||||
("OPTIMIZE", "OPTIMIZE x"),
|
||||
("CALL", "CALL system.proc()"),
|
||||
("GRANT", "GRANT SELECT ON TABLE x TO user"),
|
||||
("REVOKE", "REVOKE SELECT ON TABLE x FROM user"),
|
||||
("stacked statements", "SELECT 1; DELETE FROM x"),
|
||||
("empty", " "),
|
||||
("WITH without SELECT", "WITH x AS (DELETE FROM t)"),
|
||||
("INSERT VALUES", "INSERT INTO target VALUES (1)"),
|
||||
("INSERT OVERWRITE default blocked", "INSERT OVERWRITE target SELECT * FROM source"),
|
||||
]
|
||||
|
||||
|
||||
def expect_ok(name, sql):
|
||||
try:
|
||||
assert_select_or_insert(sql)
|
||||
print(f" [OK] {name}")
|
||||
except ValueError as e:
|
||||
raise AssertionError(f"{name} should have passed but raised: {e}")
|
||||
def main() -> int:
|
||||
print("=== EXPECT PASS ===")
|
||||
for name, sql in PASS_CASES:
|
||||
result = validate_sql(sql)
|
||||
assert result.status == "PASS", (name, sql, result)
|
||||
assert assert_allowed_sql(sql)
|
||||
print("[OK]", name)
|
||||
|
||||
print("\n=== EXPECT PASS WITH OVERWRITE ===")
|
||||
for name, sql in PASS_WITH_OVERWRITE:
|
||||
result = validate_sql(sql, allow_overwrite=True)
|
||||
assert result.status == "PASS", (name, sql, result)
|
||||
assert assert_allowed_sql(sql, allow_overwrite=True)
|
||||
print("[OK]", name)
|
||||
|
||||
print("\n=== EXPECT FAIL ===")
|
||||
for name, sql in FAIL_CASES:
|
||||
result = validate_sql(sql)
|
||||
assert result.status == "FAIL", (name, sql, result)
|
||||
print("[OK]", name, "->", result.reason)
|
||||
|
||||
print("\nALL TESTS PASSED")
|
||||
return 0
|
||||
|
||||
|
||||
def expect_raise(name, sql):
|
||||
try:
|
||||
assert_select_or_insert(sql)
|
||||
raise AssertionError(f"{name} should have raised but passed")
|
||||
except ValueError:
|
||||
print(f" [OK] {name}: raised as expected")
|
||||
|
||||
|
||||
print("=== EXPECT_OK ===")
|
||||
expect_ok("plain SELECT", "SELECT 1")
|
||||
expect_ok("WITH ... SELECT", "WITH t AS (SELECT 1 AS x) SELECT x FROM t")
|
||||
expect_ok("trailing semicolon", "SELECT 1 FROM dual;")
|
||||
expect_ok("leading -- comment", "-- comment\nSELECT 1")
|
||||
expect_ok("leading /* comment", "/* hi */ SELECT 1")
|
||||
expect_ok("INSERT ... SELECT", "INSERT INTO t SELECT 1")
|
||||
expect_ok("string 'drop'", "SELECT 'drop' AS note FROM dual")
|
||||
expect_ok("quoted identifier", 'SELECT "drop" AS note FROM dual')
|
||||
expect_ok("escaped quote", "SELECT 'don''t drop' AS note FROM dual")
|
||||
|
||||
print()
|
||||
print("=== EXPECT_RAISE ===")
|
||||
expect_raise("DROP", "DROP TABLE loan_order")
|
||||
expect_raise("UPDATE", "UPDATE loan_order SET status = 1")
|
||||
expect_raise("DELETE", "DELETE FROM loan_order")
|
||||
expect_raise("TRUNCATE", "TRUNCATE TABLE loan_order")
|
||||
expect_raise("ALTER", "ALTER TABLE loan_order ADD COLUMN x INT")
|
||||
expect_raise("MERGE", "MERGE INTO t USING s ON t.id = s.id")
|
||||
expect_raise("MSCK", "MSCK REPAIR TABLE loan_order")
|
||||
expect_raise("REFRESH", "REFRESH TABLE loan_order")
|
||||
expect_raise("VACUUM", "VACUUM TABLE loan_order")
|
||||
expect_raise("stacked ;", "SELECT 1; DROP TABLE loan_order")
|
||||
expect_raise("empty", " ")
|
||||
expect_raise("comment-hidden", "-- ok\nDROP TABLE loan_order")
|
||||
|
||||
print()
|
||||
print("ALL TESTS PASSED")
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
|
||||
Executable → Regular
+29
-46
@@ -1,60 +1,43 @@
|
||||
#!/usr/bin/env python3
|
||||
"""One-line SQL guard for the pyspark-sql-guardrails skill.
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
Usage:
|
||||
# Pipe SQL via stdin (preferred for multi-line SQL)
|
||||
echo "SELECT ..." | python3 validate_sql.py
|
||||
cat query.sql | python3 validate_sql.py
|
||||
|
||||
# Or pass SQL as the first argument (single-line, quoting-friendly)
|
||||
python3 validate_sql.py "SELECT 1"
|
||||
|
||||
Exit code:
|
||||
0 -> PASS (SQL is allowed)
|
||||
1 -> FAIL (SQL violates the allowlist; reason printed to stdout)
|
||||
|
||||
This script is the ONLY sanctioned way to validate a generated SQL
|
||||
string before showing it to the user. It is intentionally short so it
|
||||
can be invoked in a single Bash call from the assistant.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
# Allow running from anywhere; resolve bundled sql_guard.py.
|
||||
# The skill layout is:
|
||||
# .claude/skills/pyspark-sql-guardrails/
|
||||
# ├── SKILL.md
|
||||
# └── scripts/
|
||||
# ├── sql_guard.py (this script's sibling)
|
||||
# ├── validate_sql.py
|
||||
# └── test_sql_guard.py
|
||||
SCRIPTS_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
sys.path.insert(0, SCRIPTS_DIR)
|
||||
|
||||
from sql_guard import assert_select_or_insert, first_keyword # noqa: E402
|
||||
from sql_guard import validate_sql
|
||||
|
||||
|
||||
def read_sql() -> str:
|
||||
if len(sys.argv) > 1:
|
||||
return sys.argv[1]
|
||||
def read_sql(args: argparse.Namespace) -> str:
|
||||
if args.file:
|
||||
return Path(args.file).read_text(encoding="utf-8")
|
||||
if args.sql:
|
||||
return " ".join(args.sql)
|
||||
if not sys.stdin.isatty():
|
||||
return sys.stdin.read()
|
||||
print(__doc__, file=sys.stderr)
|
||||
sys.exit(2)
|
||||
raise SystemExit("Provide SQL as arguments, --file, or stdin")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
sql = read_sql()
|
||||
try:
|
||||
body = assert_select_or_insert(sql)
|
||||
except ValueError as e:
|
||||
print(f"FAIL - {e}")
|
||||
return 1
|
||||
first = first_keyword(sql)
|
||||
print(f"PASS - first keyword: {first}, body length: {len(body)}")
|
||||
return 0
|
||||
parser = argparse.ArgumentParser(description="Validate Spark SQL against a safe allowlist.")
|
||||
parser.add_argument("sql", nargs="*", help="SQL string; omitted when using --file or stdin")
|
||||
parser.add_argument("--file", help="Read SQL from a UTF-8 file")
|
||||
parser.add_argument("--json", action="store_true", help="Emit machine-readable JSON")
|
||||
parser.add_argument("--allow-overwrite", action="store_true", help="Allow INSERT OVERWRITE ... SELECT")
|
||||
args = parser.parse_args()
|
||||
|
||||
result = validate_sql(read_sql(args), allow_overwrite=args.allow_overwrite)
|
||||
if args.json:
|
||||
print(json.dumps(result.as_dict(), ensure_ascii=False))
|
||||
elif result.status == "PASS":
|
||||
print(f"PASS - first_keyword={result.first_keyword}")
|
||||
else:
|
||||
print(f"FAIL - first_keyword={result.first_keyword} reason={result.reason}")
|
||||
return 0 if result.status == "PASS" else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
raise SystemExit(main())
|
||||
|
||||
Reference in New Issue
Block a user