import os

from dotenv import load_dotenv
from sqlalchemy import create_engine
from sqlalchemy import text

load_dotenv()

engine = create_engine(
    os.getenv("DB_URL"),
    pool_pre_ping=True,
)

conn = engine.connect()


def is_select(sql: str) -> bool:
    return sql.lower().lstrip().startswith("select")


def execute_sql(sql: str):
    if not is_select(sql):
        raise ValueError("只允许 SELECT 查询")

    result = conn.execute(text(sql))
    rows = result.fetchall()
    columns = result.keys()
    return {
        "columns": list(columns),
        "rows": rows
    }