|
1 | 1 | #!/usr/bin/env python3
|
2 | 2 | import logging
|
| 3 | +import os |
| 4 | +from dotenv import load_dotenv |
| 5 | +from elasticsearch import Elasticsearch |
| 6 | +import warnings |
3 | 7 | from fastmcp import FastMCP
|
4 |
| -from .tools.index import IndexTools |
5 |
| -from .tools.document import DocumentTools |
6 |
| -from .tools.cluster import ClusterTools |
7 |
| -from .tools.alias import AliasTools |
| 8 | +from elasticsearch_mcp_server.tools.index import IndexTools |
| 9 | +from elasticsearch_mcp_server.tools.document import DocumentTools |
| 10 | +from elasticsearch_mcp_server.tools.cluster import ClusterTools |
| 11 | +from elasticsearch_mcp_server.tools.alias import AliasTools |
| 12 | +from elasticsearch_mcp_server.tools.register import ToolsRegister |
8 | 13 |
|
9 | 14 | class ElasticsearchMCPServer:
|
10 | 15 | def __init__(self):
|
11 | 16 | self.name = "elasticsearch_mcp_server"
|
12 | 17 | self.mcp = FastMCP(self.name)
|
13 |
| - |
14 |
| - # Configure logging |
15 |
| - logging.basicConfig( |
16 |
| - level=logging.INFO, |
17 |
| - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' |
18 |
| - ) |
19 | 18 | self.logger = logging.getLogger(self.name)
|
| 19 | + self.es_client = self._create_elasticsearch_client() |
20 | 20 |
|
21 | 21 | # Initialize tools
|
22 | 22 | self._register_tools()
|
| 23 | + |
| 24 | + def _get_es_config(self): |
| 25 | + """Get Elasticsearch configuration from environment variables.""" |
| 26 | + # Load environment variables from .env file |
| 27 | + load_dotenv() |
| 28 | + config = { |
| 29 | + "host": os.getenv("ELASTIC_HOST"), |
| 30 | + "username": os.getenv("ELASTIC_USERNAME"), |
| 31 | + "password": os.getenv("ELASTIC_PASSWORD") |
| 32 | + } |
| 33 | + |
| 34 | + if not all([config["username"], config["password"]]): |
| 35 | + self.logger.error("Missing required Elasticsearch configuration. Please check environment variables:") |
| 36 | + self.logger.error("ELASTIC_USERNAME and ELASTIC_PASSWORD are required") |
| 37 | + raise ValueError("Missing required Elasticsearch configuration") |
| 38 | + |
| 39 | + return config |
| 40 | + |
| 41 | + def _create_elasticsearch_client(self) -> Elasticsearch: |
| 42 | + """Create and return an Elasticsearch client using configuration from environment.""" |
| 43 | + config = self._get_es_config() |
| 44 | + |
| 45 | + # Disable SSL warnings |
| 46 | + warnings.filterwarnings("ignore", message=".*TLS with verify_certs=False is insecure.*",) |
| 47 | + |
| 48 | + return Elasticsearch( |
| 49 | + config["host"], |
| 50 | + basic_auth=(config["username"], config["password"]), |
| 51 | + verify_certs=False |
| 52 | + ) |
23 | 53 |
|
24 | 54 | def _register_tools(self):
|
25 | 55 | """Register all MCP tools."""
|
26 |
| - # Initialize tool classes |
27 |
| - index_tools = IndexTools(self.logger) |
28 |
| - document_tools = DocumentTools(self.logger) |
29 |
| - cluster_tools = ClusterTools(self.logger) |
30 |
| - alias_tools = AliasTools(self.logger) |
| 56 | + # Create a tools register |
| 57 | + register = ToolsRegister(self.logger, self.es_client, self.mcp) |
31 | 58 |
|
32 |
| - # Register tools from each module |
33 |
| - index_tools.register_tools(self.mcp) |
34 |
| - document_tools.register_tools(self.mcp) |
35 |
| - cluster_tools.register_tools(self.mcp) |
36 |
| - alias_tools.register_tools(self.mcp) |
| 59 | + # Define all tool classes to register |
| 60 | + tool_classes = [ |
| 61 | + IndexTools, |
| 62 | + DocumentTools, |
| 63 | + ClusterTools, |
| 64 | + AliasTools |
| 65 | + ] |
| 66 | + # Register all tools |
| 67 | + register.register_all_tools(tool_classes) |
37 | 68 |
|
38 | 69 | def run(self):
|
39 | 70 | """Run the MCP server."""
|
|
0 commit comments