update by wensicheng on 0629

skills/pyspark-sql-guardrails/scripts/validate_sql.py

@@ -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())