From 1196746dbe0c2394a3387ef55ffbe419b7d3f9ed Mon Sep 17 00:00:00 2001
From: Erik <erik.nihlen@gmail.com>
Date: Sat, 25 Apr 2026 23:07:00 +0200
Subject: [PATCH] fix(agent): SQL parser robust against sqlglot version drift
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The query_telemetry_db tool was crashing with AttributeError because
exp.AlterTable doesn't exist in this sqlglot version (renamed to Alter).
Made the deny-class list build defensively via getattr and dropped any
classes that the installed sqlglot doesn't expose.

Also broadened the deny list (Alter, AlterColumn, AlterDatabase, Truncate,
Grant, Revoke, Copy) and made the toplevel allowlist tolerant of missing
classes too. The walk() return shape is also normalized in case sqlglot
versions yield (node, parent, key) tuples vs. bare nodes.

Belt-and-suspenders is fine — the GRANT-SELECT-only PG role is the real
write barrier; the parser is just a faster/friendlier reject path.
---
 agent/tools.py | 48 ++++++++++++++++++++++++++++++++----------------
 1 file changed, 32 insertions(+), 16 deletions(-)

diff --git a/agent/tools.py b/agent/tools.py
index 554f587f..c7b3b8fe 100644
--- a/agent/tools.py
+++ b/agent/tools.py
@@ -82,7 +82,17 @@ async def _db() -> asyncpg.Pool:
 # ─── SQL safety ─────────────────────────────────────────────────────
 
 
-_ALLOWED_TOPLEVEL = (exp.Select, exp.With, exp.Union, exp.Subquery)
+_ALLOWED_TOPLEVEL = tuple(
+    cls for cls in (
+        getattr(exp, "Select", None),
+        getattr(exp, "With", None),
+        getattr(exp, "Union", None),
+        getattr(exp, "Subquery", None),
+        getattr(exp, "Intersect", None),
+        getattr(exp, "Except", None),
+    )
+    if cls is not None
+)
 
 
 class SqlNotAllowed(ValueError):
@@ -107,29 +117,35 @@ def assert_read_only(sql: str) -> None:
         raise SqlNotAllowed("only one statement allowed")
 
     stmt = statements[0]
+    if stmt is None:
+        raise SqlNotAllowed("empty parse result")
     if not isinstance(stmt, _ALLOWED_TOPLEVEL):
         raise SqlNotAllowed(
             f"only SELECT / WITH allowed, got {type(stmt).__name__}"
         )
 
     # Walk the tree and reject any DML/DDL hidden inside (e.g. CTE with
-    # INSERT — yes, postgres allows that).
+    # INSERT — yes, postgres allows that). Use getattr so version drift
+    # in sqlglot (renamed classes like AlterTable→Alter) doesn't crash
+    # the whole tool.
+    _DENY_NAMES = (
+        "Insert", "Update", "Delete", "Drop", "Create", "Merge",
+        "Alter", "AlterTable", "AlterColumn", "AlterDatabase",
+        "Truncate", "TruncateTable",
+        "Grant", "Revoke",
+        "Copy",  # PostgreSQL COPY can write files
+    )
+    deny_classes = tuple(
+        cls for cls in (getattr(exp, name, None) for name in _DENY_NAMES)
+        if cls is not None
+    )
     for node in stmt.walk():
-        if isinstance(
-            node,
-            (
-                exp.Insert,
-                exp.Update,
-                exp.Delete,
-                exp.Drop,
-                exp.AlterTable,
-                exp.Create,
-                exp.TruncateTable,
-                exp.Merge,
-            ),
-        ):
+        # walk() returns the node, then in some sqlglot versions a tuple of
+        # (node, parent, key). Normalize.
+        actual = node[0] if isinstance(node, tuple) else node
+        if isinstance(actual, deny_classes):
             raise SqlNotAllowed(
-                f"writes/DDL not allowed (found {type(node).__name__})"
+                f"writes/DDL not allowed (found {type(actual).__name__})"
             )