from mcp.server import Server
from mcp.types import Tool, TextContent
import json
import subprocess
import os
import asyncio
server = Server("dbtools")
# PostgreSQL Configuration - UPDATE THESE VALUES
PG_USER = "postgres"
PG_PASSWORD = "root" # Change this to your PostgreSQL password
PG_HOST = "localhost"
PG_PORT = "5432"
# Project Configuration
DB_NAME = "sample_project_db"
ENV_PATH = "/home/chaitanyaphani/monorepo/Altiushub-App-Backend/.env"
MANAGE_PY = "/home/chaitanyaphani/monorepo/Altiushub-App-Backend/manage.py"
PYTHON_EXEC = "/home/chaitanyaphani/monorepo/Altiushub-App-Backend/venv/bin/python" # Virtual environment Python
# List available tools
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="create_database",
description="Create a new PostgreSQL database",
inputSchema={
"type": "object",
"properties": {
"db_name": {
"type": "string",
"description": "Name of the database to create",
"default": "sample_project_db"
}
},
"required": []
}
),
Tool(
name="enable_hstore",
description="Enable hstore extension in the database",
inputSchema={
"type": "object",
"properties": {
"db_name": {
"type": "string",
"description": "Name of the database to enable hstore in",
"default": "sample_project_db"
}
},
"required": []
}
),
Tool(
name="update_env",
description="Update the .env file with the database name",
inputSchema={
"type": "object",
"properties": {
"db_name": {
"type": "string",
"description": "Database name to set in the .env file",
"default": "sample_project_db"
}
},
"required": []
}
),
Tool(
name="django",
description="Run a Django management command",
inputSchema={
"type": "object",
"properties": {
"cmd": {
"type": "string",
"description": "The Django command to run (e.g., 'migrate', 'makemigrations')"
}
},
"required": ["cmd"]
}
)
]
# Handle tool execution
@server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
if name == "create_database":
import psycopg2
db_name = arguments.get("db_name", DB_NAME).lower() # Convert to lowercase
conn = psycopg2.connect(dbname="postgres", user=PG_USER, password=PG_PASSWORD, host=PG_HOST)
conn.autocommit = True
cur = conn.cursor()
cur.execute(f"CREATE DATABASE {db_name};")
conn.close()
return [TextContent(type="text", text=f"Database {db_name} created.")]
elif name == "enable_hstore":
import psycopg2
db_name = arguments.get("db_name", DB_NAME).lower() # Convert to lowercase
conn = psycopg2.connect(dbname=db_name, user=PG_USER, password=PG_PASSWORD, host=PG_HOST)
conn.autocommit = True
cur = conn.cursor()
cur.execute("CREATE EXTENSION IF NOT EXISTS hstore;")
conn.close()
return [TextContent(type="text", text=f"hstore extension enabled in {db_name}.")]
elif name == "update_env":
db_name = arguments.get("db_name", DB_NAME).lower() # Convert to lowercase
with open(ENV_PATH, "r") as f:
lines = f.readlines()
new_lines = []
updated = False
for line in lines:
if line.startswith("POSTGRES_DB_NAME=") and not line.strip().startswith("#"):
new_lines.append(f"POSTGRES_DB_NAME={db_name}\n")
updated = True
else:
new_lines.append(line)
with open(ENV_PATH, "w") as f:
f.writelines(new_lines)
if updated:
return [TextContent(type="text", text=f".env updated: POSTGRES_DB_NAME={db_name}")]
else:
return [TextContent(type="text", text="Warning: POSTGRES_DB_NAME not found in .env")]
elif name == "django":
from dotenv import dotenv_values
cmd = arguments.get("cmd", "")
# Load environment variables from .env file
env = os.environ.copy()
if os.path.exists(ENV_PATH):
cfg = dotenv_values(ENV_PATH)
for k, v in cfg.items():
if v is not None:
env[str(k)] = str(v)
# Run the Django command with the loaded environment
workdir = os.path.dirname(os.path.abspath(MANAGE_PY))
result = subprocess.run(
[PYTHON_EXEC, MANAGE_PY] + cmd.split(),
cwd=workdir,
env=env,
capture_output=True,
text=True
)
# Combine stdout and stderr for complete output
output = ""
if result.stdout:
output += result.stdout
if result.stderr:
output += "\n" + result.stderr
if result.returncode != 0:
output = f"ERROR (exit code {result.returncode}):\n{output}"
return [TextContent(type="text", text=output if output else "Command completed with no output")]
else:
return [TextContent(type="text", text=f"Unknown tool: {name}")]
# Start the MCP server
async def main():
from mcp.server.stdio import stdio_server
async with stdio_server() as (read_stream, write_stream):
await server.run(
read_stream,
write_stream,
server.create_initialization_options()
)
if __name__ == "__main__":
asyncio.run(main())
