How to use PostgreSQL MCP server

pip install mcp psycopg[binary]

import os
from mcp.server import Server
from mcp.server.stdio import stdio_server
import psycopg

# PostgreSQL connection parameters
PG_CONN_INFO = {
    "host": os.environ.get("PG_HOST", "localhost"),
    "port": int(os.environ.get("PG_PORT", 5432)),
    "user": os.environ.get("PG_USER", "postgres"),
    "password": os.environ.get("PG_PASSWORD", ""),
    "dbname": os.environ.get("PG_DBNAME", "postgres")
}

class PostgresMCPServer(Server):
    def __init__(self):
        super().__init__()
        self.conn = psycopg.connect(**PG_CONN_INFO)

    async def handle_tool_call(self, tool_name, input_text):
        # Only support 'query' tool for SQL execution
        if tool_name != "query":
            return f"Error: Unsupported tool {tool_name}"
        try:
            with self.conn.cursor() as cur:
                cur.execute(input_text)
                if cur.description:  # SELECT query
                    rows = cur.fetchall()
                    columns = [desc.name for desc in cur.description]
                    # Format output as CSV-like string
                    header = ", ".join(columns)
                    data = "\n".join(", ".join(str(cell) for cell in row) for row in rows)
                    return f"{header}\n{data}"
                else:
                    self.conn.commit()
                    return "Query executed successfully."
        except Exception as e:
            return f"Error executing query: {e}"

if __name__ == "__main__":
    server = PostgresMCPServer()
    stdio_server(server)