Access modes, protocol tests, and function-prefix allowlist

README.md

@@ -247,6 +247,7 @@ Replace `postgresql://...` with your [Postgres database connection URI](https://
 Postgres MCP Pro supports multiple *access modes* to give you control over the operations that the AI agent can perform on the database:
 - **Restricted Mode (default)**: Limits operations to read-only transactions and imposes constraints on resource utilization (presently only execution time). It is suitable for production environments.
 - **Unrestricted Mode**: Allows full read/write access to modify data and schema. It is suitable for development environments. Starting with unrestricted mode active prints a startup warning, because any content the agent reads (web pages, tickets, emails) can carry prompt-injection payloads that reach `execute_sql` unfiltered.
+- **Readonly Mode**: Enforces read-only transactions at the database level without SQL validation. This allows complex queries (nested CTEs, `PERCENTILE_CONT ... WITHIN GROUP`, complex window functions) that pglast may reject, while still preventing writes via PostgreSQL's `READ ONLY` transaction mode. Note that multi-statement queries containing `COMMIT; DROP TABLE ...` will not be caught by SQL validation — protection relies solely on the database transaction.
 
 Restricted mode is the default. To allow write operations, add `--access-mode=unrestricted` to the configuration examples above explicitly.
 
@@ -670,11 +671,15 @@ We reject any SQL that contains `commit` or `rollback` statements.
 Helpfully, the popular Postgres stored procedure languages, including PL/pgSQL and PL/Python, do not allow for `COMMIT` or `ROLLBACK` statements.
 If you have unsafe stored procedure languages enabled on your database, then our read-only protections could be circumvented.
 
-At present, Postgres MCP Pro provides two levels of protection for the database, one at either extreme of the convenience/safety spectrum.
+At present, Postgres MCP Pro provides three levels of protection for the database.
 - "Unrestricted" provides maximum flexibility.
 It is suitable for development environments where speed and flexibility are paramount, and where there is no need to protect valuable or sensitive data.
-- "Restricted" provides a balance between flexibility and safety.
+- "Restricted" provides maximum safety.
 It is suitable for production environments where the database is exposed to untrusted users, and where it is important to protect valuable or sensitive data.
+- "Readonly" provides a middle ground between Unrestricted and Restricted.
+It enforces read-only transactions at the database level (via `BEGIN TRANSACTION READ ONLY`) without pglast SQL validation.
+This allows complex queries that pglast rejects, while still preventing writes.
+However, multi-statement queries like `COMMIT; DROP TABLE` are not caught by SQL validation — protection relies solely on the database transaction.
 
 Unrestricted mode aligns with the approach of [Cursor's auto-run mode](https://docs.cursor.com/chat/tools#auto-run), where the AI agent operates with limited human oversight or approvals.
 We expect auto-run to be deployed in development environments where the consequences of mistakes are low, where databases do not contain valuable or sensitive data, and where they can be recreated or restored from backups when needed.

smithery.yaml

@@ -14,7 +14,7 @@ startCommand:
         description: URI for accessing the database, e.g., postgres://user:password@host:port/database.
       accessMode:
         type: string
-        description: The access mode for the MCP, e.g., "restricted" or "unrestricted".
+        description: The access mode for the MCP, e.g., "restricted", "unrestricted", or "readonly".
   commandFunction:
     # A function that produces the CLI command to start the MCP on stdio.
     |-

src/postgres_mcp/server.py

@@ -11,6 +11,7 @@
 
 import mcp.types as types
 from mcp.server.fastmcp import FastMCP
+from mcp.server.fastmcp.exceptions import ToolError
 from mcp.types import ToolAnnotations
 from pydantic import Field
 from pydantic import validate_call
@@ -26,6 +27,7 @@
 from .index.llm_opt import LLMOptimizerTool
 from .index.presentation import TextPresentation
 from .sql import DbConnPool
+from .sql import ReadOnlySqlDriver
 from .sql import SafeSqlDriver
 from .sql import SqlDriver
 from .sql import check_hypopg_installation_status
@@ -48,21 +50,30 @@ class AccessMode(str, Enum):
 
     RESTRICTED = "restricted"  # Read-only with safety features (default)
     UNRESTRICTED = "unrestricted"  # Unrestricted access
+    READONLY = "readonly"  # Read-only at DB level, no SQL validation
 
 
 # Global variables
 db_connection = DbConnPool()
 current_access_mode = AccessMode.RESTRICTED
+current_allowed_function_prefixes: tuple[str, ...] = ()
 shutdown_in_progress = False
 
 
-async def get_sql_driver() -> SqlDriver | SafeSqlDriver:
+async def get_sql_driver() -> SqlDriver | SafeSqlDriver | ReadOnlySqlDriver:
     """Get the appropriate SQL driver based on the current access mode."""
     base_driver = SqlDriver(conn=db_connection)
 
     if current_access_mode == AccessMode.RESTRICTED:
         logger.debug("Using SafeSqlDriver with restrictions (RESTRICTED mode)")
-        return SafeSqlDriver(sql_driver=base_driver, timeout=30)  # 30 second timeout
+        return SafeSqlDriver(
+            sql_driver=base_driver,
+            timeout=30,
+            allowed_function_prefixes=current_allowed_function_prefixes,
+        )
+    elif current_access_mode == AccessMode.READONLY:
+        logger.debug("Using ReadOnlySqlDriver (READONLY mode)")
+        return ReadOnlySqlDriver(sql_driver=base_driver, timeout=30)  # 30 second timeout
     else:
         logger.debug("Using unrestricted SqlDriver (UNRESTRICTED mode)")
         return base_driver
@@ -540,6 +551,40 @@ async def execute_sql(
         return format_error_response(str(e))
 
 
+def configure_access_mode(access_mode: AccessMode) -> None:
+    """Set the access mode and register execute_sql with a matching description and annotations.
+
+    Any earlier execute_sql registration is replaced, so this can be called more than once
+    (the server calls it at startup; tests call it to switch modes).
+    """
+    global current_access_mode
+    current_access_mode = access_mode
+
+    try:
+        mcp.remove_tool("execute_sql")
+    except ToolError:
+        pass  # not registered yet (first call at startup)
+
+    if access_mode == AccessMode.UNRESTRICTED:
+        mcp.add_tool(
+            execute_sql,
+            description="Execute any SQL query",
+            annotations=ToolAnnotations(
+                title="Execute SQL",
+                destructiveHint=True,
+            ),
+        )
+    else:
+        mcp.add_tool(
+            execute_sql,
+            description="Execute a read-only SQL query",
+            annotations=ToolAnnotations(
+                title="Execute SQL (Read-Only)",
+                readOnlyHint=True,
+            ),
+        )
+
+
 @mcp.tool(
     description="Analyze frequently executed PostgreSQL queries and recommend optimal indexes.",
     annotations=ToolAnnotations(
@@ -676,7 +721,10 @@ async def main():
         type=str,
         choices=[mode.value for mode in AccessMode],
         default=AccessMode.RESTRICTED.value,
-        help="Set SQL access mode: restricted (read-only with protections, default) or unrestricted (full read/write access)",
+        help=(
+            "Set SQL access mode: restricted (read-only + SQL validation, default), "
+            "unrestricted (full read/write access), or readonly (read-only at DB level, no SQL validation)"
+        ),
     )
     parser.add_argument(
         "--transport",
@@ -718,6 +766,12 @@ async def main():
         "environment variable. Useful when defining many database configs so idle servers "
         "release their connections.",
     )
+    parser.add_argument(
+        "--allow-function-prefix",
+        action="append",
+        default=[],
+        help="Allow functions matching this lowercase prefix in restricted mode (repeatable)",
+    )
 
     args = parser.parse_args()
 
@@ -736,10 +790,15 @@ async def main():
                 DbConnPool.DEFAULT_MAX_IDLE,
             )
 
-    # Store the access mode in the global variable
-    global current_access_mode
-    current_access_mode = AccessMode(args.access_mode)
+    # Store the allowed function prefixes (used by get_sql_driver in RESTRICTED mode).
+    global current_allowed_function_prefixes
+    current_allowed_function_prefixes = tuple(p.lower() for p in args.allow_function_prefix)
+
+    # Set the access mode and register execute_sql with a matching description/annotations
+    # (configure_access_mode sets the current_access_mode global and adds the tool).
+    configure_access_mode(AccessMode(args.access_mode))
 
+    # Surface the security posture of the selected mode.
     if current_access_mode == AccessMode.UNRESTRICTED:
         logger.warning(
             "[SECURITY] UNRESTRICTED mode is active: the LLM can execute ANY SQL, "
@@ -756,26 +815,6 @@ async def main():
             "Pass --access-mode=unrestricted for write access."
         )
 
-    # Add the query tool with a description and annotations appropriate to the access mode
-    if current_access_mode == AccessMode.UNRESTRICTED:
-        mcp.add_tool(
-            execute_sql,
-            description="Execute any SQL query",
-            annotations=ToolAnnotations(
-                title="Execute SQL",
-                destructiveHint=True,
-            ),
-        )
-    else:
-        mcp.add_tool(
-            execute_sql,
-            description="Execute a read-only SQL query against the PostgreSQL database and return the results.",
-            annotations=ToolAnnotations(
-                title="Execute SQL (Read-Only)",
-                readOnlyHint=True,
-            ),
-        )
-
     logger.info(f"Starting PostgreSQL MCP Server in {current_access_mode.upper()} mode")
 
     # Get database URL from environment variable or command line

src/postgres_mcp/sql/__init__.py

@@ -9,6 +9,7 @@
 from .extension_utils import get_postgres_version
 from .extension_utils import reset_postgres_version_cache
 from .index import IndexDefinition
+from .readonly_sql import ReadOnlySqlDriver
 from .safe_sql import SafeSqlDriver
 from .sql_driver import DbConnPool
 from .sql_driver import SqlDriver
@@ -18,6 +19,7 @@
     "ColumnCollector",
     "DbConnPool",
     "IndexDefinition",
+    "ReadOnlySqlDriver",
     "SafeSqlDriver",
     "SqlBindParams",
     "SqlDriver",

src/postgres_mcp/sql/readonly_sql.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any
+
+from typing_extensions import LiteralString
+
+from .sql_driver import SqlDriver
+
+logger = logging.getLogger(__name__)
+
+
+class ReadOnlySqlDriver(SqlDriver):
+    """A wrapper around any SqlDriver that enforces read-only mode at the database level.
+
+    Unlike SafeSqlDriver, this driver does NOT perform pglast SQL validation.
+    Instead, it relies on PostgreSQL's READ ONLY transaction mode to prevent writes.
+    This allows complex but safe read-only queries (nested CTEs, PERCENTILE_CONT
+    WITHIN GROUP, complex window functions, etc.) that pglast may reject.
+    """
+
+    def __init__(self, sql_driver: SqlDriver, timeout: float | None = None):
+        """Initialize with an underlying SQL driver and optional timeout.
+
+        Args:
+            sql_driver: The underlying SQL driver to wrap
+            timeout: Optional timeout in seconds for query execution
+        """
+        self.sql_driver = sql_driver
+        self.timeout = timeout
+
+    async def execute_query(
+        self,
+        query: LiteralString,
+        params: list[Any] | None = None,
+        force_readonly: bool = True,  # do not use value passed in
+    ) -> list[SqlDriver.RowResult] | None:  # noqa: UP007
+        """Execute a query with forced read-only mode, without SQL validation."""
+        # NOTE: Always force readonly=True in ReadOnlySqlDriver regardless of what was passed
+        if self.timeout:
+            try:
+                async with asyncio.timeout(self.timeout):
+                    return await self.sql_driver.execute_query(
+                        f"/* crystaldba */ {query}",
+                        params=params,
+                        force_readonly=True,
+                    )
+            except asyncio.TimeoutError as e:
+                logger.warning(f"Query execution timed out after {self.timeout} seconds: {query[:100]}...")
+                raise ValueError(
+                    f"Query execution timed out after {self.timeout} seconds in readonly mode. "
+                    "Consider simplifying your query or increasing the timeout."
+                ) from e
+            except Exception as e:
+                logger.error(f"Error executing query: {e}")
+                raise
+        else:
+            return await self.sql_driver.execute_query(
+                f"/* crystaldba */ {query}",
+                params=params,
+                force_readonly=True,
+            )

src/postgres_mcp/sql/safe_sql.py

@@ -874,15 +874,22 @@ class SafeSqlDriver(SqlDriver):
         "postgis_topology",
     }
 
-    def __init__(self, sql_driver: SqlDriver, timeout: float | None = None):
+    def __init__(
+        self,
+        sql_driver: SqlDriver,
+        timeout: float | None = None,
+        allowed_function_prefixes: tuple[str, ...] = (),
+    ):
         """Initialize with an underlying SQL driver and optional timeout.
 
         Args:
             sql_driver: The underlying SQL driver to wrap
             timeout: Optional timeout in seconds for query execution
+            allowed_function_prefixes: Lowercase prefixes to allow beyond ALLOWED_FUNCTIONS
         """
         self.sql_driver = sql_driver
         self.timeout = timeout
+        self.allowed_function_prefixes = allowed_function_prefixes
 
     def _validate_node(self, node: Node) -> None:
         """Recursively validate a node and all its children"""
@@ -909,7 +916,8 @@ def _validate_node(self, node: Node) -> None:
             match = self.PG_CATALOG_PATTERN.match(func_name)
             unqualified_name = match.group(1) if match else func_name
             if unqualified_name not in self.ALLOWED_FUNCTIONS:
-                raise ValueError(f"Function {func_name} is not allowed")
+                if not any(unqualified_name.startswith(p) for p in self.allowed_function_prefixes):
+                    raise ValueError(f"Function {func_name} is not allowed")
 
         # Reject SELECT statements with locking clauses
         if isinstance(node, SelectStmt) and getattr(node, "lockingClause", None):