SQL: Attempt to use cratedb-sqlparse for implementing read-only mode

Author: amotl

cratedb_mcp/knowledge.py

@@ -1,6 +1,7 @@
 # ruff: noqa: E501
 import logging
 
+import cratedb_sqlparse
 import sqlparse
 
 from cratedb_mcp.settings import PERMIT_ALL_STATEMENTS
@@ -131,25 +132,55 @@ def sql_expression_permitted(expression: str) -> bool:
 
 def _sql_expression_permitted(expression: str) -> bool:
 
+    if expression:
+        expression = expression.strip()
+
     if not expression:
         return False
 
     if PERMIT_ALL_STATEMENTS:
         return True
 
     # Parse the SQL statement.
-    parsed = sqlparse.parse(expression.strip())
-    if not parsed:
-        return False
+    parsed = cratedb_sqlparse.sqlparse(expression)
 
-    # Check for multiple statements (potential SQL injection).
+    # Reject multiple statements to prevent potential SQL injections.
     if len(parsed) > 1:
         return False
 
     # Check if the expression is valid and if it's a SELECT statement,
     # also trying to consider `SELECT ... INTO ...` statements.
-    operation = parsed[0].get_type().upper()
+    operation = parsed[0].type.upper()
+    is_select = operation == 'SELECT'
+    is_rejected = _sql_is_select_into(expression) or _sql_is_evasive(expression)
+    if is_select and not is_rejected:
+        return True
+    return False
+
+
+def _sql_is_select_into(expression: str) -> bool:
+    """
+    Helper function using traditional `sqlparse` for catching `SELECT ... INTO ...` statements.
+
+    Examples:
+
+        SELECT * INTO foobar FROM bazqux
+        SELECT * FROM bazqux INTO foobar
+    """
+    parsed = sqlparse.parse(expression)
     tokens = [str(item).upper() for item in parsed[0]]
-    if operation != 'SELECT' or (operation == 'SELECT' and 'INTO' in tokens):
-        return False
-    return True
+    return "INTO" in tokens
+
+
+def _sql_is_evasive(expression: str) -> bool:
+    """
+    Helper function using traditional `sqlparse` for catching evasive SQL statements.
+
+    Reject multiple statements to prevent potential SQL injections.
+
+    Examples:
+
+        SELECT * FROM users; \uff1b DROP TABLE users
+    """
+    parsed = sqlparse.parse(expression)
+    return len(parsed) > 1

pyproject.toml

@@ -21,6 +21,7 @@ dynamic = [
 ]
 dependencies = [
   "attrs",
+  "cratedb-sqlparse==0.0.14",
   "hishel<0.2",
   "mcp[cli]>=1.5.0",
   "sqlparse<0.6",

tests/test_knowledge.py

@@ -55,6 +55,7 @@ def test_sql_expression_insert_rejected():
 def test_sql_expression_select_into_rejected():
     """SELECT+DML statements are rejected"""
     assert sql_expression_permitted("SELECT * INTO foobar FROM bazqux") is False
+    assert sql_expression_permitted("SELECT * FROM bazqux INTO foobar") is False
 
 
 def test_sql_expression_empty_rejected():