refactor: add global exception handler

Author: cr7258

src/elasticsearch_mcp_server/server.py

@@ -41,10 +41,7 @@ def _get_es_config(self):
     def _create_elasticsearch_client(self) -> Elasticsearch:
         """Create and return an Elasticsearch client using configuration from environment."""
         config = self._get_es_config()
-
-        # Disable SSL warnings
-        warnings.filterwarnings("ignore", message=".*TLS with verify_certs=False is insecure.*",)
-
+        
         return Elasticsearch(
             config["host"],
             basic_auth=(config["username"], config["password"]),

src/elasticsearch_mcp_server/tools/alias.py

@@ -1,27 +1,20 @@
 from fastmcp import FastMCP
-from mcp.types import TextContent
+from typing import Dict, List
+from elasticsearch_mcp_server.tools.base import handle_es_exceptions
 
 class AliasTools:
     def register_tools(self, mcp: FastMCP):
         @mcp.tool()
-        async def list_aliases() -> list[TextContent]:
+        async def list_aliases() -> List[Dict]:
             """List all aliases in the Elasticsearch cluster."""
-            try:
-                return self.es_client.cat.aliases(format="json")
-            except Exception as e:
-                self.logger.error(f"Error listing aliases: {e}")
-                return [TextContent(type="text", text=f"Error: {str(e)}")]
+            return self.es_client.cat.aliases()
 
         @mcp.tool()
-        async def get_alias(index: str) -> list[TextContent]:
+        async def get_alias(index: str) -> Dict:
             """
             Get alias information for a specific index.
 
             Args:
                 index: Name of the index
             """
-            try:
-                return self.es_client.indices.get_alias(index=index)
-            except Exception as e:
-                self.logger.error(f"Error getting alias information: {e}")
-                return [TextContent(type="text", text=f"Error: {str(e)}")]
+            return self.es_client.indices.get_alias(index=index)

src/elasticsearch_mcp_server/tools/base.py

@@ -0,0 +1,24 @@
+import functools
+import inspect
+import logging
+from typing import Callable, TypeVar
+from mcp.types import TextContent
+
+# Type variable for function return type
+T = TypeVar('T')
+
+def handle_es_exceptions(func: Callable[..., T]) -> Callable[..., list[TextContent]]:
+    """
+    Decorator to handle Elasticsearch exceptions in tool methods.
+    Logs the error and returns it as a TextContent object.
+    """
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        logger = logging.getLogger()
+        try:
+            return await func(*args, **kwargs)
+        except Exception as e:
+            logger.error(f"Unexpected error in {func.__name__}: {e}")
+            return [TextContent(type="text", text=f"Unexpected error in {func.__name__}: {str(e)}")]
+    
+    return wrapper

src/elasticsearch_mcp_server/tools/cluster.py

@@ -1,32 +1,25 @@
 from fastmcp import FastMCP
-from mcp.types import TextContent
+from typing import Dict
+from elasticsearch_mcp_server.tools.base import handle_es_exceptions
 
 class ClusterTools:
     def register_tools(self, mcp: FastMCP):
         """Register cluster-related tools."""
 
         @mcp.tool()
-        async def get_cluster_health() -> list[TextContent]:
+        async def get_cluster_health() -> Dict:
             """
             Get health status of the Elasticsearch cluster.
             Returns information about the number of nodes, shards, etc.
             """
-            try:
-                return self.es_client.cluster.health()
-            except Exception as e:
-                self.logger.error(f"Error getting cluster health: {e}")
-                return [TextContent(type="text", text=f"Error: {str(e)}")]
+            return self.es_client.cluster.health()
 
         @mcp.tool()
-        async def get_cluster_stats() -> list[TextContent]:
+        async def get_cluster_stats() -> Dict:
             """
             Get statistics from a cluster wide perspective. 
             The API returns basic index metrics (shard numbers, store size, memory usage) and information 
             about the current nodes that form the cluster (number, roles, os, jvm versions, memory usage, cpu and installed plugins).
             https://www.elastic.co/guide/en/elasticsearch/reference/8.17/cluster-stats.html
             """
-            try:
-                return self.es_client.cluster.stats()
-            except Exception as e:
-                self.logger.error(f"Error getting cluster stats: {e}")
-                return [TextContent(type="text", text=f"Error: {str(e)}")]
+            return self.es_client.cluster.stats()

src/elasticsearch_mcp_server/tools/document.py

@@ -1,22 +1,18 @@
 from fastmcp import FastMCP
-from mcp.types import TextContent
+from typing import Dict
+from elasticsearch_mcp_server.tools.base import handle_es_exceptions
 
 class DocumentTools:
     def register_tools(self, mcp: FastMCP):
         """Register document-related tools."""
 
         @mcp.tool()
-        async def search_documents(index: str, body: dict) -> list[TextContent]:
+        async def search_documents(index: str, body: dict) -> Dict:
             """
             Search documents in a specified index using a custom query.
             
             Args:
                 index: Name of the index to search
                 body: Elasticsearch query DSL
             """
-            self.logger.info(f"Searching in index: {index} with query: {body}")
-            try:
-                return self.es_client.search(index=index, body=body)
-            except Exception as e:
-                self.logger.error(f"Error searching documents: {e}")
-                return [TextContent(type="text", text=f"Error: {str(e)}")]
+            return self.es_client.search(index=index, body=body)

src/elasticsearch_mcp_server/tools/index.py

@@ -1,42 +1,30 @@
 from fastmcp import FastMCP
-from mcp.types import TextContent
 from typing import Dict
+from elasticsearch_mcp_server.tools.base import handle_es_exceptions
 
 class IndexTools:
     def register_tools(self, mcp: FastMCP):
         @mcp.tool()
         async def list_indices() -> Dict:
             """List all indices."""
-            try:
-                return self.es_client.cat.indices()
-            except Exception as e:
-                self.logger.error(f"Error listing indices: {e}")
-                return [TextContent(type="text", text=f"Error: {str(e)}")]
+            return self.es_client.cat.indices()
 
         @mcp.tool()
-        async def get_mapping(index: str) -> list[TextContent]:
+        async def get_mapping(index: str) -> Dict:
             """
             Get the mapping for an index.
             
             Args:
                 index: Name of the index
             """
-            try:
-                return self.es_client.indices.get_mapping(index=index)
-            except Exception as e:
-                self.logger.error(f"Error getting mapping: {e}")
-                return [TextContent(type="text", text=f"Error: {str(e)}")]
+            return self.es_client.indices.get_mapping(index=index)
 
         @mcp.tool()
-        async def get_settings(index: str) -> list[TextContent]:
+        async def get_settings(index: str) -> Dict:
             """
             Get the settings for an index.
             
             Args:
                 index: Name of the index
             """
-            try:
-                return self.es_client.indices.get_settings(index=index)
-            except Exception as e:
-                self.logger.error(f"Error getting settings: {e}")
-                return [TextContent(type="text", text=f"Error: {str(e)}")]
+            return self.es_client.indices.get_settings(index=index)

src/elasticsearch_mcp_server/tools/register.py

@@ -1,8 +1,52 @@
 import logging
-
+import functools
 from fastmcp import FastMCP
-from typing import List, Type
+from typing import List, Type, Callable, TypeVar
 from elasticsearch import Elasticsearch
+from elasticsearch_mcp_server.tools.base import handle_es_exceptions
+
+T = TypeVar('T')
+
+def with_exception_handling(tool_instance: object, mcp: FastMCP) -> None:
+    """
+    Register tools from a tool instance with automatic exception handling applied to all tools.
+    
+    This function temporarily replaces mcp.tool with a wrapped version that automatically
+    applies the handle_es_exceptions decorator to all registered tool methods.
+    
+    Args:
+        tool_instance: The tool instance that has a register_tools method
+        mcp: The FastMCP instance used for tool registration
+    """
+    # Save the original tool method
+    original_tool = mcp.tool
+    
+    @functools.wraps(original_tool)
+    def wrapped_tool(*args, **kwargs):
+        # Get the original decorator
+        decorator = original_tool(*args, **kwargs)
+        
+        # Return a new decorator that applies both the exception handler and original decorator
+        def combined_decorator(func):
+            # First apply the exception handling decorator
+            wrapped_func = handle_es_exceptions(func)
+            # Then apply the original mcp.tool decorator
+            return decorator(wrapped_func)
+        
+        return combined_decorator
+    
+    try:
+        # Temporarily replace mcp.tool with our wrapped version
+        mcp.tool = wrapped_tool
+        
+        # Call the registration method on the tool instance
+        tool_instance.register_tools(mcp)
+    finally:
+        # Restore the original mcp.tool to avoid affecting other code that might use mcp.tool
+        # This ensures that our modification is isolated to just this tool registration
+        # and prevents multiple nested decorators if register_all_tools is called multiple times
+        mcp.tool = original_tool
+
 
 class ToolsRegister:
     """
@@ -39,7 +83,7 @@ def register_all_tools(self, tool_classes: List[Type]):
             # Pass the raw Elasticsearch client instead of the wrapper
             tool_instance.es_client = self.es_client
 
-            # Register the tools from this instance
-            tool_instance.register_tools(self.mcp)
+            # Register tools with automatic exception handling
+            with_exception_handling(tool_instance, self.mcp)
 
             self.logger.info(f"Registered tools from {tool_class.__name__}")