diff --git a/README.md b/README.md index 40f1a20b..02ef74a2 100644 --- a/README.md +++ b/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. 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 3d4b88c6..0f42367e 100644 --- a/src/postgres_mcp/server.py +++ b/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 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..cb6ece20 --- /dev/null +++ b/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, + ) diff --git a/src/postgres_mcp/sql/safe_sql.py b/src/postgres_mcp/sql/safe_sql.py index 6915359c..a87ec529 100644 --- a/src/postgres_mcp/sql/safe_sql.py +++ b/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): diff --git a/tests/integration/test_mcp_protocol.py b/tests/integration/test_mcp_protocol.py new file mode 100644 index 00000000..6caf8b0c --- /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 collections.abc import AsyncGenerator +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +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) 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/sql/test_safe_sql.py b/tests/unit/sql/test_safe_sql.py index fef716e8..f2c84bbf 100644 --- a/tests/unit/sql/test_safe_sql.py +++ b/tests/unit/sql/test_safe_sql.py @@ -811,3 +811,32 @@ async def test_from_clause_dangerous_function_blocked_when_nested(safe_driver): for query in queries: with pytest.raises(ValueError, match="Error validating query"): await safe_driver.execute_query(query) + + +@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_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) diff --git a/tests/unit/test_access_mode.py b/tests/unit/test_access_mode.py index d6653ba2..dcf509d9 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): @@ -138,12 +145,10 @@ async def test_command_line_parsing_default_restricted(): patch("postgres_mcp.server.mcp.run_stdio_async", AsyncMock()), patch("postgres_mcp.server.shutdown", AsyncMock()), ): - # Reset the current_access_mode to UNRESTRICTED import postgres_mcp.server postgres_mcp.server.current_access_mode = AccessMode.UNRESTRICTED - # Run main (partially mocked to avoid actual connection) try: await main() except Exception: @@ -153,6 +158,46 @@ async def test_command_line_parsing_default_restricted(): assert postgres_mcp.server.current_access_mode == AccessMode.RESTRICTED finally: - # 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 + + 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