SQL: Stronger read-only mode, using sqlparse

Author: amotl

CHANGES.md

@@ -11,3 +11,4 @@
 - MCP documentation: Add reference to medium-sized llms.txt context file
 - Boilerplate: Added software tests and CI configuration
 - Documentation: Added development sandbox section
+- SQL: Stronger read-only mode, using `sqlparse`

README.md

@@ -40,10 +40,12 @@ LLMs can still hallucinate and give incorrect answers.
 * What is the storage consumption of my tables, give it in a graph.
 * How can I format a timestamp column to '2019 Jan 21'.
 
-# Data integrity
-We do not recommend letting the LLMs insert data or modify columns by itself, as such we tell the
-LLMs to only allow 'SELECT' statements in the `cratedb_sql` tool, you can jailbreak this against
-our recommendation.
+# Read-only access
+We do not recommend letting LLM-based agents insert or modify data by itself.
+As such, only `SELECT` statements are permitted and forwarded to the database.
+All other operations will raise a `ValueError` exception, unless the
+`CRATEDB_MCP_PERMIT_ALL_STATEMENTS` environment variable will be set
+to a truthy value.
 
 # Install
 ```shell

cratedb_mcp/__main__.py

@@ -1,7 +1,7 @@
 import httpx
 from mcp.server.fastmcp import FastMCP
 
-from .knowledge import DOCUMENTATION_INDEX, Queries
+from .knowledge import DOCUMENTATION_INDEX, Queries, sql_expression_permitted
 from .settings import HTTP_URL
 
 mcp = FastMCP("cratedb-mcp")
@@ -14,7 +14,7 @@ def query_cratedb(query: str) -> list[dict]:
 @mcp.tool(description="Send a SQL query to CrateDB, only 'SELECT' queries are allows, queries that"
                       " modify data, columns or are otherwise deemed un-safe are rejected.")
 def query_sql(query: str):
-    if 'select' not in query.lower():
+    if not sql_expression_permitted(query):
         raise ValueError('Only queries that have a SELECT statement are allowed.')
     return query_cratedb(query)
 

cratedb_mcp/knowledge.py

@@ -1,4 +1,8 @@
 # ruff: noqa: E501
+import sqlparse
+
+from cratedb_mcp.settings import PERMIT_ALL_STATEMENTS
+
 
 class Queries:
     TABLES_METADATA = """
@@ -100,3 +104,32 @@ class Queries:
         "link": "https://raw.githubusercontent.com/crate/cratedb-guide/9ab661997d7704ecbb63af9c3ee33535957e24e6/docs/performance/optimization.rst"
     }
 ]
+
+
+def sql_expression_permitted(expression: str) -> bool:
+    """
+    Validate the SQL expression, only permit read queries.
+
+    FIXME: Revisit implementation, it might be too naive and weak.
+           Issue:    https://github.com/crate/cratedb-mcp/issues/10
+           Question: Does SQLAlchemy provide a solid read-only mode, or any other library?
+    """
+
+    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
+
+    # 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()
+    tokens = [str(item).upper() for item in parsed[0]]
+    if operation != 'SELECT' or (operation == 'SELECT' and 'INTO' in tokens):
+        return False
+    return True

cratedb_mcp/settings.py

@@ -1,3 +1,8 @@
 import os
 
+from attr.converters import to_bool
+
 HTTP_URL: str = os.getenv("CRATEDB_MCP_HTTP_URL", "http://localhost:4200")
+
+# Whether to permit all statements. By default, only SELECT operations are permitted.
+PERMIT_ALL_STATEMENTS: bool = to_bool(os.getenv("CRATEDB_MCP_PERMIT_ALL_STATEMENTS", "false"))

pyproject.toml

@@ -20,7 +20,9 @@ dynamic = [
   "version",
 ]
 dependencies = [
+  "attrs",
   "mcp[cli]>=1.5.0",
+  "sqlparse<0.6",
 ]
 optional-dependencies.develop = [
   "mypy<1.16",

tests/test_knowledge.py

@@ -1,4 +1,5 @@
-from cratedb_mcp.knowledge import DOCUMENTATION_INDEX, Queries
+import cratedb_mcp
+from cratedb_mcp.knowledge import DOCUMENTATION_INDEX, Queries, sql_expression_permitted
 
 
 def test_documentation_index():
@@ -16,3 +17,32 @@ def test_queries():
     assert "sys.health" in Queries.TABLES_METADATA
     assert "WITH partitions_health" in Queries.TABLES_METADATA
     assert "LEFT JOIN" in Queries.TABLES_METADATA
+
+
+def test_sql_expression_create_forbidden():
+    assert sql_expression_permitted("CREATE TABLE foobar AS SELECT 42") is False
+
+
+def test_sql_expression_insert_forbidden():
+    assert sql_expression_permitted("INSERT INTO foobar") is False
+
+
+def test_sql_expression_select_into_forbidden():
+    assert sql_expression_permitted("SELECT * INTO foobar FROM bazqux") is False
+
+
+def test_sql_expression_empty():
+    assert sql_expression_permitted("") is False
+
+
+def test_sql_expression_almost_empty():
+    assert sql_expression_permitted(" ") is False
+
+
+def test_sql_expression_none():
+    assert sql_expression_permitted(None) is False
+
+
+def test_sql_expression_insert_allowed(mocker):
+    mocker.patch.object(cratedb_mcp.knowledge, "PERMIT_ALL_STATEMENTS", True)
+    assert sql_expression_permitted("INSERT INTO foobar") is True

tests/test_mcp.py

@@ -24,19 +24,20 @@ def test_fetch_docs_permitted():
     assert "initcap" in response
 
 
-def test_query_sql_forbidden():
+def test_query_sql_permitted():
+    assert query_sql("SELECT 42")["rows"] == [[42]]
+
+
+def test_query_sql_forbidden_easy():
     with pytest.raises(ValueError) as ex:
         assert "RelationUnknown" in str(query_sql("INSERT INTO foobar (id) VALUES (42) RETURNING id"))
     assert ex.match("Only queries that have a SELECT statement are allowed")
 
 
-def test_query_sql_permitted():
-    assert query_sql("SELECT 42")["rows"] == [[42]]
-
-
-def test_query_sql_permitted_exploit():
-    # FIXME: Read-only protection must become stronger.
-    query_sql("INSERT INTO foobar (operation) VALUES ('select')")
+def test_query_sql_forbidden_sneak_value():
+    with pytest.raises(ValueError) as ex:
+        query_sql("INSERT INTO foobar (operation) VALUES ('select')")
+    assert ex.match("Only queries that have a SELECT statement are allowed")
 
 
 def test_get_table_metadata():