fix: improve connection pool error handling

Author: alexander-zuev

supabase_mcp/logger.py

@@ -1,10 +1,43 @@
+import json
 import logging
 import logging.handlers
 from pathlib import Path
+from typing import Any
+
+
+class StructuredLogFormatter(logging.Formatter):
+    """
+    Custom formatter that outputs logs in a structured JSON format.
+    Includes error tracing and context information.
+    """
+
+    def format(self, record: logging.LogRecord) -> str:
+        log_record: dict[str, Any] = {
+            "timestamp": self.formatTime(record, self.datefmt),
+            "level": record.levelname,
+            "message": record.getMessage(),
+            "module": record.module,
+            "function": record.funcName,
+            "line": record.lineno,
+        }
+
+        # Add exception info if present with full traceback
+        if record.exc_info:
+            log_record["exception"] = self.formatException(record.exc_info)
+            # Add exception type for easier filtering
+            exc_type = record.exc_info[0].__name__ if record.exc_info[0] else "Unknown"
+            log_record["exception_type"] = exc_type
+
+        # Add extra fields from record
+        extra_data = getattr(record, "extra", None)
+        if extra_data and isinstance(extra_data, dict):
+            log_record.update(extra_data)
+
+        return json.dumps(log_record, ensure_ascii=False)
 
 
 def setup_logger() -> logging.Logger:
-    """Configure logging for the MCP server with log rotation."""
+    """Configure logging for the MCP server with structured logging and log rotation."""
     logger = logging.getLogger("supabase-mcp")
 
     # Remove existing handlers to avoid duplicate logs
@@ -13,28 +46,26 @@ def setup_logger() -> logging.Logger:
 
     # Define a consistent log directory in the user's home folder
     log_dir = Path.home() / ".local" / "share" / "supabase-mcp"
-    log_dir.mkdir(parents=True, exist_ok=True)  # Ensure the directory exists
+    log_dir.mkdir(parents=True, exist_ok=True)
 
     # Define the log file path
     log_file = log_dir / "mcp_server.log"
 
-    # Create a rotating file handler
-    # - Rotate when log reaches 5MB
-    # - Keep 3 backup files
+    # Create handlers
     file_handler = logging.handlers.RotatingFileHandler(
         log_file,
         maxBytes=5 * 1024 * 1024,  # 5MB
         backupCount=3,
         encoding="utf-8",
     )
 
-    # Create formatter
-    formatter = logging.Formatter("[%(asctime)s] %(levelname)-8s %(message)s", datefmt="%y/%m/%d %H:%M:%S")
+    # Create structured formatter
+    formatter = StructuredLogFormatter()
 
-    # Add formatter to file handler
+    # Add formatter to handlers
     file_handler.setFormatter(formatter)
 
-    # Add handler to logger
+    # Add handlers to logger
     logger.addHandler(file_handler)
 
     # Set level
@@ -43,4 +74,5 @@ def setup_logger() -> logging.Logger:
     return logger
 
 
+# Create the global logger instance
 logger = setup_logger()

supabase_mcp/services/database/postgres_client.py

@@ -203,12 +203,12 @@ async def create_pool(self) -> asyncpg.Pool[asyncpg.Record]:
                     f"4. Your Supabase project is active and the database is online\n"
                 )
 
-            logger.error(f"Failed to connect to database: {e}")
+            logger.error(f"FAILED to connect to database: {e}")
             logger.error(f"Connection details: {host_part}, Project: {self.project_ref}, Region: {self.db_region}")
 
             raise ConnectionError(error_message) from e
 
-        except OSError as e:
+        except (OSError, asyncpg.InterfaceError, asyncpg.TooManyConnectionsError) as e:
             # For network-related errors, provide a different message that clearly indicates
             # this is a network/system issue rather than a database configuration problem
             host_part = self.db_url.split("@")[1].split("/")[0] if "@" in self.db_url else "unknown"
@@ -228,6 +228,10 @@ async def create_pool(self) -> asyncpg.Pool[asyncpg.Record]:
             logger.error(f"Connection details: {host_part}")
             raise ConnectionError(error_message) from e
 
+        except Exception as e:
+            logger.error(f"Failed to connect to database: {e}")
+            raise ConnectionError(f"Failed to connect to database: {e}") from e
+
     async def ensure_pool(self) -> None:
         """Ensure a valid connection pool exists.
 
@@ -360,7 +364,6 @@ async def execute_query(
             QueryResult containing the results of all statements
 
         Raises:
-            ConnectionError: If a database connection issue occurs
             QueryError: If the query execution fails
             PermissionError: When user lacks required privileges
         """
@@ -372,24 +375,37 @@ async def execute_query(
         )
         logger.debug(f"Executing query (readonly={readonly}): {truncated_query}")
 
-        # Define the operation to execute all statements within a transaction
-        async def execute_all_statements(conn):
-            async def transaction_operation():
-                results = []
-                for statement in validated_query.statements:
-                    if statement.query:  # Skip statements with no query
-                        result = await self.execute_statement(conn, statement.query)
-                        results.append(result)
-                    else:
-                        logger.warning(f"Statement has no query, statement: {statement}")
-                return results
-
-            # Execute the operation within a transaction
-            results = await self.with_transaction(conn, transaction_operation, readonly)
-            return QueryResult(results=results)
-
-        # Execute the operation with a connection
-        return await self.with_connection(execute_all_statements)
+        try:
+            # Define the operation to execute all statements within a transaction
+            async def execute_all_statements(conn: asyncpg.Connection[Any]) -> QueryResult:
+                async def transaction_operation() -> list[StatementResult]:
+                    results = []
+                    for statement in validated_query.statements:
+                        if statement.query:  # Skip statements with no query
+                            result = await self.execute_statement(conn, statement.query)
+                            results.append(result)
+                        else:
+                            logger.warning(f"Statement has no query, statement: {statement}")
+                    return results
+
+                # Execute the operation within a transaction
+                results = await self.with_transaction(conn, transaction_operation, readonly)
+                return QueryResult(results=results)
+
+            # Execute the operation with a connection
+            return await self.with_connection(execute_all_statements)
+
+        except ConnectionError as e:
+            logger.error(f"Query execution failed because of connection error: {e}")
+            raise QueryError(f"Query execution failed because of connection error: {str(e)}") from e
+
+        except PermissionError as e:
+            logger.error(f"Query execution failed because of permission error: {e}")
+            raise PermissionError(f"Query execution failed because of permission error: {str(e)}") from e
+
+        except QueryError as e:
+            logger.error(f"Query execution failed: {e}")
+            raise QueryError(f"Query execution failed: {str(e)}") from e
 
     async def _handle_postgres_error(self, error: asyncpg.PostgresError) -> None:
         """Handle PostgreSQL errors and convert to appropriate exceptions.
@@ -408,10 +424,7 @@ async def _handle_postgres_error(self, error: asyncpg.PostgresError) -> None:
             ) from error
         elif isinstance(
             error,
-            (
-                asyncpg.exceptions.UndefinedTableError,
-                asyncpg.exceptions.UndefinedColumnError,
-            ),
+            (asyncpg.exceptions.UndefinedTableError | asyncpg.exceptions.UndefinedColumnError),
         ):
             logger.error(f"Schema error: {error}")
             raise QueryError(str(error)) from error

supabase_mcp/services/database/query_manager.py

@@ -1,4 +1,4 @@
-from supabase_mcp.exceptions import OperationNotAllowedError
+from supabase_mcp.exceptions import OperationNotAllowedError, PermissionError, QueryError
 from supabase_mcp.logger import logger
 from supabase_mcp.services.database.migration_manager import MigrationManager
 from supabase_mcp.services.database.postgres_client import PostgresClient, QueryResult
@@ -104,10 +104,17 @@ async def handle_query_execution(self, validated_query: QueryValidationResults)
         Returns:
             QueryResult: The result of the query execution
         """
-        readonly = self.check_readonly()
-        result = await self.db_client.execute_query(validated_query, readonly)
-        logger.debug(f"Query result: {result}")
-        return result
+        try:
+            readonly = self.check_readonly()
+            result = await self.db_client.execute_query(validated_query, readonly)
+            logger.debug(f"Query result: {result}")
+            return result
+        except (QueryError, PermissionError) as e:
+            logger.error(f"Query execution failed: {e}")
+            raise QueryError(f"Query execution failed: {str(e)}") from e
+        except Exception as e:
+            logger.error(f"Unexpected error during query execution: {e}")
+            raise QueryError(f"Unexpected error during query execution: {str(e)}") from e
 
     async def handle_migration(
         self, validation_result: QueryValidationResults, original_query: str, migration_name: str = ""