SQL: Use cratedb-sqlparse for implementing read-only mode

Author: amotl

cratedb_mcp/util/sql.py

@@ -2,6 +2,7 @@
 import logging
 import typing as t
 
+import cratedb_sqlparse
 import sqlparse
 
 from cratedb_mcp.settings import PERMIT_ALL_STATEMENTS
@@ -39,12 +40,21 @@ class SqlStatementClassifier:
     expression: str
     permit_all: bool = False
 
+    _parsed_cratedb: t.Any = dataclasses.field(init=False, default=None)
     _parsed_sqlparse: t.Any = dataclasses.field(init=False, default=None)
 
     def __post_init__(self) -> None:
         if self.expression:
             self.expression = self.expression.strip()
 
+    def parse_cratedb(self):
+        """
+        Parse expression using `cratedb-sqlparse` library.
+        """
+        if self._parsed_cratedb is None:
+            self._parsed_cratedb = cratedb_sqlparse.sqlparse(self.expression)
+        return self._parsed_cratedb
+
     def parse_sqlparse(self):
         """
         Parse expression using traditional `sqlparse` library.
@@ -66,7 +76,7 @@ def is_dql(self) -> bool:
             return True
 
         # Parse the SQL statement using `cratedb-sqlparse`.
-        parsed = self.parse_sqlparse()
+        parsed = self.parse_cratedb()
 
         # Reject multiple statements to prevent potential SQL injections.
         if len(parsed) > 1:
@@ -89,7 +99,7 @@ def operation(self) -> str:
         """
         The SQL operation: SELECT, INSERT, UPDATE, DELETE, CREATE, etc.
         """
-        return self._parsed_sqlparse[0].get_type().upper()
+        return self._parsed_cratedb[0].type.upper()
 
     @property
     def is_camouflage(self) -> bool:

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",