File size: 8,454 Bytes
814d95f 93166d0 814d95f 93166d0 814d95f 93166d0 814d95f 93166d0 814d95f 93166d0 814d95f 93166d0 814d95f 93166d0 814d95f 93166d0 814d95f 93166d0 814d95f 93166d0 814d95f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 | """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
|