feat: fix CTE query and new features

README.md

@@ -11,14 +11,36 @@ This server provides AI assistants with a secure and structured way to explore a
 # Project Status
 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.
 
-# Capabilities
+# Features
 
-* `list_resources` to list tables
-* `read_resource` to read table data
-* `list_tools` to list tools
-* `call_tool` to execute an SQL
-* `list_prompts` to list prompts
-* `get_prompt` to get the prompt by name
+## Resources
+- **list_resources** - List all tables in the database as browsable resources
+- **read_resource** - Read table data via `greptime://<table>/data` URIs
+
+## Tools
+
+| Tool | Description |
+|------|-------------|
+| `execute_sql` | Execute SQL queries with format (csv/json/markdown) and limit options |
+| `describe_table` | Get table schema including column names, types, and constraints |
+| `health_check` | Check database connection status and server version |
+| `execute_tql` | Execute TQL (PromQL-compatible) queries for time-series analysis |
+| `query_range` | Execute time-window aggregation queries with RANGE/ALIGN syntax |
+| `explain_query` | Analyze SQL or TQL query execution plans |
+
+## Prompts
+- **list_prompts** - List available prompt templates
+- **get_prompt** - Get a prompt template by name with argument substitution
+
+## Security
+All queries pass through a security gate that:
+- Blocks dangerous operations: DROP, DELETE, TRUNCATE, UPDATE, INSERT, ALTER, CREATE, GRANT, REVOKE
+- Prevents multiple statement execution
+- Allows read-only operations: SELECT, SHOW, DESCRIBE, TQL, EXPLAIN
+
+## Performance
+- Connection pooling for efficient database access
+- Async operations for non-blocking execution
 
 # Installation
 
@@ -51,6 +73,63 @@ Or via command-line args:
 
 # Usage
 
+## Tool Examples
+
+### execute_sql
+Execute SQL queries with optional format and limit:
+```json
+{
+ "query": "SELECT * FROM metrics WHERE host = 'server1'",
+ "format": "json",
+ "limit": 100
+}
+```
+Formats: `csv` (default), `json`, `markdown`
+
+### execute_tql
+Execute PromQL-compatible time-series queries:
+```json
+{
+ "query": "rate(http_requests_total[5m])",
+ "start": "2024-01-01T00:00:00Z",
+ "end": "2024-01-01T01:00:00Z",
+ "step": "1m",
+ "lookback": "5m"
+}
+```
+
+### query_range
+Execute time-window aggregation queries:
+```json
+{
+ "table": "metrics",
+ "select": "ts, host, avg(cpu) RANGE '5m'",
+ "align": "1m",
+ "by": "host",
+ "where": "region = 'us-east'"
+}
+```
+
+### describe_table
+Get table schema information:
+```json
+{
+ "table": "metrics"
+}
+```
+
+### explain_query
+Analyze query execution plan:
+```json
+{
+ "query": "SELECT * FROM metrics",
+ "analyze": true
+}
+```
+
+### health_check
+Check database connection (no parameters required).
+
 ## Claude Desktop Integration
 
 Configure the MCP server in Claude Desktop's configuration file:

conftest.py

@@ -7,25 +7,89 @@ class MockCursor:
  def __init__(self):
  self.query = ""
  self.rowcount = 2
+ self._results = []
+ self._fetch_index = 0
 
  def execute(self, query, args=None):
  self.query = query
-
- def fetchall(self):
+ self._fetch_index = 0
+  # Pre-populate results based on query
  if "SHOW TABLES" in self.query.upper():
- return [("users",), ("orders",)]
+ self._results = [("users",), ("orders",)]
  elif "SHOW DATABASES" in self.query.upper():
- return [("public",), ("greptime_private",)]
+ self._results = [("public",), ("greptime_private",)]
+ elif "DESCRIBE" in self.query.upper():
+ self._results = [
+ ("id", "Int64", "", "PRI", "", ""),
+ ("name", "String", "YES", "", "", ""),
+ ("ts", "TimestampMillisecond", "", "TIME INDEX", "", ""),
+ ]
+ elif "VERSION()" in self.query.upper():
+ self._results = [("GreptimeDB 0.9.0",)]
+ elif "TQL" in self.query.upper():
+ self._results = [
+ ("2024-01-01 00:00:00", "host1", 0.5),
+ ("2024-01-01 00:01:00", "host1", 0.6),
+ ]
+ elif "EXPLAIN" in self.query.upper():
+ self._results = [
+ ("GlobalLimitExec: skip=0, fetch=1000",),
+ (" SortExec: TopK(fetch=1000)",),
+ (" TableScan: users",),
+ ]
+ elif "ALIGN" in self.query.upper():
+ # Range query
+ self._results = [
+ ("2024-01-01 00:00:00", "host1", 45.5),
+ ("2024-01-01 00:05:00", "host1", 52.3),
+ ]
  elif "SELECT" in self.query.upper():
- return [(1, "John"), (2, "Jane")]
- return []
+ self._results = [(1, "John"), (2, "Jane")]
+ else:
+ self._results = []
+
+ def fetchall(self):
+ results = self._results[self._fetch_index :]
+ self._fetch_index = len(self._results)
+ return results
+
+ def fetchmany(self, size=None):
+ if size is None:
+ return self.fetchall()
+ results = self._results[self._fetch_index : self._fetch_index + size]
+ self._fetch_index += len(results)
+ return results
+
+ def fetchone(self):
+ if self._fetch_index < len(self._results):
+ result = self._results[self._fetch_index]
+ self._fetch_index += 1
+ return result
+ return None
 
  @property
  def description(self):
  if "SHOW TABLES" in self.query.upper():
  return [("table_name", None)]
  elif "SHOW DATABASES" in self.query.upper():
  return [("Databases", None)]
+ elif "DESCRIBE" in self.query.upper():
+ return [
+ ("Column", None),
+ ("Type", None),
+ ("Null", None),
+ ("Key", None),
+ ("Default", None),
+ ("Semantic Type", None),
+ ]
+ elif "VERSION()" in self.query.upper():
+ return [("version()", None)]
+ elif "TQL" in self.query.upper():
+ return [("ts", None), ("host", None), ("value", None)]
+ elif "EXPLAIN" in self.query.upper():
+ return [("plan", None)]
+ elif "ALIGN" in self.query.upper():
+ return [("ts", None), ("host", None), ("avg_cpu", None)]
  elif "SELECT" in self.query.upper():
  return [("id", None), ("name", None)]
  return []
@@ -60,6 +124,16 @@ def __exit__(self, *args):
  pass
 
 
+class MockConnectionPool:
+ """Mock for MySQL connection pool"""
+
+ def __init__(self, *args, **kwargs):
+ pass
+
+ def get_connection(self):
+ return MockConnection()
+
+
 class MockMySQLModule:
  """Mock for entire mysql.connector module"""
 
@@ -73,23 +147,37 @@ class Error(Exception):
  pass
 
 
+class MockPoolingModule:
+ """Mock for mysql.connector.pooling module"""
+
+ MySQLConnectionPool = MockConnectionPool
+
+
 def pytest_configure(config):
  """
  Called at the start of the pytest session, before tests are collected.
  This is where we apply our global patches before any imports happen.
  """
  # Create and store original modules if they exist
  original_mysql = sys.modules.get("mysql.connector")
+ original_pooling = sys.modules.get("mysql.connector.pooling")
 
- # Create mock MySQL module
+ # Create mock MySQL modules
  sys.modules["mysql.connector"] = MockMySQLModule
+ sys.modules["mysql.connector.pooling"] = MockPoolingModule
 
  # Store the original function for later import and patching
  config._mysql_original = original_mysql
+ config._pooling_original = original_pooling
 
 
 @pytest.hookimpl(trylast=True)
 def pytest_sessionfinish(session, exitstatus):
  """Restore original modules after all tests are done"""
  if hasattr(session.config, "_mysql_original") and session.config._mysql_original:
  sys.modules["mysql.connector"] = session.config._mysql_original
+ if (
+ hasattr(session.config, "_pooling_original")
+ and session.config._pooling_original
+ ):
+ sys.modules["mysql.connector.pooling"] = session.config._pooling_original