188 lines
10 KiB
Python
188 lines
10 KiB
Python
"""MCP-Server mit den einheitlichen Tools aller DB-Dienste.
|
||
|
||
Tools (alle nur lesend, gekennzeichnet mit readOnlyHint):
|
||
check_connection Version, sichtbare Tabellen, Schreibrechte des Nutzers
|
||
list_tables Tabellen und Views mit geschätzter Zeilenzahl und Kommentar
|
||
describe_table Spalten und die ersten Zeilen einer Tabelle
|
||
query genau eine lesende Abfrage, Zeilen begrenzt
|
||
|
||
Jeder Aufruf: Schlüssel prüfen (Middleware), Zugangsdaten aus Headern, Treiber in
|
||
einem Thread, Ergebnis als Text und strukturiert, eine Zeile im Audit-Log.
|
||
"""
|
||
|
||
import hmac
|
||
import time
|
||
from collections.abc import Callable
|
||
from typing import Annotated, Any
|
||
|
||
import anyio
|
||
from mcp.server.mcpserver import Context, MCPServer
|
||
from mcp.server.mcpserver.exceptions import ToolError
|
||
from mcp.server.transport_security import TransportSecuritySettings
|
||
from mcp.types import CallToolResult, ToolAnnotations
|
||
from pydantic import Field
|
||
from starlette.requests import Request
|
||
from starlette.responses import JSONResponse, PlainTextResponse
|
||
from starlette.types import ASGIApp, Receive, Scope, Send
|
||
|
||
from dbmcp.core import Invalid, audit, connection
|
||
from dbmcp.core.sql import check_query
|
||
from dbmcp.core.config import Settings
|
||
from dbmcp.core.driver import Driver
|
||
from dbmcp.core.result import Table, as_table
|
||
|
||
READ_ONLY = ToolAnnotations(readOnlyHint=True, destructiveHint=False, idempotentHint=True, openWorldHint=False)
|
||
PUBLIC = ("/live", "/ready")
|
||
|
||
|
||
def client_name(ctx: Context, settings: Settings) -> str:
|
||
"""Name des Schlüssels aus dem Authorization-Header (geprüft hat ihn schon die Middleware)."""
|
||
header = (ctx.headers or {}).get("authorization", "")
|
||
return settings.api_keys.get(header.removeprefix("Bearer ").strip(), "ohne-schluessel")
|
||
|
||
|
||
def build(driver: Driver, settings: Settings) -> MCPServer:
|
||
server = MCPServer(
|
||
f"db-mcp-{driver.key}", version="0.1.0",
|
||
instructions=(f"Lesender Zugriff auf eine {driver.name}-Datenbank. Erst list_tables und describe_table, "
|
||
"dann gezielt query. Für große Tabellen aggregieren statt alle Zeilen zu lesen; Zahlen nur "
|
||
"aus Abfrageergebnissen nennen und die verwendete Abfrage kurz angeben. "
|
||
f"SQL-Dialekt: {driver.dialect} Die Zugangsdaten setzt der Client als Header X-DB-*; sie "
|
||
"gehören nie in Tool-Argumente."))
|
||
|
||
async def call(ctx: Context, tool: str, work: Callable[[Any], tuple[Table, str]], sql_text: str = "") \
|
||
-> CallToolResult:
|
||
started, fields = time.monotonic(), {"tool": tool, "client": client_name(ctx, settings), "sql": sql_text}
|
||
try:
|
||
conn = connection.from_headers(ctx.headers, settings, driver.default_port)
|
||
fields["db"] = conn.describe()
|
||
|
||
def run() -> tuple[Table, str]:
|
||
with driver.session(conn) as session:
|
||
return work(session)
|
||
|
||
table, text = await anyio.to_thread.run_sync(run)
|
||
except Invalid as e:
|
||
audit.record(**fields, outcome="error", error=str(e), ms=round((time.monotonic() - started) * 1000))
|
||
raise ToolError(str(e)) from e
|
||
except Exception as e:
|
||
audit.record(**fields, outcome="error", error=type(e).__name__,
|
||
ms=round((time.monotonic() - started) * 1000))
|
||
raise ToolError("Unerwarteter Fehler im DB-Dienst.") from e
|
||
audit.record(**fields, outcome="ok", rows=len(table.rows), truncated=table.truncated,
|
||
ms=round((time.monotonic() - started) * 1000))
|
||
return table.result(text)
|
||
|
||
@server.tool(description="Verbindung prüfen: Version, sichtbare Tabellen und ob der Nutzer Schreibrechte hat "
|
||
"(empfohlen ist eine Rolle nur mit SELECT).", annotations=READ_ONLY)
|
||
async def check_connection(ctx: Context) -> CallToolResult:
|
||
def work(session) -> tuple[Table, str]:
|
||
found = session.check()
|
||
text = f"{driver.name} {found.version}, {found.tables} Tabellen und Views sichtbar."
|
||
if found.writable:
|
||
text += (f" Hinweis: Der Nutzer hat Schreibrechte auf {found.writable} Objekte; empfohlen ist eine "
|
||
"Rolle nur mit SELECT. Abfragen laufen trotzdem schreibgeschützt.")
|
||
if found.note:
|
||
text += f" Hinweis: {found.note}"
|
||
return Table(["version", "tables", "writable"], [[found.version, found.tables, found.writable]]), text
|
||
return await call(ctx, "check_connection", work)
|
||
|
||
@server.tool(description="Tabellen und Views der Datenbank auflisten, mit geschätzter Zeilenzahl und "
|
||
"Kommentar.", annotations=READ_ONLY)
|
||
async def list_tables(ctx: Context) -> CallToolResult:
|
||
def work(session) -> tuple[Table, str]:
|
||
table = session.tables(settings.max_tables)
|
||
if not table.rows:
|
||
return table, "Keine Tabellen sichtbar."
|
||
lines = []
|
||
for schema, name, kind, estimate, comment in table.rows:
|
||
size = f", ca. {estimate} Zeilen" if estimate is not None and estimate >= 0 else ""
|
||
lines.append(f"- {schema}.{name} ({'View' if kind == 'view' else 'Tabelle'}{size})"
|
||
+ (f": {comment}" if comment else ""))
|
||
if table.truncated:
|
||
lines.append(f"… weitere Tabellen; nur die ersten {settings.max_tables} gezeigt.")
|
||
return table, "\n".join(lines)
|
||
return await call(ctx, "list_tables", work)
|
||
|
||
@server.tool(description="Spalten, Typen und die ersten Zeilen einer Tabelle oder View zeigen.",
|
||
annotations=READ_ONLY)
|
||
async def describe_table(
|
||
ctx: Context,
|
||
table: Annotated[str, Field(min_length=1, max_length=200,
|
||
description="Tabelle oder View, gern mit Schema: schema.tabelle")],
|
||
) -> CallToolResult:
|
||
def work(session) -> tuple[Table, str]:
|
||
columns, sample = session.describe(table, settings.sample_rows)
|
||
lines = "\n".join(f"- {name}: {kind}{'' if nullable else ' (nie leer)'}" + (f" – {comment}" if comment
|
||
else "") for name, kind, nullable, comment in columns.rows)
|
||
text = (f"{table}\nSpalten:\n{lines}\n\nErste {len(sample.rows)} Zeilen:\n"
|
||
f"{as_table(sample.columns, sample.rows, settings.max_cell_chars)}")
|
||
columns.extra = {"sample": sample.structured()}
|
||
return columns, text
|
||
return await call(ctx, "describe_table", work, table)
|
||
|
||
@server.tool(description=f"Datenbank mit genau einer lesenden SQL-Abfrage ({driver.name}) auswerten. Läuft in "
|
||
"einer schreibgeschützten Transaktion; Ergebnisse sind auf wenige hundert Zeilen "
|
||
"begrenzt.", annotations=READ_ONLY)
|
||
async def query(
|
||
ctx: Context,
|
||
sql: Annotated[str, Field(min_length=1, max_length=settings.max_sql_chars,
|
||
description="Genau eine lesende Abfrage, beginnend mit "
|
||
+ " oder ".join(s.upper() for s in driver.starts) + ".")],
|
||
) -> CallToolResult:
|
||
def work(session) -> tuple[Table, str]:
|
||
checked = check_query(sql, driver.starts, settings.max_sql_chars)
|
||
table = session.query(checked, settings.max_rows)
|
||
text = as_table(table.columns, table.rows, settings.max_cell_chars) if table.columns else "(kein Ergebnis)"
|
||
text += f"\n\n{len(table.rows)} Zeile{'n' if len(table.rows) != 1 else ''}"
|
||
if table.truncated:
|
||
text += (f", gekürzt auf {settings.max_rows}; für vollständige Aussagen aggregieren (COUNT, "
|
||
"GROUP BY) statt Zeilen zu lesen")
|
||
return table, text + "."
|
||
return await call(ctx, "query", work, sql)
|
||
|
||
@server.custom_route("/live", methods=["GET"])
|
||
async def live(request: Request) -> PlainTextResponse:
|
||
return PlainTextResponse("ok")
|
||
|
||
@server.custom_route("/ready", methods=["GET"])
|
||
async def ready(request: Request) -> JSONResponse:
|
||
return JSONResponse({"status": "ok", "driver": driver.name})
|
||
|
||
return server
|
||
|
||
|
||
class RequireKey:
|
||
"""Nur Clients mit gültigem Schlüssel (Authorization: Bearer …); Health-Checks ausgenommen."""
|
||
|
||
def __init__(self, app: ASGIApp, settings: Settings):
|
||
self.app, self.settings = app, settings
|
||
|
||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||
if scope["type"] == "http" and scope["path"] not in PUBLIC and not self.settings.insecure_no_auth:
|
||
header = dict(scope["headers"]).get(b"authorization", b"").decode(errors="replace")
|
||
secret = header.removeprefix("Bearer ").strip()
|
||
if not any(hmac.compare_digest(secret, known) for known in self.settings.api_keys):
|
||
response = JSONResponse({"error": "Schlüssel fehlt oder ist ungültig."}, status_code=401,
|
||
headers={"WWW-Authenticate": "Bearer"})
|
||
await response(scope, receive, send)
|
||
return
|
||
await self.app(scope, receive, send)
|
||
|
||
|
||
def app(driver: Driver, settings: Settings) -> ASGIApp:
|
||
security = TransportSecuritySettings(enable_dns_rebinding_protection=bool(settings.http_hosts),
|
||
allowed_hosts=list(settings.http_hosts))
|
||
# Zustandslos: jeder Aufruf steht für sich, mehrere Pods brauchen keine gemeinsame Sitzung
|
||
inner = build(driver, settings).streamable_http_app(stateless_http=True, json_response=True,
|
||
transport_security=security, host=settings.bind)
|
||
return RequireKey(inner, settings)
|
||
|
||
|
||
def serve(make_driver: Callable[[Settings], Driver]) -> None:
|
||
import uvicorn
|
||
settings = Settings.from_env()
|
||
audit.setup()
|
||
uvicorn.run(app(make_driver(settings), settings), host=settings.bind, port=settings.port, log_level="warning",
|
||
proxy_headers=True)
|
||
|