z230
This commit is contained in:
+359
@@ -0,0 +1,359 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
MCP server pro PostgreSQL — používá FastMCP + psycopg v3.
|
||||
Spustit: python mcp_postgres.py
|
||||
|
||||
Host: 192.168.1.76:5432
|
||||
User: vladimir.buzalka
|
||||
Default DB: postgres (lze přepínat tool argumentem `db`)
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime, date
|
||||
from decimal import Decimal
|
||||
from typing import Optional, Union
|
||||
|
||||
import psycopg
|
||||
from psycopg.rows import dict_row
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
PG_HOST = "192.168.1.76"
|
||||
PG_PORT = 5432
|
||||
PG_USER = "vladimir.buzalka"
|
||||
PG_PASSWORD = "Vlado7309208104++"
|
||||
PG_DEFAULT_DB = "postgres"
|
||||
|
||||
WRITE_KEYWORDS = ("insert", "update", "delete", "drop", "truncate",
|
||||
"alter", "create", "grant", "revoke", "comment", "merge")
|
||||
|
||||
|
||||
def log(msg: str):
|
||||
print(msg, file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
_pool: dict[str, psycopg.Connection] = {}
|
||||
|
||||
|
||||
def get_conn(db: Optional[str] = None) -> psycopg.Connection:
|
||||
"""Lazy connection cache keyed by db name. Reconnects if closed."""
|
||||
name = db or PG_DEFAULT_DB
|
||||
conn = _pool.get(name)
|
||||
if conn is not None and not conn.closed:
|
||||
return conn
|
||||
conn = psycopg.connect(
|
||||
host=PG_HOST,
|
||||
port=PG_PORT,
|
||||
user=PG_USER,
|
||||
password=PG_PASSWORD,
|
||||
dbname=name,
|
||||
connect_timeout=5,
|
||||
autocommit=True,
|
||||
)
|
||||
_pool[name] = conn
|
||||
return conn
|
||||
|
||||
|
||||
def serialize(value):
|
||||
if isinstance(value, (datetime, date)):
|
||||
return value.isoformat()
|
||||
if isinstance(value, Decimal):
|
||||
return float(value)
|
||||
if isinstance(value, (bytes, memoryview)):
|
||||
try:
|
||||
return bytes(value).decode("utf-8", errors="replace")
|
||||
except Exception:
|
||||
return repr(value)
|
||||
if isinstance(value, dict):
|
||||
return {k: serialize(v) for k, v in value.items()}
|
||||
if isinstance(value, (list, tuple)):
|
||||
return [serialize(v) for v in value]
|
||||
return value
|
||||
|
||||
|
||||
def serialize_rows(rows):
|
||||
return [{k: serialize(v) for k, v in row.items()} for row in rows]
|
||||
|
||||
|
||||
def is_write(sql: str) -> bool:
|
||||
s = sql.lstrip().lower()
|
||||
return any(s.startswith(kw) for kw in WRITE_KEYWORDS)
|
||||
|
||||
|
||||
# Verify default DB on startup
|
||||
try:
|
||||
c = get_conn()
|
||||
with c.cursor() as cur:
|
||||
cur.execute("select version()")
|
||||
ver = cur.fetchone()[0]
|
||||
log(f"Connected to PostgreSQL ({PG_HOST}:{PG_PORT}) — {ver}")
|
||||
except Exception as e:
|
||||
log(f"PostgreSQL connection failed: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
mcp = FastMCP("janssen-postgres")
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def ping() -> dict:
|
||||
"""Check if PostgreSQL is reachable. Returns server version, latency, host."""
|
||||
try:
|
||||
start = time.monotonic()
|
||||
conn = get_conn()
|
||||
with conn.cursor() as cur:
|
||||
cur.execute("select version()")
|
||||
version = cur.fetchone()[0]
|
||||
latency_ms = round((time.monotonic() - start) * 1000, 1)
|
||||
return {
|
||||
"status": "ok",
|
||||
"version": version,
|
||||
"latency_ms": latency_ms,
|
||||
"host": PG_HOST,
|
||||
"port": PG_PORT,
|
||||
}
|
||||
except Exception as e:
|
||||
return {"status": "error", "error": str(e), "host": PG_HOST}
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def list_databases() -> dict:
|
||||
"""List all non-template databases on the server."""
|
||||
try:
|
||||
conn = get_conn()
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute("select datname from pg_database where datistemplate = false order by datname")
|
||||
dbs = [r["datname"] for r in cur.fetchall()]
|
||||
return {"count": len(dbs), "databases": dbs}
|
||||
except Exception:
|
||||
log(f"list_databases error: {traceback.format_exc()}")
|
||||
raise
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def list_schemas(db: Optional[str] = None) -> dict:
|
||||
"""List schemas in a database (excludes pg_* and information_schema)."""
|
||||
try:
|
||||
conn = get_conn(db)
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute("""
|
||||
select schema_name
|
||||
from information_schema.schemata
|
||||
where schema_name not like 'pg\\_%' escape '\\'
|
||||
and schema_name <> 'information_schema'
|
||||
order by schema_name
|
||||
""")
|
||||
schemas = [r["schema_name"] for r in cur.fetchall()]
|
||||
return {"db": db or PG_DEFAULT_DB, "count": len(schemas), "schemas": schemas}
|
||||
except Exception:
|
||||
log(f"list_schemas error: {traceback.format_exc()}")
|
||||
raise
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def list_tables(db: Optional[str] = None, schema: str = "public") -> dict:
|
||||
"""List tables (and views) in a schema."""
|
||||
try:
|
||||
conn = get_conn(db)
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute("""
|
||||
select table_name, table_type
|
||||
from information_schema.tables
|
||||
where table_schema = %s
|
||||
order by table_name
|
||||
""", (schema,))
|
||||
rows = cur.fetchall()
|
||||
return {
|
||||
"db": db or PG_DEFAULT_DB,
|
||||
"schema": schema,
|
||||
"count": len(rows),
|
||||
"tables": rows,
|
||||
}
|
||||
except Exception:
|
||||
log(f"list_tables error: {traceback.format_exc()}")
|
||||
raise
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def describe_table(table: str, db: Optional[str] = None, schema: str = "public") -> dict:
|
||||
"""Return columns, types, nullability, defaults, indexes and row count for a table."""
|
||||
try:
|
||||
conn = get_conn(db)
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute("""
|
||||
select column_name, data_type, is_nullable, column_default, character_maximum_length
|
||||
from information_schema.columns
|
||||
where table_schema = %s and table_name = %s
|
||||
order by ordinal_position
|
||||
""", (schema, table))
|
||||
columns = cur.fetchall()
|
||||
|
||||
cur.execute("""
|
||||
select indexname, indexdef
|
||||
from pg_indexes
|
||||
where schemaname = %s and tablename = %s
|
||||
order by indexname
|
||||
""", (schema, table))
|
||||
indexes = cur.fetchall()
|
||||
|
||||
qualified = f'"{schema}"."{table}"'
|
||||
try:
|
||||
cur.execute(f"select count(*) as c from {qualified}")
|
||||
row_count = cur.fetchone()["c"]
|
||||
except Exception as e:
|
||||
row_count = f"error: {e}"
|
||||
|
||||
return {
|
||||
"db": db or PG_DEFAULT_DB,
|
||||
"schema": schema,
|
||||
"table": table,
|
||||
"row_count": row_count,
|
||||
"columns": columns,
|
||||
"indexes": indexes,
|
||||
}
|
||||
except Exception:
|
||||
log(f"describe_table error: {traceback.format_exc()}")
|
||||
raise
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def query(
|
||||
sql: str,
|
||||
db: Optional[str] = None,
|
||||
params_json: Optional[Union[str, list]] = None,
|
||||
limit: int = 100,
|
||||
) -> dict:
|
||||
"""Run a READ-ONLY SQL query (SELECT / WITH / EXPLAIN / SHOW).
|
||||
For writes use `execute` (which requires `confirmed=True`).
|
||||
params_json: JSON array of positional parameters for %s placeholders, e.g. '["CZ1-01", 5]'.
|
||||
limit: cap on returned rows (max 1000). The query itself is NOT modified — limit is a Python-side cutoff.
|
||||
"""
|
||||
try:
|
||||
if is_write(sql):
|
||||
return {
|
||||
"status": "rejected",
|
||||
"reason": "Write/DDL statement detected. Use `execute` (with preview + confirmed=True).",
|
||||
}
|
||||
limit = min(max(limit, 1), 1000)
|
||||
params = json.loads(params_json) if isinstance(params_json, str) else (params_json or [])
|
||||
|
||||
conn = get_conn(db)
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(sql, params)
|
||||
if cur.description is None:
|
||||
return {"status": "ok", "rowcount": cur.rowcount, "rows": []}
|
||||
rows = cur.fetchmany(limit)
|
||||
return {
|
||||
"status": "ok",
|
||||
"db": db or PG_DEFAULT_DB,
|
||||
"row_count": len(rows),
|
||||
"limit": limit,
|
||||
"rows": serialize_rows(rows),
|
||||
}
|
||||
except Exception as e:
|
||||
log(f"query error: {traceback.format_exc()}")
|
||||
return {"status": "error", "error": str(e), "sql": sql}
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def explain(sql: str, db: Optional[str] = None, analyze: bool = False) -> dict:
|
||||
"""Return EXPLAIN (ANALYZE, BUFFERS) plan for a SELECT-style query. analyze=False is safe; True actually runs the query."""
|
||||
try:
|
||||
if is_write(sql) and analyze:
|
||||
return {"status": "rejected", "reason": "Refuse to EXPLAIN ANALYZE a write statement."}
|
||||
prefix = "explain (analyze, buffers, format json) " if analyze else "explain (format json) "
|
||||
conn = get_conn(db)
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(prefix + sql)
|
||||
plan = cur.fetchone()[0]
|
||||
return {"status": "ok", "analyze": analyze, "plan": plan}
|
||||
except Exception as e:
|
||||
log(f"explain error: {traceback.format_exc()}")
|
||||
return {"status": "error", "error": str(e)}
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def preview_execute(
|
||||
sql: str,
|
||||
db: Optional[str] = None,
|
||||
params_json: Optional[Union[str, list]] = None,
|
||||
) -> dict:
|
||||
"""Preview a write/DDL statement using a transaction + ROLLBACK. Reports rowcount that WOULD be affected.
|
||||
Always call this first and present the result to the user before calling `execute` with confirmed=True.
|
||||
"""
|
||||
try:
|
||||
params = json.loads(params_json) if isinstance(params_json, str) else (params_json or [])
|
||||
# Use a fresh, autocommit=False connection so we can ROLLBACK safely
|
||||
with psycopg.connect(
|
||||
host=PG_HOST, port=PG_PORT, user=PG_USER, password=PG_PASSWORD,
|
||||
dbname=db or PG_DEFAULT_DB, connect_timeout=5, autocommit=False,
|
||||
) as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(sql, params)
|
||||
rowcount = cur.rowcount
|
||||
conn.rollback()
|
||||
return {
|
||||
"status": "preview",
|
||||
"db": db or PG_DEFAULT_DB,
|
||||
"sql": sql,
|
||||
"params": params,
|
||||
"would_affect_rows": rowcount,
|
||||
"note": "Statement was executed inside a transaction and rolled back. Present this to the user and call `execute` with confirmed=True only after explicit approval.",
|
||||
}
|
||||
except Exception as e:
|
||||
log(f"preview_execute error: {traceback.format_exc()}")
|
||||
return {"status": "error", "error": str(e), "sql": sql}
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def execute(
|
||||
sql: str,
|
||||
db: Optional[str] = None,
|
||||
params_json: Optional[Union[str, list]] = None,
|
||||
confirmed: bool = False,
|
||||
) -> dict:
|
||||
"""Run a write/DDL statement (INSERT/UPDATE/DELETE/CREATE/ALTER/DROP/...).
|
||||
REQUIRES confirmed=True — first call `preview_execute` and obtain explicit user approval.
|
||||
"""
|
||||
if not confirmed:
|
||||
return {
|
||||
"status": "aborted",
|
||||
"reason": "confirmed=False. Call preview_execute first, present the impact to the user, then re-call with confirmed=True.",
|
||||
}
|
||||
try:
|
||||
params = json.loads(params_json) if isinstance(params_json, str) else (params_json or [])
|
||||
conn = get_conn(db)
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(sql, params)
|
||||
rowcount = cur.rowcount
|
||||
log(f"execute: db={db or PG_DEFAULT_DB} rowcount={rowcount}")
|
||||
return {
|
||||
"status": "ok",
|
||||
"db": db or PG_DEFAULT_DB,
|
||||
"rowcount": rowcount,
|
||||
"sql": sql,
|
||||
}
|
||||
except Exception as e:
|
||||
log(f"execute error: {traceback.format_exc()}")
|
||||
return {"status": "error", "error": str(e), "sql": sql}
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def list_extensions(db: Optional[str] = None) -> dict:
|
||||
"""List installed extensions in the database."""
|
||||
try:
|
||||
conn = get_conn(db)
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute("select extname, extversion from pg_extension order by extname")
|
||||
rows = cur.fetchall()
|
||||
return {"db": db or PG_DEFAULT_DB, "count": len(rows), "extensions": rows}
|
||||
except Exception:
|
||||
log(f"list_extensions error: {traceback.format_exc()}")
|
||||
raise
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
log("MCP PostgreSQL server started (FastMCP)")
|
||||
mcp.run()
|
||||
Reference in New Issue
Block a user