code.txtā¢2.03 kB
This code is only for the stdio
import psycopg2
import json
from loguru import logger
from mcp.server.fastmcp import FastMCP
# Create an MCP server
mcp = FastMCP("Postgres_Demo")
@mcp.tool()
def query_data(sql_query: str) -> str:
"""
Execute PostgreSQL queries for the litellm database.
- SELECT queries: return fetched rows
- Non-SELECT queries: commit changes and return success message
"""
logger.info(f"Executing SQL query: {sql_query}")
DB_NAME = "litellm"
DB_USER = "postgres"
DB_PASS = "postgres"
DB_HOST = "localhost"
DB_PORT = "5432"
conn = None
cursor = None
result = {}
try:
logger.info(f"Connecting to database '{DB_NAME}'...")
conn = psycopg2.connect(
dbname=DB_NAME,
user=DB_USER,
password=DB_PASS,
host=DB_HOST,
port=DB_PORT
)
logger.info("Connection successful!")
cursor = conn.cursor()
cursor.execute(sql_query)
# If it's a SELECT, fetch results
if sql_query.strip().lower().startswith("select"):
rows = cursor.fetchall()
result = {
"status": "success",
"type": "select",
"row_count": len(rows),
"data": rows
}
else:
# For CREATE, INSERT, UPDATE, DELETE, etc.
conn.commit()
result = {
"status": "success",
"type": "write",
"message": "Query executed and committed successfully."
}
except Exception as e:
logger.error(f"Database error: {e}")
result = {"status": "error", "message": str(e)}
finally:
if cursor:
cursor.close()
if conn:
conn.close()
logger.info("Database connection closed.")
return json.dumps(result, indent=2)
if __name__ == "__main__":
logger.info("Starting Postgres MCP server...")
mcp.run(transport="stdio")