Spaces:
Running
Running
File size: 1,300 Bytes
6b5f4d7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
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()
|