app2.py

import json
import os
from contextlib import contextmanager
from decimal import Decimal

import psycopg2
import requests
from dotenv import load_dotenv
from flask import Flask, jsonify, request
from flask_cors import CORS

# Load environment variables
load_dotenv()

# Flask app setup
app = Flask(__name__)
CORS(app)

# Gemini API Config
API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent"
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
headers = {"Content-Type": "application/json"}

# DB Config
DB_CONFIG = {
    "dbname": os.getenv("DB_NAME"),
    "user": os.getenv("DB_USER"),
    "password": os.getenv("DB_PASSWORD"),
    "host": os.getenv("DB_HOST"),
    "port": os.getenv("DB_PORT"),
}


# PostgreSQL connection context manager
@contextmanager
def get_connection():
    conn = None
    cur = None
    try:
        conn = psycopg2.connect(os.getenv("DB_URI"))
        cur = conn.cursor()
        yield conn, cur
    finally:
        if cur:
            cur.close()
        if conn:
            conn.close()


# Fetch database schema
def get_schema():
    try:
        with get_connection() as (conn, cur):
            cur.execute(
                """
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = 'public'
            """
            )
            tables = cur.fetchall()

            schema_text = "CREATE TABLE statements:\n\n"
            for table in tables:
                table_name = table[0]
                cur.execute(
                    f"""
                    SELECT 
                        c.column_name, 
                        c.data_type,
                        c.is_nullable,
                        c.column_default,
                        CASE 
                            WHEN tc.constraint_type = 'PRIMARY KEY' THEN 'PRIMARY KEY'
                            ELSE ''
                        END as key_type
                    FROM information_schema.columns c
                    LEFT JOIN information_schema.table_constraints tc
                        ON c.table_name = tc.table_name
                        AND tc.constraint_type = 'PRIMARY KEY'
                    WHERE c.table_name = '{table_name}'
                    ORDER BY c.ordinal_position;
                """
                )
                columns = cur.fetchall()

                create_table = f"CREATE TABLE {table_name} (\n"
                for col in columns:
                    name, data_type, nullable, default, key_type = col
                    create_table += f"    {name} {data_type}"
                    if key_type:
                        create_table += f" {key_type}"
                    if nullable == "NO":
                        create_table += " NOT NULL"
                    if default:
                        create_table += f" DEFAULT {default}"
                    create_table += ",\n"
                create_table = create_table.rstrip(",\n") + "\n);\n\n"
                schema_text += create_table

            return schema_text
    except Exception as e:
        print(f"Error fetching schema: {e}")
        return None


# Generate SQL query from Gemini
def get_sql_query(user_query, schema):
    try:
        prompt = f"""### PostgreSQL database schema:
{schema}

### Instructions:
- Generate a single, correct SQL query that answers the user's question.
- Include appropriate table aliases, JOIN conditions, and WHERE clauses.
- Provide the SQL query only, with no additional explanations or questions.

### User Query: {user_query}

### SQL Query:"""

        payload = {"contents": [{"parts": [{"text": prompt}]}]}
        response = requests.post(
            f"{API_URL}?key={GEMINI_API_KEY}",
            headers=headers,
            json=payload,
        )

        if response.status_code != 200:
            print(f"Error from API: {response.text}")
            return None

        response_data = response.json()
        if "candidates" not in response_data or not response_data["candidates"]:
            print("No candidates in response")
            return None

        sql_query = response_data["candidates"][0]["content"]["parts"][0][
            "text"
        ].strip()
        if sql_query.startswith("```sql") and sql_query.endswith("```"):
            sql_query = sql_query[7:-3].strip()

        if not sql_query.endswith(";"):
            sql_query += ";"

        return sql_query
    except Exception as e:
        print(f"Error generating SQL: {e}")
        return None


# Execute SQL query and return results
def execute_query(sql_query):
    try:
        with get_connection() as (conn, cur):
            cur.execute(sql_query)

            if sql_query.strip().upper().startswith("SELECT"):
                columns = [desc[0] for desc in cur.description]
                rows = cur.fetchall()
                results = [
                    {
                        columns[i]: (
                            float(value) if isinstance(value, Decimal) else value
                        )
                        for i, value in enumerate(row)
                    }
                    for row in rows
                ]
            else:
                conn.commit()
                results = {"message": "Query executed successfully"}

            return results
    except Exception as e:
        print(f"Error executing query: {e}")
        return None


# API endpoint for user queries
@app.route("/query", methods=["POST"])
def query():
    try:
        data = request.get_json()
        user_query = data.get("query")
        if not user_query:
            return jsonify({"error": "Query is required"}), 400

        schema = get_schema()
        if not schema:
            return jsonify({"error": "Failed to fetch schema"}), 500

        sql_query = get_sql_query(user_query, schema)
        if not sql_query:
            return jsonify({"error": "Failed to generate SQL query"}), 500

        results = execute_query(sql_query)
        if results is None:
            return jsonify({"error": "Failed to execute query"}), 500

        return jsonify({"sql_query": sql_query, "results": results})
    except Exception as e:
        return jsonify({"error": str(e)}), 500


# Health check endpoint
@app.route("/health", methods=["GET"])
def health():
    return jsonify({"status": "ok"})


# Run Flask app
if __name__ == "__main__":
    app.run(debug=True, host="0.0.0.0", port=4000)