db_utils.py

import os
import psycopg2
from sqlalchemy import create_engine, text
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
SUPABASE_CONN_STR = os.getenv("SUPABASE_CONN_STR")

def get_db_connection():
    if not SUPABASE_CONN_STR:
        raise ValueError("SUPABASE_CONN_STR not set in .env file")
    return psycopg2.connect(SUPABASE_CONN_STR)

def get_schema():
    engine = create_engine(SUPABASE_CONN_STR)
    with engine.connect() as conn:
        result = conn.execute(text("""
            SELECT table_name, column_name, data_type
            FROM information_schema.columns
            WHERE table_schema = 'public'
            ORDER BY table_name, ordinal_position;
        """))
        schema_str = "\n".join([f"{row[0]}.{row[1]} ({row[2]})" for row in result])
    return f"Database schema:\n{schema_str}"

def execute_sql(sql):
    if any(keyword in sql.lower() for keyword in ['drop', 'delete', 'truncate']):
        raise ValueError("Potentially dangerous SQL query detected")
    conn = get_db_connection()
    cur = conn.cursor()
    try:
        cur.execute(sql)
        results = cur.fetchall()
        conn.commit()
        return results
    except Exception as e:
        conn.rollback()
        raise e
    finally:
        cur.close()
        conn.close()