39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
import os
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
from dotenv import load_dotenv
|
|
|
|
# 加载 .env 文件
|
|
load_dotenv()
|
|
|
|
DATABASE_URL = os.getenv("DATABASE_URL")
|
|
|
|
# 自动将旧版的 postgres:// 替换为 SQLAlchemy 强要求的 postgresql://
|
|
if DATABASE_URL and DATABASE_URL.startswith("postgres://"):
|
|
DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql://", 1)
|
|
elif DATABASE_URL and DATABASE_URL.startswith("postgres+"):
|
|
DATABASE_URL = DATABASE_URL.replace("postgres+", "postgresql+", 1)
|
|
|
|
# 如果环境变量彻底缺失，做一个安全的报错提示
|
|
if not DATABASE_URL:
|
|
raise ValueError("DATABASE_URL environment variable is not set or empty.")
|
|
|
|
engine = create_engine(
|
|
DATABASE_URL,
|
|
pool_recycle=3600,
|
|
pool_pre_ping=True,
|
|
connect_args={
|
|
"keepalives": 1,
|
|
"keepalives_idle": 30,
|
|
"keepalives_interval": 10,
|
|
"keepalives_count": 5
|
|
}
|
|
)
|
|
SessionLocal = sessionmaker(bind=engine)
|
|
|
|
def get_db():
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close() |