init

2026-06-18 13:59:01 +08:00
parent e6552e4f43
commit 13abf8409b
14 changed files with 1676 additions and 0 deletions
@@ -0,0 +1,52 @@
+"""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)
+
+
+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(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))