32 lines
588 B
Python
32 lines
588 B
Python
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
|
|
}
|