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 is 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,12 @@
 # ruff: noqa: E501
+import logging
+
+import sqlparse
+
+from cratedb_mcp.settings import PERMIT_ALL_STATEMENTS
+
+logger = logging.getLogger(__name__)
+
 
 class Queries:
     TABLES_METADATA = """
@@ -100,3 +108,48 @@ 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 by default.
+
+    When the `CRATEDB_MCP_PERMIT_ALL_STATEMENTS` environment variable is set,
+    allow all types of statements.
+
+    FIXME: Revisit implementation, it might be too naive or weak.
+           Issue:    https://github.com/crate/cratedb-mcp/issues/10
+           Question: Does SQLAlchemy provide a solid read-only mode, or any other library?
+    """
+    outcome = _sql_expression_permitted(expression)
+    if outcome is True:
+        logger.info(f"Permitted SQL expression: {expression}")
+    else:
+        logger.warning(f"Denied SQL expression: {expression}")
+    return outcome
+
+
+def _sql_expression_permitted(expression: str) -> bool:
+
+    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 for multiple statements (potential SQL injection).
+    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()
+    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,14 @@
 import os
+import warnings
+
+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"))
+
+# TODO: Refactor into code which is not on the module level. Use OOM early.
+if PERMIT_ALL_STATEMENTS:  # pragma: no cover
+    warnings.warn("All types of SQL statements are permitted. This means the LLM "
+                  "agent can write and modify the connected database", category=UserWarning, stacklevel=2)

docs/backlog.md

@@ -0,0 +1,8 @@
+# CrateDB MCP backlog
+
+## Iteration +1
+- Docs: HTTP caching
+- Docs: Load documentation index from custom outline file
+
+## Done
+- SQL: Stronger read-only mode

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,90 @@ 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_select_permitted():
+    """Regular SQL SELECT statements are permitted"""
+    assert sql_expression_permitted("SELECT 42") is True
+    assert sql_expression_permitted(" SELECT 42") is True
+    assert sql_expression_permitted("select 42") is True
+
+
+def test_sql_expression_select_rejected():
+    """Bogus SQL SELECT statements are rejected"""
+    assert sql_expression_permitted(r"--\; select 42") is False
+
+
+def test_sql_expression_insert_allowed(mocker):
+    """When explicitly allowed, permit any kind of statement"""
+    mocker.patch.object(cratedb_mcp.knowledge, "PERMIT_ALL_STATEMENTS", True)
+    assert sql_expression_permitted("INSERT INTO foobar") is True
+
+
+def test_sql_expression_select_multiple_rejected():
+    """Multiple SQL statements are rejected"""
+    assert sql_expression_permitted("SELECT 42; SELECT 42;") is False
+
+
+def test_sql_expression_create_rejected():
+    """DDL statements are rejected"""
+    assert sql_expression_permitted("CREATE TABLE foobar AS SELECT 42") is False
+
+
+def test_sql_expression_insert_rejected():
+    """DML statements are rejected"""
+    assert sql_expression_permitted("INSERT INTO foobar") is False
+
+
+def test_sql_expression_select_into_rejected():
+    """SELECT+DML statements are rejected"""
+    assert sql_expression_permitted("SELECT * INTO foobar FROM bazqux") is False
+
+
+def test_sql_expression_empty_rejected():
+    """Empty statements are rejected"""
+    assert sql_expression_permitted("") is False
+
+
+def test_sql_expression_almost_empty_rejected():
+    """Quasi-empty statements are rejected"""
+    assert sql_expression_permitted(" ") is False
+
+
+def test_sql_expression_none_rejected():
+    """Void statements are rejected"""
+    assert sql_expression_permitted(None) is False
+
+
+def test_sql_expression_multiple_statements_rejected():
+    assert sql_expression_permitted("SELECT 42; INSERT INTO foo VALUES (1)") is False
+
+
+def test_sql_expression_with_comments_rejected():
+    assert sql_expression_permitted(
+        "/* Sneaky comment */ INSERT /* another comment */ INTO foo VALUES (1)") is False
+
+
+def test_sql_expression_update_rejected():
+    """UPDATE statements are rejected"""
+    assert sql_expression_permitted("UPDATE foobar SET column = 'value'") is False
+
+
+def test_sql_expression_delete_rejected():
+    """DELETE statements are rejected"""
+    assert sql_expression_permitted("DELETE FROM foobar") is False
+
+
+def test_sql_expression_truncate_rejected():
+    """TRUNCATE statements are rejected"""
+    assert sql_expression_permitted("TRUNCATE TABLE foobar") is False
+
+
+def test_sql_expression_drop_rejected():
+    """DROP statements are rejected"""
+    assert sql_expression_permitted("DROP TABLE foobar") is False
+
+
+def test_sql_expression_alter_rejected():
+    """ALTER statements are rejected"""
+    assert sql_expression_permitted("ALTER TABLE foobar ADD COLUMN newcol INTEGER") is False

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():