Add Limits and Maximum Payload Size

.env.example

@@ -0,0 +1,15 @@
+# PostgreSQL MCP Configuration
+
+# Maximum allowed page size for query results (default: 500)
+# This controls the upper limit of rows that can be returned in a single query
+POSTGRES_MCP_MAX_PAGE_SIZE=500
+
+# Default page size for query results (default: 100)
+# This is the default number of rows returned if no page size is specified
+POSTGRES_MCP_DEFAULT_PAGE_SIZE=100
+
+# Database connection string (required)
+DATABASE_URI=postgresql://user:password@localhost:5432/dbname
+
+# Allowed hosts for incoming connections (comma-separated, default: localhost,127.0.0.1)
+POSTGRES_MCP_ALLOWED_HOSTS=localhost:*,127.0.0.1

pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "attrs>=25.4.0",
     "psycopg-pool>=3.3.0",
     "instructor>=1.14.4",
+    "dotenv>=0.9.9",
 ]
 license = "mit"
 license-files = ["LICENSE"]

src/postgres_mcp/config.py

@@ -0,0 +1,78 @@
+"""Configuration management for postgres-mcp."""
+
+import os
+
+
+class Config:
+    """Configuration settings loaded from environment variables."""
+
+    def __init__(self):
+        """Initialize configuration from environment variables."""
+        self._load_config()
+
+    def _load_config(self):
+        """Load configuration from environment variables."""
+        # Maximum allowed page size for queries (default: 500)
+        max_page_size_str = os.getenv("POSTGRES_MCP_MAX_PAGE_SIZE", "500")
+        try:
+            self._max_page_size = int(max_page_size_str)
+            if self._max_page_size < 1:
+                raise ValueError("POSTGRES_MCP_MAX_PAGE_SIZE must be at least 1")
+        except ValueError as e:
+            raise ValueError(f"Invalid POSTGRES_MCP_MAX_PAGE_SIZE value '{max_page_size_str}': {e}") from e
+
+        # Maximum payload size in MB (default: 5)
+        max_payload_size_mb_str = os.getenv("POSTGRES_MCP_MAX_PAYLOAD_SIZE_MB", "5")
+        try:
+            self._max_payload_size_mb = int(max_payload_size_mb_str)
+            if self._max_payload_size_mb < 1:
+                raise ValueError("POSTGRES_MCP_MAX_PAYLOAD_SIZE_MB must be at least 1")
+        except ValueError as e:
+            raise ValueError(f"Invalid POSTGRES_MCP_MAX_PAYLOAD_SIZE_MB value '{max_payload_size_mb_str}': {e}") from e
+
+        # Default page size for queries (default: 100)
+        default_page_size_str = os.getenv("POSTGRES_MCP_DEFAULT_PAGE_SIZE", "100")
+        try:
+            self._default_page_size = int(default_page_size_str)
+            if self._default_page_size < 1:
+                raise ValueError("POSTGRES_MCP_DEFAULT_PAGE_SIZE must be at least 1")
+            if self._default_page_size > self._max_page_size:
+                raise ValueError(
+                    f"POSTGRES_MCP_DEFAULT_PAGE_SIZE ({self._default_page_size}) cannot exceed POSTGRES_MCP_MAX_PAGE_SIZE ({self._max_page_size})"
+                )
+        except ValueError as e:
+            raise ValueError(f"Invalid POSTGRES_MCP_DEFAULT_PAGE_SIZE value '{default_page_size_str}': {e}") from e
+
+        allowed_hosts_str = os.getenv("POSTGRES_MCP_ALLOWED_HOSTS", "localhost,localhost:*,127.0.0.1")
+        try:
+            self._allowed_hosts = [host.strip() for host in allowed_hosts_str.split(",")]
+        except Exception as e:
+            raise ValueError(f"Invalid POSTGRES_MCP_ALLOWED_HOSTS value '{allowed_hosts_str}': {e}") from e
+
+    @property
+    def max_page_size(self) -> int:
+        """Get the maximum allowed page size for queries."""
+        return self._max_page_size
+
+    @property
+    def default_page_size(self) -> int:
+        """Get the default page size for queries."""
+        return self._default_page_size
+
+    @property
+    def max_payload_size_mb(self) -> int:
+        """Get the maximum allowed payload size in MB."""
+        return self._max_payload_size_mb
+
+    @property
+    def allowed_hosts(self) -> list[str]:
+        """Get the list of allowed hosts."""
+        return self._allowed_hosts
+
+    def reload(self):
+        """Reload configuration from environment variables."""
+        self._load_config()
+
+
+# Module-level configuration instance - import this directly
+config = Config()

src/postgres_mcp/server.py

@@ -11,15 +11,22 @@
 from urllib.parse import urlparse
 
 import mcp.types as types
+from dotenv import load_dotenv
 from mcp.server.fastmcp import FastMCP
+from mcp.server.transport_security import TransportSecuritySettings
 from mcp.types import ToolAnnotations
 from pydantic import Field
 from pydantic import validate_call
 
+# Load environment variables before importing local modules that may use them
+load_dotenv()
+
+# ruff: noqa: E402
 from postgres_mcp.index.dta_calc import DatabaseTuningAdvisor
 
 from .artifacts import ErrorResult
 from .artifacts import ExplainPlanArtifact
+from .config import config
 from .database_health import DatabaseHealthTool
 from .database_health import HealthType
 from .explain import ExplainPlanTool
@@ -38,7 +45,14 @@
 from .utils.url import fix_connection_url
 
 # Initialize FastMCP with default settings
-mcp = FastMCP("postgres-mcp")
+
+
+mcp = FastMCP(
+    "postgres-mcp",
+    transport_security=TransportSecuritySettings(
+        allowed_hosts=config.allowed_hosts,
+    ),
+)
 
 # Constants
 PG_STAT_STATEMENTS = "pg_stat_statements"
@@ -270,11 +284,27 @@ async def explain_query(
 # Query function declaration without the decorator - we'll add it dynamically based on access mode
 async def execute_sql(
     sql: str = Field(description="SQL to run", default="all"),
+    page_size: int = Field(
+        description=f"Number of rows to return (1-{config.max_page_size}",
+        default=config.default_page_size,
+        ge=1,
+        le=config.max_page_size,
+    ),
+    offset: int = Field(description="Number of rows to skip for pagination", default=0, ge=0),
+    parameters: list[str | int | float | bool | None] = Field(
+        description="Optional array of parameters for parameterized queries",
+        default_factory=list,
+    ),
 ) -> ResponseType:
     """Executes a SQL query against the database."""
     try:
         sql_driver = await sql_driver_module.get_sql_driver()
-        rows = await sql_driver.execute_query(sql)  # type: ignore
+        rows = await sql_driver.execute_query(
+            sql,  # type: ignore
+            params=parameters if parameters else None,
+            page_size=page_size,
+            offset=offset,
+        )
         if rows is None:
             return format_text_response("No results")
         return format_text_response(list([r.cells for r in rows]))
@@ -460,11 +490,11 @@ async def main():
 
     # Add the query tool with a description appropriate to the access mode
     if sql_driver_module.current_access_mode == AccessMode.UNRESTRICTED:
-        mcp.add_tool(execute_sql, description="Execute any SQL query")
+        mcp.add_tool(execute_sql, description="Execute any SQL query with pagination support")
     else:
         mcp.add_tool(
             execute_sql,
-            description="Execute a read-only SQL query",
+            description="Execute a read-only SQL query with pagination support",
             annotations=ToolAnnotations(
                 title="Execute SQL (Read-Only)",
                 readOnlyHint=True,

src/postgres_mcp/sql/safe_sql.py

@@ -80,6 +80,7 @@
 from psycopg.sql import Literal
 from typing_extensions import LiteralString
 
+from ..config import config
 from .sql_driver import SqlDriver
 
 logger = logging.getLogger(__name__)
@@ -982,8 +983,14 @@ async def execute_query(
         query: LiteralString,
         params: list[Any] | None = None,
         force_readonly: bool = True,  # do not use value passed in
+        page_size: int | None = None,
+        offset: int = 0,
     ) -> Optional[list[SqlDriver.RowResult]]:  # noqa: UP007
         """Execute a query after validating it is safe"""
+        # Use configured default if page_size not specified
+        if page_size is None:
+            page_size = config.default_page_size
+
         self._validate(query)
 
         # NOTE: Always force readonly=True in SafeSqlDriver regardless of what was passed
@@ -994,6 +1001,8 @@ async def execute_query(
                         f"/* crystaldba */ {query}",
                         params=params,
                         force_readonly=True,
+                        page_size=page_size,
+                        offset=offset,
                     )
             except asyncio.TimeoutError as e:
                 logger.warning(f"Query execution timed out after {self.timeout} seconds: {query[:100]}...")
@@ -1009,6 +1018,8 @@ async def execute_query(
                 f"/* crystaldba */ {query}",
                 params=params,
                 force_readonly=True,
+                page_size=page_size,
+                offset=offset,
             )
 
     @staticmethod

src/postgres_mcp/sql/sql_driver.py

@@ -1,8 +1,12 @@
 """SQL driver adapter for PostgreSQL connections."""
 
+import io
+import json
 import logging
 import re
 from dataclasses import dataclass
+from datetime import date
+from datetime import datetime
 from typing import Any
 from typing import Dict
 from typing import List
@@ -14,9 +18,21 @@
 from psycopg_pool import AsyncConnectionPool
 from typing_extensions import LiteralString
 
+from ..config import config
+
 logger = logging.getLogger(__name__)
 
 
+def _has_limit_or_offset(query: str) -> bool:
+    """Check if SQL query has LIMIT or OFFSET (case-insensitive word boundary check)."""
+    import re
+
+    # Use word boundaries to avoid matching within identifiers
+    # This handles 95% of cases correctly
+    pattern = r"\b(LIMIT|OFFSET)\b"
+    return bool(re.search(pattern, query, re.IGNORECASE))
+
+
 def obfuscate_password(text: str | None) -> str | None:
     """
     Obfuscate password in any text containing connection information.
@@ -184,6 +200,8 @@ async def execute_query(
         query: LiteralString,
         params: list[Any] | None = None,
         force_readonly: bool = False,
+        page_size: int | None = None,
+        offset: int = 0,
     ) -> Optional[List[RowResult]]:
         """
         Execute a query and return results.
@@ -192,6 +210,8 @@ async def execute_query(
             query: SQL query to execute
             params: Query parameters
             force_readonly: Whether to enforce read-only mode
+            page_size: Number of rows to return (1-max configured), defaults to config.default_page_size
+            offset: Number of rows to skip
 
         Returns:
             List of RowResult objects or None on error
@@ -207,10 +227,24 @@ async def execute_query(
                 # For pools, get a connection from the pool
                 pool = await self.conn.pool_connect()
                 async with pool.connection() as connection:
-                    return await self._execute_with_connection(connection, query, params, force_readonly=force_readonly)
+                    return await self._execute_with_connection(
+                        connection,
+                        query,
+                        params,
+                        force_readonly=force_readonly,
+                        page_size=page_size,
+                        offset=offset,
+                    )
             else:
                 # Direct connection approach
-                return await self._execute_with_connection(self.conn, query, params, force_readonly=force_readonly)
+                return await self._execute_with_connection(
+                    self.conn,
+                    query,
+                    params,
+                    force_readonly=force_readonly,
+                    page_size=page_size,
+                    offset=offset,
+                )
         except Exception as e:
             # Mark pool as invalid if there was a connection issue
             if self.conn and self.is_pool:
@@ -221,8 +255,32 @@ async def execute_query(
 
             raise e
 
-    async def _execute_with_connection(self, connection, query, params, force_readonly) -> Optional[List[RowResult]]:
-        """Execute query with the given connection."""
+    def get_wire_size(self, data: list[dict[str, Any]]) -> int:
+        """Calculates exact bytes of the JSON-serialized data including datetimes."""
+
+        def json_serial(obj: Any) -> str:
+            """JSON serializer that converts any non-serializable object to string.
+
+            This is only used for wire size calculation, so we prioritize
+            robustness over perfect type preservation.
+            """
+            if isinstance(obj, (datetime, date)):
+                return obj.isoformat()
+
+            return str(obj)
+
+        buffer = io.StringIO()
+        json.dump(data, buffer, default=json_serial)
+        return len(buffer.getvalue().encode("utf-8"))
+
+    async def _execute_with_connection(
+        self, connection, query, params, force_readonly, page_size: int | None = None, offset: int = 0
+    ) -> Optional[List[RowResult]]:
+        """Execute query with the given connection and apply pagination."""
+
+        if page_size is None:
+            page_size = config.default_page_size
+
         transaction_started = False
         try:
             async with connection.cursor(row_factory=dict_row) as cursor:
@@ -231,10 +289,30 @@ async def _execute_with_connection(self, connection, query, params, force_readon
                     await cursor.execute("BEGIN TRANSACTION READ ONLY")
                     transaction_started = True
 
+                paginated_query = query
+                # Only apply pagination in readonly mode to avoid breaking DDL operations
+                if force_readonly and page_size > 0:
+                    # Remove trailing semicolon if present (we'll add it back later)
+                    query_trimmed = query.rstrip().rstrip(";")
+                    had_semicolon = query.rstrip().endswith(";")
+
+                    # Use proper SQL parsing to check for existing LIMIT/OFFSET
+                    if not _has_limit_or_offset(query_trimmed):
+                        # Safe to add pagination
+                        paginated_query = f"{query_trimmed} LIMIT {page_size} OFFSET {offset}"
+                        # Restore semicolon if original had one
+                        if had_semicolon:
+                            paginated_query += ";"
+                    else:
+                        # Query already has pagination, use it as-is
+                        paginated_query = query_trimmed
+                        if had_semicolon:
+                            paginated_query += ";"
+
                 if params:
-                    await cursor.execute(query, params)
+                    await cursor.execute(paginated_query, params)
                 else:
-                    await cursor.execute(query)
+                    await cursor.execute(paginated_query)
 
                 # For multiple statements, move to the last statement's results
                 while cursor.nextset():
@@ -258,7 +336,19 @@ async def _execute_with_connection(self, connection, query, params, force_readon
                     await cursor.execute("ROLLBACK")
                     transaction_started = False
 
-                return [SqlDriver.RowResult(cells=dict(row)) for row in rows]
+                result = [SqlDriver.RowResult(cells=dict(row)) for row in rows]
+
+                wire_size_bytes: int = self.get_wire_size([r.cells for r in result])
+
+                payload_size_mb = wire_size_bytes / (1024 * 1024)
+
+                if payload_size_mb > config.max_payload_size_mb:
+                    raise ValueError(
+                        f"Query result payload too large: {payload_size_mb:.2f}MB exceeds maximum allowed size of {config.max_payload_size_mb}MB. "
+                        f"Please refine your query to return less data, use pagination (LIMIT/OFFSET), or filter results."
+                    )
+
+                return result
 
         except Exception as e:
             # Try to roll back the transaction if it's still active