Sync from GitHub main @ 015473c9c5ee20c6f880c09fb6f5dfc4070596e1

nl2sql/safety.py

@@ -2,12 +2,13 @@ from __future__ import annotations
 
 import re
 import time
-from typing import List, Pattern
+from typing import List, Pattern, Any, cast
 
 import sqlglot
+from sqlglot import exp
 
 from nl2sql.types import StageResult, StageTrace
-from nl2sql.metrics import safety_blocks_total, stage_duration_ms, safety_checks_total
+from nl2sql.metrics import safety_blocks_total, safety_checks_total
 
 
 # ------------------------- Zero-width & basic regexes -------------------------
@@ -119,6 +120,54 @@ def _remove_comments(body: str) -> str:
     return body
 
 
+def _has_comments(body: str) -> bool:
+    return bool(_LINE_COMMENT_RE.search(body) or _BLOCK_COMMENT_RE.search(body))
+
+
+def _contains_forbidden_ast(root: exp.Expression) -> tuple[bool, str]:
+    """Return (blocked, reason) based on AST nodes/commands."""
+    forbidden_node_names = {
+        "insert",
+        "update",
+        "delete",
+        "drop",
+        "create",
+        "alter",
+        "truncate",
+        "merge",
+        "grant",
+        "revoke",
+        "execute",
+        "call",
+        "copy",
+        "replace",
+    }
+    forbidden_command_markers = ("pragma", "attach", "vacuum", "reindex", "analyze")
+
+    try:
+        walk = getattr(root, "walk", None)
+        if walk is None:
+            return False, ""
+        for node in root.walk():
+            name = type(node).__name__.lower()
+            if name in forbidden_node_names:
+                return True, name
+            if name == "command":
+                sql = ""
+                try:
+                    sql = node.sql(dialect="sqlite").lower()
+                except Exception:
+                    sql = str(node).lower()
+                for kw in forbidden_command_markers:
+                    if kw in sql:
+                        return True, f"command:{kw}"
+    except Exception:
+        # If AST walk fails, be conservative: do not block here (parse/root checks already ran).
+        return False, ""
+
+    return False, ""
+
+
 def _strip_strings(body: str) -> str:
     """
     Remove string literals (so forbidden keyword checks won't fire on quoted text).
@@ -160,8 +209,11 @@ class Safety:
 
     name = "safety"
 
-    def __init__(self, allow_explain: bool = True) -> None:
+    def __init__(
+        self, allow_explain: bool = True, forbid_comments: bool = False
+    ) -> None:
         self.allow_explain = allow_explain
+        self.forbid_comments = forbid_comments
 
     def check(self, sql: str) -> StageResult:
         t0 = time.perf_counter()
@@ -187,6 +239,16 @@ class Safety:
         # 1) sanitize
         body = _sanitize(sql)
 
+        # 1.5) comment policy (block if any comment tokens are present)
+        if self.forbid_comments and _has_comments(body):
+            safety_blocks_total.labels(reason="comments_not_allowed").inc()
+            safety_checks_total.labels(ok="false").inc()
+            return StageResult(
+                ok=False,
+                error=["comments_not_allowed"],
+                trace=StageTrace(stage=self.name, duration_ms=_ms(t0)),
+            )
+
         # 2) single-statement check (semicolon + parser)
         semicolon_count = _count_statements_semicolon(body)
         glot_count = _count_statements_sqlglot(body)
@@ -232,8 +294,8 @@ class Safety:
 
         # 4) read-only root kind (SELECT/EXPLAIN[/WITH])
         try:
-            trees = sqlglot.parse(body)
-            root = trees[0]
+            trees: list[Any] = sqlglot.parse(body)
+            root = cast(exp.Expression, trees[0])
         except Exception as e:
             safety_blocks_total.labels(reason="parse_error").inc()
             safety_checks_total.labels(ok="false").inc()
@@ -252,10 +314,9 @@ class Safety:
         if self.allow_explain and _EXPLAIN_HEAD_RE.match(body):
             remainder = _EXPLAIN_HEAD_RE.sub("", body, count=1).lstrip()
             try:
-                t2 = sqlglot.parse_one(remainder)
+                t2 = cast(exp.Expression, sqlglot.parse_one(remainder))
                 t2_type = type(t2).__name__.lower() if t2 else ""
                 if t2_type in {"select", "with"}:
-                    stage_duration_ms.labels("safety").observe(_ms(t0) / 1.0)
                     safety_checks_total.labels(ok="true").inc()
                     return StageResult(
                         ok=True,
@@ -292,8 +353,19 @@ class Safety:
                 trace=StageTrace(stage=self.name, duration_ms=_ms(t0)),
             )
 
+        # 4.5) AST-based forbidden nodes / commands (defense-in-depth)
+        blocked, reason = _contains_forbidden_ast(root)
+        if blocked:
+            safety_blocks_total.labels(reason="forbidden_ast").inc()
+            safety_checks_total.labels(ok="false").inc()
+            return StageResult(
+                ok=False,
+                error=[f"Forbidden AST: {reason}"],
+                trace=StageTrace(
+                    stage=self.name, duration_ms=_ms(t0), notes={"reason": reason}
+                ),
+            )
         # 5) success
-        stage_duration_ms.labels("safety").observe(_ms(t0) / 1.0)
         safety_checks_total.labels(ok="true").inc()
         return StageResult(
             ok=True,