| """MCP Database Server — main server with 7 reasoning tools.""" |
|
|
| import json |
| import sys |
| import logging |
| from mcp.server.mcpserver import MCPServer |
| from mcp_database_universal.config import DatabaseConfig |
| from mcp_database_universal.engines.base import BaseEngine |
| from mcp_database_universal.safety import SafetyValidator |
| from mcp_database_universal.schema_inspector import SchemaInspector |
| from mcp_database_universal.formatters.llm import LLMFormatter |
| from mcp_database_universal.nl2sql import translate |
| from mcp_database_universal import __version__ |
|
|
| logger = logging.getLogger("mcp-db") |
|
|
|
|
| async def create_server(config: DatabaseConfig, engine: BaseEngine) -> MCPServer: |
| server = MCPServer( |
| name="mcp-database-server", |
| version=__version__, |
| ) |
|
|
| safety = SafetyValidator( |
| read_only=config.is_effectively_read_only(), |
| max_rows=config.max_rows, |
| max_query_time=config.max_query_time, |
| ) |
| inspector = SchemaInspector(engine) |
| formatter = LLMFormatter(config) |
|
|
| @server.tool() |
| async def test_connection() -> str: |
| """ |
| Test database connection and return basic info. |
| |
| Use when: you need to verify the connection works, |
| or find out the DB type, version, and size. |
| Returns: engine type, version, database name, size. |
| """ |
| try: |
| info = await engine.get_db_info() |
| tables = await engine.get_tables() |
| result = formatter.format_db_info(info) |
| result += f"\n- **Tables:** {len(tables)}" |
| result += f"\n- **Read-only:** {'yes' if engine.is_read_only() else 'no'}" |
| return result |
| except Exception as e: |
| return f"Connection failed: {str(e)}" |
|
|
| @server.tool() |
| async def list_tables(include_stats: bool = True) -> str: |
| """ |
| List all tables in the database with metadata. |
| |
| Use when: you need an overview of what's in the database, |
| how many tables exist, and their relationships. |
| This is the first step when working with an unknown database. |
| Returns: list of tables with row counts, column counts, FK relationships. |
| """ |
| try: |
| tables = await engine.get_tables() |
| return formatter.format_table_list(tables) |
| except Exception as e: |
| return f"Error listing tables: {str(e)}" |
|
|
| @server.tool() |
| async def inspect_table(table_name: str, include_sample: bool = True, sample_size: int = 5) -> str: |
| """ |
| Inspect a table's structure: columns, types, indexes, foreign keys, sample data. |
| |
| Use when: you need to understand a specific table's structure, |
| what columns exist, their types, and relationships to other tables. |
| Returns: complete table overview with context for LLM. |
| """ |
| try: |
| detail = await engine.get_table_detail(table_name) |
| if not include_sample: |
| detail.sample_data = [] |
| return formatter.format_table_detail(detail) |
| except Exception as e: |
| return f"Error inspecting table '{table_name}': {str(e)}" |
|
|
| @server.tool() |
| async def query(sql: str, params: str = "{}") -> str: |
| """ |
| Execute a safe SQL query and return formatted results. |
| |
| Use when: you need to run a specific SQL query. |
| All queries are parametrized and pass through safety checks. |
| |
| SAFETY: |
| - Read-only by default: no INSERT/UPDATE/DELETE/DROP allowed |
| - Parametrized queries: no string formatting |
| - Max 1000 rows, 30s timeout |
| |
| Params: JSON dict for parameterized queries. |
| Example: sql="SELECT * FROM users WHERE id = :id", params='{"id": 42}' |
| """ |
| validation = safety.validate(sql) |
| if not validation.approved: |
| return f"BLOCKED: {validation.reason}" |
|
|
| try: |
| params_dict = json.loads(params) if params and params != "{}" else None |
| except json.JSONDecodeError: |
| return "BLOCKED: Invalid JSON in params" |
|
|
| safe_sql = safety.ensure_limit(sql) |
| result = await engine.execute_query(safe_sql, params_dict) |
| return formatter.format_query_result(result) |
|
|
| @server.tool() |
| async def natural_query(question: str) -> str: |
| """ |
| Ask a question in natural language and get SQL + results. |
| |
| Use when: you don't know the exact SQL, or want a quick answer |
| about the data. Examples: "How many users have orders?", |
| "What product sells the best?" |
| |
| Translation: an LLM (OpenAI/Anthropic) is used when OPENAI_API_KEY or |
| ANTHROPIC_API_KEY is configured; otherwise a built-in rules-based |
| parser handles common English question shapes. |
| |
| Returns: generated SQL + results + explanation. |
| """ |
| tables = await engine.get_tables() |
| table_names = [t.name for t in tables] |
|
|
| translation = await translate( |
| question, |
| table_names, |
| openai_key=config.openai_key, |
| anthropic_key=config.anthropic_key, |
| ) |
|
|
| sql = translation.sql |
| if not sql: |
| reason = translation.error or "no matching pattern" |
| source_note = f" (LLM: {reason})" if translation.source == "llm" else "" |
| return ( |
| "Could not automatically translate your question to SQL" |
| + source_note |
| + ".\n\n" |
| "Try using the `query` tool directly with SQL, or rephrase your question.\n" |
| f"Available tables: {', '.join(table_names)}\n\n" |
| "Examples:\n" |
| "- 'How many users are there?'\n" |
| "- 'Show me all orders'\n" |
| "- 'What products cost more than 100?'" |
| ) |
|
|
| safe_sql = safety.ensure_limit(sql) |
| validation = safety.validate(safe_sql) |
| if not validation.approved: |
| return f"Generated query was blocked: {validation.reason}" |
|
|
| result = await engine.execute_query(safe_sql) |
| output = formatter.format_query_result(result) |
| source_note = "LLM" if translation.source == "llm" else "rules-based" |
| output = f"**Question:** {question}\n**Generated SQL** ({source_note}): `{safe_sql}`\n\n{output}" |
| return output |
|
|
| @server.tool() |
| async def profile_database(table_name: str = "") -> str: |
| """ |
| Get a complete profile of the database or a specific table. |
| |
| Use when: you need to understand the data — value distributions, |
| NULL rates, sizes, relationships. Ideal first step before writing queries. |
| |
| Without parameter: database overview (table summary, relationships, sizes). |
| With parameter: detailed table profile (distributions, null rates, top values). |
| """ |
| try: |
| db_info = await engine.get_db_info() |
| tables = await engine.get_tables() |
| relationships = await inspector.discover_relationships() |
| junction_tables = inspector.detect_junction_tables(relationships) |
|
|
| profile_data = { |
| "db_info": db_info, |
| "table_count": len(tables), |
| "tables": tables, |
| "relationships": relationships, |
| "junction_tables": junction_tables, |
| } |
|
|
| if table_name: |
| table_stats = {} |
| try: |
| stats = await engine.get_table_stats(table_name) |
| table_stats[table_name] = stats |
| except Exception: |
| pass |
| profile_data["table_stats"] = table_stats |
|
|
| return formatter.format_profile(profile_data) |
| except Exception as e: |
| return f"Error profiling database: {str(e)}" |
|
|
| @server.tool() |
| async def schema_graph() -> str: |
| """ |
| Visualize table relationships as a Mermaid ER diagram. |
| |
| Use when: you need to see how tables are connected, |
| which have foreign key relationships, and the overall DB structure. |
| Returns: Mermaid diagram definition (render in markdown). |
| """ |
| try: |
| relationships = await inspector.discover_relationships() |
| mermaid = inspector.generate_mermaid(relationships) |
| return formatter.format_schema_graph(mermaid) |
| except Exception as e: |
| return f"Error generating schema graph: {str(e)}" |
|
|
| return server |
|
|
