feat: audit logging

docs/llm-instructions.md

@@ -0,0 +1,68 @@
+# LLM Instructions for GreptimeDB MCP Server
+
+Add this to your system prompt to help AI assistants work with this MCP server.
+
+## System Prompt
+
+```
+You have access to a GreptimeDB MCP server for querying and managing time-series data, logs, and metrics.
+
+## Available Tools
+- `execute_sql`: Run SQL queries (SELECT, SHOW, DESCRIBE only - read-only access)
+- `execute_tql`: Run PromQL-compatible time-series queries
+- `query_range`: Time-window aggregation with RANGE/ALIGN syntax
+- `describe_table`: Get table schema information
+- `health_check`: Check database connection status
+- `explain_query`: Analyze query execution plans
+
+### Pipeline Management
+- `list_pipelines`: View existing log pipelines
+- `create_pipeline`: Create/update pipeline with YAML config (same name creates new version)
+- `dryrun_pipeline`: Test pipeline with sample data without writing
+- `delete_pipeline`: Remove a pipeline version
+
+**Note**: All HTTP API calls (pipeline tools) require authentication. The MCP server handles auth automatically using configured credentials. When providing curl examples to users, always include `-u <username>:<password>`.
+
+## Available Prompts
+Use these prompts for specialized tasks:
+- `pipeline_creator`: Generate pipeline YAML from log samples - use when user provides log examples
+- `log_pipeline`: Log analysis with full-text search
+- `metrics_analysis`: Metrics monitoring and analysis
+- `promql_analysis`: PromQL-style queries
+- `iot_monitoring`: IoT device data analysis
+- `trace_analysis`: Distributed tracing analysis
+- `table_operation`: Table diagnostics and optimization
+
+## Workflow Tips
+1. For log pipeline creation: Get log sample → use `pipeline_creator` prompt → generate YAML → `create_pipeline` → `dryrun_pipeline` to verify
+2. For data analysis: `describe_table` first → understand schema → `execute_sql` or `execute_tql`
+3. For time-series: Prefer `query_range` for aggregations, `execute_tql` for PromQL patterns
+4. Always check `health_check` if queries fail unexpectedly
+```
+
+## Using Prompts in Claude Desktop
+
+In Claude Desktop, you need to add MCP prompts manually:
+
+1. Click the **+** button in the conversation input area
+2. Select **MCP Server**
+3. Choose **Prompt/References**
+4. Select the prompt you want to use (e.g., `pipeline_creator`)
+5. Fill in the required arguments
+
+Note: Prompts are not automatically available via `/` slash commands in Claude Desktop. You must add them through the UI as described above.
+
+## Example: Creating a Pipeline
+
+Provide your log sample and ask Claude to create a pipeline:
+
+```
+Help me create a GreptimeDB pipeline to parse this nginx log:
+127.0.0.1 - - [25/May/2024:20:16:37 +0000] "GET /index.html HTTP/1.1" 200 612 "-" "Mozilla/5.0..."
+```
+
+Claude will:
+1. Analyze your log format
+2. Generate a pipeline YAML configuration
+3. Create the pipeline using `create_pipeline` tool
+4. Test it with `dryrun_pipeline` tool

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "greptimedb-mcp-server"
-version = "0.3.1"
+version = "0.4.0"
 description = "MCP server for GreptimeDB with SQL/TQL/PromQL support, sensitive data masking, and prompt templates for observability data analysis."
 readme = "README.md"
 license = {text = "MIT"}

src/greptimedb_mcp_server/__init__.py

@@ -1,3 +1,4 @@
+import asyncio
 import sys
 
 if "-m" not in sys.argv:
@@ -6,7 +7,12 @@
 
 def main():
     """Main entry point for the package."""
-    server.main()
+    try:
+        server.main()
+    except KeyboardInterrupt:
+        print("\nReceived Ctrl+C, shutting down...")
+    except asyncio.CancelledError:
+        print("\nServer shutdown complete.")
 
 
 # Expose important items at package level

src/greptimedb_mcp_server/config.py

@@ -79,6 +79,11 @@ class Config:
     MCP HTTP server bind port (for sse/streamable-http transports)
     """
 
+    audit_enabled: bool
+    """
+    Enable audit logging for all tool calls
+    """
+
     @staticmethod
     def from_env_arguments() -> "Config":
         """
@@ -186,6 +191,13 @@ def from_env_arguments() -> "Config":
             default=int(os.getenv("GREPTIMEDB_LISTEN_PORT", "8080")),
         )
 
+        parser.add_argument(
+            "--audit-enabled",
+            type=lambda x: x.lower() not in ("false", "0", "no"),
+            help="Enable audit logging for all tool calls (default: true)",
+            default=os.getenv("GREPTIMEDB_AUDIT_ENABLED", "true"),
+        )
+
         args = parser.parse_args()
         return Config(
             host=args.host,
@@ -202,4 +214,5 @@ def from_env_arguments() -> "Config":
             transport=args.transport,
             listen_host=args.listen_host,
             listen_port=args.listen_port,
+            audit_enabled=args.audit_enabled,
         )

src/greptimedb_mcp_server/server.py

@@ -12,6 +12,7 @@
     validate_fill,
     validate_time_expression,
     format_tql_time_param,
+    audit_log,
 )
 
 import asyncio
@@ -857,11 +858,39 @@ def prompt_fn({arg_params}) -> str:
 _register_prompts()
 
 
+def _install_audit_hook():
+    """Install audit logging hook by wrapping tool manager's call_tool method."""
+    original_call_tool = mcp._tool_manager.call_tool
+
+    async def audited_call_tool(name, arguments, context=None, convert_result=False):
+        start_time = time.time()
+        try:
+            result = await original_call_tool(name, arguments, context, convert_result)
+            elapsed_ms = (time.time() - start_time) * 1000
+            audit_log(name, arguments, success=True, duration_ms=elapsed_ms)
+            return result
+        except Exception as e:
+            elapsed_ms = (time.time() - start_time) * 1000
+            audit_log(
+                name, arguments, success=False, duration_ms=elapsed_ms, error=str(e)
+            )
+            raise
+
+    mcp._tool_manager.call_tool = audited_call_tool
+
+
 def main():
     """Main entry point."""
     global _config
     _config = Config.from_env_arguments()
 
+    # Install audit logging hook if enabled
+    if _config.audit_enabled:
+        _install_audit_hook()
+        logger.info("Audit logging: enabled")
+    else:
+        logger.info("Audit logging: disabled")
+
     # Only configure HTTP server settings for non-stdio transports
     # to avoid overriding user's programmatic configuration
     if _config.transport != "stdio":

src/greptimedb_mcp_server/utils.py

@@ -2,6 +2,7 @@
 import logging
 import yaml
 import os
+from typing import Any
 
 logger = logging.getLogger("greptimedb_mcp_server")
 
@@ -157,10 +158,56 @@ def validate_time_expression(value: str, name: str) -> str:
         raise ValueError(f"{name} is required")
     if ";" in value or "--" in value:
         raise ValueError(f"Invalid characters in {name}")
-    # Guard against malformed or injected strings with unbalanced quotes
     if value.count("'") % 2 != 0:
         raise ValueError(f"Unbalanced quotes in {name}")
     is_dangerous, reason = security_gate(value)
     if is_dangerous:
         raise ValueError(f"Dangerous pattern in {name}: {reason}")
     return value
+
+
+# Audit logging
+audit_logger = logging.getLogger("greptimedb_mcp_server.audit")
+
+
+def _truncate_value(v: Any, max_len: int = 200) -> str:
+    """Truncate a value to max_len characters."""
+    v_str = str(v)
+    if len(v_str) > max_len:
+        return v_str[:max_len] + "..."
+    return v_str
+
+
+def _format_audit_params(params: dict) -> str:
+    """Format parameters for audit log."""
+    if not params:
+        return ""
+    parts = []
+    for k, v in params.items():
+        parts.append(f'{k}="{_truncate_value(v)}"')
+    return " | ".join(parts)
+
+
+def audit_log(
+    tool: str,
+    params: dict,
+    success: bool,
+    duration_ms: float,
+    error: str | None = None,
+):
+    """Record audit log for tool invocation. Never raises exceptions."""
+    try:
+        parts = [f"[AUDIT] {tool}"]
+
+        params_str = _format_audit_params(params)
+        if params_str:
+            parts.append(params_str)
+
+        parts.append(f"success={success}")
+        if error:
+            parts.append(f'error="{_truncate_value(error)}"')
+        parts.append(f"duration_ms={duration_ms:.1f}")
+
+        audit_logger.info(" | ".join(parts))
+    except Exception:
+        pass

tests/test_server.py

@@ -39,6 +39,7 @@ def setup_state():
         transport="stdio",
         listen_host="0.0.0.0",
         listen_port=8080,
+        audit_enabled=False,
     )
     # Set global config for get_config() calls
     server._config = config

tests/test_utils.py

@@ -8,6 +8,9 @@
     is_sql_time_expression,
     format_tql_time_param,
     validate_time_expression,
+    _truncate_value,
+    _format_audit_params,
+    audit_log,
 )
 from greptimedb_mcp_server.formatter import format_results
 
@@ -523,3 +526,36 @@ def test_validate_time_expression_dangerous():
     with pytest.raises(ValueError) as excinfo:
         validate_time_expression("DELETE FROM users", "start")
     assert "Dangerous pattern" in str(excinfo.value)
+
+
+# Tests for audit logging functions
+
+
+def test_truncate_value():
+    """Test _truncate_value truncates long values"""
+    assert _truncate_value("short") == "short"
+    assert _truncate_value("a" * 201) == "a" * 200 + "..."
+
+
+def test_format_audit_params():
+    """Test _format_audit_params formats params correctly"""
+    assert _format_audit_params({}) == ""
+    assert _format_audit_params({"query": "SELECT 1"}) == 'query="SELECT 1"'
+
+
+def test_audit_log(caplog):
+    """Test audit_log records tool calls"""
+    import logging
+
+    with caplog.at_level(logging.INFO, logger="greptimedb_mcp_server.audit"):
+        audit_log("execute_sql", {"query": "SELECT 1"}, success=True, duration_ms=10.5)
+
+    assert len(caplog.records) == 1
+    msg = caplog.records[0].message
+    assert "[AUDIT] execute_sql" in msg
+    assert "success=True" in msg
+
+
+def test_audit_log_never_raises():
+    """Test audit_log never raises exceptions"""
+    audit_log(None, None, None, None, None)  # Should not raise