diff --git a/skills/pyspark-sql-guardrails/scripts/sql_guard.py b/skills/pyspark-sql-guardrails/scripts/sql_guard.py index b7a7ff5..b863305 100644 --- a/skills/pyspark-sql-guardrails/scripts/sql_guard.py +++ b/skills/pyspark-sql-guardrails/scripts/sql_guard.py @@ -16,6 +16,34 @@ _FORBIDDEN_SQL = re.compile( 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: @@ -38,7 +66,7 @@ def assert_select_or_insert(sql: str) -> str: 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): + 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): diff --git a/skills/pyspark-sql-guardrails/scripts/test_sql_guard.py b/skills/pyspark-sql-guardrails/scripts/test_sql_guard.py index 30f30fa..55ca0d8 100644 --- a/skills/pyspark-sql-guardrails/scripts/test_sql_guard.py +++ b/skills/pyspark-sql-guardrails/scripts/test_sql_guard.py @@ -44,6 +44,9 @@ 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 ===") diff --git a/skills/pyspark-sql-guardrails/scripts/validate_sql.py b/skills/pyspark-sql-guardrails/scripts/validate_sql.py index 765cd33..507dad1 100755 --- a/skills/pyspark-sql-guardrails/scripts/validate_sql.py +++ b/skills/pyspark-sql-guardrails/scripts/validate_sql.py @@ -32,7 +32,7 @@ import os SCRIPTS_DIR = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, SCRIPTS_DIR) -from sql_guard import assert_select_or_insert # noqa: E402 +from sql_guard import assert_select_or_insert, first_keyword # noqa: E402 def read_sql() -> str: @@ -51,7 +51,7 @@ def main() -> int: except ValueError as e: print(f"FAIL - {e}") return 1 - first = body.split(None, 1)[0].lower() + first = first_keyword(sql) print(f"PASS - first keyword: {first}, body length: {len(body)}") return 0