server.pyā¢3.22 kB
#Hello! This script sets up a simple MCP server that can execute PostgreSQL queries on a database. It supports both SELECT and non-SELECT queries, returning results or success messages accordingly. The server can run in either stdio or HTTP mode.
# Make sure to have the required packages installed: psycopg2, loguru, and mcp via uv sync.
# You can run the server using the command: python server.py --mode http --port 9000
# Adjust the database connection parameters as needed.
# Ensure your PostgreSQL server is running and accessible.
# Enjoy querying!
#It is recommended to run this script in a virtual environment to manage dependencies effectively.
#it is created by Loki for educational purposes.
# for contact:
# email: logeshwaranshanmugam21@gmail.com
# linkedin: https://www.linkedin.com/in/logeshwaran-shanmugam
import psycopg2
import json
from loguru import logger
from mcp.server.fastmcp import FastMCP
# Create an MCP server
mcp = FastMCP("PostgresMCP", "1.0")
@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 sql_query.strip().lower().startswith("select"):
rows = cursor.fetchall()
result = {
"status": "success",
"type": "select",
"row_count": len(rows),
"data": rows
}
else:
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__":
import argparse
parser = argparse.ArgumentParser(description="Run Postgres MCP server")
parser.add_argument("--mode", choices=["stdio", "http"], default="stdio", help="Transport mode")
parser.add_argument("--port", type=int, default=8000, help="Port for HTTP mode")
args = parser.parse_args()
logger.info(f"Starting Postgres MCP server in {args.mode.upper()} mode...")
if args.mode == "http":
mcp.run(transport="streamable-http")
else:
mcp.run(transport="stdio")
# python server.py --mode http --port 9000