server.py

# ---------------- Part 1: Imports & Environment ----------------
import os
import json
import psycopg2
import logging
from loguru import logger
from typing import Optional
from dotenv import load_dotenv
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import uvicorn

# Optional OpenAI / OpenRouter support
from openai import OpenAI

# ---------------- Load environment ----------------
load_dotenv()

# OpenRouter / OpenAI Client
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
Model = os.getenv("OPENROUTER_MODEL")

client = None
if OPENROUTER_API_KEY:
    client = OpenAI(api_key=OPENROUTER_API_KEY, base_url="https://openrouter.ai/api/v1")
else:
    logger.error("❌ OpenRouter API key not configured properly")
    client = None

# ---------------- Database Config ----------------
DB_NAME = os.getenv("DB_NAME")
DB_USER = os.getenv("DB_USER")
DB_PASSWORD = os.getenv("DB_PASSWORD")
DB_HOST = os.getenv("DB_HOST")
DB_PORT = os.getenv("DB_PORT")


def get_connection():
    """Return a new PostgreSQL connection"""
    return psycopg2.connect(
        dbname=DB_NAME, user=DB_USER, password=DB_PASSWORD, host=DB_HOST, port=DB_PORT
    )


# Set up basic logging
logger = logging.getLogger(__name__)
if not logger.hasHandlers():
    logging.basicConfig(level=logging.INFO)

# ---------------- Part 2: Utility Functions ----------------


def clean_output(raw: str) -> str:
    """Remove ``` fences and extra quotes from GPT output"""
    raw = raw.strip()

    # Remove ``` fences
    if raw.startswith("```"):
        raw = raw.split("```", 1)[-1]
        raw = raw.split("```")[0]

    # Remove surrounding quotes if present
    if (raw.startswith('"') and raw.endswith('"')) or (
        raw.startswith("'") and raw.endswith("'")
    ):
        raw = raw[1:-1]

    return raw.strip()


# ---------------- Part 3: SQL Execution Functions ----------------


def execute_select_query(sql):
    """Execute SELECT queries safely"""
    if not sql.strip().lower().startswith("select"):
        return {
            "status": "error",
            "reason": "Only SELECT queries are allowed for data retrieval.",
            "query": sql,
        }

    conn, cursor = None, None
    try:
        conn = get_connection()
        cursor = conn.cursor()
        cursor.execute(sql)

        if cursor.description:
            columns = [desc[0] for desc in cursor.description]
            rows = cursor.fetchall()
            if not rows:
                return {
                    "status": "error",
                    "reason": "No matching records found.",
                    "query": sql,
                }
            return [dict(zip(columns, row)) for row in rows]
        else:
            return {
                "status": "error",
                "reason": "Query did not return rows.",
                "query": sql,
            }

    except Exception as e:
        logger.error(f"SQL execution error: {e}")
        return {
            "status": "error",
            "reason": "SQL execution failed.",
            "error_type": type(e).__name__,
            "error_message": str(e),
            "query": sql,
        }
    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()


def execute_edit_query(sql):
    """Execute INSERT/UPDATE/DELETE/DDL queries"""
    conn, cursor = None, None
    try:
        conn = get_connection()
        cursor = conn.cursor()
        cursor.execute(sql)
        conn.commit()

        query_type = sql.strip().split()[0].upper()
        return {
            "status": "success",
            "message": f"{query_type} executed successfully.",
            "rows_affected": cursor.rowcount,
            "query": sql,
        }

    except Exception as e:
        import traceback

        tb = traceback.format_exc()
        reason = "Unknown error"

        import psycopg2

        if isinstance(e, psycopg2.ProgrammingError):
            reason = "Syntax error or missing table/column"
        elif isinstance(e, psycopg2.IntegrityError):
            reason = "Constraint violation"
        elif isinstance(e, psycopg2.OperationalError):
            reason = "Connection or network issue"
        elif isinstance(e, psycopg2.DataError):
            reason = "Data type mismatch"
        elif isinstance(e, psycopg2.InternalError):
            reason = "Internal DB error"
        elif isinstance(e, psycopg2.NotSupportedError):
            reason = "Unsupported SQL operation"
        elif isinstance(e, psycopg2.errors.InsufficientPrivilege):
            reason = "Permission denied"

        return {
            "status": "error",
            "reason": reason,
            "error_type": type(e).__name__,
            "error_message": str(e),
            "traceback": tb,
            "query": sql,
        }
    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()


# ---------------- Part 4: NL-to-SQL Generation ----------------


def generate_sql(query: str):
    """Convert NL query to SQL using OpenRouter/OpenAI"""
    if client is None:
        return {"error": "OpenRouter client not configured"}

    # Fetch schema info for 'mcp'
    conn, cursor = None, None
    try:
        conn = get_connection()
        cursor = conn.cursor()
        cursor.execute(
            """
            SELECT table_schema, table_name, column_name
            FROM information_schema.columns
            WHERE table_schema = 'mcp'
            ORDER BY table_schema, table_name, ordinal_position;
            """
        )
        rows = cursor.fetchall()
        schema_info = {}
        for schema, table, column in rows:
            schema_info.setdefault(schema, {}).setdefault(table, []).append(column)
    except Exception as e:
        logger.error(f"Error fetching schema: {e}")
        return {"error": "Failed to fetch schema."}
    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()
        
    system_prompt = (
        f"You are an expert SQL generator. Use only schema 'mcp'.\n"
    )
    user_prompt = (
        "You are an expert SQL generator.\n"
        "Use ONLY the schema 'mcp'.\n"
        f"Schema information (tables, columns, relationships): {json.dumps(schema_info, indent=2)}\n"
        "Instructions:\n"
        "1. Generate valid PostgreSQL SQL for the given natural language query.\n"
        "2. Use fully qualified table and column names when necessary.\n"
        "3. Return SQL in a single line ending with a semicolon.\n"
        "4. Do NOT include quotes, markdown, explanations, or extra characters.\n"
        "5. Handle complex queries including JOINs, aggregations, GROUP BY, and WHERE clauses.\n"
        "Natural language query will follow. Generate only executable SQL.\n"
        "6. Always generate case-insensitive SQL queries by using the ILIKE operator instead of = for text comparisons (e.g., names, emails, cities, etc.).\n"
    )

    try:
        resp = client.chat.completions.create(
            model=Model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
                {"role": "user", "content": f"Natural language query: {query}\nSQL:"},
            ],
            temperature=0,
            max_tokens=800,
        )
        sql_query = resp.choices[0].message.content
        return clean_output(sql_query)

    except Exception as e:
        logger.error(f"LLM SQL generation error: {e}")
        return {"error": f"Unable to generate SQL: {str(e)}"}


# ---------------- Part 5: FastAPI App ----------------

rest_app = FastAPI(title="Postgres MCP Server - HTTP Tools")


# Request Models
class NLToSQLRequest(BaseModel):
    query: Optional[str] = None
    nl: Optional[str] = None
    natural_language_query: Optional[str] = None


class SQLRequest(BaseModel):
    sql_query: Optional[str] = None
    query: Optional[str] = None
    sql: Optional[str] = None


# Tools Endpoints
@rest_app.post("/tools/nl_to_sql_tool")
async def rest_nl_to_sql(req: NLToSQLRequest):
    text = req.query or req.nl or req.natural_language_query
    if not text:
        return {"error": "No input provided"}
    sql = generate_sql(text)
    return {"result": sql}


@rest_app.post("/tools/query_retrieve_tool")
async def rest_query_retrieve(req: SQLRequest):
    sql = req.sql_query or req.query or req.sql
    if not sql:
        return {"error": "No SQL query provided"}
    return {"result": execute_select_query(sql)}


@rest_app.post("/tools/query_edit_tool")
async def rest_query_edit(req: SQLRequest):
    sql = req.sql_query or req.query or req.sql
    if not sql:
        return {"error": "No SQL query provided"}
    return {"result": execute_edit_query(sql)}


# List available tools
@rest_app.get("/tools")
def list_tools():
    return {
        "tools": [
            {
                "name": "nl_to_sql_tool",
                "description": "Convert natural language to SQL for 'mcp' schema",
                "params": {"query": "string"},
            },
            {
                "name": "query_retrieve_tool",
                "description": "Execute SQL SELECT query on 'mcp'",
                "params": {"query": "string"},
            },
            {
                "name": "query_edit_tool",
                "description": "Execute INSERT/UPDATE/DELETE on 'mcp'",
                "params": {"query": "string"},
            },
        ]
    }


@rest_app.get("/ping")
def ping():
    return {"status": "ok"}


if __name__ == "__main__":
    uvicorn.run(rest_app, host="0.0.0.0", port=8000)