feat: support opensearch

Author: cr7258

.env

@@ -1,4 +1,9 @@
 # Elasticsearch connection settings
-ELASTIC_HOST=https://localhost:9200
-ELASTIC_USERNAME=elastic
-ELASTIC_PASSWORD=test123
+ELASTICSEARCH_HOST=https://localhost:9200
+ELASTICSEARCH_USERNAME=elastic
+ELASTICSEARCH_PASSWORD=test123
+
+# OpenSearch connection settings
+OPENSEARCH_HOST=https://localhost:9200
+OPENSEARCH_USERNAME=elastic
+OPENSEARCH_PASSWORD=test123

README.md

@@ -65,9 +65,9 @@ Using `uvx` will automatically install the package from PyPI, no need to clone t
         "elasticsearch-mcp-server"
       ],
       "env": {
-        "ELASTIC_HOST": "https://localhost:9200",
-        "ELASTIC_USERNAME": "elastic",
-        "ELASTIC_PASSWORD": "test123"
+        "ELASTICSEARCH_HOST": "https://localhost:9200",
+        "ELASTICSEARCH_USERNAME": "elastic",
+        "ELASTICSEARCH_PASSWORD": "test123"
       }
     }
   }
@@ -90,9 +90,9 @@ Using `uv` requires cloning the repository locally and specifying the path to th
         "elasticsearch-mcp-server"
       ],
       "env": {
-        "ELASTIC_HOST": "https://localhost:9200",
-        "ELASTIC_USERNAME": "elastic",
-        "ELASTIC_PASSWORD": "test123"
+        "ELASTICSEARCH_HOST": "https://localhost:9200",
+        "ELASTICSEARCH_USERNAME": "elastic",
+        "ELASTICSEARCH_PASSWORD": "test123"
       }
     }
   }

pyproject.toml

@@ -1,11 +1,12 @@
 [project]
-name = "elasticsearch-mcp-server"
+name = "search-mcp-server"
 version = "1.0.0"
-description = "MCP Server for interacting with Elasticsearch"
+description = "MCP Server for interacting with Elasticsearch and OpenSearch"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
     "elasticsearch>=8.0.0",
+    "opensearch-py>=2.0.0",
     "mcp>=1.0.0",
     "python-dotenv>=1.0.0",
     "fastmcp>=0.4.0",
@@ -15,7 +16,11 @@ dependencies = [
 file = "LICENSE"
 
 [project.scripts]
-elasticsearch-mcp-server = "elasticsearch_mcp_server:main"
+elasticsearch-mcp-server = "src.server:elasticsearch_mcp_server"
+opensearch-mcp-server = "src.server:opensearch_mcp_server"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src"]
 
 [build-system]
 requires = [

src/__init__.py

@@ -0,0 +1,6 @@
+"""
+Search MCP Server package.
+"""
+from src.server import elasticsearch_mcp_server, opensearch_mcp_server, run_search_server
+
+__all__ = ['elasticsearch_mcp_server', 'opensearch_mcp_server', 'run_search_server']

src/clients/__init__.py

@@ -0,0 +1,33 @@
+from typing import Dict
+from .elasticsearch.client import ElasticsearchClient
+from .opensearch.client import OpenSearchClient
+from .interfaces import SearchClient
+
+def create_search_client(config: Dict, engine_type: str) -> SearchClient:
+    """
+    Factory function to create the appropriate search client.
+    
+    Args:
+        config: Configuration dictionary with connection parameters
+        engine_type: Type of search engine to use ("elasticsearch" or "opensearch")
+        
+    Returns:
+        SearchClient: An instance of the appropriate search client
+        
+    Raises:
+        ValueError: If an invalid engine type is specified
+    """
+    if engine_type.lower() == "elasticsearch":
+        return ElasticsearchClient(config)
+    elif engine_type.lower() == "opensearch":
+        return OpenSearchClient(config)
+    else:
+        raise ValueError(f"Invalid engine type: {engine_type}. Must be 'elasticsearch' or 'opensearch'")
+
+__all__ = [
+    'create_search_client',
+    'handle_search_exceptions',
+    'SearchClient',
+    'ElasticsearchClient',
+    'OpenSearchClient',
+]

src/clients/base.py

@@ -0,0 +1,24 @@
+"""
+Base classes and interfaces for search clients.
+"""
+from abc import ABC, abstractmethod
+import logging
+from typing import Dict
+
+class SearchClientBase(ABC):
+    """Base abstract class for search clients with common functionality."""
+    
+    def __init__(self, config: Dict):
+        """
+        Initialize the base search client.
+        
+        Args:
+            config: Configuration dictionary with connection parameters
+        """
+        self.logger = logging.getLogger(self.__class__.__name__)
+        self.config = config
+    
+    @abstractmethod
+    def close(self):
+        """Close the client connection."""
+        pass

src/clients/elasticsearch/__init__.py

@@ -0,0 +1,3 @@
+from .client import ElasticsearchClient
+
+__all__ = ['ElasticsearchClient']

src/clients/elasticsearch/alias.py

@@ -0,0 +1,20 @@
+from typing import Dict
+from .base import BaseElasticsearchClient
+from ..interfaces.alias import AliasClientInterface
+
+class ElasticsearchAliasClient(BaseElasticsearchClient, AliasClientInterface):
+    def list_aliases(self) -> Dict:
+        """Get all aliases."""
+        return self.client.cat.aliases()
+    
+    def get_alias(self, index: str) -> Dict:
+        """Get aliases for the specified index."""
+        return self.client.indices.get_alias(index=index)
+    
+    def put_alias(self, index: str, name: str, body: Dict) -> Dict:
+        """Creates or updates an alias for the specified index."""
+        return self.client.indices.put_alias(index=index, name=name, body=body)
+    
+    def delete_alias(self, index: str, name: str) -> Dict:
+        """Delete an alias for the specified index."""
+        return self.client.indices.delete_alias(index=index, name=name)

src/clients/elasticsearch/base.py

@@ -0,0 +1,40 @@
+"""
+Base Elasticsearch client implementation.
+"""
+from elasticsearch import Elasticsearch
+from typing import Dict
+import logging
+from ..base import SearchClientBase
+
+class BaseElasticsearchClient(SearchClientBase):
+    """Base Elasticsearch client with connection management."""
+    
+    def __init__(self, config: Dict):
+        """
+        Initialize the Elasticsearch client.
+        
+        Args:
+            config: Configuration dictionary with connection parameters
+        """
+        super().__init__(config)
+        self.logger.info("Initializing Elasticsearch client")
+        
+        # Extract configuration
+        hosts = config.get("hosts", ["localhost:9200"])
+        username = config.get("username")
+        password = config.get("password")
+        verify_certs = config.get("verify_certs")
+        
+        # Create client
+        self.client = Elasticsearch(
+            hosts=hosts,
+            basic_auth=(username, password) if username and password else None,
+            verify_certs=verify_certs
+        )
+        
+        self.logger.info(f"Elasticsearch client initialized with hosts: {hosts}")
+    
+    def close(self):
+        """Close the Elasticsearch client connection."""
+        self.logger.info("Closing Elasticsearch client connection")
+        self.client.close()

src/clients/elasticsearch/client.py

@@ -0,0 +1,28 @@
+from typing import Dict
+from .index import ElasticsearchIndexClient
+from .document import ElasticsearchDocumentClient
+from .cluster import ElasticsearchClusterClient
+from .alias import ElasticsearchAliasClient
+
+class ElasticsearchClient(ElasticsearchIndexClient,
+                         ElasticsearchDocumentClient,
+                         ElasticsearchClusterClient,
+                         ElasticsearchAliasClient):
+    """
+    Elasticsearch client that implements all required interfaces.
+    
+    This class uses multiple inheritance to combine all specialized client implementations
+    into a single unified client, inheriting all methods from the specialized clients.
+    """
+    
+    def __init__(self, config: Dict):
+        """
+        Initialize the Elasticsearch client.
+        
+        Args:
+            config: Configuration dictionary with connection parameters
+        """
+        super().__init__(config)
+        
+        # Log initialization
+        self.logger.info("Initialized the Elasticsearch client")

src/clients/elasticsearch/cluster.py

@@ -0,0 +1,12 @@
+from typing import Dict
+from .base import BaseElasticsearchClient
+from ..interfaces.cluster import ClusterClientInterface
+
+class ElasticsearchClusterClient(BaseElasticsearchClient, ClusterClientInterface):
+    def get_cluster_health(self) -> Dict:
+        """Returns basic information about the health of the cluster."""
+        return self.client.cluster.health()
+    
+    def get_cluster_stats(self) -> Dict:
+        """Returns high-level overview of cluster statistics."""
+        return self.client.cluster.stats()

src/clients/elasticsearch/document.py

@@ -0,0 +1,27 @@
+from typing import Dict, Optional
+from .base import BaseElasticsearchClient
+from ..interfaces.document import DocumentClientInterface
+
+class ElasticsearchDocumentClient(BaseElasticsearchClient, DocumentClientInterface):
+    def search_documents(self, index: str, body: Dict) -> Dict:
+        """Search for documents in the index."""
+        return self.client.search(index=index, body=body)
+    
+    def index_document(self, index: str, document: Dict, id: Optional[str] = None) -> Dict:
+        """Creates or updates a document in the index. """
+        if id is not None:
+            return self.client.index(index=index, document=document, id=id)
+        else:
+            return self.client.index(index=index, document=document)
+    
+    def get_document(self, index: str, id: str) -> Dict:
+        """Get a document by ID."""
+        return self.client.get(index=index, id=id)
+    
+    def delete_document(self, index: str, id: str) -> Dict:
+        """Removes a document from the index."""
+        return self.client.delete(index=index, id=id)
+
+    def delete_by_query(self, index: str, body: Dict) -> Dict:
+        """Deletes documents matching the provided query."""
+        return self.client.delete_by_query(index=index, body=body)

src/clients/elasticsearch/index.py

@@ -0,0 +1,20 @@
+from typing import Dict, Optional
+from .base import BaseElasticsearchClient
+from ..interfaces.index import IndexClientInterface
+
+class ElasticsearchIndexClient(BaseElasticsearchClient, IndexClientInterface):
+    def list_indices(self) -> Dict:
+        """List all indices."""
+        return self.client.cat.indices()
+    
+    def get_index(self, index: str) -> Dict:
+        """Returns information (mappings, settings, aliases) about one or more indices."""
+        return self.client.indices.get(index=index)
+    
+    def create_index(self, index: str, body: Optional[Dict] = None) -> Dict:
+        """Creates an index with optional settings and mappings."""
+        return self.client.indices.create(index=index, body=body)
+    
+    def delete_index(self, index: str) -> Dict:
+        """Delete an index."""
+        return self.client.indices.delete(index=index)

src/clients/exceptions.py

@@ -0,0 +1,68 @@
+import functools
+import logging
+from typing import TypeVar, Callable
+from mcp.types import TextContent
+from fastmcp import FastMCP
+
+T = TypeVar('T')
+
+def handle_search_exceptions(func: Callable[..., T]) -> Callable[..., list[TextContent]]:
+    """
+    Decorator to handle exceptions in search client operations.
+    
+    Args:
+        func: The function to decorate
+        
+    Returns:
+        Decorated function that handles exceptions
+    """
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        logger = logging.getLogger()   
+        try:
+            return 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
+
+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_search_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_search_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

src/clients/interfaces/__init__.py

@@ -0,0 +1,23 @@
+"""
+Export all client interfaces and define a combined interface.
+"""
+from .index import IndexClientInterface
+from .document import DocumentClientInterface
+from .cluster import ClusterClientInterface
+from .alias import AliasClientInterface
+
+class SearchClient(IndexClientInterface, DocumentClientInterface, 
+                  ClusterClientInterface, AliasClientInterface):
+    """Complete search client interface combining all functionality."""
+    
+    def close(self):
+        """Close the client connection."""
+        pass
+
+__all__ = [
+    'IndexClientInterface',
+    'DocumentClientInterface',
+    'ClusterClientInterface',
+    'AliasClientInterface',
+    'SearchClient'
+]