├── .flake8 ├── .github └── workflows │ ├── publish.yml │ └── python-app.yml ├── .gitignore ├── LICENSE.md ├── README.md ├── conftest.py ├── pyproject.toml ├── pytest.ini ├── src └── greptimedb_mcp_server │ ├── __init__.py │ ├── config.py │ ├── server.py │ ├── templates │ ├── metrics_analysis │ │ ├── config.yaml │ │ └── template.md │ └── table_operation │ │ ├── config.yaml │ │ └── template.md │ └── utils.py └── tests ├── test_config.py ├── test_server.py └── test_utils.py /.flake8: -------------------------------------------------------------------------------- 1 | [flake8] 2 | ignore = E203, E266, E501, W503, F403, F401 3 | exclude = 4 | .git, 5 | __pycache__, 6 | docs/source/conf.py, 7 | build, 8 | dist 9 | max-complexity = 10 10 | max-line-length = 79 11 | -------------------------------------------------------------------------------- /.github/workflows/publish.yml: -------------------------------------------------------------------------------- 1 | name: Publish Python Package 2 | 3 | on: 4 | push: 5 | tags: 6 | - 'v*' 7 | 8 | permissions: 9 | id-token: write 10 | contents: read 11 | 12 | jobs: 13 | deploy: 14 | runs-on: ubuntu-latest 15 | environment: 'pypi' 16 | steps: 17 | - uses: actions/checkout@v4 18 | - name: Install the latest version of uv 19 | uses: astral-sh/setup-uv@v5 20 | - name: Install dependencies 21 | run: uv sync --all-groups 22 | - name: Build wheel 23 | run: uv build 24 | - name: Publish package 25 | run: uv publish 26 | -------------------------------------------------------------------------------- /.github/workflows/python-app.yml: -------------------------------------------------------------------------------- 1 | # This workflow will install Python dependencies, run tests and lint with a single version of Python 2 | # For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python 3 | 4 | name: Python application 5 | 6 | on: 7 | push: 8 | branches: [ "main" ] 9 | pull_request: 10 | branches: [ "main" ] 11 | 12 | permissions: 13 | contents: read 14 | 15 | jobs: 16 | build: 17 | 18 | runs-on: ubuntu-latest 19 | 20 | steps: 21 | - uses: actions/checkout@v4 22 | - name: Install the latest version of uv 23 | uses: astral-sh/setup-uv@v5 24 | - name: Install dependencies 25 | run: uv sync 26 | - name: Format 27 | run: uv run black --check . 28 | - name: Lint with flake8 29 | run: | 30 | # stop the build if there are Python syntax errors or undefined names 31 | uv run flake8 src --count --select=E9,F63,F7,F82 --show-source --statistics 32 | # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide 33 | uv run flake8 src --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics 34 | - name: Test with pytest 35 | run: | 36 | uv run pytest 37 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Python-generated files 2 | __pycache__/ 3 | *.py[oc] 4 | build/ 5 | dist/ 6 | venv/ 7 | wheels/ 8 | *.egg-info 9 | 10 | # Virtual environments 11 | .venv 12 | Lib/ 13 | Scripts/ 14 | 15 | # Development environment 16 | .python-version 17 | uv.lock 18 | 19 | # IDE settings (optional) 20 | .vscode/ 21 | .idea/ 22 | 23 | # Distribution directories 24 | *.dist-info/ 25 | # Mac files 26 | .DS_Store 27 | 28 | .envrc 29 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2025 Greptime Inc. 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 6 | 7 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 8 | 9 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # greptimedb-mcp-server 2 | 3 | [![PyPI - Version](https://img.shields.io/pypi/v/greptimedb-mcp-server)](https://pypi.org/project/greptimedb-mcp-server/) 4 | ![build workflow](https://github.com/GreptimeTeam/greptimedb-mcp-server/actions/workflows/python-app.yml/badge.svg) 5 | [![MIT License](https://img.shields.io/badge/license-MIT-green)](LICENSE.md) 6 | 7 | A Model Context Protocol (MCP) server implementation for [GreptimeDB](https://github.com/GreptimeTeam/greptimedb). 8 | 9 | This server provides AI assistants with a secure and structured way to explore and analyze databases. It enables them to list tables, read data, and execute SQL queries through a controlled interface, ensuring responsible database access. 10 | 11 | # Project Status 12 | This is an experimental project that is still under development. Data security and privacy issues have not been specifically addressed, so please use it with caution. 13 | 14 | # Capabilities 15 | 16 | * `list_resources` to list tables 17 | * `read_resource` to read table data 18 | * `list_tools` to list tools 19 | * `call_tool` to execute an SQL 20 | * `list_prompts` to list prompts 21 | * `get_prompt` to get the prompt by name 22 | 23 | # Installation 24 | 25 | ``` 26 | pip install greptimedb-mcp-server 27 | ``` 28 | 29 | 30 | # Configuration 31 | 32 | Set the following environment variables: 33 | 34 | ```bash 35 | GREPTIMEDB_HOST=localhost # Database host 36 | GREPTIMEDB_PORT=4002 # Optional: Database MySQL port (defaults to 4002 if not specified) 37 | GREPTIMEDB_USER=root 38 | GREPTIMEDB_PASSWORD= 39 | GREPTIMEDB_DATABASE=public 40 | GREPTIMEDB_TIMEZONE=UTC 41 | ``` 42 | 43 | Or via command-line args: 44 | 45 | * `--host` the database host, `localhost` by default, 46 | * `--port` the database port, must be MySQL protocol port, `4002` by default, 47 | * `--user` the database username, empty by default, 48 | * `--password` the database password, empty by default, 49 | * `--database` the database name, `public` by default. 50 | * `--timezone` the session time zone, empty by default(using server default time zone). 51 | 52 | # Usage 53 | 54 | ## Claude Desktop Integration 55 | 56 | Configure the MCP server in Claude Desktop's configuration file: 57 | 58 | #### MacOS 59 | 60 | Location: `~/Library/Application Support/Claude/claude_desktop_config.json` 61 | 62 | #### Windows 63 | 64 | Location: `%APPDATA%/Claude/claude_desktop_config.json` 65 | 66 | 67 | ```json 68 | { 69 | "mcpServers": { 70 | "greptimedb": { 71 | "command": "uv", 72 | "args": [ 73 | "--directory", 74 | "/path/to/greptimedb-mcp-server", 75 | "run", 76 | "-m", 77 | "greptimedb_mcp_server.server" 78 | ], 79 | "env": { 80 | "GREPTIMEDB_HOST": "localhost", 81 | "GREPTIMEDB_PORT": "4002", 82 | "GREPTIMEDB_USER": "root", 83 | "GREPTIMEDB_PASSWORD": "", 84 | "GREPTIMEDB_DATABASE": "public", 85 | "GREPTIMEDB_TIMEZONE": "" 86 | } 87 | } 88 | } 89 | } 90 | ``` 91 | 92 | # License 93 | 94 | MIT License - see LICENSE.md file for details. 95 | 96 | # Contribute 97 | 98 | ## Prerequisites 99 | - Python with `uv` package manager 100 | - GreptimeDB installation 101 | - MCP server dependencies 102 | 103 | ## Development 104 | 105 | ``` 106 | # Clone the repository 107 | git clone https://github.com/GreptimeTeam/greptimedb-mcp-server.git 108 | cd greptimedb-mcp-server 109 | 110 | # Create virtual environment 111 | uv venv 112 | source venv/bin/activate # or `venv\Scripts\activate` on Windows 113 | 114 | # Install development dependencies 115 | uv sync 116 | 117 | # Run tests 118 | pytest 119 | ``` 120 | 121 | Use [MCP Inspector](https://modelcontextprotocol.io/docs/tools/inspector) for debugging: 122 | 123 | ```bash 124 | npx @modelcontextprotocol/inspector uv \ 125 | --directory \ 126 | /path/to/greptimedb-mcp-server \ 127 | run \ 128 | -m \ 129 | greptimedb_mcp_server.server 130 | ``` 131 | 132 | # Acknowledgement 133 | This library's implementation was inspired by the following two repositories and incorporates their code, for which we express our gratitude： 134 | 135 | * [ktanaka101/mcp-server-duckdb](https://github.com/ktanaka101/mcp-server-duckdb) 136 | * [designcomputer/mysql_mcp_server](https://github.com/designcomputer/mysql_mcp_server) 137 | * [mikeskarl/mcp-prompt-templates](https://github.com/mikeskarl/mcp-prompt-templates) 138 | 139 | Thanks! 140 | -------------------------------------------------------------------------------- /conftest.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import sys 3 | 4 | 5 | # Mock classes for MySQL connection 6 | class MockCursor: 7 | def __init__(self): 8 | self.query = "" 9 | self.rowcount = 2 10 | 11 | def execute(self, query, args=None): 12 | self.query = query 13 | 14 | def fetchall(self): 15 | if "SHOW TABLES" in self.query.upper(): 16 | return [("users",), ("orders",)] 17 | elif "SHOW DATABASES" in self.query.upper(): 18 | return [("public",), ("greptime_private",)] 19 | elif "SELECT" in self.query.upper(): 20 | return [(1, "John"), (2, "Jane")] 21 | return [] 22 | 23 | @property 24 | def description(self): 25 | if "SHOW TABLES" in self.query.upper(): 26 | return [("table_name", None)] 27 | elif "SHOW DATABASES" in self.query.upper(): 28 | return [("Databases", None)] 29 | elif "SELECT" in self.query.upper(): 30 | return [("id", None), ("name", None)] 31 | return [] 32 | 33 | def close(self): 34 | pass 35 | 36 | def __enter__(self): 37 | return self 38 | 39 | def __exit__(self, *args): 40 | pass 41 | 42 | 43 | class MockConnection: 44 | def __init__(self, *args, **kwargs): 45 | pass 46 | 47 | def cursor(self): 48 | return MockCursor() 49 | 50 | def commit(self): 51 | pass 52 | 53 | def close(self): 54 | pass 55 | 56 | def __enter__(self): 57 | return self 58 | 59 | def __exit__(self, *args): 60 | pass 61 | 62 | 63 | class MockMySQLModule: 64 | """Mock for entire mysql.connector module""" 65 | 66 | @staticmethod 67 | def connect(*args, **kwargs): 68 | return MockConnection(*args, **kwargs) 69 | 70 | class Error(Exception): 71 | """Mock MySQL error class""" 72 | 73 | pass 74 | 75 | 76 | def pytest_configure(config): 77 | """ 78 | Called at the start of the pytest session, before tests are collected. 79 | This is where we apply our global patches before any imports happen. 80 | """ 81 | # Create and store original modules if they exist 82 | original_mysql = sys.modules.get("mysql.connector") 83 | 84 | # Create mock MySQL module 85 | sys.modules["mysql.connector"] = MockMySQLModule 86 | 87 | # Store the original function for later import and patching 88 | config._mysql_original = original_mysql 89 | 90 | 91 | @pytest.hookimpl(trylast=True) 92 | def pytest_sessionfinish(session, exitstatus): 93 | """Restore original modules after all tests are done""" 94 | if hasattr(session.config, "_mysql_original") and session.config._mysql_original: 95 | sys.modules["mysql.connector"] = session.config._mysql_original 96 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["hatchling"] 3 | build-backend = "hatchling.build" 4 | 5 | [project] 6 | name = "greptimedb-mcp-server" 7 | version = "0.2.2" 8 | description = "A Model Context Protocol (MCP) server that enables secure interaction with GreptimeDB databases. This server allows AI assistants to list tables, read data, and execute SQL queries through a controlled interface, making database exploration and analysis safer and more structured." 9 | readme = "README.md" 10 | license = "MIT" 11 | requires-python = ">=3.11" 12 | dependencies = [ 13 | "mcp>=1.0.0", 14 | "mysql-connector-python>=9.1.0", 15 | "pyyaml>=6.0.2", 16 | ] 17 | 18 | [[project.authors]] 19 | name = "dennis zhuang" 20 | email = "killme2008@gmail.com" 21 | 22 | [project.scripts] 23 | greptimedb-mcp-server = "greptimedb_mcp_server:main" 24 | 25 | [tool.hatch.build.targets.wheel] 26 | packages = ["src/greptimedb_mcp_server"] 27 | 28 | [tool.hatch.build] 29 | exclude = [ 30 | "venv/", 31 | ".git/", 32 | ".gitignore", 33 | "*.pyc", 34 | "__pycache__/", 35 | ".pytest_cache/", 36 | ".coverage", 37 | "tests/", 38 | "docs/", 39 | ] 40 | 41 | include = [ 42 | "src/**", 43 | "templates/**/*.md", 44 | "templates/**/*.yaml", 45 | "README.md", 46 | "LICENSE*", 47 | ] 48 | 49 | [tool.uv] 50 | dev-dependencies = [ 51 | "pyright", 52 | "black", 53 | "flake8", 54 | "pytest", 55 | "pytest-asyncio", 56 | "pytest-cov", 57 | ] 58 | -------------------------------------------------------------------------------- /pytest.ini: -------------------------------------------------------------------------------- 1 | [pytest] 2 | asyncio_mode = auto 3 | asyncio_default_fixture_loop_scope = function 4 | testpaths = tests 5 | python_files = test_*.py 6 | -------------------------------------------------------------------------------- /src/greptimedb_mcp_server/__init__.py: -------------------------------------------------------------------------------- 1 | from greptimedb_mcp_server.config import Config 2 | import sys 3 | 4 | if not "-m" in sys.argv: 5 | from . import server 6 | import asyncio 7 | 8 | 9 | def main(): 10 | """Main entry point for the package.""" 11 | config = Config.from_env_arguments() 12 | asyncio.run(server.main(config)) 13 | 14 | 15 | # Expose important items at package level 16 | __all__ = ["main", "server"] 17 | -------------------------------------------------------------------------------- /src/greptimedb_mcp_server/config.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from dataclasses import dataclass 3 | import os 4 | 5 | 6 | @dataclass 7 | class Config: 8 | """ 9 | Configuration for the greptimedb mcp server. 10 | """ 11 | 12 | host: str 13 | """ 14 | GreptimeDB host 15 | """ 16 | 17 | port: int 18 | """ 19 | GreptimeDB MySQL protocol port 20 | """ 21 | 22 | user: str 23 | """ 24 | GreptimeDB username 25 | """ 26 | 27 | password: str 28 | """ 29 | GreptimeDB password 30 | """ 31 | 32 | database: str 33 | """ 34 | GreptimeDB database name 35 | """ 36 | 37 | time_zone: str 38 | """ 39 | GreptimeDB session time zone 40 | """ 41 | 42 | @staticmethod 43 | def from_env_arguments() -> "Config": 44 | """ 45 | Parse command line arguments. 46 | """ 47 | parser = argparse.ArgumentParser(description="GreptimeDB MCP Server") 48 | 49 | parser.add_argument( 50 | "--host", 51 | type=str, 52 | help="GreptimeDB host", 53 | default=os.getenv("GREPTIMEDB_HOST", "localhost"), 54 | ) 55 | 56 | parser.add_argument( 57 | "--port", 58 | type=int, 59 | help="GreptimeDB MySQL protocol port", 60 | default=os.getenv("GREPTIMEDB_PORT", 4002), 61 | ) 62 | 63 | parser.add_argument( 64 | "--database", 65 | type=str, 66 | help="GreptimeDB connect database name", 67 | default=os.getenv("GREPTIMEDB_DATABASE", "public"), 68 | ) 69 | 70 | parser.add_argument( 71 | "--user", 72 | type=str, 73 | help="GreptimeDB username", 74 | default=os.getenv("GREPTIMEDB_USER", ""), 75 | ) 76 | 77 | parser.add_argument( 78 | "--password", 79 | type=str, 80 | help="GreptimeDB password", 81 | default=os.getenv("GREPTIMEDB_PASSWORD", ""), 82 | ) 83 | 84 | parser.add_argument( 85 | "--timezone", 86 | type=str, 87 | help="GreptimeDB session time zone", 88 | default=os.getenv("GREPTIMEDB_TIMEZONE", ""), 89 | ) 90 | 91 | args = parser.parse_args() 92 | return Config( 93 | host=args.host, 94 | port=args.port, 95 | database=args.database, 96 | user=args.user, 97 | password=args.password, 98 | time_zone=args.timezone, 99 | ) 100 | -------------------------------------------------------------------------------- /src/greptimedb_mcp_server/server.py: -------------------------------------------------------------------------------- 1 | from greptimedb_mcp_server.config import Config 2 | from greptimedb_mcp_server.utils import security_gate, templates_loader 3 | 4 | import datetime 5 | import asyncio 6 | import re 7 | import logging 8 | from logging import Logger 9 | from mysql.connector import connect, Error 10 | from mcp.server import Server 11 | from mcp.types import ( 12 | Resource, 13 | Tool, 14 | TextContent, 15 | Prompt, 16 | GetPromptResult, 17 | PromptMessage, 18 | ) 19 | from pydantic import AnyUrl 20 | 21 | # Resource URI prefix 22 | RES_PREFIX = "greptime://" 23 | # Resource query results limit 24 | RESULTS_LIMIT = 100 25 | 26 | # Configure logging 27 | logging.basicConfig( 28 | level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" 29 | ) 30 | 31 | 32 | def format_value(value): 33 | """Quote string and datetime values, leave others as-is""" 34 | if isinstance(value, (str, datetime.datetime, datetime.date, datetime.time)): 35 | return f'"{value}"' 36 | return str(value) 37 | 38 | 39 | # The GreptimeDB MCP Server 40 | class DatabaseServer: 41 | def __init__(self, logger: Logger, config: Config): 42 | """Initialize the GreptimeDB MCP server""" 43 | self.app = Server("greptimedb_mcp_server") 44 | self.logger = logger 45 | self.db_config = { 46 | "host": config.host, 47 | "port": config.port, 48 | "user": config.user, 49 | "password": config.password, 50 | "database": config.database, 51 | "time_zone": config.time_zone, 52 | } 53 | self.templates = templates_loader() 54 | 55 | self.logger.info(f"GreptimeDB Config: {self.db_config}") 56 | 57 | # Register callbacks 58 | self.app.list_resources()(self.list_resources) 59 | self.app.read_resource()(self.read_resource) 60 | self.app.list_prompts()(self.list_prompts) 61 | self.app.get_prompt()(self.get_prompt) 62 | self.app.list_tools()(self.list_tools) 63 | self.app.call_tool()(self.call_tool) 64 | 65 | async def list_resources(self) -> list[Resource]: 66 | """List GreptimeDB tables as resources.""" 67 | logger = self.logger 68 | config = self.db_config 69 | 70 | try: 71 | with connect(**config) as conn: 72 | with conn.cursor() as cursor: 73 | cursor.execute("SHOW TABLES") 74 | tables = cursor.fetchall() 75 | logger.info(f"Found tables: {tables}") 76 | 77 | resources = [] 78 | for table in tables: 79 | resources.append( 80 | Resource( 81 | uri=f"{RES_PREFIX}{table[0]}/data", 82 | name=f"Table: {table[0]}", 83 | mimeType="text/plain", 84 | description=f"Data in table: {table[0]}", 85 | ) 86 | ) 87 | return resources 88 | except Error as e: 89 | logger.error(f"Failed to list resources: {str(e)}") 90 | return [] 91 | 92 | async def read_resource(self, uri: AnyUrl) -> str: 93 | """Read table contents.""" 94 | logger = self.logger 95 | config = self.db_config 96 | 97 | uri_str = str(uri) 98 | logger.info(f"Reading resource: {uri_str}") 99 | 100 | if not uri_str.startswith(RES_PREFIX): 101 | raise ValueError(f"Invalid URI scheme: {uri_str}") 102 | 103 | parts = uri_str[len(RES_PREFIX) :].split("/") 104 | table = parts[0] 105 | if not re.match(r"^[a-zA-Z_:-][a-zA-Z0-9_:\-\.@#]*", table): 106 | raise ValueError("Invalid table name") 107 | 108 | try: 109 | with connect(**config) as conn: 110 | with conn.cursor() as cursor: 111 | cursor.execute(f"SELECT * FROM {table} LIMIT %s", (RESULTS_LIMIT,)) 112 | columns = [desc[0] for desc in cursor.description] 113 | rows = cursor.fetchall() 114 | result = [ 115 | ",".join(format_value(val) for val in row) for row in rows 116 | ] 117 | return "\n".join([",".join(columns)] + result) 118 | 119 | except Error as e: 120 | logger.error(f"Database error reading resource {uri}: {str(e)}") 121 | raise RuntimeError(f"Database error: {str(e)}") 122 | 123 | async def list_prompts(self) -> list[Prompt]: 124 | """List available GreptimeDB prompts.""" 125 | logger = self.logger 126 | 127 | logger.info("Listing prompts...") 128 | prompts = [] 129 | for name, template in self.templates.items(): 130 | logger.info(f"Found prompt: {name}") 131 | prompts.append( 132 | Prompt( 133 | name=name, 134 | description=template["config"]["description"], 135 | arguments=template["config"]["arguments"], 136 | ) 137 | ) 138 | return prompts 139 | 140 | async def get_prompt( 141 | self, name: str, arguments: dict[str, str] | None 142 | ) -> GetPromptResult: 143 | """Handle the get_prompt request.""" 144 | logger = self.logger 145 | 146 | logger.info(f"Get prompt: {name}") 147 | if name not in self.templates: 148 | logger.error(f"Unknown template: {name}") 149 | raise ValueError(f"Unknown template: {name}") 150 | 151 | template = self.templates[name] 152 | formatted_template = template["template"] 153 | 154 | # Replace placeholders with arguments 155 | if arguments: 156 | for key, value in arguments.items(): 157 | formatted_template = formatted_template.replace( 158 | f"{{{{ {key} }}}}", value 159 | ) 160 | 161 | return GetPromptResult( 162 | description=template["config"]["description"], 163 | messages=[ 164 | PromptMessage( 165 | role="user", 166 | content=TextContent(type="text", text=formatted_template), 167 | ) 168 | ], 169 | ) 170 | 171 | async def list_tools(self) -> list[Tool]: 172 | """List available GreptimeDB tools.""" 173 | logger = self.logger 174 | 175 | logger.info("Listing tools...") 176 | return [ 177 | Tool( 178 | name="execute_sql", 179 | description="Execute SQL query against GreptimeDB. Please use MySQL dialect when generating SQL queries.", 180 | inputSchema={ 181 | "type": "object", 182 | "properties": { 183 | "query": { 184 | "type": "string", 185 | "description": "The SQL query to execute (using MySQL dialect)", 186 | } 187 | }, 188 | "required": ["query"], 189 | }, 190 | ) 191 | ] 192 | 193 | async def call_tool(self, name: str, arguments: dict) -> list[TextContent]: 194 | """Execute SQL commands.""" 195 | logger = self.logger 196 | config = self.db_config 197 | 198 | logger.info(f"Calling tool: {name} with arguments: {arguments}") 199 | 200 | if name != "execute_sql": 201 | raise ValueError(f"Unknown tool: {name}") 202 | 203 | query = arguments.get("query") 204 | if not query: 205 | raise ValueError("Query is required") 206 | 207 | # Check if query is dangerous 208 | is_dangerous, reason = security_gate(query=query) 209 | if is_dangerous: 210 | return [ 211 | TextContent( 212 | type="text", 213 | text="Error: Contain dangerous operations, reason:" + reason, 214 | ) 215 | ] 216 | 217 | try: 218 | with connect(**config) as conn: 219 | with conn.cursor() as cursor: 220 | cursor.execute(query) 221 | 222 | stmt = query.strip().upper() 223 | # Special handling for SHOW DATABASES 224 | if stmt.startswith("SHOW DATABASES"): 225 | dbs = cursor.fetchall() 226 | result = ["Databases"] # Header 227 | result.extend([db[0] for db in dbs]) 228 | return [TextContent(type="text", text="\n".join(result))] 229 | # Special handling for SHOW TABLES 230 | if stmt.startswith("SHOW TABLES"): 231 | tables = cursor.fetchall() 232 | result = ["Tables_in_" + config["database"]] # Header 233 | result.extend([table[0] for table in tables]) 234 | return [TextContent(type="text", text="\n".join(result))] 235 | # Regular queries 236 | elif any( 237 | stmt.startswith(cmd) 238 | for cmd in ["SELECT", "SHOW", "DESC", "TQL", "EXPLAIN"] 239 | ): 240 | columns = [desc[0] for desc in cursor.description] 241 | rows = cursor.fetchall() 242 | result = [",".join(map(str, row)) for row in rows] 243 | return [ 244 | TextContent( 245 | type="text", 246 | text="\n".join([",".join(columns)] + result), 247 | ) 248 | ] 249 | 250 | # Non-SELECT queries 251 | else: 252 | conn.commit() 253 | return [ 254 | TextContent( 255 | type="text", 256 | text=f"Query executed successfully. Rows affected: {cursor.rowcount}", 257 | ) 258 | ] 259 | 260 | except Error as e: 261 | logger.error(f"Error executing SQL '{query}': {e}") 262 | return [TextContent(type="text", text=f"Error executing query: {str(e)}")] 263 | 264 | async def run(self): 265 | """Run the MCP server.""" 266 | logger = self.logger 267 | from mcp.server.stdio import stdio_server 268 | 269 | async with stdio_server() as (read_stream, write_stream): 270 | try: 271 | await self.app.run( 272 | read_stream, write_stream, self.app.create_initialization_options() 273 | ) 274 | except Exception as e: 275 | logger.error(f"Server error: {str(e)}", exc_info=True) 276 | raise 277 | 278 | 279 | async def main(config: Config): 280 | """Main entry point to run the MCP server.""" 281 | logger = logging.getLogger("greptimedb_mcp_server") 282 | db_server = DatabaseServer(logger, config) 283 | 284 | logger.info("Starting GreptimeDB MCP server...") 285 | 286 | await db_server.run() 287 | 288 | 289 | if __name__ == "__main__": 290 | asyncio.run(main(Config.from_env_arguments())) 291 | -------------------------------------------------------------------------------- /src/greptimedb_mcp_server/templates/metrics_analysis/config.yaml: -------------------------------------------------------------------------------- 1 | description: "Comprehensive metrics analysis template for monitoring and analyzing data from GreptimeDB" 2 | version: "1.0" 3 | arguments: 4 | - name: "topic" 5 | description: "Topic or area of metrics to analyze" 6 | required: false 7 | - name: "start_time" 8 | description: "Start time for the analysis period, example:(2025-01-01T00:00:00Z, now)" 9 | required: true 10 | - name: "end_time" 11 | description: "End time for the analysis period, example:(2025-01-01T01:00:00Z, now-1h)" 12 | required: true 13 | metadata: 14 | tags: 15 | - monitoring 16 | - metrics 17 | - performance 18 | suggested_use: "Use for system performance monitoring, resource usage analysis, and identifying performance bottlenecks" 19 | output_format: "markdown" 20 | -------------------------------------------------------------------------------- /src/greptimedb_mcp_server/templates/metrics_analysis/template.md: -------------------------------------------------------------------------------- 1 | # Monitoring Metrics Analysis Template 2 | You are a Monitoring Metrics Analysis assistant working to help users query and analyze data from GreptimeDB. Your task is to analyze the user's data needs and provide SQL queries to extract relevant information from GreptimeDB. 3 | 4 | Topic: {{ topic }} 5 | 6 | Time Range: 7 | - Start: {{ start_time }} 8 | - End: {{ end_time }} 9 | 10 | ## 1. Overview 11 | GreptimeDB is a time series database unifying metrics, logs, and events. 12 | 1. Prompts: This server provides prompts to help structure interactions with GreptimeDB. 13 | 2. Resources: You can find tables in resources with the format: "greptime:///data" 14 | 3. Tools: 15 | - "execute_sql": Execute SQL commands (MySQL syntax). 16 | 17 | ## 2. Guidelines 18 | 1. Time range is crucial - always specify a time range in your queries. 19 | 2. To explore available data, use SQL commands, such as (`DESCRIBE`, `SELECT`, `SHOW TABLES`) 20 | 3. Follow SQL best practices: 21 | - Use appropriate filtering to limit result sets 22 | - Consider using aggregation functions for time series data 23 | - Leverage GreptimeDB's built-in time functions 24 | 4. The server will block dangerous operations. Focus on read operations unless you need to modify data 25 | 5. Format your response using clean markdown with appropriate headers and bullet points. 26 | 27 | ## 3. Example Use Cases and Queries 28 | - Monitor system performance indicators 29 | - Identify performance bottlenecks 30 | - Analyze resource usage trends 31 | - Generate performance reports 32 | - Recommend alerting thresholds 33 | 34 | ```sql 35 | -- Get table schema 36 | DESCRIBE ${table}; 37 | -- Get recent data sample 38 | SELECT * FROM ${table} 39 | ${sample_queries} ORDER BY ${time_column} DESC LIMIT 100 40 | -- Current metrics summary 41 | SELECT 42 | avg(${metrics}) as avg_value, 43 | max(${metrics}) as peak_value, 44 | min(${metrics}) as min_value, 45 | percentile_cont(0.95) WITHIN GROUP (ORDER BY ${metrics}) as p95 46 | FROM ${table} 47 | WHERE ${time_column} >= NOW() - INTERVAL '${time_range}'; 48 | -- Detect anomalies (values outside 2 standard deviations) 49 | WITH stats AS ( 50 | SELECT 51 | avg(${metrics}) as avg_value, 52 | stddev(${metrics}) as stddev_value 53 | FROM ${table} 54 | WHERE ${time_column} >= NOW() - INTERVAL '${time_range}' 55 | ) 56 | SELECT 57 | ${time_column}, 58 | ${metrics}, 59 | 'Anomaly' as status 60 | FROM ${table}, stats 61 | WHERE 62 | ${time_column} >= NOW() - INTERVAL '${time_range}' 63 | AND (${metrics} > avg_value + 2 * stddev_value 64 | OR ${metrics} < avg_value - 2 * stddev_value); 65 | ``` 66 | 67 | ## 4. Additional Notes 68 | 1. If you don't know how to answer a specific question, suggest exploring the schema first to understand the available data structure. 69 | 2. Explain query results in a clear, informative way. 70 | 3. Help them analyze time series data effectively. 71 | -------------------------------------------------------------------------------- /src/greptimedb_mcp_server/templates/table_operation/config.yaml: -------------------------------------------------------------------------------- 1 | description: "Table operations such as querying region metadata, backing up/restoring data for table in GreptimeDB." 2 | version: "1.0" 3 | arguments: 4 | - name: "table" 5 | description: "The table to operate on." 6 | required: false 7 | metadata: 8 | tags: 9 | - operation 10 | - backup 11 | - restore 12 | - metadata 13 | suggested_use: "Used for table data backup and restore, or querying metadata information such as partitions, regions, etc." 14 | output_format: "markdown" 15 | -------------------------------------------------------------------------------- /src/greptimedb_mcp_server/templates/table_operation/template.md: -------------------------------------------------------------------------------- 1 | # Table Operation Template 2 | You are a Table Operation assistant working to help users manage table data in **GreptimeDB**. Your task is to assist users with backup/restoration and querying region metadata, utilizing familiar SQL operations reminiscent of MySQL's **information_schema**. 3 | 4 | ## Table: **{{ table }}** 5 | 6 | --- 7 | 8 | ## 1. Overview 9 | GreptimeDB provides efficient tools for managing table data and distributed region metadata. By using SQL syntax derived from MySQL's **information_schema**, users can easily observe table schema and distribution details or perform data operations like backups and restores. 10 | 11 | ### Key Operations: 12 | 1. **Backup and Restore Table Data**: Export table contents using `COPY TO` and restore them with `COPY FROM`. 13 | 2. **Query Region Metadata**: Use system tables like `region_peers` to retrieve region details and peer states. 14 | 15 | ### Supported Features: 16 | - **File Formats**: Parquet, CSV, JSON for backup/restore. 17 | - **Metadata Inspection**: Query the schema and region details using structured SQL commands. 18 | - **Cloud Storage Integration**: Enable operations with services like AWS S3, provided proper credentials. 19 | 20 | --- 21 | 22 | ## 2. Guidelines for Table Operations 23 | 1. **Specify Table**: Always use table name (`{{ table }}`) in operations and queries. 24 | 2. **Backup/Restore Commands**: Provide accurate file paths and formats to ensure compatibility. 25 | 3. **Region Queries**: Use the `region_peers` table for distribution and state monitoring. 26 | 4. **Metadata Tables**: Query metadata like `region_peers` or standard schema views (similar to MySQL's **information_schema**). 27 | 5. **Time Range Filtering**: Backup and restoration commands can include time constraints using `START_TIME` and `END_TIME`. 28 | 29 | --- 30 | 31 | ## 3. Example Operations 32 | 33 | ### **Backup Table Data** 34 | Export table `{{ table }}` data to file with optional filtering: 35 | 36 | ```sql 37 | COPY TO 's3://my-backup-bucket/{{ table }}.parquet' 38 | FROM {{ table }} 39 | WITH ( 40 | FORMAT = 'parquet', -- Change format if needed (e.g., 'csv', 'json') 41 | START_TIME = '2023-01-01T00:00:00Z', -- Optional 42 | END_TIME = '2025-01-01T00:00:00Z', -- Optional 43 | CONNECTION = { 44 | URL = 's3://my-backup-bucket', 45 | REGION = 'us-west-1', 46 | ACCESS_KEY = 'your-access-key', 47 | SECRET_KEY = 'your-secret-key' 48 | } 49 | ); 50 | ``` 51 | 52 | ### **Restore Table Data** 53 | Import data from a file into table `{{ table }}`: 54 | 55 | ```sql 56 | COPY FROM 's3://my-backup-bucket/{{ table }}.parquet' 57 | INTO {{ table }} 58 | WITH ( 59 | FORMAT = 'parquet', -- Change format if needed (e.g., 'csv', 'json') 60 | CONNECTION = { 61 | URL = 's3://my-backup-bucket', 62 | REGION = 'us-west-1', 63 | ACCESS_KEY = 'your-access-key', 64 | SECRET_KEY = 'your-secret-key' 65 | } 66 | ); 67 | ``` 68 | 69 | --- 70 | 71 | ### **Query Region Metadata** 72 | 73 | #### View Region Peers Metadata 74 | Query peer distribution across regions of table `{{ table }}`: 75 | 76 | ```sql 77 | SELECT * 78 | FROM information_schema.region_peers 79 | WHERE table_name = '{{ table }}'; 80 | ``` 81 | 82 | #### Inspect Peer States 83 | Find regions with problematic peer states (e.g., not `"RUNNING"`): 84 | 85 | ```sql 86 | SELECT region_id, peer_id, role, state 87 | FROM information_schema.region_peers 88 | WHERE table_name = '{{ table }}' 89 | AND state != 'RUNNING'; 90 | ``` 91 | 92 | --- 93 | 94 | ### Additional Metadata Queries 95 | #### Describe Table Schema 96 | Gain full details on table `{{ table }}` via **information_schema**: 97 | 98 | ```sql 99 | DESCRIBE {{ table }}; 100 | ``` 101 | 102 | #### View Available Tables 103 | List all user tables from the database: 104 | 105 | ```sql 106 | SHOW TABLES; 107 | ``` 108 | 109 | #### Check Available Regions 110 | View region configurations for any table: 111 | 112 | ```sql 113 | SELECT region_id, partition_key 114 | FROM information_schema.region_metadata 115 | WHERE table_name = '{{ table }}'; 116 | ``` 117 | 118 | --- 119 | 120 | ## 4. Additional Notes 121 | 1. For **cloud storage**, ensure setup includes access credentials and correct bucket URI. 122 | 2. Leverage **information_schema**-style queries for metadata, mirroring traditional MySQL layouts. 123 | 3. Use **Parquet** format where possible for efficient storage/restore. 124 | 4. Paths for table backup/restoration should always be valid (adjust for Windows compatibility with `/` instead of `\`). 125 | 5. Focus on filtering and specifying time ranges to limit large data operations. 126 | 127 | This template integrates familiar MySQL-style **information_schema** syntax to make GreptimeDB operations seamless for users transitioning from relational databases while also covering advanced distributed table and region queries. 128 | -------------------------------------------------------------------------------- /src/greptimedb_mcp_server/utils.py: -------------------------------------------------------------------------------- 1 | import re 2 | import logging 3 | import yaml 4 | import os 5 | 6 | logger = logging.getLogger("greptimedb_mcp_server") 7 | 8 | 9 | def security_gate(query: str) -> tuple[bool, str]: 10 | """ 11 | Simple security check for SQL queries. 12 | Args: 13 | query: The SQL query to check 14 | Returns: 15 | tuple: A boolean indicating if the query is dangerous, and a reason message 16 | """ 17 | if not query or not query.strip(): 18 | return True, "Empty query not allowed" 19 | 20 | # Remove comments and normalize whitespace 21 | clean_query = re.sub(r"/\*.*?\*/", " ", query, flags=re.DOTALL) # Remove /* */ 22 | clean_query = re.sub(r"--.*", "", clean_query) # Remove -- 23 | clean_query = re.sub(r"\s+", " ", clean_query).strip().upper() # Normalize spaces 24 | 25 | # Check for dangerous patterns 26 | dangerous_patterns = [ 27 | (r"\bDROP\b", "Forbided `DROP` operation"), 28 | (r"\bDELETE\b", "Forbided `DELETE` operation"), 29 | (r"\bREVOKE\b", "Forbided `REVOKE` operation"), 30 | (r"\bTRUNCATE\b", "Forbided `TRUNCATE` operation"), 31 | (r"\bUPDATE\b", "Forbided `UPDATE` operation"), 32 | (r"\bINSERT\b", "Forbided `INSERT` operation"), 33 | (r"\bALTER\b", "Forbided `ALTER` operation"), 34 | (r"\bCREATE\b", "Forbided `CREATE` operation"), 35 | (r"\bGRANT\b", "Forbided `GRANT` operation"), 36 | (r";\s*\w+", "Forbided multiple statements"), 37 | ] 38 | 39 | for pattern, reason in dangerous_patterns: 40 | if re.search(pattern, clean_query): 41 | logger.warning(f"Dangerous pattern detected: {query[:50]}...") 42 | return True, reason 43 | 44 | return False, "" 45 | 46 | 47 | def templates_loader() -> dict[str, dict[str, str]]: 48 | templates = {} 49 | template_dir = os.path.join(os.path.dirname(__file__), "templates") 50 | 51 | for category in os.listdir(template_dir): 52 | category_path = os.path.join(template_dir, category) 53 | if os.path.isdir(category_path): 54 | # Load config 55 | with open(os.path.join(category_path, "config.yaml"), "r") as f: 56 | config = yaml.safe_load(f) 57 | 58 | # Load template 59 | with open(os.path.join(category_path, "template.md"), "r") as f: 60 | template = f.read() 61 | 62 | templates[category] = {"config": config, "template": template} 63 | 64 | return templates 65 | -------------------------------------------------------------------------------- /tests/test_config.py: -------------------------------------------------------------------------------- 1 | import os 2 | from unittest.mock import patch 3 | from greptimedb_mcp_server.config import Config 4 | 5 | 6 | def test_config_default_values(): 7 | """ 8 | Test default configuration values 9 | """ 10 | with patch.dict(os.environ, {}, clear=True): 11 | with patch("sys.argv", ["script_name"]): 12 | config = Config.from_env_arguments() 13 | 14 | assert config.host == "localhost" 15 | assert config.port == 4002 16 | assert config.database == "public" 17 | assert config.user == "" 18 | assert config.password == "" 19 | assert config.time_zone == "" 20 | 21 | 22 | def test_config_env_variables(): 23 | """ 24 | Test configuration via environment variables 25 | """ 26 | env_vars = { 27 | "GREPTIMEDB_HOST": "test-host", 28 | "GREPTIMEDB_PORT": "5432", 29 | "GREPTIMEDB_DATABASE": "test_db", 30 | "GREPTIMEDB_USER": "test_user", 31 | "GREPTIMEDB_PASSWORD": "test_password", 32 | "GREPTIMEDB_TIMEZONE": "test_tz", 33 | } 34 | 35 | with patch.dict(os.environ, env_vars): 36 | with patch("sys.argv", ["script_name"]): 37 | config = Config.from_env_arguments() 38 | 39 | assert config.host == "test-host" 40 | assert config.port == 5432 41 | assert config.database == "test_db" 42 | assert config.user == "test_user" 43 | assert config.password == "test_password" 44 | assert config.time_zone == "test_tz" 45 | 46 | 47 | def test_config_cli_arguments(): 48 | """ 49 | Test configuration via command-line arguments 50 | """ 51 | cli_args = [ 52 | "script_name", 53 | "--host", 54 | "cli-host", 55 | "--port", 56 | "9999", 57 | "--database", 58 | "cli_db", 59 | "--user", 60 | "cli_user", 61 | "--password", 62 | "cli_password", 63 | "--timezone", 64 | "cli_tz", 65 | ] 66 | 67 | with patch.dict(os.environ, {}, clear=True): 68 | with patch("sys.argv", cli_args): 69 | config = Config.from_env_arguments() 70 | 71 | assert config.host == "cli-host" 72 | assert config.port == 9999 73 | assert config.database == "cli_db" 74 | assert config.user == "cli_user" 75 | assert config.password == "cli_password" 76 | assert config.time_zone == "cli_tz" 77 | 78 | 79 | def test_config_precedence(): 80 | """ 81 | Test configuration precedence (CLI arguments override environment variables) 82 | """ 83 | env_vars = { 84 | "GREPTIMEDB_HOST": "env-host", 85 | "GREPTIMEDB_PORT": "6666", 86 | "GREPTIMEDB_DATABASE": "env_db", 87 | "GREPTIMEDB_USER": "env_user", 88 | "GREPTIMEDB_PASSWORD": "env_password", 89 | "GREPTIMEDB_TIMEZONE": "env_tz", 90 | } 91 | 92 | cli_args = [ 93 | "script_name", 94 | "--host", 95 | "cli-host", 96 | "--port", 97 | "9999", 98 | "--database", 99 | "cli_db", 100 | "--user", 101 | "cli_user", 102 | "--password", 103 | "cli_password", 104 | "--timezone", 105 | "cli_tz", 106 | ] 107 | 108 | with patch.dict(os.environ, env_vars): 109 | with patch("sys.argv", cli_args): 110 | config = Config.from_env_arguments() 111 | 112 | assert config.host == "cli-host" 113 | assert config.port == 9999 114 | assert config.database == "cli_db" 115 | assert config.user == "cli_user" 116 | assert config.password == "cli_password" 117 | assert config.time_zone == "cli_tz" 118 | 119 | 120 | def test_config_object_creation(): 121 | """ 122 | Test direct creation of Config object 123 | """ 124 | config = Config( 125 | host="manual-host", 126 | port=1234, 127 | database="manual_db", 128 | user="manual_user", 129 | password="manual_password", 130 | time_zone="manual_timezone", 131 | ) 132 | 133 | assert config.host == "manual-host" 134 | assert config.port == 1234 135 | assert config.database == "manual_db" 136 | assert config.user == "manual_user" 137 | assert config.password == "manual_password" 138 | assert config.time_zone == "manual_timezone" 139 | -------------------------------------------------------------------------------- /tests/test_server.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import logging 3 | 4 | # Now we can safely import these 5 | from greptimedb_mcp_server.server import DatabaseServer 6 | from greptimedb_mcp_server.config import Config 7 | 8 | 9 | @pytest.fixture 10 | def config(): 11 | """Create a test configuration""" 12 | return Config( 13 | host="localhost", 14 | port=4002, 15 | user="testuser", 16 | password="testpassword", 17 | database="testdb", 18 | time_zone="", 19 | ) 20 | 21 | 22 | @pytest.fixture 23 | def logger(): 24 | """Create a test logger""" 25 | return logging.getLogger("test_logger") 26 | 27 | 28 | @pytest.mark.asyncio 29 | async def test_list_resources(logger, config): 30 | """Test listing database resources""" 31 | server = DatabaseServer(logger, config) 32 | resources = await server.list_resources() 33 | 34 | # Verify the results 35 | assert len(resources) == 2 36 | assert resources[0].name == "Table: users" 37 | assert str(resources[0].uri) == "greptime://users/data" 38 | 39 | 40 | @pytest.mark.asyncio 41 | async def test_read_resource(logger, config): 42 | """Test reading a specific database resource""" 43 | server = DatabaseServer(logger, config) 44 | result = await server.read_resource("greptime://users/data") 45 | 46 | # Verify the results contain expected data 47 | assert "id,name" in result 48 | assert '1,"John"' in result 49 | assert '2,"Jane"' in result 50 | 51 | 52 | @pytest.mark.asyncio 53 | async def test_list_tools(logger, config): 54 | """Test listing available database tools""" 55 | server = DatabaseServer(logger, config) 56 | tools = await server.list_tools() 57 | 58 | # Verify the tool list 59 | assert len(tools) == 1 60 | assert tools[0].name == "execute_sql" 61 | assert "query" in tools[0].inputSchema["properties"] 62 | 63 | 64 | @pytest.mark.asyncio 65 | async def test_call_tool_select_query(logger, config): 66 | """Test executing a SELECT query via tool""" 67 | server = DatabaseServer(logger, config) 68 | result = await server.call_tool("execute_sql", {"query": "SELECT * FROM users"}) 69 | 70 | # Verify the results 71 | assert len(result) == 1 72 | assert "id,name" in result[0].text 73 | assert "1,John" in result[0].text 74 | 75 | 76 | @pytest.mark.asyncio 77 | async def test_security_gate_dangerous_query(logger, config): 78 | """Test security gate blocking dangerous queries""" 79 | server = DatabaseServer(logger, config) 80 | 81 | result = await server.call_tool("execute_sql", {"query": "DROP TABLE users"}) 82 | 83 | # Verify that the security gate blocked the query 84 | assert "Error: Contain dangerous operations" in result[0].text 85 | assert "Forbided `DROP` operation" in result[0].text 86 | 87 | 88 | @pytest.mark.asyncio 89 | async def test_show_tables_query(logger, config): 90 | """Test SHOW TABLES query execution""" 91 | server = DatabaseServer(logger, config) 92 | result = await server.call_tool("execute_sql", {"query": "SHOW TABLES"}) 93 | 94 | # Verify the results 95 | assert len(result) == 1 96 | assert "Tables_in_testdb" in result[0].text 97 | assert "users" in result[0].text 98 | assert "orders" in result[0].text 99 | 100 | 101 | @pytest.mark.asyncio 102 | async def test_show_dbs_query(logger, config): 103 | """Test SHOW DATABASES query execution""" 104 | server = DatabaseServer(logger, config) 105 | result = await server.call_tool("execute_sql", {"query": "SHOW DATABASES"}) 106 | 107 | # Verify the results 108 | assert len(result) == 1 109 | assert "Databases" in result[0].text 110 | print(result[0].text) 111 | assert "public" in result[0].text 112 | assert "greptime_private" in result[0].text 113 | 114 | 115 | @pytest.mark.asyncio 116 | async def test_list_prompts(logger, config): 117 | """Test listing available prompts""" 118 | server = DatabaseServer(logger, config) 119 | prompts = await server.list_prompts() 120 | 121 | # Verify the results 122 | assert len(prompts) > 0 123 | # Check that each prompt has the expected properties 124 | for prompt in prompts: 125 | assert hasattr(prompt, "name") 126 | assert hasattr(prompt, "description") 127 | assert hasattr(prompt, "arguments") 128 | 129 | 130 | @pytest.mark.asyncio 131 | async def test_get_prompt_without_args(logger, config): 132 | """Test getting a prompt without arguments""" 133 | server = DatabaseServer(logger, config) 134 | # Get the first prompt from the list to test with 135 | prompts = await server.list_prompts() 136 | if not prompts: 137 | pytest.skip("No prompts available for testing") 138 | 139 | test_prompt_name = prompts[0].name 140 | result = await server.get_prompt(test_prompt_name, {}) 141 | 142 | # Verify the result has the expected structure 143 | assert hasattr(result, "messages") 144 | assert len(result.messages) > 0 145 | for message in result.messages: 146 | assert hasattr(message, "role") 147 | assert hasattr(message, "content") 148 | 149 | 150 | @pytest.mark.asyncio 151 | async def test_get_prompt_with_args(logger, config): 152 | """Test getting a prompt with argument substitution""" 153 | server = DatabaseServer(logger, config) 154 | # Assume there's a prompt with arguments 155 | prompts = await server.list_prompts() 156 | prompt_with_args = None 157 | 158 | # Find a prompt that has arguments 159 | for prompt in prompts: 160 | if prompt.arguments and len(prompt.arguments) > 0: 161 | prompt_with_args = prompt 162 | break 163 | 164 | if not prompt_with_args: 165 | pytest.skip("No prompts with arguments available for testing") 166 | 167 | # Create args dictionary with test values for each required argument 168 | args = {} 169 | for arg in prompt_with_args.arguments: 170 | args[arg.name] = f"test_{arg.name}" 171 | 172 | result = await server.get_prompt(prompt_with_args.name, args) 173 | 174 | # Verify result structure and argument substitution 175 | assert hasattr(result, "messages") 176 | assert len(result.messages) > 0 177 | 178 | # Check that at least one message contains our test values 179 | substitution_found = False 180 | for message in result.messages: 181 | for arg_name, arg_value in args.items(): 182 | if arg_value in message.content.text: 183 | substitution_found = True 184 | break 185 | if substitution_found: 186 | break 187 | 188 | assert substitution_found, "Argument substitution not found in prompt messages" 189 | 190 | 191 | @pytest.mark.asyncio 192 | async def test_get_prompt_nonexistent(logger, config): 193 | """Test getting a non-existent prompt""" 194 | server = DatabaseServer(logger, config) 195 | 196 | # Try to get a prompt that doesn't exist 197 | with pytest.raises(ValueError) as excinfo: 198 | await server.get_prompt("non_existent_prompt", {}) 199 | 200 | # Verify the error message 201 | assert "Unknown template: non_existent_prompt" in str(excinfo.value) 202 | 203 | 204 | def test_server_initialization(logger, config): 205 | """Test server initialization with configuration""" 206 | server = DatabaseServer(logger, config) 207 | 208 | # Verify the server was initialized correctly 209 | assert server.logger == logger 210 | assert server.db_config["host"] == "localhost" 211 | assert server.db_config["port"] == 4002 212 | assert server.db_config["user"] == "testuser" 213 | assert server.db_config["password"] == "testpassword" 214 | assert server.db_config["database"] == "testdb" 215 | -------------------------------------------------------------------------------- /tests/test_utils.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | from greptimedb_mcp_server.utils import templates_loader, security_gate 3 | 4 | 5 | def test_templates_loader_basic(): 6 | """Test that templates_loader can load existing templates""" 7 | # Call the function under test 8 | templates = templates_loader() 9 | 10 | # Basic validation that we got something back 11 | assert templates is not None 12 | 13 | # Check if templates is a dictionary 14 | assert isinstance(templates, dict), "Expected templates to be a dictionary" 15 | assert len(templates) > 0, "Expected templates dictionary to have items" 16 | 17 | # Check if the metrics_analysis template is in the dictionary 18 | assert "metrics_analysis" in templates, "metrics_analysis template not found" 19 | 20 | # Get the metrics_analysis template 21 | metrics_template = templates["metrics_analysis"] 22 | 23 | # Get the metrics_analysis template config 24 | config = metrics_template["config"] 25 | 26 | # Check that the config has the expected structure 27 | assert isinstance(config, dict), "Expected template to be a dictionary" 28 | assert "description" in config, "Template config missing 'description' field" 29 | assert "arguments" in config, "Template config missing 'arguments' field" 30 | assert "metadata" in config, "Template config missing 'metadata' field" 31 | 32 | # Check that the template has the expected arguments 33 | arguments = config["arguments"] 34 | assert isinstance(arguments, list), "Expected arguments to be a list" 35 | 36 | arg_names = [ 37 | arg.get("name") for arg in arguments if isinstance(arg, dict) and "name" in arg 38 | ] 39 | expected_args = ["topic", "start_time", "end_time"] 40 | for arg in expected_args: 41 | assert ( 42 | arg in arg_names 43 | ), f"Expected argument '{arg}' not found in metrics_analysis template" 44 | 45 | # Check template content 46 | tpl = metrics_template["template"] 47 | assert "{{ topic }}" in tpl 48 | assert "{{ start_time }}" in tpl 49 | assert "{{ end_time }}" in tpl 50 | 51 | 52 | def test_empty_queries(): 53 | """Test empty queries""" 54 | assert security_gate("") == (True, "Empty query not allowed") 55 | assert security_gate(" ") == (True, "Empty query not allowed") 56 | assert security_gate(None) == (True, "Empty query not allowed") 57 | 58 | 59 | def test_safe_queries(): 60 | """Test safe queries""" 61 | assert security_gate("SELECT * FROM users") == (False, "") 62 | assert security_gate("select id from products") == (False, "") 63 | 64 | 65 | def test_dangerous_operations(): 66 | """Test dangerous operations""" 67 | assert security_gate("DROP TABLE users") == (True, "Forbided `DROP` operation") 68 | assert security_gate("DELETE FROM users") == (True, "Forbided `DELETE` operation") 69 | assert security_gate("UPDATE users SET name='test'") == ( 70 | True, 71 | "Forbided `UPDATE` operation", 72 | ) 73 | assert security_gate("INSERT INTO users VALUES (1)") == ( 74 | True, 75 | "Forbided `INSERT` operation", 76 | ) 77 | 78 | 79 | def test_multiple_statements(): 80 | """Test multiple statements""" 81 | assert security_gate("SELECT * FROM users; SELECT * FROM test") == ( 82 | True, 83 | "Forbided multiple statements", 84 | ) 85 | 86 | 87 | def test_comment_bypass(): 88 | """Test comment bypass attempts""" 89 | assert security_gate("DROP/**/TABLE users") == (True, "Forbided `DROP` operation") 90 | assert security_gate("DROP--comment\nTABLE users") == ( 91 | True, 92 | "Forbided `DROP` operation", 93 | ) 94 | 95 | 96 | @pytest.mark.parametrize( 97 | "query,expected", 98 | [ 99 | ("SELECT * FROM users", (False, "")), 100 | ("DROP TABLE users", (True, "Forbided `DROP` operation")), 101 | ("DELETE FROM users", (True, "Forbided `DELETE` operation")), 102 | ("", (True, "Empty query not allowed")), 103 | ], 104 | ) 105 | def test_parametrized_queries(query, expected): 106 | """Parametrized test for multiple queries""" 107 | assert security_gate(query) == expected 108 | --------------------------------------------------------------------------------