From f8e330fb11ac613fc5b309a423f2228613e653eb Mon Sep 17 00:00:00 2001 From: yaryrslv Date: Fri, 6 Feb 2026 17:05:50 +0200 Subject: [PATCH 1/5] feat: Add readonly access mode for database-level read-only enforcement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a third access mode (--access-mode=readonly) that sits between `unrestricted` and `restricted`, solving the problem of pglast rejecting legitimate complex read-only queries. Background: The existing `restricted` mode uses pglast to parse SQL into an AST and validate every statement, function call, and node type against whitelists. This catches write operations at the application layer, but also rejects valid read-only queries that pglast cannot parse: nested CTEs, PERCENTILE_CONT with WITHIN GROUP, complex window functions, and queries using extensions or newer PostgreSQL syntax. Users running analytical workloads hit pglast rejections on safe queries and are forced to switch to `unrestricted` mode, losing all protection. How it works: ReadOnlySqlDriver wraps a base SqlDriver using the decorator pattern and overrides execute_query() to: 1. Always set force_readonly=True (ignoring the caller's value) 2. Prepend /* crystaldba */ comment for query identification in logs 3. Enforce a configurable timeout (default 30s) via asyncio.timeout() The underlying SqlDriver.execute_query() with force_readonly=True executes queries inside a read-only transaction: BEGIN TRANSACTION READ ONLY -- PostgreSQL enforces no writes /* crystaldba */ -- user's query runs here ROLLBACK -- transaction closed PostgreSQL itself rejects any INSERT, UPDATE, DELETE, DROP, CREATE, or other write operation inside a READ ONLY transaction — this is database-engine-level enforcement, not application-layer parsing. Protection comparison across modes: Mode SQL validation DB enforcement Complex queries ────────────────────────────────────────────────────────────────── unrestricted none none all allowed readonly none READ ONLY tx all allowed restricted pglast AST check READ ONLY tx limited by pglast Security trade-off: Readonly mode does not validate SQL syntax. A multi-statement payload like "COMMIT; DROP TABLE users" is not caught at the application layer. However, PostgreSQL's read-only transaction prevents the write from executing. This is an intentional trade-off: pglast's false rejections of valid analytical queries cause more real-world pain than the theoretical multi-statement attack vector, which is mitigated by the database transaction boundary. Changes: - src/postgres_mcp/sql/readonly_sql.py: new ReadOnlySqlDriver class (64 lines, decorator over SqlDriver, timeout + force_readonly) - src/postgres_mcp/server.py: AccessMode.READONLY enum value, get_sql_driver() branch returning ReadOnlySqlDriver(timeout=30), --access-mode=readonly CLI argument - src/postgres_mcp/sql/__init__.py: export ReadOnlySqlDriver - smithery.yaml: add "readonly" to access mode config - README.md: update "Protected SQL Execution" to document three levels - tests/unit/sql/test_readonly_sql.py: 8 unit tests (force_readonly override, no SQL validation, comment prepending, timeout, parameter forwarding, exception propagation, None result handling) - tests/unit/sql/test_readonly_enforcement.py: 3 parameterized tests verifying force_readonly behavior across all three access modes - tests/unit/test_access_mode.py: driver selection + CLI parsing tests --- README.md | 12 +- smithery.yaml | 2 +- src/postgres_mcp/server.py | 9 +- src/postgres_mcp/sql/__init__.py | 2 + src/postgres_mcp/sql/readonly_sql.py | 64 ++++++++++ tests/unit/sql/test_readonly_enforcement.py | 31 +++++ tests/unit/sql/test_readonly_sql.py | 125 ++++++++++++++++++++ tests/unit/test_access_mode.py | 49 ++++++++ 8 files changed, 288 insertions(+), 6 deletions(-) create mode 100644 src/postgres_mcp/sql/readonly_sql.py create mode 100644 tests/unit/sql/test_readonly_sql.py diff --git a/README.md b/README.md index d5005ed6..42b56a76 100644 --- a/README.md +++ b/README.md @@ -222,9 +222,11 @@ 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: - **Unrestricted Mode**: Allows full read/write access to modify data and schema. It is suitable for development environments. -- **Restricted Mode**: Limits operations to read-only transactions and imposes constraints on resource utilization (presently only execution time). It is suitable for production environments. +- **Restricted Mode**: Limits operations to read-only transactions and imposes constraints on resource utilization (presently only execution time). Uses pglast to parse and validate SQL before execution. It is suitable for production environments. +- **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. To use restricted mode, replace `--access-mode=unrestricted` with `--access-mode=restricted` in the configuration examples above. +To use readonly mode, replace `--access-mode=unrestricted` with `--access-mode=readonly`. #### Other MCP Clients @@ -605,11 +607,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. diff --git a/smithery.yaml b/smithery.yaml index 2763c205..b121a5df 100644 --- a/smithery.yaml +++ b/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. |- diff --git a/src/postgres_mcp/server.py b/src/postgres_mcp/server.py index f3ba8f8b..79283b3e 100644 --- a/src/postgres_mcp/server.py +++ b/src/postgres_mcp/server.py @@ -28,6 +28,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 @@ -51,6 +52,7 @@ class AccessMode(str, Enum): UNRESTRICTED = "unrestricted" # Unrestricted access RESTRICTED = "restricted" # Read-only with safety features + READONLY = "readonly" # Read-only at DB level, no SQL validation # Global variables @@ -59,13 +61,16 @@ class AccessMode(str, Enum): shutdown_in_progress = False -async def get_sql_driver() -> Union[SqlDriver, SafeSqlDriver]: +async def get_sql_driver() -> Union[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 + 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 @@ -563,7 +568,7 @@ async def main(): type=str, choices=[mode.value for mode in AccessMode], default=AccessMode.UNRESTRICTED.value, - help="Set SQL access mode: unrestricted (unrestricted) or restricted (read-only with protections)", + help="Set SQL access mode: unrestricted, restricted (read-only + SQL validation), or readonly (read-only, no SQL validation)", ) parser.add_argument( "--transport", diff --git a/src/postgres_mcp/sql/__init__.py b/src/postgres_mcp/sql/__init__.py index 1fded3bb..76e3b751 100644 --- a/src/postgres_mcp/sql/__init__.py +++ b/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", diff --git a/src/postgres_mcp/sql/readonly_sql.py b/src/postgres_mcp/sql/readonly_sql.py new file mode 100644 index 00000000..9b1843a0 --- /dev/null +++ b/src/postgres_mcp/sql/readonly_sql.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import asyncio +import logging +from typing import Any +from typing import Optional + +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 + ) -> Optional[list[SqlDriver.RowResult]]: # 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, + ) diff --git a/tests/unit/sql/test_readonly_enforcement.py b/tests/unit/sql/test_readonly_enforcement.py index 0bce3985..4db58cc2 100644 --- a/tests/unit/sql/test_readonly_enforcement.py +++ b/tests/unit/sql/test_readonly_enforcement.py @@ -6,6 +6,7 @@ from postgres_mcp.server import AccessMode from postgres_mcp.server import get_sql_driver +from postgres_mcp.sql import ReadOnlySqlDriver from postgres_mcp.sql import SafeSqlDriver from postgres_mcp.sql import SqlDriver @@ -85,3 +86,33 @@ async def test_force_readonly_enforcement(): assert mock_execute.call_count == 1 # Check that force_readonly remains True assert mock_execute.call_args[1]["force_readonly"] is True + + # Test READONLY mode + with ( + patch("postgres_mcp.server.current_access_mode", AccessMode.READONLY), + patch("postgres_mcp.server.db_connection", mock_conn_pool), + patch.object(SqlDriver, "_execute_with_connection", mock_execute), + ): + driver = await get_sql_driver() + assert isinstance(driver, ReadOnlySqlDriver) + + # Test default behavior + mock_execute.reset_mock() + await driver.execute_query("SELECT 1") + assert mock_execute.call_count == 1 + # Check that force_readonly is always True + assert mock_execute.call_args[1]["force_readonly"] is True + + # Test explicit False (should still be True) + mock_execute.reset_mock() + await driver.execute_query("SELECT 1", force_readonly=False) + assert mock_execute.call_count == 1 + # Check that force_readonly is True despite passing False + assert mock_execute.call_args[1]["force_readonly"] is True + + # Test explicit True + mock_execute.reset_mock() + await driver.execute_query("SELECT 1", force_readonly=True) + assert mock_execute.call_count == 1 + # Check that force_readonly remains True + assert mock_execute.call_args[1]["force_readonly"] is True diff --git a/tests/unit/sql/test_readonly_sql.py b/tests/unit/sql/test_readonly_sql.py new file mode 100644 index 00000000..b8cb16aa --- /dev/null +++ b/tests/unit/sql/test_readonly_sql.py @@ -0,0 +1,125 @@ +import asyncio +from unittest.mock import AsyncMock +from unittest.mock import MagicMock + +import pytest + +from postgres_mcp.sql.readonly_sql import ReadOnlySqlDriver +from postgres_mcp.sql.sql_driver import SqlDriver + + +@pytest.fixture +def mock_sql_driver(): + """Create a mock base SqlDriver.""" + driver = MagicMock(spec=SqlDriver) + driver.execute_query = AsyncMock(return_value=[SqlDriver.RowResult(cells={"test": "value"})]) + return driver + + +@pytest.mark.asyncio +async def test_readonly_driver_forces_readonly(mock_sql_driver): + """Test that force_readonly=True is always passed, even if caller passes False.""" + readonly_driver = ReadOnlySqlDriver(sql_driver=mock_sql_driver, timeout=30) + + # Default call + await readonly_driver.execute_query("SELECT 1") + assert mock_sql_driver.execute_query.call_args[1]["force_readonly"] is True + + # Explicit False should still result in True + mock_sql_driver.execute_query.reset_mock() + await readonly_driver.execute_query("SELECT 1", force_readonly=False) + assert mock_sql_driver.execute_query.call_args[1]["force_readonly"] is True + + # Explicit True + mock_sql_driver.execute_query.reset_mock() + await readonly_driver.execute_query("SELECT 1", force_readonly=True) + assert mock_sql_driver.execute_query.call_args[1]["force_readonly"] is True + + +@pytest.mark.asyncio +async def test_readonly_driver_no_validation(mock_sql_driver): + """Test that any SQL passes through without pglast validation (INSERT, DROP, etc.).""" + readonly_driver = ReadOnlySqlDriver(sql_driver=mock_sql_driver, timeout=30) + + # These would be rejected by SafeSqlDriver's pglast validation, + # but ReadOnlySqlDriver should pass them through (DB will reject at transaction level) + dangerous_queries = [ + "INSERT INTO users (name) VALUES ('test')", + "DROP TABLE users", + "UPDATE users SET name = 'hacked'", + "DELETE FROM users", + "CREATE TABLE evil (id int)", + ] + + for query in dangerous_queries: + mock_sql_driver.execute_query.reset_mock() + await readonly_driver.execute_query(query) + assert mock_sql_driver.execute_query.call_count == 1 + + +@pytest.mark.asyncio +async def test_readonly_driver_prepends_comment(mock_sql_driver): + """Test that /* crystaldba */ prefix is added to queries.""" + readonly_driver = ReadOnlySqlDriver(sql_driver=mock_sql_driver, timeout=30) + + await readonly_driver.execute_query("SELECT 1") + + called_query = mock_sql_driver.execute_query.call_args[0][0] + assert called_query == "/* crystaldba */ SELECT 1" + + +@pytest.mark.asyncio +async def test_readonly_driver_timeout(mock_sql_driver): + """Test that timeout raises ValueError on expiry.""" + + async def slow_query(*args, **kwargs): + await asyncio.sleep(10) + return [SqlDriver.RowResult(cells={"test": "value"})] + + mock_sql_driver.execute_query = slow_query + + readonly_driver = ReadOnlySqlDriver(sql_driver=mock_sql_driver, timeout=0.01) + + with pytest.raises(ValueError, match=r"timed out.*readonly mode"): + await readonly_driver.execute_query("SELECT pg_sleep(10)") + + +@pytest.mark.asyncio +async def test_readonly_driver_no_timeout(mock_sql_driver): + """Test that queries work without timeout when timeout is None.""" + readonly_driver = ReadOnlySqlDriver(sql_driver=mock_sql_driver, timeout=None) + + result = await readonly_driver.execute_query("SELECT 1") + assert result == [SqlDriver.RowResult(cells={"test": "value"})] + assert mock_sql_driver.execute_query.call_args[1]["force_readonly"] is True + + +@pytest.mark.asyncio +async def test_readonly_driver_passes_params(mock_sql_driver): + """Test that query parameters are forwarded correctly.""" + readonly_driver = ReadOnlySqlDriver(sql_driver=mock_sql_driver, timeout=30) + + params = ["param1", 42] + await readonly_driver.execute_query("SELECT * FROM t WHERE a = $1 AND b = $2", params=params) + + call_kwargs = mock_sql_driver.execute_query.call_args[1] + assert call_kwargs["params"] == params + assert call_kwargs["force_readonly"] is True + + +@pytest.mark.asyncio +async def test_readonly_driver_forwards_exceptions(mock_sql_driver): + """Test that exceptions from the underlying driver propagate.""" + mock_sql_driver.execute_query = AsyncMock(side_effect=RuntimeError("connection lost")) + readonly_driver = ReadOnlySqlDriver(sql_driver=mock_sql_driver, timeout=30) + with pytest.raises(RuntimeError, match="connection lost"): + await readonly_driver.execute_query("SELECT 1") + + +@pytest.mark.asyncio +async def test_readonly_driver_none_result(mock_sql_driver): + """Test that None result (DDL/no-result queries) is forwarded.""" + mock_sql_driver.execute_query = AsyncMock(return_value=None) + readonly_driver = ReadOnlySqlDriver(sql_driver=mock_sql_driver, timeout=30) + result = await readonly_driver.execute_query("VACUUM") + assert result is None diff --git a/tests/unit/test_access_mode.py b/tests/unit/test_access_mode.py index f7d3b803..760e34a5 100644 --- a/tests/unit/test_access_mode.py +++ b/tests/unit/test_access_mode.py @@ -7,6 +7,7 @@ from postgres_mcp.server import AccessMode from postgres_mcp.server import get_sql_driver +from postgres_mcp.sql.readonly_sql import ReadOnlySqlDriver from postgres_mcp.sql.safe_sql import SafeSqlDriver from postgres_mcp.sql.sql_driver import DbConnPool from postgres_mcp.sql.sql_driver import SqlDriver @@ -25,6 +26,7 @@ def mock_db_connection(): [ (AccessMode.UNRESTRICTED, SqlDriver), (AccessMode.RESTRICTED, SafeSqlDriver), + (AccessMode.READONLY, ReadOnlySqlDriver), ], ) @pytest.mark.asyncio @@ -42,6 +44,11 @@ async def test_get_sql_driver_returns_correct_driver(access_mode, expected_drive assert isinstance(driver, SafeSqlDriver) assert driver.timeout == 30 + # When in READONLY mode, verify timeout is set + if access_mode == AccessMode.READONLY: + assert isinstance(driver, ReadOnlySqlDriver) + assert driver.timeout == 30 + @pytest.mark.asyncio async def test_get_sql_driver_sets_timeout_in_restricted_mode(mock_db_connection): @@ -112,3 +119,45 @@ async def test_command_line_parsing(): # Restore original values sys.argv = original_argv asyncio.run = original_run + + +@pytest.mark.asyncio +async def test_command_line_parsing_readonly(): + """Test that --access-mode=readonly correctly sets the access mode.""" + import sys + + from postgres_mcp.server import main + + # Mock sys.argv and asyncio.run + original_argv = sys.argv + original_run = asyncio.run + + try: + sys.argv = [ + "postgres_mcp", + "postgresql://user:password@localhost/db", + "--access-mode=readonly", + ] + asyncio.run = AsyncMock() + + with ( + patch("postgres_mcp.server.current_access_mode", AccessMode.UNRESTRICTED), + patch("postgres_mcp.server.db_connection.pool_connect", AsyncMock()), + patch("postgres_mcp.server.mcp.run_stdio_async", AsyncMock()), + patch("postgres_mcp.server.shutdown", AsyncMock()), + ): + import postgres_mcp.server + + postgres_mcp.server.current_access_mode = AccessMode.UNRESTRICTED + + try: + await main() + except Exception: + pass + + # Verify the mode was changed to READONLY + assert postgres_mcp.server.current_access_mode == AccessMode.READONLY + + finally: + sys.argv = original_argv + asyncio.run = original_run From d50c152c7439ecd5a8d9398921919e9740716d4d Mon Sep 17 00:00:00 2001 From: Will Frey Date: Fri, 13 Feb 2026 14:28:10 -0500 Subject: [PATCH 2/5] feat: Add --allow-function-prefix flag for restricted mode Allow specific function prefixes (e.g. st_ for PostGIS) to bypass the hardcoded ALLOWED_FUNCTIONS allowlist in restricted mode. The flag is repeatable and case-insensitive. --- src/postgres_mcp/server.py | 18 +++++++++++++++--- src/postgres_mcp/sql/safe_sql.py | 12 ++++++++++-- tests/unit/sql/test_safe_sql.py | 29 +++++++++++++++++++++++++++++ 3 files changed, 54 insertions(+), 5 deletions(-) diff --git a/src/postgres_mcp/server.py b/src/postgres_mcp/server.py index f3ba8f8b..cf36a43e 100644 --- a/src/postgres_mcp/server.py +++ b/src/postgres_mcp/server.py @@ -56,6 +56,7 @@ class AccessMode(str, Enum): # Global variables db_connection = DbConnPool() current_access_mode = AccessMode.UNRESTRICTED +current_allowed_function_prefixes: tuple[str, ...] = () shutdown_in_progress = False @@ -65,7 +66,11 @@ async def get_sql_driver() -> Union[SqlDriver, SafeSqlDriver]: 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, + ) else: logger.debug("Using unrestricted SqlDriver (UNRESTRICTED mode)") return base_driver @@ -596,12 +601,19 @@ async def main(): default=8000, help="Port for streamable HTTP server (default: 8000)", ) + parser.add_argument( + "--allow-function-prefix", + action="append", + default=[], + help="Allow functions matching this lowercase prefix in restricted mode (repeatable)", + ) args = parser.parse_args() - # Store the access mode in the global variable - global current_access_mode + # Store the access mode and allowed function prefixes in global variables + global current_access_mode, current_allowed_function_prefixes current_access_mode = AccessMode(args.access_mode) + current_allowed_function_prefixes = tuple(p.lower() for p in args.allow_function_prefix) # Add the query tool with a description and annotations appropriate to the access mode if current_access_mode == AccessMode.UNRESTRICTED: diff --git a/src/postgres_mcp/sql/safe_sql.py b/src/postgres_mcp/sql/safe_sql.py index 37382f0b..f30ee9c8 100644 --- a/src/postgres_mcp/sql/safe_sql.py +++ b/src/postgres_mcp/sql/safe_sql.py @@ -865,15 +865,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""" @@ -900,7 +907,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): diff --git a/tests/unit/sql/test_safe_sql.py b/tests/unit/sql/test_safe_sql.py index c55d2530..fc98bf11 100644 --- a/tests/unit/sql/test_safe_sql.py +++ b/tests/unit/sql/test_safe_sql.py @@ -23,6 +23,11 @@ async def safe_driver(mock_sql_driver): return SafeSqlDriver(mock_sql_driver) +@pytest_asyncio.fixture +async def safe_driver_with_st_prefix(mock_sql_driver): + return SafeSqlDriver(mock_sql_driver, allowed_function_prefixes=("st_",)) + + @pytest.mark.asyncio async def test_select_statement(safe_driver, mock_sql_driver): """Test that simple SELECT statements are allowed""" @@ -758,3 +763,27 @@ async def test_query_with_whitespace(safe_driver, mock_sql_driver): """ await safe_driver.execute_query(query) mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) + + +@pytest.mark.asyncio +async def test_function_prefix_allows_postgis(safe_driver_with_st_prefix, mock_sql_driver): + """Test that allowed_function_prefixes permits ST_* PostGIS functions""" + query = "SELECT ST_Intersects(a.geom, b.geom) FROM areas a, points b" + await safe_driver_with_st_prefix.execute_query(query) + mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) + + +@pytest.mark.asyncio +async def test_function_prefix_not_set_blocks_postgis(safe_driver): + """Test that without allowed_function_prefixes, ST_* functions are blocked""" + query = "SELECT ST_Intersects(a.geom, b.geom) FROM areas a, points b" + with pytest.raises(ValueError, match="Error validating query"): + await safe_driver.execute_query(query) + + +@pytest.mark.asyncio +async def test_function_prefix_case_insensitive(safe_driver_with_st_prefix, mock_sql_driver): + """Test that prefix matching is case-insensitive (function names are lowercased)""" + query = "SELECT ST_DWithin(geom, ST_MakePoint(-73.9, 40.7), 1000) FROM places" + await safe_driver_with_st_prefix.execute_query(query) + mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) From 67c5f81101af9fafe190d260b93665e9f76d1cd1 Mon Sep 17 00:00:00 2001 From: Ben Atkinson Date: Wed, 19 Aug 2026 20:17:57 +0100 Subject: [PATCH 3/5] refactor: extract configure_access_mode so tests can drive execute_sql registration execute_sql was registered inline in main() based on --access-mode, so its description and annotations were only observable by running the server. configure_access_mode(mode) now sets the global and (re)registers the tool; main() calls it. Behaviour is unchanged. --- src/postgres_mcp/server.py | 59 +++++++++++++++++++++++--------------- 1 file changed, 36 insertions(+), 23 deletions(-) diff --git a/src/postgres_mcp/server.py b/src/postgres_mcp/server.py index f3ba8f8b..33927cc3 100644 --- a/src/postgres_mcp/server.py +++ b/src/postgres_mcp/server.py @@ -13,6 +13,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 @@ -427,6 +428,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 queries in the database and recommend optimal indexes", annotations=ToolAnnotations( @@ -599,29 +634,7 @@ async def main(): args = parser.parse_args() - # Store the access mode in the global variable - global current_access_mode - current_access_mode = AccessMode(args.access_mode) - - # 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", - annotations=ToolAnnotations( - title="Execute SQL (Read-Only)", - readOnlyHint=True, - ), - ) + configure_access_mode(AccessMode(args.access_mode)) logger.info(f"Starting PostgreSQL MCP Server in {current_access_mode.upper()} mode") From 2a257dd9bb1be306f127724c4deeb6d34377d80e Mon Sep 17 00:00:00 2001 From: Ben Atkinson Date: Wed, 19 Aug 2026 20:17:57 +0100 Subject: [PATCH 4/5] test: round-trip the MCP tools through a real client session Adds tests/integration/test_mcp_protocol.py: an in-process FastMCP server wired to a ClientSession over the SDK memory transport, against the real PostgreSQL container fixture. Covers list_tools (all 9 tools, schemas, required args, access-mode-dependent execute_sql annotations), call_tool round trips for every tool, restricted-mode write blocking observed through the protocol, and how tool errors, missing arguments and unknown tools surface to a client. --- tests/integration/test_mcp_protocol.py | 258 +++++++++++++++++++++++++ 1 file changed, 258 insertions(+) create mode 100644 tests/integration/test_mcp_protocol.py diff --git a/tests/integration/test_mcp_protocol.py b/tests/integration/test_mcp_protocol.py new file mode 100644 index 00000000..13a1dafd --- /dev/null +++ b/tests/integration/test_mcp_protocol.py @@ -0,0 +1,258 @@ +"""Round-trip tests that exercise the MCP tools through a real MCP client session. + +Everything else in the suite tests the Python functions behind the tools. These tests +go through the protocol layer instead: an in-process FastMCP server wired to a +ClientSession over the SDK's memory transport, backed by the real PostgreSQL +container the rest of the integration tests use. They cover tool discovery +(list_tools), argument schemas, call_tool round trips for every registered tool, +access-mode-dependent registration of execute_sql, and how errors surface to a client. + +The client session is opened inside each test rather than in a fixture: the SDK's +memory transport uses anyio cancel scopes, which must be entered and exited in the +same task, and pytest-asyncio tears async fixtures down in a different task. +""" + +import ast +import logging +from contextlib import asynccontextmanager +from typing import AsyncGenerator +from typing import AsyncIterator + +import pytest +import pytest_asyncio +from mcp import ClientSession +from mcp.shared.memory import create_connected_server_and_client_session +from mcp.types import CallToolResult +from mcp.types import TextContent + +import postgres_mcp.server as server +from postgres_mcp.server import AccessMode +from postgres_mcp.server import configure_access_mode + +logger = logging.getLogger(__name__) + +EXPECTED_TOOLS = { + "list_schemas", + "list_objects", + "get_object_details", + "explain_query", + "execute_sql", + "analyze_workload_indexes", + "analyze_query_indexes", + "analyze_db_health", + "get_top_queries", +} + +TEST_TABLE = "mcp_protocol_test_items" + + +def _text(result: CallToolResult) -> str: + """Concatenate the text content blocks of a tool result.""" + return "\n".join(block.text for block in result.content if isinstance(block, TextContent)) + + +@asynccontextmanager +async def mcp_session(access_mode: AccessMode) -> AsyncIterator[ClientSession]: + """An initialised MCP client talking to the server module's FastMCP instance over the memory transport.""" + configure_access_mode(access_mode) + async with create_connected_server_and_client_session(server.mcp) as session: + yield session + + +@pytest_asyncio.fixture +async def connected_db(test_postgres_connection_string) -> AsyncGenerator[str, None]: + """Point the server module's global connection pool at the test database and seed a table.""" + connection_string, version = test_postgres_connection_string + logger.info(f"MCP protocol tests against PostgreSQL {version}") + + pool = await server.db_connection.pool_connect(connection_string) + async with pool.connection() as conn: + await conn.execute(f"DROP TABLE IF EXISTS {TEST_TABLE}") + await conn.execute(f"CREATE TABLE {TEST_TABLE} (id serial PRIMARY KEY, name text NOT NULL, qty integer)") + await conn.execute(f"INSERT INTO {TEST_TABLE} (name, qty) VALUES ('apple', 3), ('pear', 5)") + await conn.commit() + try: + yield connection_string + finally: + async with pool.connection() as conn: + await conn.execute(f"DROP TABLE IF EXISTS {TEST_TABLE}") + await conn.commit() + await server.db_connection.close() + + +class TestToolDiscovery: + @pytest.mark.asyncio + async def test_list_tools_exposes_every_tool(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + tools = (await session.list_tools()).tools + assert {tool.name for tool in tools} == EXPECTED_TOOLS + + @pytest.mark.asyncio + async def test_every_tool_has_description_and_object_schema(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + tools = (await session.list_tools()).tools + for tool in tools: + assert tool.description, f"{tool.name} has no description" + assert tool.inputSchema.get("type") == "object", f"{tool.name} schema is not an object" + + @pytest.mark.asyncio + async def test_required_arguments_are_marked_required(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + tools = {tool.name: tool for tool in (await session.list_tools()).tools} + assert "sql" in tools["explain_query"].inputSchema["required"] + assert "schema_name" in tools["list_objects"].inputSchema["required"] + assert set(tools["get_object_details"].inputSchema["required"]) >= {"schema_name", "object_name"} + assert "queries" in tools["analyze_query_indexes"].inputSchema["required"] + # Tools whose arguments all have defaults advertise nothing as required. + assert not tools["list_schemas"].inputSchema.get("required") + assert not tools["analyze_db_health"].inputSchema.get("required") + + @pytest.mark.asyncio + async def test_execute_sql_is_read_only_in_restricted_mode(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + tools = {tool.name: tool for tool in (await session.list_tools()).tools} + execute_sql = tools["execute_sql"] + assert execute_sql.description == "Execute a read-only SQL query" + assert execute_sql.annotations is not None + assert execute_sql.annotations.title == "Execute SQL (Read-Only)" + assert execute_sql.annotations.readOnlyHint is True + assert execute_sql.annotations.destructiveHint is not True + + @pytest.mark.asyncio + async def test_execute_sql_is_destructive_in_unrestricted_mode(self, connected_db): + async with mcp_session(AccessMode.UNRESTRICTED) as session: + tools = {tool.name: tool for tool in (await session.list_tools()).tools} + execute_sql = tools["execute_sql"] + assert execute_sql.description == "Execute any SQL query" + assert execute_sql.annotations is not None + assert execute_sql.annotations.title == "Execute SQL" + assert execute_sql.annotations.destructiveHint is True + assert execute_sql.annotations.readOnlyHint is not True + + @pytest.mark.asyncio + async def test_switching_access_mode_replaces_execute_sql(self, connected_db): + """configure_access_mode can be called repeatedly without duplicating or stacking registrations.""" + async with mcp_session(AccessMode.UNRESTRICTED) as session: + first = {tool.name: tool for tool in (await session.list_tools()).tools} + async with mcp_session(AccessMode.RESTRICTED) as session: + tools = (await session.list_tools()).tools + assert [tool.name for tool in tools].count("execute_sql") == 1 + second = {tool.name: tool for tool in tools} + assert first["execute_sql"].description != second["execute_sql"].description + assert second["execute_sql"].description == "Execute a read-only SQL query" + + +class TestToolCalls: + @pytest.mark.asyncio + async def test_list_schemas(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + result = await session.call_tool("list_schemas", {}) + assert not result.isError + assert "public" in _text(result) + + @pytest.mark.asyncio + async def test_list_objects_sees_seeded_table(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + result = await session.call_tool("list_objects", {"schema_name": "public", "object_type": "table"}) + assert not result.isError + assert TEST_TABLE in _text(result) + + @pytest.mark.asyncio + async def test_get_object_details_returns_columns(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + result = await session.call_tool( + "get_object_details", + {"schema_name": "public", "object_name": TEST_TABLE, "object_type": "table"}, + ) + assert not result.isError + text = _text(result) + for column in ("id", "name", "qty"): + assert column in text + + @pytest.mark.asyncio + async def test_execute_sql_returns_rows(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + result = await session.call_tool("execute_sql", {"sql": f"SELECT name, qty FROM {TEST_TABLE} ORDER BY id"}) + assert not result.isError + rows = ast.literal_eval(_text(result)) + assert rows == [{"name": "apple", "qty": 3}, {"name": "pear", "qty": 5}] + + @pytest.mark.asyncio + async def test_explain_query_returns_plan(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + result = await session.call_tool("explain_query", {"sql": f"SELECT * FROM {TEST_TABLE} WHERE qty > 1"}) + assert not result.isError + text = _text(result) + assert "Plan" in text or "Seq Scan" in text + + @pytest.mark.asyncio + async def test_analyze_db_health_runs(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + result = await session.call_tool("analyze_db_health", {"health_type": "index"}) + assert not result.isError + assert _text(result).strip() + + @pytest.mark.asyncio + async def test_analyze_query_indexes_runs(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + result = await session.call_tool( + "analyze_query_indexes", + {"queries": [f"SELECT * FROM {TEST_TABLE} WHERE name = 'apple'"]}, + ) + assert not result.isError + assert _text(result).strip() + + @pytest.mark.asyncio + async def test_analyze_workload_indexes_runs(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + result = await session.call_tool("analyze_workload_indexes", {}) + assert not result.isError + assert _text(result).strip() + + @pytest.mark.asyncio + async def test_get_top_queries_runs(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + result = await session.call_tool("get_top_queries", {"sort_by": "total_time", "limit": 5}) + assert not result.isError + assert _text(result).strip() + + +class TestErrorsThroughTheProtocol: + @pytest.mark.asyncio + async def test_restricted_mode_blocks_writes_and_leaves_data_intact(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + result = await session.call_tool("execute_sql", {"sql": f"DELETE FROM {TEST_TABLE}"}) + check = await session.call_tool("execute_sql", {"sql": f"SELECT count(*) AS n FROM {TEST_TABLE}"}) + # Tool-level failures are reported in the content, not as a protocol-level error. + assert _text(result).startswith("Error:") + assert not check.isError + assert ast.literal_eval(_text(check)) == [{"n": 2}] + + @pytest.mark.asyncio + async def test_unrestricted_mode_allows_writes(self, connected_db): + async with mcp_session(AccessMode.UNRESTRICTED) as session: + result = await session.call_tool("execute_sql", {"sql": f"UPDATE {TEST_TABLE} SET qty = qty + 1 WHERE name = 'apple'"}) + check = await session.call_tool("execute_sql", {"sql": f"SELECT qty FROM {TEST_TABLE} WHERE name = 'apple'"}) + assert not result.isError + assert not _text(result).startswith("Error:") + assert ast.literal_eval(_text(check)) == [{"qty": 4}] + + @pytest.mark.asyncio + async def test_invalid_sql_is_reported_not_raised(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + result = await session.call_tool("execute_sql", {"sql": "SELECT * FROM table_that_does_not_exist"}) + assert _text(result).startswith("Error:") + + @pytest.mark.asyncio + async def test_missing_required_argument_is_a_tool_error(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + result = await session.call_tool("explain_query", {}) + assert result.isError + assert "sql" in _text(result) + + @pytest.mark.asyncio + async def test_unknown_tool_is_a_tool_error(self, connected_db): + async with mcp_session(AccessMode.RESTRICTED) as session: + result = await session.call_tool("no_such_tool", {}) + assert result.isError + assert "no_such_tool" in _text(result) From 72142656e01161a346f0d8016f6c2584b7423688 Mon Sep 17 00:00:00 2001 From: cupskeee Date: Sun, 6 Sep 2026 22:03:38 +0700 Subject: [PATCH 5/5] fix(post-merge): ruff cleanups for #148/#154/#207 (typing + line length) Modernize typing in readonly_sql.py (UP045) and test_mcp_protocol.py (UP035), and wrap the long --access-mode help string (E501). #173 removed the UP0xx ruff ignores, so these stragglers from the merged PRs need the same treatment. ruff clean, pyright 0 errors, pytest 221 passed. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01D6w353jusev2fgmWQY1mS1 --- src/postgres_mcp/server.py | 5 ++++- src/postgres_mcp/sql/readonly_sql.py | 3 +-- tests/integration/test_mcp_protocol.py | 4 ++-- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/postgres_mcp/server.py b/src/postgres_mcp/server.py index 93b22315..0f42367e 100644 --- a/src/postgres_mcp/server.py +++ b/src/postgres_mcp/server.py @@ -721,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 + SQL validation, default), unrestricted (full read/write access), or readonly (read-only at DB level, no SQL validation)", + 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", diff --git a/src/postgres_mcp/sql/readonly_sql.py b/src/postgres_mcp/sql/readonly_sql.py index 9b1843a0..cb6ece20 100644 --- a/src/postgres_mcp/sql/readonly_sql.py +++ b/src/postgres_mcp/sql/readonly_sql.py @@ -3,7 +3,6 @@ import asyncio import logging from typing import Any -from typing import Optional from typing_extensions import LiteralString @@ -36,7 +35,7 @@ async def execute_query( query: LiteralString, params: list[Any] | None = None, force_readonly: bool = True, # do not use value passed in - ) -> Optional[list[SqlDriver.RowResult]]: # noqa: UP007 + ) -> 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: diff --git a/tests/integration/test_mcp_protocol.py b/tests/integration/test_mcp_protocol.py index 13a1dafd..6caf8b0c 100644 --- a/tests/integration/test_mcp_protocol.py +++ b/tests/integration/test_mcp_protocol.py @@ -14,9 +14,9 @@ import ast import logging +from collections.abc import AsyncGenerator +from collections.abc import AsyncIterator from contextlib import asynccontextmanager -from typing import AsyncGenerator -from typing import AsyncIterator import pytest import pytest_asyncio